From b0f5eff550faeb9653fe4c7a9ce8fcb597013b2f Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Thu, 23 Apr 2026 15:39:51 -0700 Subject: [PATCH] feat: surface active and cached modal nodes --- README.md | 2 + comfyui_modal_sync_cloud.py | 23 +++++- docs/LEARNINGS.md | 1 + remote/modal_app.py | 48 ++++++++++++ tests/test_modal_executor_node.py | 79 +++++++++++++++++++ tests/test_web_modal_toggle.py | 31 +++++++- web/modal_toggle.js | 125 ++++++++++++++++++++++++++++-- 7 files changed, 299 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 52c79ad..9bfa927 100644 --- a/README.md +++ b/README.md @@ -311,6 +311,8 @@ Current mapped-execution rules: - mapped node highlighting now derives its purple/active state from live per-node lane or numeric progress, so overlapping `A -> B` work across different containers can keep both remote nodes active at once while ordinary non-mapped runs still use the older single-active-node behavior - worker lanes no longer emit a fake execution-progress bar on the component representative node before real node progress arrives, which avoids leaving stray lane bars behind on the first node in the remote island - mapped lane setup now emits its own `setup_only` lane events before a worker starts real node execution, and the frontend renders those as translucent pulsing lane bars. That makes cold-lane provisioning visible without pretending the worker is already doing sampler progress. +- ready-state node outlines now distinguish component membership: bright green means the node belongs to the component that is currently running, while dark green means the node is remote-ready but belongs to some other component waiting its turn +- persisted Modal node-cache hits now emit lightweight `cached_hit` UI markers, and the frontend renders those nodes with a slower green pulse so users can see when outputs were reused instead of recomputed - the proxy node itself now opts into ComfyUI `INPUT_IS_LIST` handling and unwraps singleton list wrappers on both `original_node_data` and ordinary inputs before dispatch, so list-valued mapped inputs reach the internal Modal scheduler without causing the whole proxy node to be auto-mapped once per item - static and per-item sub-runs prune the rewritten prompt to only the nodes they actually depend on before handing it to `PromptExecutor`, so shared-upstream hybrid components do not validate the other branch's batched inputs during the wrong sub-run diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index d5c4412..62641e7 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -24,7 +24,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, Iterator, Mapping +from typing import Any, Callable, Iterator, Mapping, Sequence _REPO_ROOT = Path(__file__).resolve().parent _REMOTE_REPO_ROOT = Path("/root/comfyui_modal_sync_repo") @@ -2522,6 +2522,26 @@ def _install_prompt_executor_persisted_cache_restore( ) +def _emit_restored_node_cache_events( + status_callback: Callable[[dict[str, Any]], None] | None, + restored_node_ids: Sequence[str], +) -> None: + """Publish one status event per node whose outputs were restored from the persisted cache.""" + if status_callback is None: + return + + for node_id in restored_node_ids: + safe_node_id = str(node_id) + status_callback( + { + "event_type": "node_cached", + "node_id": safe_node_id, + "display_node_id": safe_node_id, + "real_node_id": safe_node_id, + } + ) + + def _persist_node_output_cache_entries( executor: Any, *, @@ -3283,6 +3303,7 @@ def _execute_subgraph_prompt( component_id, restored_node_ids, ) + _emit_restored_node_cache_events(status_callback, restored_node_ids) if not executor.success: _log_prompt_executor_failure_details( component_id=component_id, diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 1ee90eb..620151b 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -210,4 +210,5 @@ - The in-process "already auto-deployed" cache cannot be treated as authoritative once operators can delete the Modal app out-of-band. A later `app not found` lookup has to invalidate that cache entry and redeploy instead of assuming the earlier successful deploy still exists remotely. - With deployed `modal.Cls` handles, "app missing" is not only a lookup-time failure. A stale handle can also survive lookup and then fail later during `hydrate()` or the first `remote()` call. The invoke path has to treat that the same way as a missing-app lookup, auto-deploy, and retry before surfacing an error locally. - Mapped lane observability needs a distinct setup state, not just "first real progress wins." If the backend emits a dedicated `setup_only` lane progress event as soon as a worker starts seeding, the frontend can render a translucent pulsing placeholder bar while the lane is provisioning and then seamlessly replace it with real node progress once execution begins. +- Remote node UI needs two separate non-executing signals to stay legible under Modal fan-out: "ready but not in the active component" and "reused from the persisted node cache." The former should be derived locally from component membership and `active_node_id`, while the latter needs an explicit backend event because cached nodes often never emit ordinary execution progress at all. - Combined first-run and second-run performance logs are now clear enough to prioritize work. In this workflow, uploads and volume reload were already cheap, while the recurring non-sampling bottleneck was static-phase model loading (`UNETLoader` + `CLIPLoader` at roughly `24-25s` total per run). The second run also showed that local proxy caching can skip the static proxy dispatch while still leaving a large replay cost inside the mapped phase, which makes two no-brainer optimizations obvious: reuse heavy loader outputs on the same worker, and avoid replaying the whole static producer branch just to reconstruct session-backed bridge values. diff --git a/remote/modal_app.py b/remote/modal_app.py index 96295d8..2b50491 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -2078,6 +2078,7 @@ def _emit_local_modal_progress( item_index: int | None = None, aggregate_only: bool = False, setup_only: bool = False, + cached_hit: bool = False, ) -> None: """Forward remote numeric node progress into the local ComfyUI websocket stream.""" if client_id is None: @@ -2107,6 +2108,8 @@ def _emit_local_modal_progress( payload["aggregate_only"] = True if setup_only: payload["setup_only"] = True + if cached_hit: + payload["cached_hit"] = True prompt_server.send_sync("modal_progress", payload, client_id) @@ -2674,6 +2677,51 @@ def _consume_remote_payload_stream( image_value=deserialize_value(stream_event.get("value")), ) continue + if event_type == "node_cached": + progress_metadata = _progress_stream_event_metadata(stream_event) + filter_node_id = ( + progress_metadata["filter_node_id"] if progress_metadata is not None else None + ) + if ( + suppress_status_stream + and not _should_forward_suppressed_stream_event(payload, filter_node_id) + ): + logger.debug( + "Suppressing streamed Modal cached-node marker for component=%s node_id=%s real_node_id=%s because it does not belong to this mapped/static payload.", + payload.get("component_id"), + stream_event.get("node_id"), + stream_event.get("real_node_id"), + ) + continue + reported_node_id = ( + progress_metadata["node_id"] if progress_metadata is not None else None + ) + if reported_node_id is not None: + display_node_id = ( + progress_metadata["display_node_id"] + if progress_metadata is not None + else str(reported_node_id) + ) + real_node_id = ( + progress_metadata["real_node_id"] if progress_metadata is not None else None + ) + logger.debug( + "Forwarding streamed Modal cached-node marker for component=%s node_id=%s real_node_id=%s.", + payload.get("component_id"), + reported_node_id, + real_node_id, + ) + _emit_local_modal_progress( + prompt_id=prompt_id, + client_id=client_id, + node_id=str(reported_node_id), + value=0.0, + max_value=1.0, + display_node_id=display_node_id, + real_node_id=real_node_id, + cached_hit=True, + ) + continue logger.info( "Forwarding streamed Modal progress for component=%s phase=%s active_node_id=%s.", payload.get("component_id"), diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 1c108e5..dc8502d 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -4163,6 +4163,30 @@ def test_modal_cloud_tracing_prompt_server_emits_executed_outputs( ] +def test_modal_cloud_emits_restored_node_cache_events( + modal_cloud_module: Any, +) -> None: + """Persisted node-cache restores should surface one cached-node marker per restored node.""" + observed_updates: list[dict[str, Any]] = [] + + modal_cloud_module._emit_restored_node_cache_events(observed_updates.append, ["7", "9"]) + + assert observed_updates == [ + { + "event_type": "node_cached", + "node_id": "7", + "display_node_id": "7", + "real_node_id": "7", + }, + { + "event_type": "node_cached", + "node_id": "9", + "display_node_id": "9", + "real_node_id": "9", + }, + ] + + def test_modal_cloud_tracing_prompt_server_emits_boundary_image_outputs( modal_cloud_module: Any, ) -> None: @@ -7303,6 +7327,61 @@ def test_consume_remote_payload_stream_filters_static_sibling_ui_events_from_map ] +def test_consume_remote_payload_stream_forwards_cached_node_markers( + remote_modal_app_module: Any, + serialization_module: Any, + monkeypatch: Any, +) -> None: + """Streamed node-cache hits should mark the matching UI node as cached without fake progress.""" + progress_calls: list[dict[str, Any]] = [] + + monkeypatch.setattr( + remote_modal_app_module, + "_emit_local_modal_progress", + lambda **kwargs: progress_calls.append(kwargs), + ) + + payload = { + "component_id": "1::static", + "prompt_id": "prompt-1", + "component_node_ids": ["1", "2", "12", "4"], + "execute_node_ids": ["12", "4"], + "extra_data": {"client_id": "client-1"}, + "suppress_status_stream": True, + } + stream_events = iter( + [ + { + "kind": "progress", + "event_type": "node_cached", + "node_id": "1", + "display_node_id": "1", + "real_node_id": "12", + }, + { + "kind": "result", + "outputs": serialization_module.serialize_node_outputs(("done",)), + }, + ] + ) + + response = remote_modal_app_module._consume_remote_payload_stream(payload, stream_events) + + assert serialization_module.deserialize_node_outputs(response) == ("done",) + assert progress_calls == [ + { + "prompt_id": "prompt-1", + "client_id": "client-1", + "node_id": "1", + "value": 0.0, + "max_value": 1.0, + "display_node_id": "1", + "real_node_id": "12", + "cached_hit": True, + } + ] + + def test_modal_cloud_initializes_remote_comfy_runtime_once_per_custom_node_root( modal_cloud_module: Any, monkeypatch: Any, diff --git a/tests/test_web_modal_toggle.py b/tests/test_web_modal_toggle.py index aa51cb4..146dfe4 100644 --- a/tests/test_web_modal_toggle.py +++ b/tests/test_web_modal_toggle.py @@ -66,7 +66,8 @@ def test_remote_modal_uses_distinct_ready_active_and_complete_colors() -> None: """The frontend should distinguish ready, active, and completed remote nodes visually.""" source = _modal_toggle_source() - assert 'const READY_BORDER_COLOR = "#22c55e";' in source + assert 'const READY_ACTIVE_COMPONENT_BORDER_COLOR = "#22c55e";' in source + assert 'const READY_INACTIVE_COMPONENT_BORDER_COLOR = "#166534";' in source assert 'const ACTIVE_BORDER_COLOR = "#a855f7";' in source assert 'const COMPLETE_BORDER_COLOR = "#16a34a";' in source assert 'const STATE_WAITING = "waiting";' in source @@ -138,6 +139,8 @@ def test_streamed_modal_node_progress_updates_active_overlay() -> None: assert "setNodeProgressLane(" in source assert "function deriveRemoteNodePhase(phase, hasLiveProgress)" in source assert "isActiveRemoteNode: hasLiveProgress || promptState?.activeNodeId === nodeId(node)," in source + assert "isActiveComponentMember: isNodeInActiveComponent(state.promptId, nodeId(node))," in source + assert "isCachedRemoteNode: Boolean(cachedState)," in source assert 'const panelY = node.size[1] + 6 / scale;' in source assert 'ctx.roundRect(-borderWidth, panelY, barWidth, panelHeight, 10 / scale);' in source assert 'const headerText = hasAggregateProgress' in source @@ -178,6 +181,32 @@ def test_mapped_lane_setup_events_render_placeholder_lane_bars() -> None: assert "const hasVisibleLaneProgress = visibleLaneProgress.length > 0;" in source +def test_ready_nodes_distinguish_active_component_members_and_cached_nodes() -> None: + """Ready nodes should visually distinguish active-component membership and cache hits.""" + source = _modal_toggle_source() + + assert "function isNodeInActiveComponent(promptId, nodeIdValue)" in source + assert "const activeComponentNodeIds = resolveComponentNodeIds(promptId, activeNodeId);" in source + assert "const descendantNodeIds = promptState.descendantNodeIdsByAncestor.get(safeNodeIdValue);" in source + assert "const pulseRate = state?.isCachedRemoteNode ? 2 : 6;" in source + assert "borderColor = state?.isActiveComponentMember" in source + assert "READY_INACTIVE_COMPONENT_BORDER_COLOR" in source + assert "READY_ACTIVE_COMPONENT_BORDER_COLOR" in source + + +def test_cached_node_hits_are_marked_without_fake_progress() -> None: + """Cache-hit markers should be tracked separately from numeric progress.""" + source = _modal_toggle_source() + + assert "const modalNodeCachedStates = new Map();" in source + assert "function markNodeCached(nodeIdValue, promptId)" in source + assert "function clearNodeCached(nodeIdValue, promptId)" in source + assert "function nodeCachedState(nodeIdValue, promptId)" in source + assert "if (detail.cached_hit) {" in source + assert "markNodeCached(progressNodeId, promptId);" in source + assert "const hasCachedPulse = Array.from(modalNodeCachedStates.values()).length > 0;" in source + + def test_modal_context_menu_can_expand_required_upstream_nodes() -> None: """Right-clicking a node should offer a dry-run expansion action backed by the backend.""" source = _modal_toggle_source() diff --git a/web/modal_toggle.js b/web/modal_toggle.js index 20bae69..12320af 100644 --- a/web/modal_toggle.js +++ b/web/modal_toggle.js @@ -8,7 +8,8 @@ const INTERNAL_NODE_PREFIX = "ModalUniversalExecutor"; const IDLE_BORDER_COLOR = "#1d9bf0"; const SETUP_BORDER_COLOR = "#f59e0b"; -const READY_BORDER_COLOR = "#22c55e"; +const READY_ACTIVE_COMPONENT_BORDER_COLOR = "#22c55e"; +const READY_INACTIVE_COMPONENT_BORDER_COLOR = "#166534"; const ACTIVE_BORDER_COLOR = "#a855f7"; const COMPLETE_BORDER_COLOR = "#16a34a"; const ERROR_BORDER_COLOR = "#ef4444"; @@ -27,6 +28,7 @@ const modalNodeStates = new Map(); const modalNodeProgress = new Map(); const modalNodeProgressLanes = new Map(); const modalNodeBatchProgress = new Map(); +const modalNodeCachedStates = new Map(); const modalNodeClearTimers = new Map(); const modalPromptStates = new Map(); const syntheticPromptUiStates = new Map(); @@ -448,7 +450,8 @@ function refreshCanvasAnimation() { const hasAnimatedState = Array.from(modalNodeStates.values()).some((state) => [STATE_SETUP, STATE_READY, STATE_ACTIVE, STATE_ERROR].includes(state.phase), ); - if (!hasAnimatedState) { + const hasCachedPulse = Array.from(modalNodeCachedStates.values()).length > 0; + if (!hasAnimatedState && !hasCachedPulse) { animationFrameHandle = null; return; } @@ -521,6 +524,9 @@ function setNodesPhase(nodeIds, phase, promptId, errorMessage) { if (phase === STATE_ERROR) { scheduleNodeClear(currentNodeId, promptId, ERROR_CLEAR_DELAY_MS); } + if (phase === STATE_SETUP || phase === STATE_ERROR) { + clearNodeCached(currentNodeId, promptId); + } if (phase === STATE_COMPLETE || phase === STATE_ERROR) { clearNodeProgress(currentNodeId, promptId); } @@ -788,6 +794,56 @@ function clearNodeProgress(nodeIdValue, promptId) { } } +/** + * Return one cached-node marker when it belongs to the prompt. + * @param {string} nodeIdValue + * @param {string} promptId + * @returns {{ promptId: string, cachedAt: number } | null} + */ +function nodeCachedState(nodeIdValue, promptId) { + const cachedState = modalNodeCachedStates.get(String(nodeIdValue)) ?? null; + return cachedState?.promptId === promptId ? cachedState : null; +} + +/** + * Mark one node and its visible ancestors as restored from the remote Modal node-output cache. + * @param {string} nodeIdValue + * @param {string} promptId + */ +function markNodeCached(nodeIdValue, promptId) { + const progressNodeIds = [String(nodeIdValue), ...ancestorNodeIds(nodeIdValue)]; + for (const progressNodeId of progressNodeIds) { + if (!shouldApplyPromptState(progressNodeId, promptId)) { + continue; + } + modalNodeCachedStates.set(progressNodeId, { + promptId, + cachedAt: nowMs(), + }); + } + ensureAnimationLoop(); + app.graph?.setDirtyCanvas(true, true); +} + +/** + * Remove one cached-node marker from a node and its visible ancestors. + * @param {string} nodeIdValue + * @param {string | undefined} promptId + */ +function clearNodeCached(nodeIdValue, promptId) { + const progressNodeIds = [String(nodeIdValue), ...ancestorNodeIds(nodeIdValue)]; + for (const progressNodeId of progressNodeIds) { + const cachedState = modalNodeCachedStates.get(progressNodeId); + if (!cachedState) { + continue; + } + if (promptId && cachedState.promptId !== promptId) { + continue; + } + modalNodeCachedStates.delete(progressNodeId); + } +} + /** * Record aggregate mapped-batch progress on the representative remote node only. * @param {string} nodeIdValue @@ -1278,18 +1334,52 @@ function getRemoteVisualState(node) { const promptState = modalPromptStates.get(state.promptId); const progressState = nodeProgressState(nodeId(node), state.promptId); const progressLanes = nodeProgressLanes(nodeId(node), state.promptId); + const cachedState = nodeCachedState(nodeId(node), state.promptId); const batchProgressState = modalNodeBatchProgress.get(nodeId(node)) ?? null; const hasLiveProgress = hasLiveNodeProgress(nodeId(node), state.promptId); return { ...state, phase: deriveRemoteNodePhase(state.phase, hasLiveProgress), isActiveRemoteNode: hasLiveProgress || promptState?.activeNodeId === nodeId(node), + isActiveComponentMember: isNodeInActiveComponent(state.promptId, nodeId(node)), + isCachedRemoteNode: Boolean(cachedState), + cachedAt: cachedState?.cachedAt ?? null, progress: progressState, batchProgress: batchProgressState?.promptId === state.promptId ? batchProgressState : null, progressLanes, }; } +/** + * Return whether one visible node belongs to the component that is currently executing. + * @param {string} promptId + * @param {string} nodeIdValue + * @returns {boolean} + */ +function isNodeInActiveComponent(promptId, nodeIdValue) { + const promptState = modalPromptStates.get(promptId); + const activeNodeId = promptState?.activeNodeId ?? null; + if (!promptState || !activeNodeId) { + return false; + } + + const activeComponentNodeIds = resolveComponentNodeIds(promptId, activeNodeId); + if (!activeComponentNodeIds || activeComponentNodeIds.length === 0) { + return activeNodeId === String(nodeIdValue); + } + + const safeNodeIdValue = String(nodeIdValue); + if (activeComponentNodeIds.includes(safeNodeIdValue)) { + return true; + } + + const descendantNodeIds = promptState.descendantNodeIdsByAncestor.get(safeNodeIdValue); + if (!descendantNodeIds || descendantNodeIds.size === 0) { + return false; + } + return activeComponentNodeIds.some((componentNodeId) => descendantNodeIds.has(componentNodeId)); +} + /** * Draw the remote execution border and shading for a node. * @param {LGraphNode} node @@ -1317,10 +1407,17 @@ function drawRemoteNodeDecoration(node, ctx) { .padStart(2, "0")}`; shadowColor = `rgba(245, 158, 11, ${0.25 + pulse * 0.35})`; } else if (state?.phase === STATE_READY) { - const pulse = (Math.sin(elapsed * 6) + 1) / 2; - borderColor = READY_BORDER_COLOR; - shadowColor = "rgba(34, 197, 94, 0.35)"; - fillColor = `rgba(134, 239, 172, ${0.12 + pulse * 0.08})`; + const pulseRate = state?.isCachedRemoteNode ? 2 : 6; + const pulse = (Math.sin(elapsed * pulseRate) + 1) / 2; + borderColor = state?.isActiveComponentMember + ? READY_ACTIVE_COMPONENT_BORDER_COLOR + : READY_INACTIVE_COMPONENT_BORDER_COLOR; + shadowColor = state?.isActiveComponentMember + ? `rgba(34, 197, 94, ${0.24 + pulse * 0.18})` + : `rgba(22, 101, 52, ${0.18 + pulse * 0.12})`; + fillColor = state?.isActiveComponentMember + ? `rgba(134, 239, 172, ${0.1 + pulse * 0.07})` + : `rgba(74, 222, 128, ${0.06 + pulse * 0.04})`; } else if (state?.phase === STATE_ACTIVE) { const pulse = (Math.sin(elapsed * 7) + 1) / 2; borderColor = `${ACTIVE_BORDER_COLOR}${Math.round((0.7 + pulse * 0.3) * 255) @@ -1329,9 +1426,15 @@ function drawRemoteNodeDecoration(node, ctx) { shadowColor = `rgba(168, 85, 247, ${0.28 + pulse * 0.32})`; fillColor = `rgba(216, 180, 254, ${0.16 + pulse * 0.1})`; } else if (state?.phase === STATE_COMPLETE) { + const pulseRate = state?.isCachedRemoteNode ? 2 : 0; + const pulse = pulseRate > 0 ? (Math.sin(elapsed * pulseRate) + 1) / 2 : 0; borderColor = COMPLETE_BORDER_COLOR; - shadowColor = "rgba(34, 197, 94, 0.28)"; - fillColor = "rgba(134, 239, 172, 0.14)"; + shadowColor = state?.isCachedRemoteNode + ? `rgba(34, 197, 94, ${0.16 + pulse * 0.12})` + : "rgba(34, 197, 94, 0.28)"; + fillColor = state?.isCachedRemoteNode + ? `rgba(134, 239, 172, ${0.1 + pulse * 0.06})` + : "rgba(134, 239, 172, 0.14)"; } else if (state?.phase === STATE_ERROR) { const pulse = (Math.sin(elapsed * 6) + 1) / 2; borderColor = `${ERROR_BORDER_COLOR}${Math.round((0.7 + pulse * 0.3) * 255) @@ -1638,6 +1741,10 @@ function handleModalProgress(event) { setGlobalStatusPhase(promptId, EXECUTION_PHASE, promptState.remoteNodeIds.length || 1); return; } + if (detail.cached_hit) { + markNodeCached(progressNodeId, promptId); + return; + } setGlobalStatusPhase(promptId, EXECUTION_PHASE, promptState.remoteNodeIds.length || 1); if (detail.lane_id != null) { setPromptActiveNode(promptId, null); @@ -1741,6 +1848,7 @@ function clearPromptRemoteStates(promptId) { for (const remoteNodeId of promptState.remoteNodeIds) { clearNodeTimer(remoteNodeId); clearNodeProgress(remoteNodeId, promptId); + clearNodeCached(remoteNodeId, promptId); const currentState = modalNodeStates.get(remoteNodeId); if (currentState?.promptId === promptId) { modalNodeStates.delete(remoteNodeId); @@ -1749,6 +1857,7 @@ function clearPromptRemoteStates(promptId) { for (const ancestorNodeId of promptState.descendantNodeIdsByAncestor.keys()) { clearNodeTimer(ancestorNodeId); clearNodeProgress(ancestorNodeId, promptId); + clearNodeCached(ancestorNodeId, promptId); const currentState = modalNodeStates.get(ancestorNodeId); if (currentState?.promptId === promptId) { modalNodeStates.delete(ancestorNodeId);