From 6908fe7ea2e60310bc26dcb6cf258dead2498c7e Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Sun, 23 Aug 2026 09:52:19 -0700 Subject: [PATCH] Show multi-environment execution locations --- README.md | 2 +- api_intercept.py | 109 ++++++--- docs/LEARNINGS.md | 2 + remote/modal_app.py | 55 +++++ tests/modal_toggle_frontend.test.mjs | 166 ++++++++++++++ tests/test_api_intercept.py | 8 + tests/test_modal_stream_progress.py | 79 +++++++ tests/test_web_modal_toggle.py | 16 ++ web/modal_toggle.js | 330 +++++++++++++++++++++++++-- 9 files changed, 715 insertions(+), 52 deletions(-) diff --git a/README.md b/README.md index 69e198b..8bb7f86 100644 --- a/README.md +++ b/README.md @@ -257,7 +257,7 @@ These state colours work with both the legacy LiteGraph renderer and ComfyUI's e Subgraph nodes inherit a derived decoration whenever any nested descendant is set to run on Modal, including through multiple levels of nested subgraphs. The container remains a local organizational node: the derived highlight does not change its `Run on Modal` setting. During execution, a `Σ` badge identifies the decoration as a descendant summary and its tooltip lists the descendant phase counts. When descendants have conflicting phases, the border uses the most operationally important state in this order: error, cancelling, active, starting, setup, ready, finalizing, then complete. For example, a subgraph containing one active node and one starting node is purple for active work while the tooltip preserves both counts. -Remote sampler-style progress is rendered in a small temporary panel near the node. Each active progress bar includes a smoothed iterations-per-second (`it/s`) rate derived from successive streamed progress samples; LLM token streams use engine-reported `tok/s`, show token count and TTFT, and reuse the header for compact startup-stage text. Parallel mapped lanes report their rates independently. Static progress redraws happen only when progress events arrive, while indeterminate LLM work, pulsing node phases, setup lane placeholders, and short fade-outs use a throttled canvas animation loop. Preview images and ComfyUI UI payloads emitted by remote nodes are streamed back into the local PromptServer while the remote component is still running. When the browser regains focus, Modal-Sync replays recent UI events and reconciles them against ComfyUI queue/history state so cancelled or completed prompts do not leave stale progress bars behind. +Remote sampler-style progress is rendered in a small temporary panel near the node. Each active progress bar includes a smoothed iterations-per-second (`it/s`) rate derived from successive streamed progress samples; LLM token streams use engine-reported `tok/s`, show token count and TTFT, and reuse the header for compact startup-stage text. Parallel mapped lanes report their rates independently. When one prompt is actually scheduled across more than one remote environment, the panel adds a compact execution-location row beneath its bars: the Modal mark precedes the exact running container ID once known, while the SSH mark precedes the configured host name. Long identities are clipped to the panel width. Nodes 2.0 renders the same identity in its DOM layer so the label remains crisp and visible outside the node at every zoom level; legacy canvas nodes retain the progress-panel row. Runtime status events carry the identity before numeric sampler progress begins. Static progress redraws happen only when progress events arrive, while indeterminate LLM work, pulsing node phases, setup lane placeholders, and short fade-outs use a throttled canvas animation loop. Preview images and ComfyUI UI payloads emitted by remote nodes are streamed back into the local PromptServer while the remote component is still running. When the browser regains focus, Modal-Sync replays recent UI events and reconciles them against ComfyUI queue/history state so cancelled or completed prompts do not leave stale progress bars behind. Cancelling a local prompt propagates a targeted interrupt to the active Modal work. If Modal is still deploying, provisioning, or slow to observe the interrupt, the local proxy releases the ComfyUI prompt after the configured grace period while remote cleanup continues. diff --git a/api_intercept.py b/api_intercept.py index da28c2a..0f1a921 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -1097,11 +1097,14 @@ def _stamp_execution_assignment( assignment: ExecutionAssignment, worker_index: int = 0, execution_history_signature: str | None = None, + execution_location: str | None = None, ) -> None: """Attach provider placement to a payload and every nested proxy phase.""" payload["execution_provider"] = assignment.provider.value payload["execution_environment_id"] = assignment.environment_id payload["execution_worker_index"] = worker_index + if execution_location: + payload["execution_location"] = execution_location if execution_history_signature: payload["execution_history_signature"] = execution_history_signature split_payloads = payload.get("split_proxy_payloads") @@ -1124,9 +1127,29 @@ def _stamp_execution_assignment( assignment, worker_index, execution_history_signature, + execution_location, ) +def _ssh_hostname(ssh_target: str) -> str: + """Return the host portion of one validated OpenSSH destination.""" + host = ssh_target.rsplit("@", 1)[-1].strip() + if host.startswith("[") and "]" in host: + return host[1 : host.index("]")] + return host + + +def _execution_location_for_assignment( + assignment: ExecutionAssignment, + ssh_hosts_by_id: Mapping[str, SshHostConfig], +) -> str | None: + """Return the runtime location label known before remote dispatch.""" + if assignment.provider is ExecutionProvider.MODAL: + return None + host = ssh_hosts_by_id.get(assignment.environment_id) + return _ssh_hostname(host.ssh_target) if host is not None else assignment.environment_id + + def _ensure_remote_sync_backend( settings: ModalSyncSettings, sync_engine: ModalAssetSyncEngine, @@ -6216,6 +6239,7 @@ def rewrite_prompt_for_modal( 0, ), _component_execution_signature(component, prompt), + _execution_location_for_assignment(assignment, ssh_hosts_by_id), ) implicitly_mapped_output_sources = _implicitly_mapped_boundary_output_sources( component=component, @@ -6233,6 +6257,14 @@ def rewrite_prompt_for_modal( payload=payload, nodes_module=resolved_nodes_module, ) + for proxy_node_id in proxy_node_ids: + summary.execution_assignments_by_representative[proxy_node_id] = assignment + summary.execution_worker_indices_by_representative[proxy_node_id] = ( + summary.execution_worker_indices_by_representative.get( + component.representative_node_id, + 0, + ) + ) split_proxy_payloads = payload.get("split_proxy_payloads") if isinstance(split_proxy_payloads, dict): static_proxy_node_id, mapped_proxy_node_id = proxy_node_ids @@ -6374,6 +6406,39 @@ async def rewrite_prompt_for_modal_async( ) +def _execution_assignments_payload( + summary: RewriteSummary, + settings: ModalSyncSettings, +) -> dict[str, dict[str, Any]]: + """Return planner assignments with member nodes and safe runtime labels.""" + ssh_hosts_by_id = { + host.environment_id: host for host in _configured_ssh_hosts(settings) + } + return { + component_id: { + "provider": assignment.provider.value, + "environment_id": assignment.environment_id, + "execution_location": _execution_location_for_assignment( + assignment, + ssh_hosts_by_id, + ), + "node_ids": list( + summary.component_node_ids_by_representative.get(component_id, []) + ), + "predicted_cost_usd": assignment.predicted_cost_usd, + "predicted_completion_seconds": assignment.predicted_completion_seconds, + "worker_index": summary.execution_worker_indices_by_representative.get( + component_id, + 0, + ), + "reasons": list(assignment.reasons), + } + for component_id, assignment in sorted( + summary.execution_assignments_by_representative.items() + ) + } + + async def _queue_prompt_json( prompt_server: Any, json_data: dict[str, Any], @@ -7206,24 +7271,10 @@ def setup_modal_queue_route( asset.remote_path for asset in summary.synced_assets ] json_data["extra_data"]["remote_execution"] = { - "assignments": { - component_id: { - "provider": assignment.provider.value, - "environment_id": assignment.environment_id, - "predicted_cost_usd": assignment.predicted_cost_usd, - "predicted_completion_seconds": ( - assignment.predicted_completion_seconds - ), - "worker_index": summary.execution_worker_indices_by_representative.get( - component_id, - 0, - ), - "reasons": list(assignment.reasons), - } - for component_id, assignment in sorted( - summary.execution_assignments_by_representative.items() - ) - } + "assignments": _execution_assignments_payload( + summary, + request_settings, + ) } if summary.custom_nodes_bundle is not None: json_data["extra_data"]["modal"][ @@ -7252,24 +7303,10 @@ def setup_modal_queue_route( "modal_parallel_local_branch_node_ids": list( summary.parallel_local_branch_node_ids ), - "remote_execution_assignments": { - component_id: { - "provider": assignment.provider.value, - "environment_id": assignment.environment_id, - "predicted_cost_usd": assignment.predicted_cost_usd, - "predicted_completion_seconds": ( - assignment.predicted_completion_seconds - ), - "worker_index": summary.execution_worker_indices_by_representative.get( - component_id, - 0, - ), - "reasons": list(assignment.reasons), - } - for component_id, assignment in sorted( - summary.execution_assignments_by_representative.items() - ) - }, + "remote_execution_assignments": _execution_assignments_payload( + summary, + request_settings, + ), "modal_components": [ { "representative_node_id": representative_node_id, diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 628c5c3..5b4903d 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -481,3 +481,5 @@ - Replacing a provider image's entrypoint does not replace its Docker health check. The llama-server base probes HTTP port 8080, while the managed SSH worker initially serves only its Unix socket; explicitly emit `HEALTHCHECK NONE` in the derived image so an otherwise ready worker is not trapped in an unhealthy startup state. - Copying an executable from a provider image does not make its sibling shared libraries discoverable after the derived image overrides `LD_LIBRARY_PATH`. Preserve `/app` in the SSH image library path and also prefix the configured llama-server binary directory in the spawned process environment; validate with `ldd` or an actual server launch because image construction alone does not reveal the missing dynamic dependency. - Provider-aware placement must happen before remote-session bridging. Cut explicit backend boundaries during component formation, preserve an acyclic coarse graph when later expensive-edge merges would wrap around that cut, and materialize cross-environment values through ComfyUI; otherwise session co-location silently overwrites the scheduler's provider decisions. Modal speculative prewarm and local-gap keepalive metadata must carry and match both provider and environment identity, because recursively embedded SSH payloads can otherwise stage self-hosted-only models on Modal. +- A node execution-location badge needs both queue-time and runtime identity. Return the planner's provider/environment assignment keyed to rewritten proxy representatives so every original member node can render immediately, then replace the Modal environment label with the exact task/container ID learned from the remote log stream. For SSH, derive the display hostname from the configured target instead of exposing the scheduler's environment key. Gate the row on the count of distinct environments actually assigned to that prompt, not the number configured, so a one-environment plan stays visually quiet. +- Nodes 2.0 node bodies can cover legacy canvas foreground drawings, so execution-location UI needs a DOM-layer equivalent rather than relying only on `onDrawForeground`. Attach the exact runtime identity to streamed status events as well as numeric progress; otherwise loader and other non-numeric phases know the active node but cannot display its Modal task ID. diff --git a/remote/modal_app.py b/remote/modal_app.py index 95c1be8..4e9d65c 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -3560,6 +3560,9 @@ def _emit_local_modal_status( status_current: int | None = None, status_total: int | None = None, completed_ancestor_node_ids: list[str] | None = None, + execution_provider: str | None = None, + execution_environment_id: str | None = None, + execution_location: str | None = None, ) -> None: """Forward remote execution progress into the local ComfyUI websocket stream.""" if client_id is None: @@ -3592,6 +3595,12 @@ def _emit_local_modal_status( payload["status_total"] = int(status_total) if completed_ancestor_node_ids: payload["completed_ancestor_node_ids"] = list(completed_ancestor_node_ids) + if execution_provider: + payload["execution_provider"] = execution_provider + if execution_environment_id: + payload["execution_environment_id"] = execution_environment_id + if execution_location: + payload["execution_location"] = execution_location _record_local_modal_ui_event("modal_status", payload, client_id) prompt_server.send_sync("modal_status", payload, client_id) @@ -3607,6 +3616,31 @@ def _remote_execution_destination(payload: Mapping[str, Any]) -> str: ) +def _remote_execution_identity( + payload: Mapping[str, Any], + modal_task_id: str | None = None, +) -> dict[str, str]: + """Return non-empty provider, environment, and runtime location fields.""" + provider = str( + payload.get("execution_provider") or ("modal" if modal_task_id else "") + ).strip() + environment_id = str(payload.get("execution_environment_id") or "").strip() + location = ( + modal_task_id + if provider == "modal" and modal_task_id + else str(payload.get("execution_location") or "").strip() + ) + return { + key: value + for key, value in { + "execution_provider": provider, + "execution_environment_id": environment_id, + "execution_location": location, + }.items() + if value + } + + def _emit_local_remote_dispatch_status(payload: dict[str, Any]) -> None: """Tell the local UI a remote component was dispatched before progress streams.""" _emit_local_remote_startup_status( @@ -3652,6 +3686,7 @@ def _emit_local_remote_startup_status( str(payload["modal_gpu"]) if payload.get("modal_gpu") is not None else None ), status_message=status_message, + **_remote_execution_identity(payload), ) @@ -3679,6 +3714,9 @@ def _emit_local_modal_progress( time_to_first_token_seconds: float | None = None, tokens_per_second: float | None = None, pre_gpu: bool = False, + execution_provider: str | None = None, + execution_environment_id: str | None = None, + execution_location: str | None = None, ) -> None: """Forward remote numeric and stage progress into the local websocket stream.""" if client_id is None: @@ -3730,6 +3768,12 @@ def _emit_local_modal_progress( payload["tokens_per_second"] = float(tokens_per_second) if pre_gpu: payload["pre_gpu"] = True + if execution_provider: + payload["execution_provider"] = execution_provider + if execution_environment_id: + payload["execution_environment_id"] = execution_environment_id + if execution_location: + payload["execution_location"] = execution_location _record_local_modal_ui_event("modal_progress", payload, client_id) prompt_server.send_sync("modal_progress", payload, client_id) @@ -4435,6 +4479,7 @@ def _consume_remote_payload_stream( suppress_status_stream = bool(payload.get("suppress_status_stream")) result_payload: bytes | bytearray | None = None suppressed_progress_node_metadata: dict[str, dict[str, str | None]] = {} + active_remote_task_id: str | None = None active_remote_log_task_id: str | None = None should_close_stream = False component_id = str(payload.get("component_id") or "payload") @@ -4466,6 +4511,8 @@ def _consume_remote_payload_stream( event_kind = str(stream_event.get("kind", "")) if event_kind == "remote_logs": task_id = _coerce_modal_task_id(stream_event.get("task_id")) + if task_id is not None and active_remote_task_id is None: + active_remote_task_id = task_id if ( task_id is not None and active_remote_log_task_id is None @@ -4577,6 +4624,9 @@ def _consume_remote_payload_stream( ), "aggregate_only": aggregate_only, } + progress_kwargs.update( + _remote_execution_identity(payload, active_remote_task_id) + ) for string_field in ("stage", "message", "unit"): if stream_event.get(string_field) is not None: progress_kwargs[string_field] = str( @@ -4749,6 +4799,9 @@ def _consume_remote_payload_stream( display_node_id=display_node_id, real_node_id=real_node_id, cached_hit=True, + **_remote_execution_identity( + payload, active_remote_task_id + ), ) continue logger.info( @@ -4791,6 +4844,7 @@ def _consume_remote_payload_stream( "Receiving remote outputs from " f"{_remote_execution_destination(payload)}" ), + **_remote_execution_identity(payload, active_remote_task_id), ) continue _emit_local_modal_status( @@ -4821,6 +4875,7 @@ def _consume_remote_payload_stream( if stream_event.get("active_node_role") is not None else None ), + **_remote_execution_identity(payload, active_remote_task_id), ) continue if event_kind == "result": diff --git a/tests/modal_toggle_frontend.test.mjs b/tests/modal_toggle_frontend.test.mjs index 7c1b413..36c53c3 100644 --- a/tests/modal_toggle_frontend.test.mjs +++ b/tests/modal_toggle_frontend.test.mjs @@ -42,6 +42,7 @@ const transformedSource = `${[ "globalThis.__modalToggleExports = {", " ensurePromptState,", " registerPromptComponents,", + " registerPromptExecutionAssignments,", " resolveComponentNodeIds,", " handleModalProgress,", " handleModalStatus,", @@ -69,6 +70,7 @@ const transformedSource = `${[ " hasRemoteDescendants,", " remoteContainerTooltip,", " drawModalNodeDecoration,", + " updateVueExecutionLocation,", " currentGlobalStatus,", " formatIterationRate,", " fadeNodeProgress,", @@ -213,6 +215,170 @@ const progressCanvasContext = { assert.doesNotThrow(() => remoteProgressNode.onDrawForeground(progressCanvasContext)); assert.match(progressPanelStrokeStyles.at(-1), /^rgba\(168, 85, 247, 0\.44/); +resetFrontendState(); +const mixedLocationNode = { + id: 177, + comfyClass: "KSampler", + properties: { is_modal_remote: true }, + size: [132, 80], +}; +modalToggle.registerPromptComponents("prompt-mixed-location", ["177", "178"], [ + { representative_node_id: "177", node_ids: ["177"] }, + { representative_node_id: "178", node_ids: ["178"] }, +]); +modalToggle.registerPromptExecutionAssignments("prompt-mixed-location", { + "177": { + provider: "modal", + environment_id: "modal:B300", + node_ids: ["177"], + }, + "178": { + provider: "ssh_docker", + environment_id: "spark-one", + execution_location: "spark-one.internal.example", + node_ids: ["178"], + }, +}); +modalToggle.handleModalProgress({ + detail: { + prompt_id: "prompt-mixed-location", + node_id: "177", + value: 2, + max: 10, + execution_provider: "modal", + execution_environment_id: "modal:B300", + execution_location: "ta-01K3M0DALCONTAINERIDENTITY", + }, +}); +const mixedLocationState = modalToggle.getRemoteVisualState(mixedLocationNode); +assert.equal(mixedLocationState?.scheduledEnvironmentCount, 2); +assert.deepEqual(mixedLocationState?.executionLocation, { + provider: "modal", + label: "ta-01K3M0DALCONTAINERIDENTITY", +}); +const vueLocationIcon = { + hidden: true, + attributes: new Map(), + getAttribute(name) { + return this.attributes.get(name) ?? null; + }, + setAttribute(name, value) { + this.attributes.set(name, value); + }, +}; +const vueLocationLabel = { textContent: "" }; +const vueLocationElement = { + hidden: true, + dataset: {}, + title: "", + removeAttribute(name) { + if (name === "title") { + this.title = ""; + } + }, + querySelector(selector) { + if (selector === ".comfy-modal-vue-execution-location-icon") { + return vueLocationIcon; + } + if (selector === ".comfy-modal-vue-execution-location-label") { + return vueLocationLabel; + } + return null; + }, +}; +const vueLocationDecoration = { + querySelector(selector) { + return selector === ".comfy-modal-vue-execution-location" + ? vueLocationElement + : null; + }, +}; +modalToggle.updateVueExecutionLocation(vueLocationDecoration, mixedLocationState); +assert.equal(vueLocationElement.hidden, false); +assert.equal(vueLocationElement.dataset.provider, "modal"); +assert.equal(vueLocationElement.title, "ta-01K3M0DALCONTAINERIDENTITY"); +assert.equal(vueLocationLabel.textContent, "ta-01K3M0DALCONTAINERIDENTITY"); +assert.match(vueLocationIcon.getAttribute("src"), /^data:image\/svg\+xml/); +const mixedLocationLabels = []; +const mixedLocationCanvasContext = { + globalAlpha: 1, + save() {}, + restore() {}, + beginPath() {}, + roundRect() {}, + arc() {}, + fill() {}, + fillRect() {}, + stroke() {}, + fillText(text) { + mixedLocationLabels.push(String(text)); + }, + measureText(text) { + return { width: String(text).length * 6 }; + }, +}; +modalToggle.drawModalNodeDecoration(mixedLocationNode, mixedLocationCanvasContext); +assert.equal(mixedLocationLabels.some((label) => label.endsWith("…")), true); + +modalToggle.handleModalProgress({ + detail: { + prompt_id: "prompt-mixed-location", + node_id: "178", + value: 1, + max: 4, + execution_provider: "ssh_docker", + execution_environment_id: "spark-one", + execution_location: "spark-one.internal.example", + }, +}); +assert.deepEqual(modalToggle.getRemoteVisualState({ id: 178 })?.executionLocation, { + provider: "ssh_docker", + label: "spark-one.internal.example", +}); +modalToggle.updateVueExecutionLocation( + vueLocationDecoration, + modalToggle.getRemoteVisualState({ id: 178 }), +); +assert.equal(vueLocationElement.dataset.provider, "ssh_docker"); +assert.equal(vueLocationLabel.textContent, "spark-one.internal.example"); +assert.match(vueLocationIcon.getAttribute("src"), /^data:image\/svg\+xml/); +assert.match(decodeURIComponent(vueLocationIcon.getAttribute("src")), /viewBox="0 0 73 73"/); + +resetFrontendState(); +modalToggle.registerPromptComponents("prompt-single-location", ["179"], [ + { representative_node_id: "179", node_ids: ["179"] }, +]); +modalToggle.registerPromptExecutionAssignments("prompt-single-location", { + "179": { + provider: "ssh_docker", + environment_id: "spark-one", + execution_location: "spark-one.internal.example", + node_ids: ["179"], + }, +}); +modalToggle.handleModalProgress({ + detail: { + prompt_id: "prompt-single-location", + node_id: "179", + value: 1, + max: 4, + }, +}); +const singleLocationLabels = []; +modalToggle.drawModalNodeDecoration( + { id: 179, size: [132, 80] }, + { + ...mixedLocationCanvasContext, + fillText(text) { + singleLocationLabels.push(String(text)); + }, + }, +); +assert.equal( + singleLocationLabels.some((label) => label.includes("spark-one")), + false, +); + resetFrontendState(); modalToggle.registerPromptComponents("prompt-a", ["10", "11", "12"], [ { diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index 916225e..2d8ee0d 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -13,6 +13,14 @@ from typing import Any import pytest +def test_ssh_hostname_extracts_safe_runtime_badge_label( + api_intercept_module: Any, +) -> None: + """Planner UI metadata should show the host rather than an SSH user target.""" + assert api_intercept_module._ssh_hostname("worker@example.internal") == "example.internal" + assert api_intercept_module._ssh_hostname("[2001:db8::17]") == "2001:db8::17" + + class _FakeRemoteModelNode: """Fake node that produces a non-transportable MODEL output.""" diff --git a/tests/test_modal_stream_progress.py b/tests/test_modal_stream_progress.py index d8bd693..624046c 100644 --- a/tests/test_modal_stream_progress.py +++ b/tests/test_modal_stream_progress.py @@ -196,3 +196,82 @@ def test_llm_stream_progress_preserves_stage_and_token_metrics( "tokens_per_second": 9.75, } ] + + +def test_stream_progress_reports_exact_modal_container_identity(monkeypatch: Any) -> None: + """A Modal task id should follow its component progress into the node badge.""" + modal_app = _load_modal_app_module() + emitted_progress: list[dict[str, Any]] = [] + monkeypatch.setattr( + modal_app, + "_emit_local_modal_progress", + lambda **kwargs: emitted_progress.append(kwargs), + ) + payload = { + "prompt_id": "prompt-location", + "component_id": "node-162", + "component_node_ids": ["node-162"], + "execution_provider": "modal", + "execution_environment_id": "modal:B300", + "extra_data": {"client_id": "client-1"}, + } + + modal_app._consume_remote_payload_stream( + payload, + iter( + [ + {"kind": "remote_logs", "task_id": "ta-01K3MODAL"}, + { + "kind": "progress", + "event_type": "node_progress", + "node_id": "node-162", + "value": 1, + "max": 8, + }, + {"kind": "result", "outputs": ["done"]}, + ] + ), + ) + + assert emitted_progress[0]["execution_provider"] == "modal" + assert emitted_progress[0]["execution_environment_id"] == "modal:B300" + assert emitted_progress[0]["execution_location"] == "ta-01K3MODAL" + + +def test_stream_status_reports_exact_modal_container_identity(monkeypatch: Any) -> None: + """Node status should expose a Modal task id before numeric progress begins.""" + modal_app = _load_modal_app_module() + emitted_status: list[dict[str, Any]] = [] + monkeypatch.setattr( + modal_app, + "_emit_local_modal_status", + lambda **kwargs: emitted_status.append(kwargs), + ) + payload = { + "prompt_id": "prompt-location", + "component_id": "node-162", + "component_node_ids": ["node-162"], + "execution_provider": "modal", + "execution_environment_id": "modal:B300", + "extra_data": {"client_id": "client-1"}, + } + + modal_app._consume_remote_payload_stream( + payload, + iter( + [ + {"kind": "remote_logs", "task_id": "ta-01K3MODAL"}, + { + "kind": "progress", + "event_type": "status", + "phase": "executing", + "active_node_id": "node-162", + }, + {"kind": "result", "outputs": ["done"]}, + ] + ), + ) + + assert emitted_status[0]["execution_provider"] == "modal" + assert emitted_status[0]["execution_environment_id"] == "modal:B300" + assert emitted_status[0]["execution_location"] == "ta-01K3MODAL" diff --git a/tests/test_web_modal_toggle.py b/tests/test_web_modal_toggle.py index 9032d19..b195900 100644 --- a/tests/test_web_modal_toggle.py +++ b/tests/test_web_modal_toggle.py @@ -417,6 +417,22 @@ def test_streamed_modal_node_progress_updates_active_overlay() -> None: assert "function fadeNodeProgress(nodeIdValue, promptId)" in source assert "function progressVisualOpacity(progressState)" in source assert 'api.addEventListener("modal_progress", handleModalProgress);' in source + + +def test_mixed_remote_environments_render_clipped_runtime_locations() -> None: + """Mixed planner assignments should add provider identity beneath progress bars.""" + source = _modal_toggle_source() + + assert "function registerPromptExecutionAssignments(promptId, assignments)" in source + assert "promptState.scheduledEnvironmentCount = environmentIds.size;" in source + assert "updateNodeExecutionLocation(promptId, progressNodeId, detail);" in source + assert "updateNodeExecutionLocation(promptId, nextActiveNodeId, detail);" in source + assert "function drawRemoteExecutionLocation(" in source + assert "function updateVueExecutionLocation(decoration, state)" in source + assert 'text-overflow: ellipsis;' in source + assert "Number(state?.scheduledEnvironmentCount ?? 0) > 1" in source + assert "fitCanvasText(ctx, label, availableTextWidth)" in source + assert "REMOTE_LOCATION_ICON_SOURCES" in source assert "function handleModalProgress(event)" in source assert "if (detail.aggregate_only) {" in source assert "setNodeBatchProgress(" in source diff --git a/web/modal_toggle.js b/web/modal_toggle.js index 10cc86e..41952a6 100644 --- a/web/modal_toggle.js +++ b/web/modal_toggle.js @@ -56,6 +56,14 @@ const FINALIZING_NODE_BORDER_COLOR = "#00358A"; const ERROR_BORDER_COLOR = "#ef4444"; const LOCAL_BOTTLENECK_BADGE_BORDER_COLOR = "rgba(148, 163, 184, 0.72)"; const LOCAL_BOTTLENECK_TOOLTIP = "Did you mean to make this node execute on Modal?"; +const REMOTE_LOCATION_ICON_SOURCES = { + modal: `data:image/svg+xml;charset=utf-8,${encodeURIComponent( + '', + )}`, + ssh_docker: `data:image/svg+xml;charset=utf-8,${encodeURIComponent( + '', + )}`, +}; const STATE_SETUP = "setup"; const STATE_STARTING = "starting"; @@ -92,6 +100,7 @@ const modalSandwichedLocalNodeIds = new Set(); const modalRemoteDescendantNodeIdsByAncestor = new Map(); const syntheticPromptUiStates = new Map(); const modalGlobalStatusStates = new Map(); +const remoteLocationIconImages = new Map(); let animationFrameHandle = null; let modalLastAnimationRedrawAt = 0; @@ -296,6 +305,27 @@ function fitCanvasText(ctx, text, maxWidth) { return `${safeText.slice(0, end)}…`; } +/** + * Return a lazily loaded provider icon for the canvas location line. + * @param {string | null | undefined} provider + * @returns {HTMLImageElement | null} + */ +function remoteLocationIcon(provider) { + const safeProvider = String(provider ?? ""); + const source = REMOTE_LOCATION_ICON_SOURCES[safeProvider]; + if (!source || typeof Image !== "function") { + return null; + } + if (!remoteLocationIconImages.has(safeProvider)) { + const icon = new Image(); + icon.decoding = "async"; + icon.src = source; + icon.addEventListener?.("load", refreshNodeDecorations, { once: true }); + remoteLocationIconImages.set(safeProvider, icon); + } + return remoteLocationIconImages.get(safeProvider) ?? null; +} + /** * Draw an iteration-rate label over a progress bar with a tight opaque backing. * @param {CanvasRenderingContext2D} ctx @@ -1755,7 +1785,7 @@ function refreshModalUiAfterVisibilityChange() { /** * Return the prompt metadata bucket, creating it if needed. * @param {string} promptId - * @returns {{ startedAt: number, modalGpu: string | null, executionLabel: string | null, remoteNodeIds: string[], componentsByRepresentative: Map, componentNodeIdsByMember: Map, representativeNodeIdByMember: Map, componentLabelByMember: Map, laneNodeIdsByLane: Map }} + * @returns {{ startedAt: number, modalGpu: string | null, executionLabel: string | null, scheduledEnvironmentCount: number, remoteNodeIds: string[], componentsByRepresentative: Map, componentNodeIdsByMember: Map, representativeNodeIdByMember: Map, componentLabelByMember: Map, executionAssignmentsByRepresentative: Map, executionAssignmentByMember: Map, laneNodeIdsByLane: Map }} */ function ensurePromptState(promptId) { if (isPromptTerminal(promptId)) { @@ -1767,11 +1797,14 @@ function ensurePromptState(promptId) { startedAt, modalGpu: null, executionLabel: null, + scheduledEnvironmentCount: 0, remoteNodeIds: [], componentsByRepresentative: new Map(), componentNodeIdsByMember: new Map(), representativeNodeIdByMember: new Map(), componentLabelByMember: new Map(), + executionAssignmentsByRepresentative: new Map(), + executionAssignmentByMember: new Map(), descendantNodeIdsByAncestor: new Map(), laneNodeIdsByLane: new Map(), activeNodeId: null, @@ -2094,6 +2127,82 @@ function registerPromptComponents(promptId, remoteNodeIds, components) { } } +/** + * Store planner assignments by both representative and visible member node. + * @param {string} promptId + * @param {Record>} assignments + */ +function registerPromptExecutionAssignments(promptId, assignments) { + const promptState = ensurePromptState(promptId); + if (!promptState) { + return; + } + promptState.executionAssignmentsByRepresentative.clear(); + promptState.executionAssignmentByMember.clear(); + const environmentIds = new Set(); + + for (const [representativeNodeIdValue, rawAssignment] of Object.entries(assignments ?? {})) { + if (!rawAssignment || typeof rawAssignment !== "object") { + continue; + } + const representativeNodeId = String(representativeNodeIdValue); + const environmentId = String(rawAssignment.environment_id ?? "").trim(); + const assignment = { + provider: String(rawAssignment.provider ?? "modal"), + environmentId, + location: String(rawAssignment.execution_location ?? "").trim() || null, + }; + if (environmentId) { + environmentIds.add(environmentId); + } + promptState.executionAssignmentsByRepresentative.set(representativeNodeId, assignment); + const assignedNodeIds = Array.isArray(rawAssignment.node_ids) && rawAssignment.node_ids.length > 0 + ? rawAssignment.node_ids.map((nodeIdValue) => String(nodeIdValue)) + : (promptState.componentNodeIdsByMember.get(representativeNodeId) ?? [representativeNodeId]); + for (const assignedNodeId of assignedNodeIds) { + promptState.executionAssignmentByMember.set(assignedNodeId, assignment); + } + promptState.executionAssignmentByMember.set(representativeNodeId, assignment); + } + promptState.scheduledEnvironmentCount = environmentIds.size; +} + +/** + * Update one component's runtime identity from its streamed progress metadata. + * @param {string} promptId + * @param {string} nodeIdValue + * @param {Record} metadata + */ +function updateNodeExecutionLocation(promptId, nodeIdValue, metadata) { + const promptState = modalPromptStates.get(promptId); + const location = String(metadata.execution_location ?? "").trim(); + if (!promptState || !location) { + return; + } + const safeNodeIdValue = String(nodeIdValue); + const representativeNodeId = + promptState.representativeNodeIdByMember.get(safeNodeIdValue) ?? safeNodeIdValue; + const existingAssignment = + promptState.executionAssignmentByMember.get(safeNodeIdValue) ?? + promptState.executionAssignmentsByRepresentative.get(representativeNodeId) ?? + {}; + const assignment = { + ...existingAssignment, + provider: String(metadata.execution_provider ?? existingAssignment.provider ?? "modal"), + environmentId: String( + metadata.execution_environment_id ?? existingAssignment.environmentId ?? "", + ), + location, + }; + promptState.executionAssignmentsByRepresentative.set(representativeNodeId, assignment); + const componentNodeIds = + promptState.componentNodeIdsByMember.get(safeNodeIdValue) ?? [safeNodeIdValue]; + for (const componentNodeId of componentNodeIds) { + promptState.executionAssignmentByMember.set(String(componentNodeId), assignment); + } + promptState.executionAssignmentByMember.set(representativeNodeId, assignment); +} + /** * Return ancestor node ids for a composed subgraph prompt id like `24:23`. * @param {string} nodeIdValue @@ -3490,6 +3599,8 @@ function getRemoteVisualState(node) { const cachedState = nodeCachedState(visualNodeId, state.promptId); const batchProgressState = modalNodeBatchProgress.get(visualNodeId) ?? null; const hasLiveProgress = hasLiveNodeProgress(visualNodeId, state.promptId); + const executionAssignment = + promptState?.executionAssignmentByMember.get(visualNodeId) ?? null; return { ...state, phase: deriveRemoteNodePhase(state.phase, hasLiveProgress), @@ -3503,6 +3614,13 @@ function getRemoteVisualState(node) { progress: progressState, batchProgress: batchProgressState?.promptId === state.promptId ? batchProgressState : null, progressLanes, + scheduledEnvironmentCount: promptState?.scheduledEnvironmentCount ?? 0, + executionLocation: executionAssignment?.location + ? { + provider: executionAssignment.provider, + label: executionAssignment.location, + } + : null, }; } @@ -3740,19 +3858,72 @@ function remoteContainerTooltip(state) { */ function ensureVueNodeDecoration(nodeElement) { let decoration = nodeElement.querySelector(":scope > .comfy-modal-vue-node-decoration"); - if (decoration) { + if (!decoration) { + decoration = document.createElement("div"); + decoration.className = "comfy-modal-vue-node-decoration"; + decoration.setAttribute("aria-hidden", "true"); + const badge = document.createElement("span"); + badge.className = "comfy-modal-vue-node-badge"; + decoration.appendChild(badge); + nodeElement.appendChild(decoration); + } + if (decoration.querySelector(".comfy-modal-vue-execution-location")) { return decoration; } - decoration = document.createElement("div"); - decoration.className = "comfy-modal-vue-node-decoration"; - decoration.setAttribute("aria-hidden", "true"); - const badge = document.createElement("span"); - badge.className = "comfy-modal-vue-node-badge"; - decoration.appendChild(badge); - nodeElement.appendChild(decoration); + const executionLocation = document.createElement("div"); + executionLocation.className = "comfy-modal-vue-execution-location"; + executionLocation.hidden = true; + const executionLocationIcon = document.createElement("img"); + executionLocationIcon.className = "comfy-modal-vue-execution-location-icon"; + executionLocationIcon.alt = ""; + executionLocationIcon.setAttribute("aria-hidden", "true"); + const executionLocationLabel = document.createElement("span"); + executionLocationLabel.className = "comfy-modal-vue-execution-location-label"; + executionLocation.append(executionLocationIcon, executionLocationLabel); + decoration.appendChild(executionLocation); return decoration; } +/** + * Update the execution-location badge rendered by the Nodes 2.0 DOM layer. + * @param {HTMLDivElement} decoration + * @param {Record | null} state + */ +function updateVueExecutionLocation(decoration, state) { + const executionLocation = decoration.querySelector( + ".comfy-modal-vue-execution-location", + ); + if (!executionLocation) { + return; + } + const location = + Number(state?.scheduledEnvironmentCount ?? 0) > 1 && + (state?.isActiveRemoteNode || state?.isActiveComponentMember) + ? state?.executionLocation + : null; + const label = String(location?.label ?? ""); + const provider = String(location?.provider ?? ""); + const iconSource = REMOTE_LOCATION_ICON_SOURCES[provider] ?? ""; + executionLocation.hidden = !label; + if (!label) { + executionLocation.removeAttribute("title"); + return; + } + executionLocation.dataset.provider = provider; + executionLocation.title = label; + const icon = executionLocation.querySelector(".comfy-modal-vue-execution-location-icon"); + const labelElement = executionLocation.querySelector(".comfy-modal-vue-execution-location-label"); + if (icon) { + icon.hidden = !iconSource; + if (iconSource && icon.getAttribute("src") !== iconSource) { + icon.setAttribute("src", iconSource); + } + } + if (labelElement) { + labelElement.textContent = label; + } +} + /** * Apply Modal's current visual state to one Nodes 2.0 DOM node. * @param {HTMLElement} nodeElement @@ -3798,6 +3969,7 @@ function updateVueNodeDecoration(nodeElement, node, timestamp) { ? LOCAL_BOTTLENECK_BADGE_BORDER_COLOR : palette.borderColor; badge.style.boxShadow = localBottleneck ? "none" : `0 0 8px ${palette.shadowColor}`; + updateVueExecutionLocation(decoration, state); } /** @@ -3939,6 +4111,51 @@ function updateLegacyModalTooltip(graphCanvas, tooltip) { } } +/** + * Draw a clipped provider icon and runtime identity beneath node progress bars. + * @param {CanvasRenderingContext2D} ctx + * @param {{ provider?: string, label?: string }} location + * @param {number} rowY + * @param {number} rowHeight + * @param {number} nodeWidth + * @param {number} paddingX + * @param {number} scale + */ +function drawRemoteExecutionLocation( + ctx, + location, + rowY, + rowHeight, + nodeWidth, + paddingX, + scale, +) { + const provider = String(location?.provider ?? ""); + const label = String(location?.label ?? ""); + if (!label) { + return; + } + const icon = remoteLocationIcon(provider); + const iconHeight = 10 / scale; + const iconWidth = (provider === "modal" ? 19 : 10) / scale; + const iconX = paddingX; + const iconY = rowY + (rowHeight - iconHeight) / 2; + if (icon?.complete && Number(icon.naturalWidth) > 0 && typeof ctx.drawImage === "function") { + ctx.drawImage(icon, iconX, iconY, iconWidth, iconHeight); + } + const textX = iconX + iconWidth + 5 / scale; + const availableTextWidth = Math.max(0, nodeWidth - paddingX - textX); + ctx.fillStyle = "#cbd5e1"; + ctx.font = `${Math.max(8 / scale, 6.5)}px ui-sans-serif, system-ui, sans-serif`; + ctx.textAlign = "left"; + ctx.textBaseline = "middle"; + ctx.fillText( + fitCanvasText(ctx, label, availableTextWidth), + textX, + rowY + rowHeight / 2, + ); +} + /** * Draw the Modal execution decoration or neutral local-bottleneck badge for a node. * @param {LGraphNode} node @@ -4065,13 +4282,25 @@ function drawModalNodeDecoration(node, ctx) { const visibleLaneProgress = hasLaneProgress ? visibleActiveProgressLanes : setupProgressLanes; const hasVisibleLaneProgress = visibleLaneProgress.length > 0; const hasIterationRateLabels = hasAggregateProgress || hasLaneProgress; + const executionLocation = + Number(state?.scheduledEnvironmentCount ?? 0) > 1 ? state?.executionLocation : null; + const hasExecutionLocation = Boolean(executionLocation?.label); const progressBarWidth = barWidth; const laneBlockHeight = hasVisibleLaneProgress ? visibleLaneProgress.length * laneHeight + (visibleLaneProgress.length - 1) * laneGap : 0; - const bodyHeight = + const progressBodyHeight = (hasVisibleLaneProgress ? laneBlockHeight + laneGap : 0) + (hasAggregateProgress ? aggregateHeight : 0); - const panelHeight = panelPaddingY * 2 + headerHeight + bodyHeight; + const progressTopGap = progressBodyHeight > 0 ? laneGap : 0; + const locationRowHeight = 13 / scale; + const locationTopGap = hasExecutionLocation ? 5 / scale : 0; + const panelHeight = + panelPaddingY * 2 + + headerHeight + + progressTopGap + + progressBodyHeight + + locationTopGap + + (hasExecutionLocation ? locationRowHeight : 0); const laneColors = [ "rgba(196, 181, 253, 0.94)", "rgba(147, 197, 253, 0.94)", @@ -4154,7 +4383,7 @@ function drawModalNodeDecoration(node, ctx) { ctx.fillText(badgeText, badgeX + badgeWidth / 2, badgeY + badgeHeight / 2); } - const barY = panelY + panelPaddingY + headerHeight + laneGap; + const barY = panelY + panelPaddingY + headerHeight + progressTopGap; if (hasIterationRateLabels) { ctx.font = `${Math.max(9 / scale, 7)}px ui-sans-serif, system-ui, sans-serif`; @@ -4219,6 +4448,24 @@ function drawModalNodeDecoration(node, ctx) { state.progress.unit, ); } + if (hasExecutionLocation) { + const locationY = + panelY + + panelPaddingY + + headerHeight + + progressTopGap + + progressBodyHeight + + locationTopGap; + drawRemoteExecutionLocation( + ctx, + executionLocation, + locationY, + locationRowHeight, + node.size[0], + panelPaddingX, + scale, + ); + } ctx.restore(); } @@ -4452,6 +4699,7 @@ function handleModalStatus(event) { nextActiveNodeId, ); if (nextActiveNodeId) { + updateNodeExecutionLocation(promptId, nextActiveNodeId, detail); setNodesPhase([nextActiveNodeId], STATE_ACTIVE, promptId); } setPromptActiveNode(promptId, nextActiveNodeId); @@ -4513,6 +4761,7 @@ function handleModalProgress(event) { if (!detail.cached_hit && !detail.pre_gpu) { promptState.hasRemoteExecutionStarted = true; } + updateNodeExecutionLocation(promptId, progressNodeId, detail); const componentNodeIds = resolveComponentNodeIds(promptId, progressNodeId); const readyNodeIds = (componentNodeIds ?? []).filter((nodeIdValue) => nodeIdValue !== progressNodeId); promptState.hasStreamedProgress = true; @@ -4999,9 +5248,9 @@ function patchQueuePrompt() { const acceptedModalGpu = MODAL_GPU_TYPES.includes(responsePayload.modal_gpu) ? responsePayload.modal_gpu : modalGpu; - const executionAssignments = Object.values( - responsePayload.remote_execution_assignments ?? {}, - ); + const executionAssignmentsByRepresentative = + responsePayload.remote_execution_assignments ?? {}; + const executionAssignments = Object.values(executionAssignmentsByRepresentative); const usesSshExecution = executionAssignments.some( (assignment) => assignment?.provider === "ssh_docker", ); @@ -5023,6 +5272,7 @@ function patchQueuePrompt() { if (resolvedRemoteNodeIds.length > 0 || resolvedComponents.length > 0) { registerPromptComponents(promptId, resolvedRemoteNodeIds, resolvedComponents); } + registerPromptExecutionAssignments(promptId, executionAssignmentsByRepresentative); const promptState = ensurePromptState(promptId); if (!promptState) { return responsePayload; @@ -5194,6 +5444,56 @@ function installGlobalStatusStyles() { .comfy-modal-vue-node-badge[hidden] { display: none; } + + .comfy-modal-vue-execution-location { + position: absolute; + top: calc(100% + 6px); + left: 6px; + right: 6px; + display: flex; + min-width: 0; + height: 24px; + align-items: center; + gap: 6px; + box-sizing: border-box; + overflow: hidden; + padding: 5px 7px; + border: 1px solid rgba(148, 163, 184, 0.34); + border-radius: 8px; + background: rgba(15, 23, 42, 0.94); + box-shadow: 0 3px 10px rgba(2, 6, 23, 0.42); + color: #cbd5e1; + font: 500 10px/1 ui-sans-serif, system-ui, sans-serif; + letter-spacing: 0.01em; + white-space: nowrap; + } + + .comfy-modal-vue-execution-location[hidden] { + display: none; + } + + .comfy-modal-vue-execution-location-icon { + width: 12px; + height: 12px; + flex: 0 0 auto; + object-fit: contain; + } + + .comfy-modal-vue-execution-location[data-provider="modal"] + .comfy-modal-vue-execution-location-icon { + width: 22px; + } + + .comfy-modal-vue-execution-location-icon[hidden] { + display: none; + } + + .comfy-modal-vue-execution-location-label { + min-width: 0; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; + } `; document.head.appendChild(style); }