feat: mirror remote node progress locally

This commit is contained in:
Ken Simpson
2026-04-18 11:40:16 -07:00
parent 0fbd3a44eb
commit 39a9e2eecd
7 changed files with 344 additions and 3 deletions
+2
View File
@@ -164,6 +164,8 @@ The frontend shows remote state directly on the canvas:
- steady green border: finished for the current run
- red border: queue-time or execution failure
When a remote node reports numeric progress through ComfyUI's progress hooks, the active Modal-marked node now also shows an in-node progress bar and percentage locally while the remote work is running. That is especially useful for sampler-style nodes such as `KSampler`.
The extension also keeps a global activity badge visible during queue-time sync and remote execution, including the period before the prompt is formally queued.
### 6. Asset sync expectations
+90 -3
View File
@@ -82,6 +82,7 @@ class _NullPromptServer:
"""Initialize the no-op prompt server state."""
self.client_id: str | None = None
self.last_node_id: str | None = None
self.last_prompt_id: str | None = None
def send_sync(self, event: str, data: dict[str, Any], client_id: str | None) -> None:
"""Discard PromptExecutor progress and status events."""
@@ -131,6 +132,7 @@ class _TracingPromptServer(_NullPromptServer):
self._status_callback = status_callback
self._active_node_id: str | None = None
self._active_node_started_at: float | None = None
self.last_prompt_id = prompt_id
def _classify_node_role(self, class_type: str) -> str:
"""Return a coarse role name for a node class."""
@@ -185,6 +187,7 @@ class _TracingPromptServer(_NullPromptServer):
)
self._active_node_id = str(next_node_id)
self._active_node_started_at = time.perf_counter()
self.last_node_id = self._active_node_id
_emit_cloud_info(
"Remote node %s class_type=%s role=%s started",
self._active_node_id,
@@ -193,6 +196,49 @@ class _TracingPromptServer(_NullPromptServer):
)
return
if event == "progress_state":
if self._status_callback is None:
return
nodes_payload = data.get("nodes")
if not isinstance(nodes_payload, dict):
return
tracked_node_id = self._active_node_id
tracked_node_state: dict[str, Any] | None = None
if tracked_node_id is not None:
candidate_state = nodes_payload.get(tracked_node_id)
if isinstance(candidate_state, dict):
tracked_node_state = candidate_state
if tracked_node_state is None:
for node_state in nodes_payload.values():
if isinstance(node_state, dict) and node_state.get("state") == "running":
tracked_node_state = node_state
break
if tracked_node_state is None:
return
display_node_id = tracked_node_state.get("display_node_id")
real_node_id = tracked_node_state.get("real_node_id")
reported_node_id = display_node_id or real_node_id or tracked_node_state.get("node_id")
if reported_node_id is None:
return
self._status_callback(
{
"event_type": "node_progress",
"node_id": str(reported_node_id),
"display_node_id": (
str(display_node_id) if display_node_id is not None else None
),
"value": float(tracked_node_state.get("value", 0.0)),
"max": float(tracked_node_state.get("max", 1.0)),
}
)
return
if event == "executed":
executed_node_id = data.get("node")
if executed_node_id is not None and str(executed_node_id) == self._active_node_id:
@@ -427,6 +473,44 @@ def _ensure_headless_prompt_server_instance() -> None:
logger.info("Installed headless PromptServer.instance for remote custom-node initialization.")
@contextmanager
def _temporary_progress_hook(prompt_server: _NullPromptServer) -> Iterator[None]:
"""Install a ComfyUI progress hook so remote samplers emit numeric progress updates."""
import comfy.utils
from comfy_execution.progress import get_progress_state
from comfy_execution.utils import get_executing_context
previous_hook = comfy.utils.PROGRESS_BAR_HOOK
def hook(
value: float,
total: float,
preview_image: Any,
prompt_id: str | None = None,
node_id: str | None = None,
) -> None:
"""Mirror ComfyUI progress-bar updates into the headless progress registry."""
executing_context = get_executing_context()
if prompt_id is None and executing_context is not None:
prompt_id = executing_context.prompt_id
if node_id is None and executing_context is not None:
node_id = executing_context.node_id
if prompt_id is None:
prompt_id = prompt_server.last_prompt_id
if node_id is None:
node_id = prompt_server.last_node_id
if node_id is None:
return
get_progress_state().update_progress(str(node_id), value, total, preview_image)
comfy.utils.set_progress_bar_global_hook(hook)
try:
yield
finally:
comfy.utils.set_progress_bar_global_hook(previous_hook)
def _ensure_default_custom_nodes_dir() -> Path | None:
"""Create the default ComfyUI custom_nodes directory when the image omits its contents."""
comfyui_root = _active_comfyui_root()
@@ -848,12 +932,15 @@ def _execute_subgraph_prompt(
execution = _load_execution_module()
cache_type, cache_args = _prompt_executor_cache_config(execution)
with _temporary_node_mapping(None), _patched_folder_paths_absolute_lookup():
prompt_server = _TracingPromptServer(
with (
_temporary_node_mapping(None),
_patched_folder_paths_absolute_lookup(),
_temporary_progress_hook(prompt_server := _TracingPromptServer(
component_id,
prompt,
status_callback=status_callback,
)
)),
):
with _timed_phase("create_prompt_executor", component=component_id):
executor_state = _get_or_create_prompt_executor_state(
execution=execution,
+1
View File
@@ -5,6 +5,7 @@
- The README had drifted into release-note territory, especially in the old `Important distinction:` block. For this project, the maintainable structure is: short overview, explicit local-vs-remote execution sections, concise operational caveats, and a separate learnings log for historical implementation details.
- ComfyUI registry setup is mostly metadata, not code. The repo needs a real `[tool.comfy]` block in `pyproject.toml`, a stable immutable node-pack id, repository URLs, and an optional GitHub publish workflow; choosing a software license should stay explicit instead of being guessed during packaging work.
- `comfy node pack` is not git-aware by default; it uses `.gitignore` as its practical packaging filter. For a clean registry artifact, this repo needs an allowlist-style `.gitignore` so untracked local files like logs, caches, virtualenvs, and scratch docs cannot leak into `node.zip`.
- Remote per-node phase events are not enough to mirror sampler progress in the local UI. The Modal worker also needs a headless equivalent of ComfyUI's global progress hook so `ProgressBar` updates feed `progress_state`, and the local relay should forward those numeric updates as a separate high-frequency UI event rather than overloading coarse execution-phase messages.
## 2026-04-17
+51
View File
@@ -523,6 +523,34 @@ def _emit_local_modal_status(
prompt_server.send_sync("modal_status", payload, client_id)
def _emit_local_modal_progress(
*,
prompt_id: str | None,
client_id: str | None,
node_id: str,
value: float,
max_value: float,
display_node_id: str | None = None,
) -> None:
"""Forward remote numeric node progress into the local ComfyUI websocket stream."""
if client_id is None:
return
prompt_server = _lookup_local_prompt_server()
if prompt_server is None:
return
payload: dict[str, Any] = {
"prompt_id": prompt_id,
"node_id": node_id,
"value": float(value),
"max": float(max_value),
}
if display_node_id is not None:
payload["display_node_id"] = display_node_id
prompt_server.send_sync("modal_progress", payload, client_id)
def _should_stream_remote_progress(payload: dict[str, Any]) -> bool:
"""Return whether the local client has enough context to mirror remote node progress."""
extra_data = payload.get("extra_data") or {}
@@ -551,6 +579,29 @@ def _consume_remote_payload_stream(
for stream_event in stream_events:
event_kind = str(stream_event.get("kind", ""))
if event_kind == "progress":
if stream_event.get("event_type") == "node_progress":
reported_node_id = stream_event.get("node_id")
if reported_node_id is not None:
logger.debug(
"Forwarding streamed Modal node progress for component=%s node_id=%s value=%s max=%s.",
payload.get("component_id"),
reported_node_id,
stream_event.get("value"),
stream_event.get("max"),
)
_emit_local_modal_progress(
prompt_id=prompt_id,
client_id=client_id,
node_id=str(reported_node_id),
value=float(stream_event.get("value", 0.0)),
max_value=float(stream_event.get("max", 1.0)),
display_node_id=(
str(stream_event["display_node_id"])
if stream_event.get("display_node_id") is not None
else None
),
)
continue
logger.info(
"Forwarding streamed Modal progress for component=%s phase=%s active_node_id=%s.",
payload.get("component_id"),
+76
View File
@@ -266,6 +266,15 @@ def test_modal_cloud_streams_progress_and_result_events(
"active_node_role": "model_load",
}
)
status_callback(
{
"event_type": "node_progress",
"node_id": "7",
"display_node_id": "7",
"value": 3,
"max": 10,
}
)
status_callback({"phase": "execution_success"})
progress_callbacks.append({"component_id": payload["component_id"], "kwargs": kwargs_payload})
return b"serialized-outputs"
@@ -288,6 +297,14 @@ def test_modal_cloud_streams_progress_and_result_events(
"active_node_class_type": "UNETLoader",
"active_node_role": "model_load",
},
{
"kind": "progress",
"event_type": "node_progress",
"node_id": "7",
"display_node_id": "7",
"value": 3,
"max": 10,
},
{
"kind": "progress",
"phase": "execution_success",
@@ -748,6 +765,14 @@ def test_remote_modal_consumes_streamed_progress_and_result(
"active_node_class_type": "UNETLoader",
"active_node_role": "model_load",
},
{
"kind": "progress",
"event_type": "node_progress",
"node_id": "7",
"display_node_id": "7",
"value": 4,
"max": 20,
},
{
"kind": "progress",
"phase": "execution_success",
@@ -774,6 +799,17 @@ def test_remote_modal_consumes_streamed_progress_and_result(
},
"client-1",
),
(
"modal_progress",
{
"prompt_id": "prompt-1",
"node_id": "7",
"display_node_id": "7",
"value": 4.0,
"max": 20.0,
},
"client-1",
),
(
"modal_status",
{
@@ -786,6 +822,46 @@ def test_remote_modal_consumes_streamed_progress_and_result(
]
def test_modal_cloud_tracing_prompt_server_emits_numeric_node_progress(
modal_cloud_module: Any,
) -> None:
"""The cloud tracing prompt server should forward active-node numeric progress updates."""
observed_updates: list[dict[str, Any]] = []
server = modal_cloud_module._TracingPromptServer(
"component-1",
{"7": {"class_type": "KSampler", "inputs": {}}},
status_callback=observed_updates.append,
)
server.send_sync("executing", {"node": "7"}, None)
server.send_sync(
"progress_state",
{
"prompt_id": "component-1",
"nodes": {
"7": {
"node_id": "7",
"display_node_id": "7",
"real_node_id": "7",
"state": "running",
"value": 5,
"max": 20,
}
},
},
None,
)
assert observed_updates[0]["phase"] == "executing"
assert observed_updates[1] == {
"event_type": "node_progress",
"node_id": "7",
"display_node_id": "7",
"value": 5.0,
"max": 20.0,
}
def test_remote_modal_consumes_streamed_tensor_result_payload(
remote_modal_app_module: Any,
serialization_module: Any,
+12
View File
@@ -75,6 +75,18 @@ def test_streamed_modal_progress_takes_precedence_over_proxy_events() -> None:
assert "if (promptState.hasStreamedProgress && phase !== STATE_ERROR)" in source
def test_streamed_modal_node_progress_updates_active_overlay() -> None:
"""The frontend should listen for numeric Modal node progress and render it on the node."""
source = _modal_toggle_source()
assert 'api.addEventListener("modal_progress", handleModalProgress);' in source
assert "function handleModalProgress(event)" in source
assert "function setNodeProgress(nodeIdValue, promptId, value, maxValue)" in source
assert "function clearNodeProgress(nodeIdValue, promptId)" in source
assert "state?.progress" in source
assert 'ctx.fillText(`${Math.round(progressRatio * 100)}%`' in source
def test_prompt_cleanup_prunes_orphaned_global_status_entries() -> None:
"""Workflow cleanup should remove stale global badge states once prompt state is gone."""
source = _modal_toggle_source()
+112
View File
@@ -22,6 +22,7 @@ const STATE_ERROR = "error";
const ERROR_CLEAR_DELAY_MS = 5000;
const modalNodeStates = new Map();
const modalNodeProgress = new Map();
const modalNodeClearTimers = new Map();
const modalPromptStates = new Map();
const syntheticPromptUiStates = new Map();
@@ -431,6 +432,9 @@ function setNodesPhase(nodeIds, phase, promptId, errorMessage) {
if (phase === STATE_ERROR) {
scheduleNodeClear(currentNodeId, promptId, ERROR_CLEAR_DELAY_MS);
}
if (phase === STATE_COMPLETE || phase === STATE_ERROR) {
clearNodeProgress(currentNodeId, promptId);
}
for (const ancestorNodeId of ancestorNodeIds(currentNodeId)) {
affectedAncestorNodeIds.add(ancestorNodeId);
}
@@ -563,6 +567,53 @@ function setPromptActiveNode(promptId, activeNodeId) {
promptState.activeNodeId = activeNodeId ? String(activeNodeId) : null;
}
/**
* Remove stored numeric progress for one node and its visible ancestors.
* @param {string} nodeIdValue
* @param {string | undefined} promptId
*/
function clearNodeProgress(nodeIdValue, promptId) {
const progressNodeIds = [String(nodeIdValue), ...ancestorNodeIds(nodeIdValue)];
for (const progressNodeId of progressNodeIds) {
const progressState = modalNodeProgress.get(progressNodeId);
if (!progressState) {
continue;
}
if (promptId && progressState.promptId !== promptId) {
continue;
}
modalNodeProgress.delete(progressNodeId);
}
}
/**
* Record numeric progress for one node and its visible ancestors.
* @param {string} nodeIdValue
* @param {string} promptId
* @param {number} value
* @param {number} maxValue
*/
function setNodeProgress(nodeIdValue, promptId, value, maxValue) {
const safeMaxValue = Math.max(1, Number(maxValue) || 1);
const safeValue = Math.max(0, Math.min(safeMaxValue, Number(value) || 0));
const progressNodeIds = [String(nodeIdValue), ...ancestorNodeIds(nodeIdValue)];
for (const progressNodeId of progressNodeIds) {
if (!shouldApplyPromptState(progressNodeId, promptId)) {
continue;
}
modalNodeProgress.set(progressNodeId, {
promptId,
value: safeValue,
max: safeMaxValue,
updatedAt: nowMs(),
});
}
ensureAnimationLoop();
app.graph?.setDirtyCanvas(true, true);
}
/**
* Return remote component node ids for an executing proxy node event.
* @param {string} promptId
@@ -614,9 +665,11 @@ function getRemoteVisualState(node) {
return state;
}
const promptState = modalPromptStates.get(state.promptId);
const progressState = modalNodeProgress.get(nodeId(node)) ?? null;
return {
...state,
isActiveRemoteNode: promptState?.activeNodeId === nodeId(node),
progress: progressState?.promptId === state.promptId ? progressState : null,
};
}
@@ -691,6 +744,26 @@ function drawRemoteNodeDecoration(node, ctx) {
node.size[1] + titleHeight + borderWidth,
);
ctx.restore();
if (!state?.progress || state.phase !== STATE_ACTIVE) {
return;
}
const progressRatio = Math.max(0, Math.min(1, state.progress.value / state.progress.max));
const progressWidth = Math.max(0, (node.size[0] + borderWidth * 2) * progressRatio);
const progressY = node.size[1] - Math.max(10 / scale, borderWidth * 2);
ctx.save();
ctx.fillStyle = "rgba(15, 23, 42, 0.72)";
ctx.fillRect(-borderWidth, progressY, node.size[0] + borderWidth * 2, 8 / scale);
ctx.fillStyle = "rgba(216, 180, 254, 0.92)";
ctx.fillRect(-borderWidth, progressY, progressWidth, 8 / scale);
ctx.fillStyle = "#f8fafc";
ctx.font = `${Math.max(11 / scale, 8)}px ui-sans-serif, system-ui, sans-serif`;
ctx.textAlign = "right";
ctx.textBaseline = "bottom";
ctx.fillText(`${Math.round(progressRatio * 100)}%`, node.size[0], progressY - 2 / scale);
ctx.restore();
}
/**
@@ -758,6 +831,9 @@ function handleModalStatus(event) {
setGlobalStatusPhase(promptId, STATE_ERROR, nodeIds.length);
setTimeout(() => clearGlobalStatusPhase(promptId), ERROR_CLEAR_DELAY_MS);
setPromptActiveNode(promptId, null);
for (const nodeIdValue of nodeIds) {
clearNodeProgress(nodeIdValue, promptId);
}
setNodesPhase(nodeIds, STATE_ERROR, promptId, detail.error_message);
return;
}
@@ -770,6 +846,7 @@ function handleModalStatus(event) {
setGlobalStatusPhase(promptId, EXECUTION_PHASE, nodeIds.length);
setNodesPhase(nodeIds, STATE_READY, promptId);
if (previousActiveNodeId && previousActiveNodeId !== nextActiveNodeId) {
clearNodeProgress(previousActiveNodeId, promptId);
setNodesPhase([previousActiveNodeId], STATE_COMPLETE, promptId);
}
if (nextActiveNodeId) {
@@ -782,6 +859,7 @@ function handleModalStatus(event) {
if (detail.phase === "execution_success") {
promptState.hasStreamedProgress = true;
if (promptState.activeNodeId) {
clearNodeProgress(promptState.activeNodeId, promptId);
setNodesPhase([promptState.activeNodeId], STATE_COMPLETE, promptId);
}
setPromptActiveNode(promptId, null);
@@ -790,6 +868,31 @@ function handleModalStatus(event) {
}
}
/**
* Apply a streamed numeric Modal node-progress event.
* @param {CustomEvent} event
*/
function handleModalProgress(event) {
const detail = eventDetail(event);
const promptId = String(detail.prompt_id ?? "");
const progressNodeId = String(detail.display_node_id ?? detail.node_id ?? "");
if (!promptId || !progressNodeId) {
return;
}
const promptState = ensurePromptState(promptId);
promptState.hasStreamedProgress = true;
setPromptActiveNode(promptId, progressNodeId);
setGlobalStatusPhase(promptId, EXECUTION_PHASE, promptState.remoteNodeIds.length || 1);
setNodesPhase([progressNodeId], STATE_ACTIVE, promptId);
setNodeProgress(
progressNodeId,
promptId,
Number(detail.value ?? 0),
Number(detail.max ?? 1),
);
}
/**
* Update remote component visuals from a native ComfyUI execution event.
* @param {CustomEvent} event
@@ -824,11 +927,17 @@ function handleExecutionPhase(event, phase) {
setGlobalStatusPhase(promptId, STATE_ERROR, componentNodeIds.length);
setTimeout(() => clearGlobalStatusPhase(promptId), ERROR_CLEAR_DELAY_MS);
setPromptActiveNode(promptId, null);
for (const nodeIdValue of componentNodeIds) {
clearNodeProgress(nodeIdValue, promptId);
}
setNodesPhase(componentNodeIds, STATE_ERROR, promptId, detail.exception_message);
return;
}
if (phase === STATE_COMPLETE) {
setPromptActiveNode(promptId, null);
for (const nodeIdValue of componentNodeIds) {
clearNodeProgress(nodeIdValue, promptId);
}
setNodesPhase(componentNodeIds, STATE_COMPLETE, promptId, detail.exception_message);
}
}
@@ -846,6 +955,7 @@ function clearPromptRemoteStates(promptId) {
}
for (const remoteNodeId of promptState.remoteNodeIds) {
clearNodeTimer(remoteNodeId);
clearNodeProgress(remoteNodeId, promptId);
const currentState = modalNodeStates.get(remoteNodeId);
if (currentState?.promptId === promptId) {
modalNodeStates.delete(remoteNodeId);
@@ -853,6 +963,7 @@ function clearPromptRemoteStates(promptId) {
}
for (const ancestorNodeId of promptState.descendantNodeIdsByAncestor.keys()) {
clearNodeTimer(ancestorNodeId);
clearNodeProgress(ancestorNodeId, promptId);
const currentState = modalNodeStates.get(ancestorNodeId);
if (currentState?.promptId === promptId) {
modalNodeStates.delete(ancestorNodeId);
@@ -974,6 +1085,7 @@ function registerExecutionListeners() {
}
api.addEventListener("modal_status", handleModalStatus);
api.addEventListener("modal_progress", handleModalProgress);
api.addEventListener("executing", (event) => handleExecutionPhase(event, EXECUTION_PHASE));
api.addEventListener("executed", (event) => {
endSyntheticExecutionUi(String(eventDetail(event).prompt_id ?? ""));