diff --git a/.gitignore b/.gitignore index 078af1e..7d5b497 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/api_intercept.py b/api_intercept.py index 101ad0e..5a541ea 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -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 diff --git a/modal_ui_events.py b/modal_ui_events.py new file mode 100644 index 0000000..2dbe98d --- /dev/null +++ b/modal_ui_events.py @@ -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() diff --git a/tests/conftest.py b/tests/conftest.py index 1d657e6..b57ff52 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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.""" diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index 8587bb7..5dfc337 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -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: