Report remote CPU and GPU memory telemetry
This commit is contained in:
@@ -82,6 +82,7 @@
|
||||
!/remote_executor_router.py
|
||||
!/remote_protocol.py
|
||||
!/remote_plan_types.py
|
||||
!/resource_telemetry.py
|
||||
!/r2_cache.py
|
||||
!/r2_credentials.py
|
||||
!/runtime_environment.py
|
||||
|
||||
@@ -242,7 +242,7 @@ Remote state is shown directly on the canvas, in both the legacy LiteGraph rende
|
||||
|
||||
Subgraph containers inherit a derived summary decoration (marked with a `Σ` badge) whenever any nested descendant runs remotely, using the most operationally important descendant state for the border color.
|
||||
|
||||
Active remote nodes show a small progress panel with a Modal, SSH terminal, or Vast.ai icon beside every bar, per-bar `it/s` rates, human-readable transferred/total byte counts and `B/s` through `TB/s` rates for large input and output transfers, and `tok/s`, token counts, and TTFT for LLM nodes. Modal uploads pulse until the SDK reports its first remote event; large Modal results travel through the durable volume so their download has exact byte progress. SSH and Vast input/output frames stream through their relays in bounded chunks and report exact byte progress rather than buffering the complete payload first. LLM preparation is reported on the actual LLM node—not the remote component's representative node—with distinct metadata inspection, snapshot lookup, concurrent-download wait, download preparation, transfer, cached, and ready states. When a prompt is scheduled across more than one remote environment, the panel also names where each component is running — the Modal container id or the SSH host name. Preview images and UI payloads from remote nodes stream back into the local session while the component is still running, and progress state is reconciled against queue history when the browser regains focus so cancelled prompts do not leave stale bars behind.
|
||||
Active remote nodes show a small progress panel with a Modal, SSH terminal, or Vast.ai icon beside every bar, per-bar `it/s` rates, human-readable transferred/total byte counts and `B/s` through `TB/s` rates for large input and output transfers, and `tok/s`, token counts, and TTFT for LLM nodes. Modal uploads pulse until the SDK reports its first remote event; large Modal results travel through the durable volume so their download has exact byte progress. SSH and Vast input/output frames stream through their relays in bounded chunks and report exact byte progress rather than buffering the complete payload first. During execution, every Modal, SSH, and Vast worker samples its process tree's CPU RAM and the memory consumed across all visible CUDA devices twice per second. The Configurator reports current and peak GPU/CPU memory for each environment, aggregates concurrently active execution sources, and leaves the peak readout visible when the workflow finishes. LLM preparation is reported on the actual LLM node—not the remote component's representative node—with distinct metadata inspection, snapshot lookup, concurrent-download wait, download preparation, transfer, cached, and ready states. When a prompt is scheduled across more than one remote environment, the panel also names where each component is running — the Modal container id or the SSH host name. Preview images and UI payloads from remote nodes stream back into the local session while the component is still running, and progress state is reconciled against queue history when the browser regains focus so cancelled prompts do not leave stale bars behind.
|
||||
|
||||
Remote planning and capacity acquisition count as active queue work even though they happen before the rewritten prompt enters ComfyUI's native execution queue. During Vast startup, the canvas receives readable updates for rental, worker-image download, SSH readiness, and runtime initialization. A provider-side terminal image-pull failure is reported immediately rather than appearing as an empty queue until the startup timeout expires. A rental that never produces an SSH-usable worker is marked unusable and destroyed after the bounded initialization wait so a later workflow cannot silently reuse it.
|
||||
|
||||
|
||||
@@ -40,6 +40,7 @@ try:
|
||||
serialize_mapping,
|
||||
serialize_node_outputs,
|
||||
)
|
||||
from .resource_telemetry import RemoteResourceTelemetrySampler
|
||||
from .settings import get_settings
|
||||
except ImportError: # pragma: no cover - flat Modal-container import.
|
||||
from cloud_comfy_bootstrap import (
|
||||
@@ -70,6 +71,7 @@ except ImportError: # pragma: no cover - flat Modal-container import.
|
||||
serialize_mapping,
|
||||
serialize_node_outputs,
|
||||
)
|
||||
from resource_telemetry import RemoteResourceTelemetrySampler
|
||||
from settings import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -263,14 +265,18 @@ def _stream_remote_payload_events(
|
||||
def execute_payload() -> None:
|
||||
"""Run compute in a worker thread while deferring Modal volume commits."""
|
||||
pending_batch: DurableObjectCommitBatch | None = None
|
||||
telemetry_sampler = RemoteResourceTelemetrySampler(publish_status)
|
||||
telemetry_sampler.start()
|
||||
try:
|
||||
with object_store.batch_commits(commit_on_exit=False) as pending_batch:
|
||||
outputs = _execute_payload_with_output_capture(payload, execute_once)
|
||||
except (
|
||||
Exception
|
||||
) as exc: # pragma: no cover - exercised through generator consumer tests.
|
||||
telemetry_sampler.stop()
|
||||
event_buffer.publish_terminal("error", (exc, pending_batch))
|
||||
else:
|
||||
telemetry_sampler.stop()
|
||||
result_bytes = len(outputs)
|
||||
pending_object_write = bool(pending_batch and pending_batch.wrote_object)
|
||||
logger.info(
|
||||
|
||||
@@ -1,5 +1,12 @@
|
||||
# Learnings
|
||||
|
||||
## 2026-08-30: Resource telemetry belongs inside the shared execution kernel
|
||||
|
||||
- Modal, SSH, and Vast converge on the same streamed cloud execution surface inside their workers. Sampling there produces one telemetry schema and cadence for every provider while provider clients remain transport adapters.
|
||||
- CPU memory must include recursive child processes: inference engines frequently move substantial allocations into spawned workers, so sampling only the Python coordinator's RSS materially understates use. Visible-device used memory (`total - free`) likewise captures CUDA allocations outside PyTorch's allocator.
|
||||
- Environment-level telemetry needs a source identity. Aggregate current usage only across active component/location sources, retain the maximum observed aggregate and per-source peak, and mark the terminal sample inactive so a finished worker does not remain in the live total.
|
||||
- Periodic telemetry can crowd important status events out of a bounded browser-refocus buffer. Coalesce samples by prompt, environment, location, and component while retaining the newest cumulative peak for each source.
|
||||
|
||||
## 2026-08-30: Progress requires ownership of the transfer loop
|
||||
|
||||
- A provider SDK can expose a remote function as finished while its client is still materializing a large return value. Status events around the call cannot measure that hidden transfer; return a small durable object reference and read the object through an explicit chunk iterator instead.
|
||||
|
||||
@@ -110,10 +110,44 @@ def record_modal_ui_event(
|
||||
client_id,
|
||||
deque(maxlen=_MODAL_UI_EVENT_LIMIT_PER_CLIENT),
|
||||
)
|
||||
_discard_replaced_telemetry_event_locked(client_events, event_record)
|
||||
client_events.append(event_record)
|
||||
_prune_modal_ui_events_locked(client_events)
|
||||
|
||||
|
||||
def _discard_replaced_telemetry_event_locked(
|
||||
client_events: deque[dict[str, Any]],
|
||||
event_record: Mapping[str, Any],
|
||||
) -> None:
|
||||
"""Retain only the newest replay sample for one prompt execution source."""
|
||||
if event_record.get("event") != "modal_telemetry":
|
||||
return
|
||||
payload = event_record.get("payload")
|
||||
if not isinstance(payload, Mapping):
|
||||
return
|
||||
replacement_key = (
|
||||
str(payload.get("prompt_id") or ""),
|
||||
str(payload.get("execution_environment_id") or ""),
|
||||
str(payload.get("execution_location") or ""),
|
||||
str(payload.get("component_id") or ""),
|
||||
)
|
||||
for existing_record in tuple(client_events):
|
||||
existing_payload = existing_record.get("payload")
|
||||
if existing_record.get("event") != "modal_telemetry" or not isinstance(
|
||||
existing_payload, Mapping
|
||||
):
|
||||
continue
|
||||
existing_key = (
|
||||
str(existing_payload.get("prompt_id") or ""),
|
||||
str(existing_payload.get("execution_environment_id") or ""),
|
||||
str(existing_payload.get("execution_location") or ""),
|
||||
str(existing_payload.get("component_id") or ""),
|
||||
)
|
||||
if existing_key == replacement_key:
|
||||
client_events.remove(existing_record)
|
||||
return
|
||||
|
||||
|
||||
def modal_ui_events_for_client(client_id: str | None) -> list[dict[str, Any]]:
|
||||
"""Return recent Modal UI events for one websocket client."""
|
||||
if not client_id:
|
||||
|
||||
@@ -150,6 +150,57 @@ def _emit_local_remote_dispatch_status(payload: dict[str, Any]) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _emit_local_remote_resource_telemetry(
|
||||
payload: Mapping[str, Any],
|
||||
telemetry: Mapping[str, Any],
|
||||
*,
|
||||
modal_task_id: str | None = None,
|
||||
) -> None:
|
||||
"""Forward one remote memory sample to the Configurator UI."""
|
||||
extra_data = payload.get("extra_data") or {}
|
||||
client_id = (
|
||||
str(extra_data["client_id"])
|
||||
if isinstance(extra_data, Mapping) and extra_data.get("client_id") is not None
|
||||
else None
|
||||
)
|
||||
if client_id is None:
|
||||
return
|
||||
prompt_server = _lookup_local_prompt_server()
|
||||
if prompt_server is None:
|
||||
return
|
||||
event_payload: dict[str, Any] = {
|
||||
"event_type": "resource_telemetry",
|
||||
"prompt_id": (
|
||||
str(payload["prompt_id"])
|
||||
if payload.get("prompt_id") is not None
|
||||
else None
|
||||
),
|
||||
"component_id": str(payload.get("component_id") or "payload"),
|
||||
"node_ids": [
|
||||
str(node_id) for node_id in payload.get("component_node_ids", [])
|
||||
],
|
||||
"active": bool(telemetry.get("active", True)),
|
||||
**_remote_execution_identity(payload, modal_task_id),
|
||||
}
|
||||
for field_name in (
|
||||
"sample_sequence",
|
||||
"cpu_memory_bytes",
|
||||
"cpu_memory_peak_bytes",
|
||||
"cpu_memory_total_bytes",
|
||||
"gpu_memory_bytes",
|
||||
"gpu_memory_peak_bytes",
|
||||
"gpu_memory_total_bytes",
|
||||
):
|
||||
field_value = telemetry.get(field_name)
|
||||
if isinstance(field_value, int) and not isinstance(field_value, bool):
|
||||
event_payload[field_name] = max(0, field_value)
|
||||
sampled_at = telemetry.get("sampled_at")
|
||||
if isinstance(sampled_at, int | float) and not isinstance(sampled_at, bool):
|
||||
event_payload["sampled_at"] = float(sampled_at)
|
||||
_record_local_modal_ui_event("modal_telemetry", event_payload, client_id)
|
||||
prompt_server.send_sync("modal_telemetry", event_payload, client_id)
|
||||
|
||||
|
||||
def _emit_local_remote_startup_status(
|
||||
payload: Mapping[str, Any],
|
||||
*,
|
||||
|
||||
@@ -20,6 +20,7 @@ from .local_ui_events import (
|
||||
_emit_local_modal_status,
|
||||
_emit_local_preview_boundary_output,
|
||||
_emit_local_preview_image,
|
||||
_emit_local_remote_resource_telemetry,
|
||||
_progress_stream_event_metadata,
|
||||
_remote_execution_destination,
|
||||
_remote_execution_identity,
|
||||
@@ -509,6 +510,7 @@ def _handle_progress_stream_event(
|
||||
"preview": _handle_preview_event,
|
||||
"boundary_output": _handle_boundary_output_event,
|
||||
"node_cached": _handle_cached_node_event,
|
||||
"resource_telemetry": _handle_resource_telemetry_event,
|
||||
}
|
||||
handler = handlers.get(event_type)
|
||||
if handler is not None:
|
||||
@@ -517,6 +519,18 @@ def _handle_progress_stream_event(
|
||||
_handle_status_event(state, stream_event)
|
||||
|
||||
|
||||
def _handle_resource_telemetry_event(
|
||||
state: _PayloadStreamState,
|
||||
stream_event: dict[str, Any],
|
||||
) -> None:
|
||||
"""Forward one provider-neutral remote memory sample to the local UI."""
|
||||
_emit_local_remote_resource_telemetry(
|
||||
state.payload,
|
||||
stream_event,
|
||||
modal_task_id=state.active_remote_task_id,
|
||||
)
|
||||
|
||||
|
||||
def _handle_result_event(
|
||||
state: _PayloadStreamState,
|
||||
stream_event: dict[str, Any],
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
"""Backend-neutral CPU and GPU memory telemetry for remote execution."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import logging
|
||||
import resource
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Callable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_RESOURCE_SAMPLE_INTERVAL_SECONDS = 0.5
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResourceMemorySnapshot:
|
||||
"""Capture one process-RAM and visible-GPU memory observation."""
|
||||
|
||||
cpu_memory_bytes: int
|
||||
cpu_memory_total_bytes: int
|
||||
gpu_memory_bytes: int
|
||||
gpu_memory_total_bytes: int
|
||||
|
||||
|
||||
class RemoteResourceTelemetrySampler:
|
||||
"""Sample remote memory on a daemon thread and publish cumulative peaks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
callback: Callable[[dict[str, Any]], None],
|
||||
*,
|
||||
interval_seconds: float = _RESOURCE_SAMPLE_INTERVAL_SECONDS,
|
||||
sample_reader: Callable[[], ResourceMemorySnapshot] | None = None,
|
||||
) -> None:
|
||||
"""Initialize a sampler with an injectable reader for deterministic tests."""
|
||||
self._callback = callback
|
||||
self._interval_seconds = max(0.05, float(interval_seconds))
|
||||
self._sample_reader = sample_reader or read_resource_memory_snapshot
|
||||
self._stop_event = threading.Event()
|
||||
self._thread: threading.Thread | None = None
|
||||
self._sample_sequence = 0
|
||||
self._cpu_memory_peak_bytes = 0
|
||||
self._gpu_memory_peak_bytes = 0
|
||||
self._sampling_error_logged = False
|
||||
|
||||
def start(self) -> None:
|
||||
"""Publish an initial sample and start periodic collection."""
|
||||
if self._thread is not None:
|
||||
return
|
||||
self._publish_sample(active=True)
|
||||
self._thread = threading.Thread(
|
||||
target=self._sample_until_stopped,
|
||||
name="remote-resource-telemetry",
|
||||
daemon=True,
|
||||
)
|
||||
self._thread.start()
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop periodic collection and publish one terminal peak sample."""
|
||||
self._stop_event.set()
|
||||
if self._thread is not None:
|
||||
self._thread.join(timeout=max(1.0, self._interval_seconds * 2.0))
|
||||
self._publish_sample(active=False)
|
||||
|
||||
def _sample_until_stopped(self) -> None:
|
||||
"""Publish observations at a bounded cadence until execution finishes."""
|
||||
while not self._stop_event.wait(self._interval_seconds):
|
||||
self._publish_sample(active=True)
|
||||
|
||||
def _publish_sample(self, *, active: bool) -> None:
|
||||
"""Read, peak, and publish one transport-safe telemetry event."""
|
||||
try:
|
||||
snapshot = self._sample_reader()
|
||||
except (OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
if not self._sampling_error_logged:
|
||||
logger.warning("Remote resource telemetry sampling failed: %s", exc)
|
||||
self._sampling_error_logged = True
|
||||
return
|
||||
self._cpu_memory_peak_bytes = max(
|
||||
self._cpu_memory_peak_bytes,
|
||||
snapshot.cpu_memory_bytes,
|
||||
)
|
||||
self._gpu_memory_peak_bytes = max(
|
||||
self._gpu_memory_peak_bytes,
|
||||
snapshot.gpu_memory_bytes,
|
||||
)
|
||||
self._sample_sequence += 1
|
||||
try:
|
||||
self._callback(
|
||||
{
|
||||
"event_type": "resource_telemetry",
|
||||
"active": active,
|
||||
"sample_sequence": self._sample_sequence,
|
||||
"sampled_at": time.time(),
|
||||
"cpu_memory_bytes": snapshot.cpu_memory_bytes,
|
||||
"cpu_memory_peak_bytes": self._cpu_memory_peak_bytes,
|
||||
"cpu_memory_total_bytes": snapshot.cpu_memory_total_bytes,
|
||||
"gpu_memory_bytes": snapshot.gpu_memory_bytes,
|
||||
"gpu_memory_peak_bytes": self._gpu_memory_peak_bytes,
|
||||
"gpu_memory_total_bytes": snapshot.gpu_memory_total_bytes,
|
||||
}
|
||||
)
|
||||
except (RuntimeError, TypeError, ValueError) as exc:
|
||||
logger.debug("Remote resource telemetry callback was unavailable: %s", exc)
|
||||
|
||||
|
||||
def read_resource_memory_snapshot() -> ResourceMemorySnapshot:
|
||||
"""Read current process RSS and aggregate visible CUDA device memory."""
|
||||
cpu_memory_bytes, cpu_memory_total_bytes = _read_cpu_memory()
|
||||
gpu_memory_bytes, gpu_memory_total_bytes = _read_gpu_memory()
|
||||
return ResourceMemorySnapshot(
|
||||
cpu_memory_bytes=cpu_memory_bytes,
|
||||
cpu_memory_total_bytes=cpu_memory_total_bytes,
|
||||
gpu_memory_bytes=gpu_memory_bytes,
|
||||
gpu_memory_total_bytes=gpu_memory_total_bytes,
|
||||
)
|
||||
|
||||
|
||||
def _read_cpu_memory() -> tuple[int, int]:
|
||||
"""Return current process RSS and host/container RAM capacity."""
|
||||
try:
|
||||
import psutil
|
||||
except ModuleNotFoundError:
|
||||
peak_rss = int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)
|
||||
if sys.platform != "darwin":
|
||||
peak_rss *= 1024
|
||||
return max(0, peak_rss), 0
|
||||
try:
|
||||
root_process = psutil.Process()
|
||||
processes = [root_process, *root_process.children(recursive=True)]
|
||||
process_rss = 0
|
||||
for process in processes:
|
||||
try:
|
||||
process_rss += int(process.memory_info().rss)
|
||||
except (OSError, psutil.Error):
|
||||
continue
|
||||
total_memory = int(psutil.virtual_memory().total)
|
||||
except (OSError, psutil.Error) as exc:
|
||||
raise RuntimeError("Unable to read process memory usage.") from exc
|
||||
return max(0, process_rss), max(0, total_memory)
|
||||
|
||||
|
||||
def _read_gpu_memory() -> tuple[int, int]:
|
||||
"""Return used and total memory across every CUDA device visible to the worker."""
|
||||
try:
|
||||
import torch
|
||||
except ModuleNotFoundError:
|
||||
return 0, 0
|
||||
try:
|
||||
if not torch.cuda.is_available():
|
||||
return 0, 0
|
||||
used_bytes = 0
|
||||
total_bytes = 0
|
||||
for device_index in range(torch.cuda.device_count()):
|
||||
free_device_bytes, total_device_bytes = torch.cuda.mem_get_info(device_index)
|
||||
used_bytes += max(0, int(total_device_bytes) - int(free_device_bytes))
|
||||
total_bytes += max(0, int(total_device_bytes))
|
||||
return used_bytes, total_bytes
|
||||
except (AssertionError, RuntimeError) as exc:
|
||||
raise RuntimeError("Unable to read CUDA memory usage.") from exc
|
||||
@@ -631,6 +631,12 @@ def runtime_environment_module(extension_package: object) -> object:
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.runtime_environment")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def resource_telemetry_module(extension_package: object) -> object:
|
||||
"""Return the backend-neutral resource telemetry module."""
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.resource_telemetry")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def sync_engine_module(extension_package: object) -> object:
|
||||
"""Return the sync engine module."""
|
||||
|
||||
@@ -4,6 +4,17 @@ from __future__ import annotations
|
||||
|
||||
from modal_executor_test_support import * # noqa: F401,F403
|
||||
|
||||
|
||||
def _without_resource_telemetry(
|
||||
events: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return stream events unrelated to periodic memory telemetry."""
|
||||
return [
|
||||
event
|
||||
for event in events
|
||||
if event.get("event_type") != "resource_telemetry"
|
||||
]
|
||||
|
||||
def test_modal_cloud_streams_progress_and_result_events(
|
||||
modal_cloud_module: Any,
|
||||
monkeypatch: Any,
|
||||
@@ -52,7 +63,11 @@ def test_modal_cloud_streams_progress_and_result_events(
|
||||
)
|
||||
|
||||
assert progress_callbacks == [{"component_id": "component-1", "kwargs": b"{}"}]
|
||||
assert events == [
|
||||
telemetry_events = [
|
||||
event for event in events if event.get("event_type") == "resource_telemetry"
|
||||
]
|
||||
assert [event["active"] for event in telemetry_events] == [True, False]
|
||||
assert _without_resource_telemetry(events) == [
|
||||
{
|
||||
"kind": "progress",
|
||||
"phase": "executing",
|
||||
@@ -108,7 +123,7 @@ def test_modal_cloud_streams_remote_log_task_id_before_progress(
|
||||
)
|
||||
|
||||
assert events[0] == {"kind": "remote_logs", "task_id": "ta-remote-123"}
|
||||
assert events[1:] == [
|
||||
assert _without_resource_telemetry(events[1:]) == [
|
||||
{
|
||||
"kind": "progress",
|
||||
"phase": "executing",
|
||||
@@ -165,7 +180,9 @@ def test_modal_cloud_returns_large_durable_result_by_reference(
|
||||
)
|
||||
|
||||
assert begin_calls == [True]
|
||||
assert events == [{"kind": "result", "output_ref": result_ref.to_payload()}]
|
||||
assert _without_resource_telemetry(events) == [
|
||||
{"kind": "result", "output_ref": result_ref.to_payload()}
|
||||
]
|
||||
|
||||
def test_modal_cloud_streams_tensor_safe_progress_and_result_events(
|
||||
modal_cloud_module: Any,
|
||||
@@ -204,15 +221,18 @@ def test_modal_cloud_streams_tensor_safe_progress_and_result_events(
|
||||
)
|
||||
)
|
||||
|
||||
assert events[0]["kind"] == "progress"
|
||||
assert events[0]["phase"] == "executing"
|
||||
execution_events = _without_resource_telemetry(events)
|
||||
assert execution_events[0]["kind"] == "progress"
|
||||
assert execution_events[0]["phase"] == "executing"
|
||||
assert torch.equal(
|
||||
serialization_module.deserialize_value(events[0]["preview_tensor"]),
|
||||
serialization_module.deserialize_value(execution_events[0]["preview_tensor"]),
|
||||
tensor,
|
||||
)
|
||||
|
||||
assert events[1]["kind"] == "result"
|
||||
decoded_outputs = serialization_module.deserialize_node_outputs(events[1]["outputs"])
|
||||
assert execution_events[1]["kind"] == "result"
|
||||
decoded_outputs = serialization_module.deserialize_node_outputs(
|
||||
execution_events[1]["outputs"]
|
||||
)
|
||||
assert len(decoded_outputs) == 1
|
||||
assert torch.equal(decoded_outputs[0], tensor)
|
||||
|
||||
@@ -524,6 +544,89 @@ def test_large_transfer_reporter_emits_byte_ui_metadata(
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def test_remote_resource_telemetry_reaches_local_websocket(
|
||||
remote_modal_app_module: Any,
|
||||
local_ui_events_module: Any,
|
||||
monkeypatch: Any,
|
||||
) -> None:
|
||||
"""Provider-neutral memory samples should retain identity and peak fields."""
|
||||
|
||||
class PromptServer:
|
||||
"""Capture telemetry forwarded to the ComfyUI client."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize the captured message list."""
|
||||
self.messages: list[tuple[str, dict[str, Any], str]] = []
|
||||
|
||||
def send_sync(self, event: str, data: dict[str, Any], sid: str) -> None:
|
||||
"""Record one websocket event."""
|
||||
self.messages.append((event, data, sid))
|
||||
|
||||
prompt_server = PromptServer()
|
||||
monkeypatch.setattr(
|
||||
local_ui_events_module,
|
||||
"_lookup_local_prompt_server",
|
||||
lambda: prompt_server,
|
||||
)
|
||||
payload = {
|
||||
"prompt_id": "prompt-memory",
|
||||
"component_id": "170",
|
||||
"component_node_ids": ["170", "169"],
|
||||
"execution_provider": "modal",
|
||||
"execution_environment_id": "modal:gpu-a",
|
||||
"extra_data": {"client_id": "client-memory"},
|
||||
}
|
||||
|
||||
result = remote_modal_app_module._consume_remote_payload_stream(
|
||||
payload,
|
||||
iter(
|
||||
[
|
||||
{"kind": "remote_logs", "task_id": "ta-memory"},
|
||||
{
|
||||
"kind": "progress",
|
||||
"event_type": "resource_telemetry",
|
||||
"active": False,
|
||||
"sample_sequence": 4,
|
||||
"sampled_at": 123.5,
|
||||
"cpu_memory_bytes": 200,
|
||||
"cpu_memory_peak_bytes": 350,
|
||||
"cpu_memory_total_bytes": 1000,
|
||||
"gpu_memory_bytes": 500,
|
||||
"gpu_memory_peak_bytes": 800,
|
||||
"gpu_memory_total_bytes": 2000,
|
||||
},
|
||||
{"kind": "result", "outputs": b"outputs"},
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
assert result == b"outputs"
|
||||
assert prompt_server.messages == [
|
||||
(
|
||||
"modal_telemetry",
|
||||
{
|
||||
"event_type": "resource_telemetry",
|
||||
"prompt_id": "prompt-memory",
|
||||
"component_id": "170",
|
||||
"node_ids": ["170", "169"],
|
||||
"active": False,
|
||||
"execution_provider": "modal",
|
||||
"execution_environment_id": "modal:gpu-a",
|
||||
"execution_location": "ta-memory",
|
||||
"sample_sequence": 4,
|
||||
"cpu_memory_bytes": 200,
|
||||
"cpu_memory_peak_bytes": 350,
|
||||
"cpu_memory_total_bytes": 1000,
|
||||
"gpu_memory_bytes": 500,
|
||||
"gpu_memory_peak_bytes": 800,
|
||||
"gpu_memory_total_bytes": 2000,
|
||||
"sampled_at": 123.5,
|
||||
},
|
||||
"client-memory",
|
||||
)
|
||||
]
|
||||
|
||||
def test_emit_local_mapped_lane_progress_start_marks_lane_as_setup_only(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
|
||||
@@ -28,6 +28,17 @@ def _cloud_streaming_owner() -> Any:
|
||||
return importlib.import_module("cloud_streaming")
|
||||
|
||||
|
||||
def _without_resource_telemetry(
|
||||
events: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return durable-stream events unrelated to memory sampling."""
|
||||
return [
|
||||
event
|
||||
for event in events
|
||||
if event.get("event_type") != "resource_telemetry"
|
||||
]
|
||||
|
||||
|
||||
def _patch_cloud_durable_state(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
name: str,
|
||||
@@ -244,7 +255,7 @@ def test_cloud_stream_commits_durable_outputs_on_consumer_thread(
|
||||
completed_record = invocation_store.get_record("RIV_stream")
|
||||
assert completed_record.state == "completed"
|
||||
assert completed_record.result_object is not None
|
||||
assert events == [
|
||||
assert _without_resource_telemetry(events) == [
|
||||
{
|
||||
"kind": "result",
|
||||
"output_ref": completed_record.result_object.to_payload(),
|
||||
@@ -290,7 +301,9 @@ def test_cloud_stream_logs_result_persistence_and_transport_boundaries(
|
||||
modal_cloud_module._stream_remote_payload_events(payload, b"inputs")
|
||||
)
|
||||
|
||||
assert events == [{"kind": "result", "outputs": serialized_outputs}]
|
||||
assert _without_resource_telemetry(events) == [
|
||||
{"kind": "result", "outputs": serialized_outputs}
|
||||
]
|
||||
assert "Remote stream worker produced result component=component-timing" in caplog.text
|
||||
assert "Finished publishing remote stream result to event buffer" in caplog.text
|
||||
assert "Remote stream consumer received buffered result" in caplog.text
|
||||
@@ -466,7 +479,8 @@ def test_cloud_stream_close_after_result_preserves_completed_record(
|
||||
b"inputs",
|
||||
)
|
||||
|
||||
assert next(stream) == {"kind": "result", "outputs": serialized_outputs}
|
||||
result_event = next(event for event in stream if event.get("kind") == "result")
|
||||
assert result_event == {"kind": "result", "outputs": serialized_outputs}
|
||||
stream.close()
|
||||
|
||||
completed_record = invocation_store.get_record(invocation_id)
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Tests for backend-neutral remote memory telemetry."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def test_sampler_publishes_current_and_peak_memory(
|
||||
resource_telemetry_module: Any,
|
||||
) -> None:
|
||||
"""The terminal sample should preserve peaks observed during execution."""
|
||||
snapshots = iter(
|
||||
(
|
||||
resource_telemetry_module.ResourceMemorySnapshot(100, 1000, 200, 2000),
|
||||
resource_telemetry_module.ResourceMemorySnapshot(350, 1000, 800, 2000),
|
||||
resource_telemetry_module.ResourceMemorySnapshot(250, 1000, 500, 2000),
|
||||
)
|
||||
)
|
||||
events: list[dict[str, Any]] = []
|
||||
sampler = resource_telemetry_module.RemoteResourceTelemetrySampler(
|
||||
events.append,
|
||||
interval_seconds=60.0,
|
||||
sample_reader=lambda: next(snapshots),
|
||||
)
|
||||
|
||||
sampler.start()
|
||||
sampler._publish_sample(active=True)
|
||||
sampler.stop()
|
||||
|
||||
assert [event["active"] for event in events] == [True, True, False]
|
||||
assert events[-1]["cpu_memory_bytes"] == 250
|
||||
assert events[-1]["cpu_memory_peak_bytes"] == 350
|
||||
assert events[-1]["gpu_memory_bytes"] == 500
|
||||
assert events[-1]["gpu_memory_peak_bytes"] == 800
|
||||
assert events[-1]["cpu_memory_total_bytes"] == 1000
|
||||
assert events[-1]["gpu_memory_total_bytes"] == 2000
|
||||
|
||||
|
||||
def test_snapshot_reader_returns_non_negative_memory(
|
||||
resource_telemetry_module: Any,
|
||||
) -> None:
|
||||
"""The production reader should work on CPU-only local test hosts."""
|
||||
snapshot = resource_telemetry_module.read_resource_memory_snapshot()
|
||||
|
||||
assert snapshot.cpu_memory_bytes > 0
|
||||
assert snapshot.cpu_memory_total_bytes >= snapshot.cpu_memory_bytes
|
||||
assert snapshot.gpu_memory_bytes >= 0
|
||||
assert snapshot.gpu_memory_total_bytes >= snapshot.gpu_memory_bytes
|
||||
@@ -137,6 +137,52 @@ def test_modal_ui_event_replay_is_client_scoped(modal_ui_events_module: Any) ->
|
||||
}
|
||||
assert modal_ui_events_module.modal_ui_events_for_client(None) == []
|
||||
|
||||
|
||||
def test_modal_telemetry_replay_coalesces_each_execution_source(
|
||||
modal_ui_events_module: Any,
|
||||
) -> None:
|
||||
"""Periodic samples must not evict durable status events from replay history."""
|
||||
modal_ui_events_module._MODAL_UI_EVENTS_BY_CLIENT.clear()
|
||||
base_payload = {
|
||||
"prompt_id": "prompt-memory",
|
||||
"execution_environment_id": "modal:gpu-a",
|
||||
"execution_location": "ta-one",
|
||||
"component_id": "170",
|
||||
}
|
||||
modal_ui_events_module.record_modal_ui_event(
|
||||
"modal_telemetry",
|
||||
{**base_payload, "cpu_memory_peak_bytes": 100},
|
||||
"client-memory",
|
||||
)
|
||||
modal_ui_events_module.record_modal_ui_event(
|
||||
"modal_status",
|
||||
{"prompt_id": "prompt-memory", "phase": "executing"},
|
||||
"client-memory",
|
||||
)
|
||||
modal_ui_events_module.record_modal_ui_event(
|
||||
"modal_telemetry",
|
||||
{**base_payload, "cpu_memory_peak_bytes": 300},
|
||||
"client-memory",
|
||||
)
|
||||
modal_ui_events_module.record_modal_ui_event(
|
||||
"modal_telemetry",
|
||||
{**base_payload, "component_id": "171", "cpu_memory_peak_bytes": 200},
|
||||
"client-memory",
|
||||
)
|
||||
|
||||
replay_events = modal_ui_events_module.modal_ui_events_for_client("client-memory")
|
||||
|
||||
assert [event["event"] for event in replay_events] == [
|
||||
"modal_status",
|
||||
"modal_telemetry",
|
||||
"modal_telemetry",
|
||||
]
|
||||
assert [
|
||||
event["payload"]["cpu_memory_peak_bytes"]
|
||||
for event in replay_events
|
||||
if event["event"] == "modal_telemetry"
|
||||
] == [300, 200]
|
||||
|
||||
def test_progress_state_route_is_queue_route_sibling(api_intercept_module: Any) -> None:
|
||||
"""The frontend should have a stable sibling route for Modal UI replay."""
|
||||
assert api_intercept_module._progress_state_route_path("/modal/queue_prompt") == (
|
||||
@@ -166,4 +212,3 @@ def test_modal_reset_route_paths_are_queue_route_siblings(api_intercept_module:
|
||||
assert api_intercept_module._delete_modal_volume_route_path("/custom/modal") == (
|
||||
"/custom/modal/delete_volume"
|
||||
)
|
||||
|
||||
|
||||
@@ -606,6 +606,22 @@ def test_mixed_remote_environments_render_clipped_runtime_locations() -> None:
|
||||
assert "const llmTimingLabel = aggregateProgress?.timeToFirstTokenSeconds" in source
|
||||
|
||||
|
||||
def test_configurator_renders_current_and_peak_remote_memory() -> None:
|
||||
"""The Configurator should aggregate telemetry and retain terminal peaks."""
|
||||
source = _modal_toggle_source()
|
||||
|
||||
assert 'api.addEventListener("modal_telemetry", handleModalTelemetry);' in source
|
||||
assert 'eventName === "modal_telemetry"' in source
|
||||
assert "function remoteConfiguratorTelemetryPatch(existing, detail)" in source
|
||||
assert "resourceTelemetryBySource: samples" in source
|
||||
assert "resourceActiveCount: activeSamples.length" in source
|
||||
assert "function remoteConfiguratorResourceUsageLabel(state)" in source
|
||||
assert "CPU RAM ${remoteStorageSizeLabel" in source
|
||||
assert "Peak CPU RAM ${peak}" in source
|
||||
assert "Peak GPU ${peak}" in source
|
||||
assert 'resourceText.className = "remote-configurator-environment-resources";' in source
|
||||
|
||||
|
||||
def test_modal_canvas_animation_loop_is_selective_and_throttled() -> None:
|
||||
"""Static progress data should not keep an unbounded canvas rAF loop alive."""
|
||||
source = _modal_toggle_source()
|
||||
|
||||
+127
-4
@@ -1747,6 +1747,8 @@ function replayModalUiEvent(eventRecord) {
|
||||
handleModalStatus({ detail: payload });
|
||||
} else if (eventName === "modal_progress") {
|
||||
handleModalProgress({ detail: payload });
|
||||
} else if (eventName === "modal_telemetry") {
|
||||
handleModalTelemetry({ detail: payload });
|
||||
}
|
||||
} finally {
|
||||
modalReplayedEventUpdatedAtMs = null;
|
||||
@@ -3584,6 +3586,9 @@ function createRemoteConfiguratorEnvironmentRow(environmentId, label) {
|
||||
statusText.className = "remote-configurator-environment-status";
|
||||
statusText.textContent = "Waiting";
|
||||
heading.append(dot, labelElement, statusText);
|
||||
const resourceText = document.createElement("div");
|
||||
resourceText.className = "remote-configurator-environment-resources";
|
||||
resourceText.hidden = true;
|
||||
const progress = document.createElement("div");
|
||||
progress.className = "remote-configurator-environment-progress";
|
||||
const progressFill = document.createElement("span");
|
||||
@@ -3591,18 +3596,42 @@ function createRemoteConfiguratorEnvironmentRow(environmentId, label) {
|
||||
const progressValue = document.createElement("span");
|
||||
progressValue.className = "remote-configurator-environment-progress-value";
|
||||
progress.append(progressFill, progressValue);
|
||||
root.append(heading, progress);
|
||||
root.append(heading, resourceText, progress);
|
||||
return {
|
||||
environmentId,
|
||||
root,
|
||||
label: labelElement,
|
||||
statusText,
|
||||
resourceText,
|
||||
progress,
|
||||
progressFill,
|
||||
progressValue,
|
||||
};
|
||||
}
|
||||
|
||||
/** Return the Configurator readout for current and peak remote memory usage. */
|
||||
function remoteConfiguratorResourceUsageLabel(state) {
|
||||
const labels = [];
|
||||
const hasActiveSamples = Number(state?.resourceActiveCount ?? 0) > 0;
|
||||
if (Boolean(state?.gpuMemoryObserved)) {
|
||||
const peak = remoteStorageSizeLabel(Number(state?.gpuMemoryPeakBytes ?? 0));
|
||||
labels.push(
|
||||
hasActiveSamples
|
||||
? `GPU ${remoteStorageSizeLabel(Number(state?.gpuMemoryBytes ?? 0))} now · ${peak} peak`
|
||||
: `Peak GPU ${peak}`,
|
||||
);
|
||||
}
|
||||
if (Boolean(state?.cpuMemoryObserved)) {
|
||||
const peak = remoteStorageSizeLabel(Number(state?.cpuMemoryPeakBytes ?? 0));
|
||||
labels.push(
|
||||
hasActiveSamples
|
||||
? `CPU RAM ${remoteStorageSizeLabel(Number(state?.cpuMemoryBytes ?? 0))} now · ${peak} peak`
|
||||
: `Peak CPU RAM ${peak}`,
|
||||
);
|
||||
}
|
||||
return labels.join(" | ");
|
||||
}
|
||||
|
||||
/**
|
||||
* Render status for one remote environment without affecting another lane.
|
||||
* @param {Record<string, any>} row
|
||||
@@ -3615,6 +3644,9 @@ function renderRemoteConfiguratorEnvironmentStatus(row, state) {
|
||||
const phase = String(state?.phase ?? "idle");
|
||||
row.root.dataset.phase = phase;
|
||||
row.statusText.textContent = remoteConfiguratorStatusMessage(state);
|
||||
const resourceLabel = remoteConfiguratorResourceUsageLabel(state);
|
||||
row.resourceText.hidden = !resourceLabel;
|
||||
row.resourceText.textContent = resourceLabel;
|
||||
const total = Math.max(0, Number(state?.statusTotal ?? 0));
|
||||
const current = Math.max(0, Math.min(total, Number(state?.statusCurrent ?? 0)));
|
||||
row.progress.hidden = total <= 1;
|
||||
@@ -3624,6 +3656,72 @@ function renderRemoteConfiguratorEnvironmentStatus(row, state) {
|
||||
: "";
|
||||
}
|
||||
|
||||
/** Merge one component/location memory sample into environment-level telemetry. */
|
||||
function remoteConfiguratorTelemetryPatch(existing, detail) {
|
||||
if (String(detail?.event_type ?? "") !== "resource_telemetry") {
|
||||
return {};
|
||||
}
|
||||
const samples = new Map(
|
||||
existing?.resourceTelemetryBySource instanceof Map
|
||||
? existing.resourceTelemetryBySource
|
||||
: [],
|
||||
);
|
||||
const sourceKey = [
|
||||
String(detail?.execution_location ?? ""),
|
||||
String(detail?.component_id ?? "payload"),
|
||||
].join("|");
|
||||
samples.set(sourceKey, {
|
||||
active: detail?.active !== false,
|
||||
cpuMemoryBytes: Math.max(0, Number(detail?.cpu_memory_bytes ?? 0)),
|
||||
cpuMemoryPeakBytes: Math.max(0, Number(detail?.cpu_memory_peak_bytes ?? 0)),
|
||||
cpuMemoryTotalBytes: Math.max(0, Number(detail?.cpu_memory_total_bytes ?? 0)),
|
||||
gpuMemoryBytes: Math.max(0, Number(detail?.gpu_memory_bytes ?? 0)),
|
||||
gpuMemoryPeakBytes: Math.max(0, Number(detail?.gpu_memory_peak_bytes ?? 0)),
|
||||
gpuMemoryTotalBytes: Math.max(0, Number(detail?.gpu_memory_total_bytes ?? 0)),
|
||||
});
|
||||
const activeSamples = Array.from(samples.values()).filter((sample) => sample.active);
|
||||
const cpuMemoryBytes = activeSamples.reduce(
|
||||
(total, sample) => total + sample.cpuMemoryBytes,
|
||||
0,
|
||||
);
|
||||
const gpuMemoryBytes = activeSamples.reduce(
|
||||
(total, sample) => total + sample.gpuMemoryBytes,
|
||||
0,
|
||||
);
|
||||
const sourceCpuPeakBytes = Math.max(
|
||||
0,
|
||||
...Array.from(samples.values(), (sample) => sample.cpuMemoryPeakBytes),
|
||||
);
|
||||
const sourceGpuPeakBytes = Math.max(
|
||||
0,
|
||||
...Array.from(samples.values(), (sample) => sample.gpuMemoryPeakBytes),
|
||||
);
|
||||
return {
|
||||
resourceTelemetryBySource: samples,
|
||||
resourceActiveCount: activeSamples.length,
|
||||
cpuMemoryBytes,
|
||||
cpuMemoryPeakBytes: Math.max(
|
||||
Number(existing?.cpuMemoryPeakBytes ?? 0),
|
||||
sourceCpuPeakBytes,
|
||||
cpuMemoryBytes,
|
||||
),
|
||||
cpuMemoryObserved:
|
||||
Boolean(existing?.cpuMemoryObserved) ||
|
||||
Array.from(samples.values()).some(
|
||||
(sample) => sample.cpuMemoryBytes > 0 || sample.cpuMemoryTotalBytes > 0,
|
||||
),
|
||||
gpuMemoryBytes,
|
||||
gpuMemoryPeakBytes: Math.max(
|
||||
Number(existing?.gpuMemoryPeakBytes ?? 0),
|
||||
sourceGpuPeakBytes,
|
||||
gpuMemoryBytes,
|
||||
),
|
||||
gpuMemoryObserved:
|
||||
Boolean(existing?.gpuMemoryObserved) ||
|
||||
Array.from(samples.values()).some((sample) => sample.gpuMemoryTotalBytes > 0),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Rebuild the independent environment rows for a placement plan.
|
||||
* @param {Record<string, any>} panel
|
||||
@@ -4386,7 +4484,7 @@ function renderRemoteConfiguratorPlan(panel, assignments, configurations) {
|
||||
panel.minHeight =
|
||||
118 +
|
||||
Math.max(31, planHeight) +
|
||||
panel.environmentRows.size * 52 +
|
||||
panel.environmentRows.size * 66 +
|
||||
panel.storageHeight +
|
||||
panel.capacityHeight +
|
||||
panel.keychainUnlockHeight;
|
||||
@@ -4525,14 +4623,17 @@ function updateRemoteConfiguratorEnvironmentStatus(promptId, detail) {
|
||||
promoteProvisionalVastEnvironment(panel, promptState, environmentId);
|
||||
}
|
||||
const existing = promptState.remoteEnvironmentStatuses.get(environmentId) ?? {};
|
||||
const telemetryPatch = remoteConfiguratorTelemetryPatch(existing, detail);
|
||||
const state = {
|
||||
...existing,
|
||||
...telemetryPatch,
|
||||
phase: String(detail?.phase ?? existing.phase ?? STATE_SETUP),
|
||||
statusMessage:
|
||||
detail?.status_message ?? detail?.message ?? detail?.stage ?? existing.statusMessage ?? null,
|
||||
statusCurrent:
|
||||
detail?.status_current ?? detail?.value ?? existing.statusCurrent ?? null,
|
||||
statusTotal: detail?.status_total ?? detail?.max ?? existing.statusTotal ?? null,
|
||||
updatedAt: nowMs(),
|
||||
updatedAt: progressEventTimestampMs(),
|
||||
};
|
||||
promptState.remoteEnvironmentStatuses.set(environmentId, state);
|
||||
if (!panel || panel.promptId !== promptId) {
|
||||
@@ -4548,7 +4649,7 @@ function updateRemoteConfiguratorEnvironmentStatus(promptId, detail) {
|
||||
panel.environmentRows.set(environmentId, row);
|
||||
panel.environments.appendChild(row.root);
|
||||
panel.environments.hidden = false;
|
||||
panel.minHeight += 52;
|
||||
panel.minHeight += 66;
|
||||
panel.node.setSize?.([
|
||||
Math.max(Number(panel.node.size?.[0]) || 0, 500),
|
||||
Math.max(Number(panel.node.size?.[1]) || 0, panel.minHeight),
|
||||
@@ -6718,6 +6819,19 @@ function handleModalProgress(event) {
|
||||
);
|
||||
}
|
||||
|
||||
/** Apply remote CPU/GPU memory telemetry without changing execution phase. */
|
||||
function handleModalTelemetry(event) {
|
||||
const detail = eventDetail(event);
|
||||
const promptId = String(detail?.prompt_id ?? "");
|
||||
if (!promptId || isPromptTerminal(promptId)) {
|
||||
return;
|
||||
}
|
||||
if (!ensurePromptState(promptId)) {
|
||||
return;
|
||||
}
|
||||
updateRemoteConfiguratorEnvironmentStatus(promptId, detail);
|
||||
}
|
||||
|
||||
/**
|
||||
* Update remote component visuals from a native ComfyUI execution event.
|
||||
* @param {CustomEvent} event
|
||||
@@ -7002,6 +7116,7 @@ function registerExecutionListeners() {
|
||||
|
||||
api.addEventListener("modal_status", handleModalStatus);
|
||||
api.addEventListener("modal_progress", handleModalProgress);
|
||||
api.addEventListener("modal_telemetry", handleModalTelemetry);
|
||||
api.addEventListener("executing", (event) => handleExecutionPhase(event, EXECUTION_PHASE));
|
||||
api.addEventListener("executed", (event) => {
|
||||
handleExecutionPhase(event, STATE_COMPLETE);
|
||||
@@ -7584,6 +7699,14 @@ function installGlobalStatusStyles() {
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.remote-configurator-environment-resources {
|
||||
overflow: hidden;
|
||||
color: #cbd5e1;
|
||||
font: 600 8px/11px ui-monospace, SFMono-Regular, monospace;
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.remote-configurator-environment-progress {
|
||||
position: relative;
|
||||
height: 10px;
|
||||
|
||||
Reference in New Issue
Block a user