feat: show mapped modal progress lanes

This commit is contained in:
Ken Simpson
2026-04-18 21:30:33 -07:00
parent ddf6617165
commit d4c819ae99
6 changed files with 312 additions and 45 deletions
+1
View File
@@ -102,6 +102,7 @@ Remote runtime behavior:
- Remote proxy nodes now execute through ComfyUI's async node path, so independent Modal-backed components can overlap instead of being forced through one blocking local proxy at a time.
- The local Modal call executor keeps multiple worker threads available, which removes the previous `max_workers=1` bottleneck when several remote components are ready at once.
- `ModalMapInput` can turn one remote component boundary into a locally scheduled mapped execution. List inputs and batched tensors fan out into multiple per-item Modal subgraph calls, and the local scheduler refills that queue up to the configured `COMFY_MODAL_MAX_CONTAINERS`.
- When a mapped remote component has several Modal workers running at once, the local node overlay now shows one progress lane per active worker plus the aggregate completion bar, instead of letting concurrent runs overwrite a single progress bar.
- Remote cancellation now uses a shared Modal `Dict` control store instead of a second RPC lane into the execution class, so per-container execution concurrency can stay at `1` without losing interrupt propagation.
- If a run is cancelled and restarted quickly, the remote worker now gives Modal volume reload a short bounded retry window so recently released model files can close before the next request needs a fresh `vol.reload()`.
- Mapped remote execution now carries a per-request volume reload marker, so one warm container only performs `vol.reload()` once for that uploaded asset snapshot even if the local scheduler fans the component out into many per-item Modal calls.
+1
View File
@@ -75,6 +75,7 @@
- Once warm-container caches keep models and executors alive, unconditional `vol.reload()` becomes wrong: Modal can reject reloads while files from the mounted volume are still open. The queue side already knows whether a request actually uploaded anything new, so that fact should travel in the payload and only requests with fresh uploads should reload the mounted volume.
- Fast cancel-and-restart sequences can still race the container's file-handle cleanup even after warm caches are cleared. For the specific Modal `"open files preventing the operation"` reload failure, the remote worker needs a short bounded retry window after cache/model cleanup rather than a single immediate retry.
- That queue-time `requires_volume_reload` bit is not enough once one logical request fans out into many per-item Modal calls. Mapped execution also needs a stable per-request reload marker so each warm container can reload the mounted volume once for that asset snapshot and skip the duplicate reloads from later mapped items.
- Mapped Modal fanout needs stable local worker lanes, not just a semaphore over one task per item. The UI can only render sensible concurrent progress if each in-flight mapped run carries a durable lane id; otherwise parallel item updates all target one node-progress slot and the bar flickers incoherently.
- The streamed Modal generator path needs the same tensor-aware serialization rules as ordinary node input/output transport. If a progress or result envelope leaks a raw `torch.Tensor`, Modal's JSON encoding fails before the local client ever sees it, so the stream layer should normalize event payloads through the shared serializer too.
- Modal volume metadata calls can become a bottleneck on their own. Repeated `VolumeListFiles` checks for the same marker path are worth memoizing for the life of the queue request, and transient `ResourceExhaustedError` responses should be retried with backoff instead of surfacing immediately as hard prompt failures.
- Modal volume SDK calls should not run directly on the aiohttp request event loop. Running them in a worker thread avoids `AsyncUsageWarning`, and a missing path from `Volume.listdir()` should be treated as an ordinary "not uploaded yet" cache miss instead of a hard failure.
+114 -23
View File
@@ -549,6 +549,9 @@ def _emit_local_modal_progress(
value: float,
max_value: float,
display_node_id: str | None = None,
lane_id: str | None = None,
clear: bool = False,
item_index: int | None = None,
) -> None:
"""Forward remote numeric node progress into the local ComfyUI websocket stream."""
if client_id is None:
@@ -566,6 +569,12 @@ def _emit_local_modal_progress(
}
if display_node_id is not None:
payload["display_node_id"] = display_node_id
if lane_id is not None:
payload["lane_id"] = lane_id
if clear:
payload["clear"] = True
if item_index is not None:
payload["item_index"] = int(item_index)
prompt_server.send_sync("modal_progress", payload, client_id)
@@ -852,16 +861,22 @@ def _consume_remote_payload_stream(
if event_kind == "progress":
event_type = str(stream_event.get("event_type", ""))
if event_type == "node_progress":
if suppress_status_stream:
lane_id = (
str(payload["mapped_progress_lane_id"])
if payload.get("mapped_progress_lane_id") is not None
else None
)
if suppress_status_stream and lane_id is None:
continue
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.",
"Forwarding streamed Modal node progress for component=%s node_id=%s value=%s max=%s lane_id=%s.",
payload.get("component_id"),
reported_node_id,
stream_event.get("value"),
stream_event.get("max"),
lane_id,
)
_emit_local_modal_progress(
prompt_id=prompt_id,
@@ -870,8 +885,18 @@ def _consume_remote_payload_stream(
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
str(payload["mapped_progress_display_node_id"])
if payload.get("mapped_progress_display_node_id") is not None
else (
str(stream_event["display_node_id"])
if stream_event.get("display_node_id") is not None
else None
)
),
lane_id=lane_id,
item_index=(
int(payload["map_item_index"])
if payload.get("map_item_index") is not None
else None
),
)
@@ -1012,7 +1037,11 @@ def _mapped_execution_parallelism(item_count: int) -> int:
return max(1, min(item_count, configured_limit))
def _build_mapped_item_payload(payload: dict[str, Any], item_index: int) -> dict[str, Any]:
def _build_mapped_item_payload(
payload: dict[str, Any],
item_index: int,
lane_index: int,
) -> dict[str, Any]:
"""Return one per-item subgraph payload derived from a mapped remote component payload."""
item_payload = copy.deepcopy(payload)
item_payload["payload_kind"] = "subgraph"
@@ -1020,6 +1049,10 @@ def _build_mapped_item_payload(payload: dict[str, Any], item_index: int) -> dict
item_payload["mapped_input"] = None
item_payload["suppress_status_stream"] = True
item_payload["map_item_index"] = item_index
item_payload["mapped_progress_lane_id"] = str(lane_index)
item_payload["mapped_progress_display_node_id"] = str(
payload.get("component_id", "modal-subgraph")
)
return item_payload
@@ -1071,6 +1104,55 @@ def _emit_local_mapped_progress(
)
def _clear_local_mapped_lane_progress(
payload: dict[str, Any],
lane_index: int,
item_index: int,
) -> None:
"""Remove one mapped worker lane from the local node overlay."""
prompt_id = str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None
extra_data = payload.get("extra_data") or {}
client_id = str(extra_data.get("client_id")) if extra_data.get("client_id") is not None else None
display_node_id = str(payload.get("component_id") or "")
if not prompt_id or not client_id or not display_node_id:
return
_emit_local_modal_progress(
prompt_id=prompt_id,
client_id=client_id,
node_id=display_node_id,
value=0.0,
max_value=1.0,
display_node_id=display_node_id,
lane_id=str(lane_index),
clear=True,
item_index=item_index,
)
def _emit_local_mapped_lane_progress_start(
payload: dict[str, Any],
lane_index: int,
item_index: int,
) -> None:
"""Create or reset one mapped worker lane before remote progress begins arriving."""
prompt_id = str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None
extra_data = payload.get("extra_data") or {}
client_id = str(extra_data.get("client_id")) if extra_data.get("client_id") is not None else None
display_node_id = str(payload.get("component_id") or "")
if not prompt_id or not client_id or not display_node_id:
return
_emit_local_modal_progress(
prompt_id=prompt_id,
client_id=client_id,
node_id=display_node_id,
value=0.0,
max_value=1.0,
display_node_id=display_node_id,
lane_id=str(lane_index),
item_index=item_index,
)
async def _invoke_mapped_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
"""Split one mapped boundary input, run per-item remote executions, and reassemble outputs."""
mapped_input = payload.get("mapped_input") or {}
@@ -1101,31 +1183,40 @@ async def _invoke_mapped_remote_engine_async(payload: dict[str, Any], kwargs_pay
)
_emit_local_mapped_progress(payload, 0, total_items)
semaphore = asyncio.Semaphore(parallelism)
per_item_outputs: list[tuple[Any, ...] | None] = [None] * total_items
completed_items = 0
item_queue: asyncio.Queue[tuple[int, Any] | None] = asyncio.Queue()
for item_index, item_value in enumerate(mapped_items):
item_queue.put_nowait((item_index, item_value))
for _ in range(parallelism):
item_queue.put_nowait(None)
async def run_item(item_index: int, item_value: Any) -> None:
"""Execute one mapped item through the ordinary remote subgraph path."""
async def run_worker(lane_index: int) -> None:
"""Execute queued mapped items through one stable local worker lane."""
nonlocal completed_items
async with semaphore:
while True:
queued_item = await item_queue.get()
if queued_item is None:
return
item_index, item_value = queued_item
if _local_processing_interrupted():
_raise_local_interrupt()
item_payload = _build_mapped_item_payload(payload, item_index)
item_inputs = dict(broadcast_inputs)
item_inputs[mapped_input_name] = item_value
item_response = await invoke_remote_engine_async(
item_payload,
serialize_node_inputs(item_inputs),
)
per_item_outputs[item_index] = deserialize_node_outputs(item_response)
completed_items += 1
_emit_local_mapped_progress(payload, completed_items, total_items)
_emit_local_mapped_lane_progress_start(payload, lane_index, item_index)
try:
item_payload = _build_mapped_item_payload(payload, item_index, lane_index)
item_inputs = dict(broadcast_inputs)
item_inputs[mapped_input_name] = item_value
item_response = await invoke_remote_engine_async(
item_payload,
serialize_node_inputs(item_inputs),
)
per_item_outputs[item_index] = deserialize_node_outputs(item_response)
completed_items += 1
_emit_local_mapped_progress(payload, completed_items, total_items)
finally:
_clear_local_mapped_lane_progress(payload, lane_index, item_index)
tasks = [
asyncio.create_task(run_item(item_index, item_value))
for item_index, item_value in enumerate(mapped_items)
]
tasks = [asyncio.create_task(run_worker(lane_index)) for lane_index in range(parallelism)]
try:
await asyncio.gather(*tasks)
except Exception:
+26 -6
View File
@@ -1548,12 +1548,13 @@ def test_invoke_mapped_remote_engine_async_uses_bounded_parallelism(
"""Mapped remote execution should refill a local queue up to COMFY_MODAL_MAX_CONTAINERS."""
in_flight = 0
max_in_flight = 0
progress_updates: list[float] = []
progress_updates: list[dict[str, Any]] = []
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
nonlocal in_flight, max_in_flight
assert payload["payload_kind"] == "subgraph"
assert payload["suppress_status_stream"] is True
assert payload["mapped_progress_lane_id"] in {"0", "1"}
mapped_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
item_value = mapped_inputs["remote_input_0"]
in_flight += 1
@@ -1572,7 +1573,7 @@ def test_invoke_mapped_remote_engine_async_uses_bounded_parallelism(
monkeypatch.setattr(
remote_modal_app_module,
"_emit_local_modal_progress",
lambda **kwargs: progress_updates.append(float(kwargs["value"])),
lambda **kwargs: progress_updates.append(kwargs),
)
try:
payload = {
@@ -1598,8 +1599,13 @@ def test_invoke_mapped_remote_engine_async_uses_bounded_parallelism(
["done:a", "done:b", "done:c", "done:d"],
)
assert max_in_flight == 2
assert progress_updates[0] == 0.0
assert progress_updates[-1] == 4.0
assert progress_updates[0]["value"] == 0.0
assert progress_updates[0].get("lane_id") is None
aggregate_updates = [update for update in progress_updates if update.get("lane_id") is None]
lane_updates = [update for update in progress_updates if update.get("lane_id") is not None]
assert aggregate_updates[-1]["value"] == 4.0
assert {update["lane_id"] for update in lane_updates} == {"0", "1"}
assert any(update.get("clear") is True for update in lane_updates)
def test_consume_remote_payload_stream_suppresses_status_but_keeps_boundary_previews(
@@ -1607,7 +1613,7 @@ def test_consume_remote_payload_stream_suppresses_status_but_keeps_boundary_prev
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Mapped per-item remote calls should suppress node-status chatter but still forward previews."""
"""Mapped per-item remote calls should suppress status chatter but still forward previews and lane progress."""
progress_calls: list[dict[str, Any]] = []
status_calls: list[dict[str, Any]] = []
preview_calls: list[dict[str, Any]] = []
@@ -1634,6 +1640,9 @@ def test_consume_remote_payload_stream_suppresses_status_but_keeps_boundary_prev
"component_node_ids": ["6", "7"],
"extra_data": {"client_id": "client-1"},
"suppress_status_stream": True,
"mapped_progress_lane_id": "1",
"mapped_progress_display_node_id": "6",
"map_item_index": 0,
}
stream_events = iter(
[
@@ -1666,7 +1675,18 @@ def test_consume_remote_payload_stream_suppresses_status_but_keeps_boundary_prev
response = remote_modal_app_module._consume_remote_payload_stream(payload, stream_events)
assert serialization_module.deserialize_node_outputs(response) == ("done",)
assert progress_calls == []
assert progress_calls == [
{
"prompt_id": "prompt-1",
"client_id": "client-1",
"node_id": "7",
"value": 1.0,
"max_value": 4.0,
"display_node_id": "6",
"lane_id": "1",
"item_index": 0,
}
]
assert status_calls == []
assert preview_calls == [
{
+16
View File
@@ -92,11 +92,27 @@ def test_streamed_modal_node_progress_updates_active_overlay() -> None:
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 setNodeProgressLane(nodeIdValue, promptId, laneId, value, maxValue, itemIndex)" in source
assert "function clearNodeProgressLane(nodeIdValue, promptId, laneId)" in source
assert "function clearNodeProgress(nodeIdValue, promptId)" in source
assert "state?.progress" in source
assert "state?.progressLanes" in source
assert "detail.lane_id != null" in source
assert "clearNodeProgressLane(progressNodeId, promptId, String(detail.lane_id));" in source
assert "setNodeProgressLane(" in source
assert 'ctx.fillText(`${Math.round(progressRatio * 100)}%`' in source
def test_mapped_parallel_modal_progress_renders_multiple_lane_bars() -> None:
"""Parallel mapped Modal runs should render one local progress lane per active worker."""
source = _modal_toggle_source()
assert "const modalNodeProgressLanes = new Map();" in source
assert "progressLanes.length > 0" in source
assert "const laneColors = [" in source
assert 'ctx.fillText(`${progressLanes.length}x`' 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()
+154 -16
View File
@@ -24,6 +24,7 @@ const ERROR_CLEAR_DELAY_MS = 5000;
const modalNodeStates = new Map();
const modalNodeProgress = new Map();
const modalNodeProgressLanes = new Map();
const modalNodeClearTimers = new Map();
const modalPromptStates = new Map();
const syntheticPromptUiStates = new Map();
@@ -578,12 +579,18 @@ function clearNodeProgress(nodeIdValue, promptId) {
for (const progressNodeId of progressNodeIds) {
const progressState = modalNodeProgress.get(progressNodeId);
if (!progressState) {
modalNodeProgress.delete(progressNodeId);
} else if (!promptId || progressState.promptId === promptId) {
modalNodeProgress.delete(progressNodeId);
}
const laneState = modalNodeProgressLanes.get(progressNodeId);
if (!laneState) {
continue;
}
if (promptId && progressState.promptId !== promptId) {
if (promptId && laneState.promptId !== promptId) {
continue;
}
modalNodeProgress.delete(progressNodeId);
modalNodeProgressLanes.delete(progressNodeId);
}
}
@@ -615,6 +622,78 @@ function setNodeProgress(nodeIdValue, promptId, value, maxValue) {
app.graph?.setDirtyCanvas(true, true);
}
/**
* Record numeric progress for one worker lane on a node and its visible ancestors.
* @param {string} nodeIdValue
* @param {string} promptId
* @param {string} laneId
* @param {number} value
* @param {number} maxValue
* @param {number | null | undefined} itemIndex
*/
function setNodeProgressLane(nodeIdValue, promptId, laneId, value, maxValue, itemIndex) {
const safeLaneId = String(laneId ?? "");
if (!safeLaneId) {
return;
}
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;
}
const existingLaneState = modalNodeProgressLanes.get(progressNodeId);
const laneState =
existingLaneState?.promptId === promptId
? existingLaneState
: {
promptId,
lanes: new Map(),
};
laneState.lanes.set(safeLaneId, {
laneId: safeLaneId,
value: safeValue,
max: safeMaxValue,
itemIndex: Number.isFinite(Number(itemIndex)) ? Number(itemIndex) : null,
updatedAt: nowMs(),
});
modalNodeProgressLanes.set(progressNodeId, laneState);
}
ensureAnimationLoop();
app.graph?.setDirtyCanvas(true, true);
}
/**
* Remove one worker lane progress bar from a node and its visible ancestors.
* @param {string} nodeIdValue
* @param {string} promptId
* @param {string} laneId
*/
function clearNodeProgressLane(nodeIdValue, promptId, laneId) {
const safeLaneId = String(laneId ?? "");
if (!safeLaneId) {
return;
}
const progressNodeIds = [String(nodeIdValue), ...ancestorNodeIds(nodeIdValue)];
for (const progressNodeId of progressNodeIds) {
const laneState = modalNodeProgressLanes.get(progressNodeId);
if (!laneState || laneState.promptId !== promptId) {
continue;
}
laneState.lanes.delete(safeLaneId);
if (laneState.lanes.size === 0) {
modalNodeProgressLanes.delete(progressNodeId);
}
}
ensureAnimationLoop();
app.graph?.setDirtyCanvas(true, true);
}
/**
* Return remote component node ids for an executing proxy node event.
* @param {string} promptId
@@ -942,10 +1021,17 @@ function getRemoteVisualState(node) {
}
const promptState = modalPromptStates.get(state.promptId);
const progressState = modalNodeProgress.get(nodeId(node)) ?? null;
const progressLaneState = modalNodeProgressLanes.get(nodeId(node)) ?? null;
return {
...state,
isActiveRemoteNode: promptState?.activeNodeId === nodeId(node),
progress: progressState?.promptId === state.promptId ? progressState : null,
progressLanes:
progressLaneState?.promptId === state.promptId
? Array.from(progressLaneState.lanes.values()).sort(
(left, right) => Number(left.laneId) - Number(right.laneId),
)
: [],
};
}
@@ -1021,24 +1107,61 @@ function drawRemoteNodeDecoration(node, ctx) {
);
ctx.restore();
if (!state?.progress || state.phase !== STATE_ACTIVE) {
const progressLanes = Array.isArray(state?.progressLanes) ? state.progressLanes : [];
const hasAggregateProgress = Boolean(state?.progress && state.phase === STATE_ACTIVE);
const hasLaneProgress = progressLanes.length > 0 && state?.phase === STATE_ACTIVE;
if (!hasAggregateProgress && !hasLaneProgress) {
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);
const barWidth = node.size[0] + borderWidth * 2;
const aggregateHeight = 8 / scale;
const laneHeight = 5 / scale;
const laneGap = 2 / scale;
const progressY = node.size[1] - Math.max(10 / scale, borderWidth * 2);
const laneColors = [
"rgba(196, 181, 253, 0.94)",
"rgba(147, 197, 253, 0.94)",
"rgba(110, 231, 183, 0.94)",
"rgba(253, 224, 71, 0.94)",
"rgba(251, 146, 60, 0.94)",
"rgba(244, 114, 182, 0.94)",
];
if (hasLaneProgress) {
const laneBlockHeight = progressLanes.length * laneHeight + (progressLanes.length - 1) * laneGap;
let laneY = progressY - laneGap - laneBlockHeight;
for (const [laneIndex, laneProgress] of progressLanes.entries()) {
const laneRatio = Math.max(0, Math.min(1, laneProgress.value / laneProgress.max));
const laneWidth = Math.max(0, barWidth * laneRatio);
ctx.fillStyle = "rgba(15, 23, 42, 0.66)";
ctx.fillRect(-borderWidth, laneY, barWidth, laneHeight);
ctx.fillStyle = laneColors[laneIndex % laneColors.length];
ctx.fillRect(-borderWidth, laneY, laneWidth, laneHeight);
laneY += laneHeight + laneGap;
}
}
if (hasAggregateProgress) {
const progressRatio = Math.max(0, Math.min(1, state.progress.value / state.progress.max));
const progressWidth = Math.max(0, barWidth * progressRatio);
ctx.fillStyle = "rgba(15, 23, 42, 0.72)";
ctx.fillRect(-borderWidth, progressY, barWidth, aggregateHeight);
ctx.fillStyle = "rgba(216, 180, 254, 0.92)";
ctx.fillRect(-borderWidth, progressY, progressWidth, aggregateHeight);
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);
} else if (hasLaneProgress) {
ctx.fillStyle = "#f8fafc";
ctx.font = `${Math.max(10 / scale, 8)}px ui-sans-serif, system-ui, sans-serif`;
ctx.textAlign = "right";
ctx.textBaseline = "bottom";
ctx.fillText(`${progressLanes.length}x`, node.size[0], progressY - 2 / scale);
}
ctx.restore();
}
@@ -1163,6 +1286,21 @@ function handleModalProgress(event) {
setPromptActiveNode(promptId, progressNodeId);
setGlobalStatusPhase(promptId, EXECUTION_PHASE, promptState.remoteNodeIds.length || 1);
setNodesPhase([progressNodeId], STATE_ACTIVE, promptId);
if (detail.lane_id != null) {
if (detail.clear) {
clearNodeProgressLane(progressNodeId, promptId, String(detail.lane_id));
return;
}
setNodeProgressLane(
progressNodeId,
promptId,
String(detail.lane_id),
Number(detail.value ?? 0),
Number(detail.max ?? 1),
detail.item_index,
);
return;
}
setNodeProgress(
progressNodeId,
promptId,