nodes/ttulttul-ComfyUI-Modal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Extract Modal UI event buffering

Browse Source
This commit is contained in:
Ken Simpson
2026-08-29 10:31:04 -07:00
parent 85ebb3b934
commit 4c648ab9d4
5 changed files with 161 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -32,6 +32,7 @@
!/modal_llm_node.py
!/modal_llm_runtime.py
!/modal_sdk.py
!/modal_ui_events.py
!/output_artifacts.py
!/remote_hosts.py
!/remote_configuration_nodes.py
api_intercept.py
+10 -121
View File
@@ -147,6 +147,11 @@ if __package__:
delete_modal_cache_dicts,
delete_modal_volume,
)
from .modal_ui_events import (
_emit_modal_status,
modal_ui_events_for_client,
record_modal_ui_event,
)
else: # pragma: no cover - flat import inside the Modal container.
from modal_executor_node import (
MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS,
@@ -274,6 +279,11 @@ else: # pragma: no cover - flat import inside the Modal container.
delete_modal_cache_dicts,
delete_modal_volume,
)
from modal_ui_events import (
_emit_modal_status,
modal_ui_events_for_client,
record_modal_ui_event,
)
logger = logging.getLogger(__name__)
SetupStatusCallback = Callable[[str, int | None, int | None], None]
@@ -285,10 +295,6 @@ ExecutionPlanStatusCallback = Callable[
]
_ROUTE_REGISTERED = False
_MODAL_UI_EVENT_RETENTION_SECONDS = 2 * 60 * 60
_MODAL_UI_EVENT_LIMIT_PER_CLIENT = 512
_MODAL_UI_EVENTS_LOCK = threading.Lock()
_MODAL_UI_EVENTS_BY_CLIENT: dict[str, deque[dict[str, Any]]] = {}
_MODAL_INTERRUPT_QUEUE_BRIDGE_ATTR = "__comfy_modal_interrupt_queue_bridge_installed"
_REMOTE_PREPARATION_PROMPTS_ATTR = "__comfy_modal_remote_preparation_prompts"
_REMOTE_PREPARATION_CANCELLATIONS_ATTR = (
@@ -2332,123 +2338,6 @@ def _ensure_remote_sync_backend(
)
def _emit_modal_status(
prompt_server: Any,
phase: str,
*,
client_id: str | None,
prompt_id: str | None,
node_ids: list[str],
configurator_node_id: str | None = None,
modal_gpu: str | None = None,
component_node_ids_by_representative: dict[str, list[str]] | None = None,
active_node_id: str | None = None,
active_node_class_type: str | None = None,
active_node_role: str | None = None,
error_message: str | None = None,
status_message: str | None = None,
status_current: int | None = None,
status_total: int | None = None,
execution_environment_id: str | None = None,
remote_execution_assignments: Mapping[str, Mapping[str, Any]] | None = None,
remote_execution_configurations: Sequence[Mapping[str, Any]] | None = None,
) -> None:
"""Send a Modal execution status event to the active websocket client."""
if client_id is None:
logger.debug(
"Skipping Modal status event %s because no client id is available.", phase
)
return
payload: dict[str, Any] = {
"phase": phase,
"prompt_id": prompt_id,
"node_ids": list(node_ids),
}
if modal_gpu is not None:
payload["modal_gpu"] = modal_gpu
if configurator_node_id is not None:
payload["configurator_node_id"] = configurator_node_id
if component_node_ids_by_representative:
payload["components"] = [
{
"representative_node_id": representative_node_id,
"node_ids": list(component_node_ids),
}
for representative_node_id, component_node_ids in sorted(
component_node_ids_by_representative.items()
)
]
if active_node_id is not None:
payload["active_node_id"] = active_node_id
if active_node_class_type is not None:
payload["active_node_class_type"] = active_node_class_type
if active_node_role is not None:
payload["active_node_role"] = active_node_role
if error_message is not None:
payload["error_message"] = error_message
if status_message is not None:
payload["status_message"] = status_message
if status_current is not None:
payload["status_current"] = int(status_current)
if status_total is not None:
payload["status_total"] = int(status_total)
if execution_environment_id is not None:
payload["execution_environment_id"] = execution_environment_id
if remote_execution_assignments is not None:
payload["remote_execution_assignments"] = copy.deepcopy(
dict(remote_execution_assignments)
)
if remote_execution_configurations is not None:
payload["remote_execution_configurations"] = copy.deepcopy(
list(remote_execution_configurations)
)
record_modal_ui_event("modal_status", payload, client_id)
prompt_server.send_sync("modal_status", payload, client_id)
def record_modal_ui_event(
event: str, payload: Mapping[str, Any], client_id: str | None
) -> None:
"""Store one client-scoped Modal UI event so the browser can replay it after refocus."""
if client_id is None:
return
event_record = {
"event": event,
"payload": copy.deepcopy(dict(payload)),
"updated_at": time.time(),
}
with _MODAL_UI_EVENTS_LOCK:
client_events = _MODAL_UI_EVENTS_BY_CLIENT.setdefault(
client_id,
deque(maxlen=_MODAL_UI_EVENT_LIMIT_PER_CLIENT),
)
client_events.append(event_record)
_prune_modal_ui_events_locked(client_events)
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:
return []
with _MODAL_UI_EVENTS_LOCK:
client_events = _MODAL_UI_EVENTS_BY_CLIENT.get(client_id)
if client_events is None:
return []
_prune_modal_ui_events_locked(client_events)
return [copy.deepcopy(event_record) for event_record in client_events]
def _prune_modal_ui_events_locked(client_events: deque[dict[str, Any]]) -> None:
"""Discard stale Modal UI events while the caller holds the event lock."""
cutoff = time.time() - _MODAL_UI_EVENT_RETENTION_SECONDS
while client_events and float(client_events[0].get("updated_at", 0.0)) < cutoff:
client_events.popleft()
def _get_nodes_module() -> Any:
"""Import the ComfyUI nodes module lazily."""
import nodes
modal_ui_events.py
+134
View File
@@ -0,0 +1,134 @@
"""Client-scoped status emission and replay for remote workflow UI state."""
from __future__ import annotations
import copy
import logging
import threading
import time
from collections import deque
from typing import Any, Mapping, Sequence
logger = logging.getLogger(__name__)
_MODAL_UI_EVENT_RETENTION_SECONDS = 2 * 60 * 60
_MODAL_UI_EVENT_LIMIT_PER_CLIENT = 512
_MODAL_UI_EVENTS_LOCK = threading.Lock()
_MODAL_UI_EVENTS_BY_CLIENT: dict[str, deque[dict[str, Any]]] = {}
def _emit_modal_status(
prompt_server: Any,
phase: str,
*,
client_id: str | None,
prompt_id: str | None,
node_ids: list[str],
configurator_node_id: str | None = None,
modal_gpu: str | None = None,
component_node_ids_by_representative: dict[str, list[str]] | None = None,
active_node_id: str | None = None,
active_node_class_type: str | None = None,
active_node_role: str | None = None,
error_message: str | None = None,
status_message: str | None = None,
status_current: int | None = None,
status_total: int | None = None,
execution_environment_id: str | None = None,
remote_execution_assignments: Mapping[str, Mapping[str, Any]] | None = None,
remote_execution_configurations: Sequence[Mapping[str, Any]] | None = None,
) -> None:
"""Send a Modal execution status event to the active websocket client."""
if client_id is None:
logger.debug(
"Skipping Modal status event %s because no client id is available.", phase
)
return
payload: dict[str, Any] = {
"phase": phase,
"prompt_id": prompt_id,
"node_ids": list(node_ids),
}
if modal_gpu is not None:
payload["modal_gpu"] = modal_gpu
if configurator_node_id is not None:
payload["configurator_node_id"] = configurator_node_id
if component_node_ids_by_representative:
payload["components"] = [
{
"representative_node_id": representative_node_id,
"node_ids": list(component_node_ids),
}
for representative_node_id, component_node_ids in sorted(
component_node_ids_by_representative.items()
)
]
if active_node_id is not None:
payload["active_node_id"] = active_node_id
if active_node_class_type is not None:
payload["active_node_class_type"] = active_node_class_type
if active_node_role is not None:
payload["active_node_role"] = active_node_role
if error_message is not None:
payload["error_message"] = error_message
if status_message is not None:
payload["status_message"] = status_message
if status_current is not None:
payload["status_current"] = int(status_current)
if status_total is not None:
payload["status_total"] = int(status_total)
if execution_environment_id is not None:
payload["execution_environment_id"] = execution_environment_id
if remote_execution_assignments is not None:
payload["remote_execution_assignments"] = copy.deepcopy(
dict(remote_execution_assignments)
)
if remote_execution_configurations is not None:
payload["remote_execution_configurations"] = copy.deepcopy(
list(remote_execution_configurations)
)
record_modal_ui_event("modal_status", payload, client_id)
prompt_server.send_sync("modal_status", payload, client_id)
def record_modal_ui_event(
event: str, payload: Mapping[str, Any], client_id: str | None
) -> None:
"""Store one client-scoped Modal UI event for replay after browser refocus."""
if client_id is None:
return
event_record = {
"event": event,
"payload": copy.deepcopy(dict(payload)),
"updated_at": time.time(),
}
with _MODAL_UI_EVENTS_LOCK:
client_events = _MODAL_UI_EVENTS_BY_CLIENT.setdefault(
client_id,
deque(maxlen=_MODAL_UI_EVENT_LIMIT_PER_CLIENT),
)
client_events.append(event_record)
_prune_modal_ui_events_locked(client_events)
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:
return []
with _MODAL_UI_EVENTS_LOCK:
client_events = _MODAL_UI_EVENTS_BY_CLIENT.get(client_id)
if client_events is None:
return []
_prune_modal_ui_events_locked(client_events)
return [copy.deepcopy(event_record) for event_record in client_events]
def _prune_modal_ui_events_locked(client_events: deque[dict[str, Any]]) -> None:
"""Discard stale Modal UI events while the caller holds the event lock."""
cutoff = time.time() - _MODAL_UI_EVENT_RETENTION_SECONDS
while client_events and float(client_events[0].get("updated_at", 0.0)) < cutoff:
client_events.popleft()
tests/conftest.py
+6
View File
@@ -315,6 +315,12 @@ def modal_admin_ops_module(extension_package: object) -> object:
return importlib.import_module(f"{PACKAGE_NAME}.modal_admin_ops")
@pytest.fixture(scope="session")
def modal_ui_events_module(extension_package: object) -> object:
"""Return the client-scoped Modal UI event module."""
return importlib.import_module(f"{PACKAGE_NAME}.modal_ui_events")
@pytest.fixture(scope="session")
def modal_endpoint_module(extension_package: object) -> object:
"""Return the Modal hosted-model endpoint node module."""
tests/test_api_intercept.py
+10 -9
View File
@@ -5961,9 +5961,10 @@ def test_rewrite_keeps_nested_remote_nodes_remote_when_root_ids_collide(
def test_emit_modal_status_targets_prompt_client(
api_intercept_module: Any,
modal_ui_events_module: Any,
) -> None:
"""Modal status events should preserve prompt and component metadata for the UI."""
api_intercept_module._MODAL_UI_EVENTS_BY_CLIENT.clear()
modal_ui_events_module._MODAL_UI_EVENTS_BY_CLIENT.clear()
class FakePromptServer:
"""Capture websocket events emitted by the queue route."""
@@ -5977,7 +5978,7 @@ def test_emit_modal_status_targets_prompt_client(
self.messages.append((event, data, sid))
prompt_server = FakePromptServer()
api_intercept_module._emit_modal_status(
modal_ui_events_module._emit_modal_status(
prompt_server=prompt_server,
phase="executing",
client_id="client-1",
@@ -6030,7 +6031,7 @@ def test_emit_modal_status_targets_prompt_client(
"client-1",
)
]
replay_events = api_intercept_module.modal_ui_events_for_client("client-1")
replay_events = modal_ui_events_module.modal_ui_events_for_client("client-1")
assert replay_events == [
{
"event": "modal_status",
@@ -6325,22 +6326,22 @@ def test_workflow_ssh_metadata_preserves_probed_gpu_capabilities(
assert queued_host.last_error is None
def test_modal_ui_event_replay_is_client_scoped(api_intercept_module: Any) -> None:
def test_modal_ui_event_replay_is_client_scoped(modal_ui_events_module: Any) -> None:
"""Refocus replay should only return events for the requesting ComfyUI client."""
api_intercept_module._MODAL_UI_EVENTS_BY_CLIENT.clear()
modal_ui_events_module._MODAL_UI_EVENTS_BY_CLIENT.clear()
api_intercept_module.record_modal_ui_event(
modal_ui_events_module.record_modal_ui_event(
"modal_progress",
{"prompt_id": "prompt-1", "node_id": "4", "value": 2.0, "max": 10.0},
"client-1",
)
api_intercept_module.record_modal_ui_event(
modal_ui_events_module.record_modal_ui_event(
"modal_status",
{"prompt_id": "prompt-2", "phase": "executing", "node_ids": ["9"]},
"client-2",
)
replay_events = api_intercept_module.modal_ui_events_for_client("client-1")
replay_events = modal_ui_events_module.modal_ui_events_for_client("client-1")
assert len(replay_events) == 1
assert replay_events[0]["event"] == "modal_progress"
@@ -6350,7 +6351,7 @@ def test_modal_ui_event_replay_is_client_scoped(api_intercept_module: Any) -> No
"value": 2.0,
"max": 10.0,
}
assert api_intercept_module.modal_ui_events_for_client(None) == []
assert modal_ui_events_module.modal_ui_events_for_client(None) == []
def test_progress_state_route_is_queue_route_sibling(api_intercept_module: Any) -> None: