From 093b2bdeecd7f341eb731f1e936bd3d93db03fb5 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Wed, 19 Aug 2026 20:48:18 -0700 Subject: [PATCH] Run local previews alongside remote continuations --- README.md | 2 +- api_intercept.py | 169 ++++++++++++++++++--------- docs/LEARNINGS.md | 3 +- modal_executor_node.py | 185 +++++++++++++++++++++++------ tests/test_modal_executor_node.py | 188 +++++++++++++++++++++++++++--- 5 files changed, 437 insertions(+), 110 deletions(-) diff --git a/README.md b/README.md index 595e226..075b300 100644 --- a/README.md +++ b/README.md @@ -292,7 +292,7 @@ When the planner detects a remote → local → remote gap, those remote proxies Independent Modal-backed components can overlap through ComfyUI's async proxy path and the local Modal call executor. Large tensor and media edges are deliberately co-located first, so this parallelism is reserved for genuinely independent branches, inexpensive scalar boundaries, mapped execution, and graph shapes that require a split for correctness. Ordinary components with several remote execution targets remain one proxy unless a local re-entry dependency would create a scheduler cycle. Each Modal GPU container handles one active workflow execution at a time, so parallel ready components can scale out across containers instead of multiplexing several active executions onto one worker. `COMFY_MODAL_MAX_INFLIGHT_CALLS` bounds local dispatch independently from the local CPU count and the remote autoscaler. -If a remote output also feeds a local branch that cannot transitively reach any later remote component, Modal-Sync moves that local-only branch behind completion of the downstream remote work. The local node still executes and produces the same output, but ComfyUI cannot choose it ahead of the ready remote continuation. Local branches that do feed remote execution remain ordinary dependencies and are never deferred. +If a remote output also feeds a local branch that cannot transitively reach any later remote component, Modal-Sync inserts an asynchronous dispatch gate on that local-only branch. The gate first lets every nearest downstream Modal continuation begin dispatch, then immediately releases the local branch while that remote work remains in flight. Mid-workflow previews and saves therefore still appear as soon as their source value has crossed the remote-to-local boundary, without placing their local rendering or display work on the continuing remote critical path. Local branches that do feed remote execution remain ordinary dependencies and preserve their original ordering. When a component output is consumed exclusively by other remote components, Modal-Sync keeps it remote. The producer returns a small durable bridge reference through the local ComfyUI proxy instead of returning the underlying `IMAGE`, `AUDIO`, `LATENT`, `MASK`, `SIGMAS`, `VIDEO`, or scalar value. A downstream remote component resolves that reference from the worker's warm bridge cache when possible, or restores it from shared Modal storage when Modal schedules the components on different containers. If any consumer is local, including an interim preview or final `SaveVideo`, the planner leaves that boundary materialized so the local node receives the ordinary ComfyUI value it expects. diff --git a/api_intercept.py b/api_intercept.py index 82e7c90..27f03fb 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -20,13 +20,14 @@ from .modal_executor_node import ( MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS, MODAL_ARTIFACT_FINALIZER_NODE_ID, MODAL_COMPONENT_COMPLETION_OUTPUT_NAME, - MODAL_DEFERRED_LOCAL_PASSTHROUGH_NODE_ID, MODAL_MAP_INPUT_NODE_ID, + MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, ensure_modal_artifact_finalizer_registered, ensure_modal_component_proxy_node_registered, - ensure_modal_deferred_local_passthrough_registered, + ensure_modal_parallel_local_passthrough_registered, register_cache_friendly_proxy_payload, register_modal_map_input_warmup_context, + update_registered_proxy_payload_fields, ) from .session_state import RemoteSessionHandle from .settings import ( @@ -184,7 +185,7 @@ class RewriteSummary: uploaded_volume_paths: list[str] = field(default_factory=list) rewritten_node_id_map: dict[str, str] = field(default_factory=dict) sandwiched_local_node_ids: list[str] = field(default_factory=list) - deferred_local_branch_node_ids: list[str] = field(default_factory=list) + parallel_local_branch_node_ids: list[str] = field(default_factory=list) synced_assets: list[SyncedAsset] = field(default_factory=list) custom_nodes_bundle: SyncedAsset | None = None artifact_finalizer_node_id: str | None = None @@ -549,8 +550,8 @@ def _modal_rewritten_prompt_diagnostics( diagnostics["component_execution_stages"] = copy.deepcopy( summary.component_execution_stages ) - diagnostics["deferred_local_branch_node_ids"] = list( - summary.deferred_local_branch_node_ids + diagnostics["parallel_local_branch_node_ids"] = list( + summary.parallel_local_branch_node_ids ) diagnostics["rewritten_node_id_map"] = copy.deepcopy(summary.rewritten_node_id_map) return diagnostics @@ -3918,20 +3919,64 @@ def _attach_modal_artifact_finalizer( return finalizer_node_id -def _defer_non_returning_local_branches( +def _nearest_downstream_remote_component_ids( + *, + rewritten_prompt: dict[str, Any], + consumers: dict[LinkedOutputRef, list[InputTarget]], + seed_targets: list[InputTarget], + remote_component_id_set: set[str], + nodes_module: Any, +) -> list[str]: + """Return the first remote proxies reached along every downstream branch.""" + nearest_remote_component_ids: set[str] = set() + visited_node_ids: set[str] = set() + pending_node_ids = deque(str(target.node_id) for target in seed_targets) + while pending_node_ids: + node_id = pending_node_ids.popleft() + if node_id in visited_node_ids or node_id not in rewritten_prompt: + continue + visited_node_ids.add(node_id) + if node_id in remote_component_id_set: + nearest_remote_component_ids.add(node_id) + continue + for output_ref in _node_output_refs( + rewritten_prompt, + node_id, + nodes_module, + ): + for downstream_target in consumers.get(output_ref, []): + pending_node_ids.append(str(downstream_target.node_id)) + return sorted(nearest_remote_component_ids) + + +def _parallelize_non_returning_local_branches( *, rewritten_prompt: dict[str, Any], remote_component_ids: list[str], nodes_module: Any, ) -> list[str]: - """Keep local-only taps from winning scheduler priority over downstream remote work.""" + """Release local-only taps once the nearest remote continuations are in flight.""" remote_component_id_set = set(remote_component_ids) if len(remote_component_id_set) < 2: return [] consumers = _build_consumer_map(rewritten_prompt) - ensure_modal_deferred_local_passthrough_registered(nodes_module) - barrier_node_ids: list[str] = [] + ensure_modal_parallel_local_passthrough_registered(nodes_module) + passthrough_node_ids: list[str] = [] + dispatch_group_id = uuid.uuid4().hex + for component_id in remote_component_ids: + embedded_payload = ( + rewritten_prompt.get(component_id, {}) + .get("inputs", {}) + .get("original_node_data") + ) + if not isinstance(embedded_payload, Mapping): + continue + payload_prompt_id = embedded_payload.get("prompt_id") + if payload_prompt_id is not None and str(payload_prompt_id).strip(): + dispatch_group_id = str(payload_prompt_id) + break + for source_component_id in remote_component_ids: for source in _node_output_refs( rewritten_prompt, @@ -3941,16 +3986,18 @@ def _defer_non_returning_local_branches( output_consumers = list(consumers.get(source, [])) if not output_consumers: continue - downstream_node_ids = _downstream_node_ids_from_targets( - prompt=rewritten_prompt, + downstream_remote_component_ids = _nearest_downstream_remote_component_ids( + rewritten_prompt=rewritten_prompt, consumers=consumers, seed_targets=output_consumers, + remote_component_id_set=remote_component_id_set, nodes_module=nodes_module, ) - downstream_remote_component_ids = sorted( - (downstream_node_ids & remote_component_id_set) - - {source_component_id} - ) + downstream_remote_component_ids = [ + component_id + for component_id in downstream_remote_component_ids + if component_id != source_component_id + ] if not downstream_remote_component_ids: continue @@ -3970,46 +4017,54 @@ def _defer_non_returning_local_branches( if not local_only_consumers: continue - barrier_node_id = ( - f"__{MODAL_DEFERRED_LOCAL_PASSTHROUGH_NODE_ID}__" - f"{len(barrier_node_ids)}" + passthrough_node_id = ( + f"__{MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID}__" + f"{len(passthrough_node_ids)}" ) - while barrier_node_id in rewritten_prompt: - barrier_node_id = f"{barrier_node_id}_proxy" - barrier_inputs: dict[str, Any] = { + while passthrough_node_id in rewritten_prompt: + passthrough_node_id = f"{passthrough_node_id}_proxy" + passthrough_inputs: dict[str, Any] = { "value": [source.node_id, source.output_index], + "dispatch_context": { + "dispatch_group_id": dispatch_group_id, + "component_ids": downstream_remote_component_ids, + }, } - for component_index, component_id in enumerate( - downstream_remote_component_ids - ): - barrier_inputs[f"components.component_{component_index}"] = [ - component_id, - _modal_component_completion_output_index( - proxy_node_id=component_id, - rewritten_prompt=rewritten_prompt, - nodes_module=nodes_module, - ), - ] - rewritten_prompt[barrier_node_id] = { - "class_type": MODAL_DEFERRED_LOCAL_PASSTHROUGH_NODE_ID, - "inputs": barrier_inputs, - "_meta": {"title": "Modal Deferred Local Branch"}, + rewritten_prompt[passthrough_node_id] = { + "class_type": MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, + "inputs": passthrough_inputs, + "_meta": {"title": "Modal Parallel Local Branch"}, } for target in local_only_consumers: rewritten_prompt[target.node_id]["inputs"][target.input_name] = [ - barrier_node_id, + passthrough_node_id, 0, ] - barrier_node_ids.append(barrier_node_id) + for component_id in downstream_remote_component_ids: + prompt_inputs = rewritten_prompt[component_id].get("inputs", {}) + embedded_payload = prompt_inputs.get("original_node_data") + if not isinstance(embedded_payload, Mapping): + continue + prompt_inputs["original_node_data"] = ( + update_registered_proxy_payload_fields( + component_id, + embedded_payload, + { + "parallel_local_dispatch_group_id": dispatch_group_id, + "signal_parallel_local_dispatch": True, + }, + ) + ) + passthrough_node_ids.append(passthrough_node_id) logger.info( - "Deferred local-only consumers %s of remote output %s:%d until downstream Modal components %s complete via barrier=%s.", + "Parallelized local-only consumers %s of remote output %s:%d after downstream Modal dispatches %s via passthrough=%s.", [target.node_id for target in local_only_consumers], source.node_id, source.output_index, downstream_remote_component_ids, - barrier_node_id, + passthrough_node_id, ) - return barrier_node_ids + return passthrough_node_ids def _configure_local_gap_keepalive_payloads( @@ -4022,7 +4077,7 @@ def _configure_local_gap_keepalive_payloads( if not sandwiched_local_node_ids: return - payloads_by_component_id: dict[str, dict[str, Any]] = {} + payloads_by_component_id: dict[str, Mapping[str, Any]] = {} for component_id in remote_component_ids: payload = ( rewritten_prompt.get(component_id, {}) @@ -4031,9 +4086,6 @@ def _configure_local_gap_keepalive_payloads( ) if not isinstance(payload, dict): continue - payload["remote_local_gap_pool"] = True - payload.pop("keepalive_after_remote_component", None) - payload.pop("stop_local_gap_keepalive_before_remote_component", None) payloads_by_component_id[component_id] = payload keepalive_component_ids: set[str] = set() @@ -4053,14 +4105,19 @@ def _configure_local_gap_keepalive_payloads( ): continuation_component_ids.add(component_id) - for component_id in sorted(keepalive_component_ids): - payload = payloads_by_component_id.get(component_id) - if payload is not None: - payload["keepalive_after_remote_component"] = True - for component_id in sorted(continuation_component_ids): - payload = payloads_by_component_id.get(component_id) - if payload is not None: - payload["stop_local_gap_keepalive_before_remote_component"] = True + for component_id, embedded_payload in payloads_by_component_id.items(): + payload_fields: dict[str, Any] = {"remote_local_gap_pool": True} + if component_id in keepalive_component_ids: + payload_fields["keepalive_after_remote_component"] = True + if component_id in continuation_component_ids: + payload_fields["stop_local_gap_keepalive_before_remote_component"] = True + rewritten_prompt[component_id]["inputs"]["original_node_data"] = ( + update_registered_proxy_payload_fields( + component_id, + embedded_payload, + payload_fields, + ) + ) logger.info( "Configured local-gap Modal pool for components=%s keepalive_producers=%s continuations=%s.", sorted(payloads_by_component_id), @@ -4271,7 +4328,7 @@ def rewrite_prompt_for_modal( remote_component_ids=summary.remote_component_ids, sandwiched_local_node_ids=sandwiched_local_node_id_set, ) - summary.deferred_local_branch_node_ids = _defer_non_returning_local_branches( + summary.parallel_local_branch_node_ids = _parallelize_non_returning_local_branches( rewritten_prompt=rewritten_prompt, remote_component_ids=summary.remote_component_ids, nodes_module=resolved_nodes_module, @@ -4920,8 +4977,8 @@ def setup_modal_queue_route( "modal_sandwiched_local_node_ids": list( summary.sandwiched_local_node_ids ), - "modal_deferred_local_branch_node_ids": list( - summary.deferred_local_branch_node_ids + "modal_parallel_local_branch_node_ids": list( + summary.parallel_local_branch_node_ids ), "modal_components": [ { diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 425e58a..c6481c7 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -4,7 +4,8 @@ - A reusable worker-pool affinity key is ineffective if it appears only in diagnostics. Modal class parameters define the independently autoscaled instance identity, so the stable slot key must be passed as a `RemoteEngine` parameter on every real invocation and keepalive pulse; prompt-scoped session identity still belongs only in the payload. - Modal's `scaledown_window` is an idle upper bound, not a reservation that guarantees the same loaded GPU worker survives a long local phase. For a known remote → local → remote gap, combine an extended per-call class option with a bounded lightweight heartbeat, and stop that heartbeat as soon as the next remote component dispatches. -- ComfyUI gives ready output nodes priority, which can put an independent local preview or save branch on the remote critical path even when its result cannot affect later remote work. Insert a completion-token dependency barrier only on branches whose full downstream closure contains no remote proxy; any branch that transitively re-enters remote execution must retain its original dependency ordering. +- ComfyUI gives ready output and async nodes priority, which can put an independent local preview or save branch on the remote critical path even when its result cannot affect later remote work. Insert an async gate only on branches whose full downstream closure contains no remote proxy; wait for the nearest remote continuation to begin dispatch, then release the local branch while that remote call is in flight. Any branch that transitively re-enters remote execution must retain its original dependency ordering. +- Cache-friendly proxy inputs are only a stable cache surface; execution rehydrates the full run-scoped payload from the proxy context registry. Any rewrite-time control fields added after initial registration must update and re-register that full payload as well as the embedded prompt value, or the live proxy silently executes without them. A parallel-preview dispatch group must remain part of the cache surface so its signalling proxy executes on every prompt instead of being skipped by ComfyUI cache reuse. - Modal Dict operator assignment is always synchronous, even inside an async helper. When a cache path becomes asynchronous for ComfyUI compatibility, use `await dictionary.put.aio(key, value)` for remote writes and retain operator assignment only as a local-dictionary compatibility fallback; otherwise Modal 1.3.3 and newer warn while the network call blocks the active event loop. - Hugging Face Xet acceleration is a transport capability of `huggingface_hub`, not a separate model-download path. Pin and import-validate `hf-xet` in the CPU staging image, keep passing the environment token through the existing Hub calls, and let the Hub select Xet for repositories that use it. Adaptive concurrency is the safe default for the 16 GiB stager; high-performance mode raises buffers for machines with at least 64 GiB of RAM and must remain an explicit operator choice. diff --git a/modal_executor_node.py b/modal_executor_node.py index bf00062..ef0a419 100644 --- a/modal_executor_node.py +++ b/modal_executor_node.py @@ -8,6 +8,7 @@ import inspect import json import logging import threading +from collections import OrderedDict from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass @@ -20,7 +21,7 @@ from .serialization import deserialize_node_outputs, serialize_node_inputs, spli logger = logging.getLogger(__name__) MODAL_MAP_INPUT_NODE_ID = "ModalMapInput" MODAL_ARTIFACT_FINALIZER_NODE_ID = "ModalArtifactFinalizer" -MODAL_DEFERRED_LOCAL_PASSTHROUGH_NODE_ID = "ModalDeferredLocalPassthrough" +MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID = "ModalParallelLocalPassthrough" MODAL_COMPONENT_COMPLETION_OUTPUT_NAME = "modal_component_complete" MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS = 100 _PROXY_CACHE_CONTEXT_ID_KEY = "__comfy_modal_proxy_cache_context_id__" @@ -83,6 +84,11 @@ _MODAL_WORKFLOW_EXECUTION_GATE = threading.Condition() _MODAL_WORKFLOW_ACTIVE_PROMPT_ID: str | None = None _MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS = 0 _MODAL_WORKFLOW_ABANDONED_RELEASES_BY_PROMPT_ID: dict[str, int] = {} +_MODAL_PARALLEL_DISPATCH_EVENTS_LOCK = threading.Lock() +_MODAL_PARALLEL_DISPATCH_EVENTS: OrderedDict[ + tuple[str, str], asyncio.Event +] = OrderedDict() +_MODAL_PARALLEL_DISPATCH_EVENT_LIMIT = 512 @dataclass(frozen=True) @@ -197,6 +203,8 @@ def abandon_modal_workflow_execution_prompt(prompt_id: str | None, reason: str) if normalized_prompt_id is None: return + _clear_parallel_dispatch_events(normalized_prompt_id) + global _MODAL_WORKFLOW_ACTIVE_PROMPT_ID global _MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS global _MODAL_WORKFLOW_ABANDONED_RELEASES_BY_PROMPT_ID @@ -257,6 +265,112 @@ async def _modal_workflow_execution_slot(payload: Mapping[str, Any]) -> AsyncIte ) +def _parallel_dispatch_key( + dispatch_group_id: Any, + component_id: Any, +) -> tuple[str, str] | None: + """Return one normalized parallel-preview dispatch key.""" + normalized_group_id = _normalize_prompt_id(dispatch_group_id) + normalized_component_id = _normalize_prompt_id(component_id) + if normalized_group_id is None or normalized_component_id is None: + return None + return (normalized_group_id, normalized_component_id) + + +def _parallel_dispatch_event( + dispatch_group_id: Any, + component_id: Any, +) -> asyncio.Event | None: + """Return the bounded dispatch event for one remote continuation.""" + dispatch_key = _parallel_dispatch_key(dispatch_group_id, component_id) + if dispatch_key is None: + return None + with _MODAL_PARALLEL_DISPATCH_EVENTS_LOCK: + dispatch_event = _MODAL_PARALLEL_DISPATCH_EVENTS.get(dispatch_key) + if dispatch_event is None: + dispatch_event = asyncio.Event() + _MODAL_PARALLEL_DISPATCH_EVENTS[dispatch_key] = dispatch_event + while len(_MODAL_PARALLEL_DISPATCH_EVENTS) > ( + _MODAL_PARALLEL_DISPATCH_EVENT_LIMIT + ): + completed_dispatch_key = next( + ( + candidate_key + for candidate_key, candidate_event in ( + _MODAL_PARALLEL_DISPATCH_EVENTS.items() + ) + if candidate_event.is_set() + ), + None, + ) + if completed_dispatch_key is None: + break + _MODAL_PARALLEL_DISPATCH_EVENTS.pop(completed_dispatch_key, None) + else: + _MODAL_PARALLEL_DISPATCH_EVENTS.move_to_end(dispatch_key) + return dispatch_event + + +def _signal_parallel_local_dispatch(payload: Mapping[str, Any]) -> bool: + """Signal that one remote continuation acquired its local dispatch slot.""" + if not bool(payload.get("signal_parallel_local_dispatch")): + return False + dispatch_event = _parallel_dispatch_event( + payload.get("parallel_local_dispatch_group_id"), + payload.get("component_id"), + ) + if dispatch_event is None: + return False + dispatch_event.set() + logger.info( + "Released parallel local branches after remote dispatch started group=%s component=%s.", + payload.get("parallel_local_dispatch_group_id"), + payload.get("component_id"), + ) + return True + + +async def _wait_for_parallel_local_dispatches( + dispatch_context: Mapping[str, Any], +) -> None: + """Wait until every nearest downstream remote continuation has started.""" + dispatch_group_id = dispatch_context.get("dispatch_group_id") + component_ids = dispatch_context.get("component_ids") + if not isinstance(component_ids, Sequence) or isinstance( + component_ids, + str | bytes | bytearray, + ): + raise TypeError("Parallel local dispatch context must define component_ids.") + dispatch_events = [ + dispatch_event + for component_id in component_ids + if ( + dispatch_event := _parallel_dispatch_event( + dispatch_group_id, + component_id, + ) + ) + is not None + ] + if not dispatch_events: + raise RuntimeError( + "Parallel local dispatch context did not identify a remote continuation." + ) + await asyncio.gather(*(dispatch_event.wait() for dispatch_event in dispatch_events)) + + +def _clear_parallel_dispatch_events(dispatch_group_id: str) -> None: + """Forget dispatch events for one completed or abandoned prompt group.""" + with _MODAL_PARALLEL_DISPATCH_EVENTS_LOCK: + dispatch_keys = [ + dispatch_key + for dispatch_key in _MODAL_PARALLEL_DISPATCH_EVENTS + if dispatch_key[0] == dispatch_group_id + ] + for dispatch_key in dispatch_keys: + _MODAL_PARALLEL_DISPATCH_EVENTS.pop(dispatch_key, None) + + async def _execute_payload_async( client: RemoteExecutorClient, payload: Mapping[str, Any], @@ -418,6 +532,21 @@ def register_cache_friendly_proxy_payload( return sanitized_payload +def update_registered_proxy_payload_fields( + node_id: str, + embedded_payload: Mapping[str, Any], + fields: Mapping[str, Any], +) -> Mapping[str, Any]: + """Update both the embedded cache surface and run-scoped execution payload.""" + with _PROXY_EXECUTION_CONTEXTS_LOCK: + context = _PROXY_EXECUTION_CONTEXTS.get(str(node_id)) + execution_payload = dict( + context.execution_payload if context is not None else embedded_payload + ) + execution_payload.update(fields) + return register_cache_friendly_proxy_payload(node_id, execution_payload) + + def register_modal_map_input_warmup_context( node_id: str, payload: Mapping[str, Any], @@ -584,6 +713,7 @@ def _build_proxy_node_class( ) async with _modal_workflow_execution_slot(payload): + _signal_parallel_local_dispatch(payload) outputs = _normalize_scheduler_list_outputs( payload, await _execute_payload_async( @@ -688,14 +818,14 @@ def ensure_modal_artifact_finalizer_registered(nodes_module: Any) -> None: ) -def ensure_modal_deferred_local_passthrough_registered(nodes_module: Any) -> None: - """Register the internal local-branch scheduling barrier.""" - nodes_module.NODE_CLASS_MAPPINGS[MODAL_DEFERRED_LOCAL_PASSTHROUGH_NODE_ID] = ( - ModalDeferredLocalPassthrough +def ensure_modal_parallel_local_passthrough_registered(nodes_module: Any) -> None: + """Register the internal parallel local-branch dispatch gate.""" + nodes_module.NODE_CLASS_MAPPINGS[MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID] = ( + ModalParallelLocalPassthrough ) nodes_module.NODE_DISPLAY_NAME_MAPPINGS[ - MODAL_DEFERRED_LOCAL_PASSTHROUGH_NODE_ID - ] = "Modal Deferred Local Passthrough" + MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID + ] = "Modal Parallel Local Passthrough" class ModalUniversalExecutor(io.ComfyNode): @@ -811,29 +941,23 @@ class ModalArtifactFinalizer(io.ComfyNode): return io.NodeOutput() -class ModalDeferredLocalPassthrough(io.ComfyNode): - """Delay a non-returning local branch until continuing remote work completes.""" +class ModalParallelLocalPassthrough(io.ComfyNode): + """Release a local-only branch once continuing remote work is in flight.""" @classmethod def define_schema(cls) -> io.Schema: - """Accept one arbitrary value plus completion tokens from remote components.""" - completion_template = io.Autogrow.TemplatePrefix( - input=io.Boolean.Input("completion", force_input=True), - prefix="component_", - min=1, - max=MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS, - ) + """Accept one arbitrary value plus remote dispatch metadata.""" return io.Schema( - node_id=MODAL_DEFERRED_LOCAL_PASSTHROUGH_NODE_ID, - display_name="Modal Deferred Local Passthrough", + node_id=MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, + display_name="Modal Parallel Local Passthrough", category="Modal", description=( - "Internal scheduling barrier that prevents an independent local tap " - "from delaying downstream Modal execution." + "Internal async gate that releases an independent local tap once " + "downstream Modal execution has started." ), inputs=[ io.AnyType.Input("value"), - io.Autogrow.Input("components", template=completion_template), + io.AnyType.Input("dispatch_context"), ], outputs=[io.AnyType.Output(display_name="value")], is_dev_only=True, @@ -841,16 +965,11 @@ class ModalDeferredLocalPassthrough(io.ComfyNode): ) @classmethod - def execute(cls, value: Any, components: io.Autogrow.Type) -> io.NodeOutput: - """Return the original value after every required remote component completes.""" - incomplete_component_names = [ - str(component_name) - for component_name, completed in components.items() - if completed is not True - ] - if incomplete_component_names: - raise RuntimeError( - "Deferred local branch received incomplete Modal component tokens: " - f"{incomplete_component_names}." - ) + async def execute( + cls, + value: Any, + dispatch_context: Mapping[str, Any], + ) -> io.NodeOutput: + """Return the original value once the remote continuation is in flight.""" + await _wait_for_parallel_local_dispatches(dispatch_context) return io.NodeOutput(value) diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index a902d82..57cdbc5 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -794,22 +794,161 @@ def test_modal_artifact_finalizer_requires_completed_components( ) -def test_modal_deferred_local_passthrough_requires_completed_components( +def test_modal_parallel_local_passthrough_releases_after_remote_dispatch( modal_executor_module: Any, ) -> None: - """The local scheduling barrier should preserve values after remote completion.""" - schema = modal_executor_module.ModalDeferredLocalPassthrough.GET_SCHEMA() + """A local preview should unblock when its remote continuation starts.""" + schema = modal_executor_module.ModalParallelLocalPassthrough.GET_SCHEMA() + dispatch_context = { + "dispatch_group_id": "prompt-parallel", + "component_ids": ["component-b"], + } + + async def run_scenario() -> tuple[Any, ...]: + """Prove the passthrough remains pending until remote dispatch is signalled.""" + passthrough_task = asyncio.create_task( + modal_executor_module.ModalParallelLocalPassthrough.execute( + "image-value", + dispatch_context, + ) + ) + await asyncio.sleep(0) + assert passthrough_task.done() is False + assert modal_executor_module._signal_parallel_local_dispatch( + { + "parallel_local_dispatch_group_id": "prompt-parallel", + "component_id": "component-b", + "signal_parallel_local_dispatch": True, + } + ) is True + result = await asyncio.wait_for(passthrough_task, timeout=1.0) + return result.result + + try: + result = asyncio.run(run_scenario()) + finally: + modal_executor_module._clear_parallel_dispatch_events("prompt-parallel") assert schema.is_output_node is False - assert modal_executor_module.ModalDeferredLocalPassthrough.execute( - "image-value", - {"component_0": True, "component_1": True}, - ).result == ("image-value",) - with pytest.raises(RuntimeError, match="component_1"): - modal_executor_module.ModalDeferredLocalPassthrough.execute( - "image-value", - {"component_0": True, "component_1": False}, + assert result == ("image-value",) + + +def test_modal_parallel_local_passthrough_runs_before_remote_result_returns( + modal_executor_module: Any, +) -> None: + """A local preview should run while its downstream remote call is in flight.""" + fake_nodes_module = types.SimpleNamespace( + NODE_CLASS_MAPPINGS={}, + NODE_DISPLAY_NAME_MAPPINGS={}, + ) + proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered( + output_types=("STRING",), + output_names=("value",), + output_is_list=(False,), + nodes_module=fake_nodes_module, + is_output_node=False, + ) + proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] + remote_started = asyncio.Event() + release_remote = asyncio.Event() + + class FakeClient: + """Fake async remote client that remains in flight until explicitly released.""" + + async def execute_payload_async( + self, + payload: dict[str, Any], + kwargs: dict[str, Any], + ) -> tuple[str]: + """Record dispatch, then keep the remote result pending.""" + remote_started.set() + await release_remote.wait() + return (str(payload["component_id"]),) + + async def run_scenario() -> tuple[tuple[Any, ...], tuple[Any, ...]]: + """Prove local passthrough completion precedes the remote result.""" + dispatch_group_id = "prompt-parallel-in-flight" + local_task = asyncio.create_task( + modal_executor_module.ModalParallelLocalPassthrough.execute( + "preview-image", + { + "dispatch_group_id": dispatch_group_id, + "component_ids": ["component-b"], + }, + ) ) + remote_task = asyncio.create_task( + proxy_class.execute( + original_node_data={ + "payload_kind": "subgraph", + "prompt_id": dispatch_group_id, + "component_id": "component-b", + "parallel_local_dispatch_group_id": dispatch_group_id, + "signal_parallel_local_dispatch": True, + }, + unique_id="component-b", + ) + ) + await asyncio.wait_for(remote_started.wait(), timeout=1.0) + local_result = await asyncio.wait_for(local_task, timeout=1.0) + assert remote_task.done() is False + release_remote.set() + remote_result = await asyncio.wait_for(remote_task, timeout=1.0) + return local_result.result, remote_result.result + + modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) + try: + local_result, remote_result = asyncio.run(run_scenario()) + finally: + modal_executor_module.set_remote_executor_client_factory(None) + modal_executor_module._clear_parallel_dispatch_events( + "prompt-parallel-in-flight" + ) + + assert local_result == ("preview-image",) + assert remote_result == ("component-b",) + + +def test_parallel_local_dispatch_frontier_stops_at_nearest_remote_component( + api_intercept_module: Any, +) -> None: + """A preview gate should not wait for remote descendants of a remote frontier.""" + fake_nodes_module = types.SimpleNamespace( + NODE_CLASS_MAPPINGS={ + "RemoteProxy": _FakeRewriteRemoteSamplerNode, + "LocalFeedback": _FakeRewriteLocalFeedbackNode, + }, + ) + rewritten_prompt = { + "source": {"class_type": "RemoteProxy", "inputs": {}}, + "feedback": { + "class_type": "LocalFeedback", + "inputs": {"value": ["source", 0]}, + }, + "nearest": { + "class_type": "RemoteProxy", + "inputs": {"value": ["feedback", 0]}, + }, + "later": { + "class_type": "RemoteProxy", + "inputs": {"value": ["nearest", 0]}, + }, + } + consumers = api_intercept_module._build_consumer_map(rewritten_prompt) + + nearest_component_ids = ( + api_intercept_module._nearest_downstream_remote_component_ids( + rewritten_prompt=rewritten_prompt, + consumers=consumers, + seed_targets=consumers[ + api_intercept_module.LinkedOutputRef("source", 0) + ], + remote_component_id_set={"source", "nearest", "later"}, + nodes_module=fake_nodes_module, + ) + ) + + assert nearest_component_ids == ["nearest"] def test_proxy_execution_uses_injected_remote_client( @@ -10394,6 +10533,7 @@ def test_split_hybrid_proxies_allow_local_downstream_work_before_mapped_completi def test_mapped_component_with_local_reentry_rewrites_to_ordered_acyclic_proxies( api_intercept_module: Any, + modal_executor_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, @@ -10507,15 +10647,16 @@ def test_mapped_component_with_local_reentry_rewrites_to_ordered_acyclic_proxies assert summary.component_execution_stages == [["3"], ["7"]] assert rewritten_prompt["4"]["inputs"]["image"] == ["3", 0] assert rewritten_prompt["7"]["inputs"]["remote_input_2"] == ["4", 0] - assert len(summary.deferred_local_branch_node_ids) == 1 - barrier_node_id = summary.deferred_local_branch_node_ids[0] - barrier_node = rewritten_prompt[barrier_node_id] - assert barrier_node["class_type"] == ( - api_intercept_module.MODAL_DEFERRED_LOCAL_PASSTHROUGH_NODE_ID + assert len(summary.parallel_local_branch_node_ids) == 1 + passthrough_node_id = summary.parallel_local_branch_node_ids[0] + passthrough_node = rewritten_prompt[passthrough_node_id] + assert passthrough_node["class_type"] == ( + api_intercept_module.MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID ) - assert barrier_node["inputs"]["value"] == ["3", 0] - assert barrier_node["inputs"]["components.component_0"][0] == "7" - assert rewritten_prompt["9"]["inputs"]["image"] == [barrier_node_id, 0] + assert passthrough_node["inputs"]["value"] == ["3", 0] + dispatch_context = passthrough_node["inputs"]["dispatch_context"] + assert dispatch_context["component_ids"] == ["7"] + assert rewritten_prompt["9"]["inputs"]["image"] == [passthrough_node_id, 0] first_payload = rewritten_prompt["3"]["inputs"]["original_node_data"] mapped_payload = rewritten_prompt["7"]["inputs"]["original_node_data"] assert first_payload["component_node_ids"] == ["1", "3"] @@ -10526,6 +10667,15 @@ def test_mapped_component_with_local_reentry_rewrites_to_ordered_acyclic_proxies assert mapped_payload["remote_local_gap_pool"] is True assert "keepalive_after_remote_component" not in mapped_payload assert mapped_payload["stop_local_gap_keepalive_before_remote_component"] is True + assert mapped_payload["parallel_local_dispatch_group_id"] == ( + dispatch_context["dispatch_group_id"] + ) + assert mapped_payload["signal_parallel_local_dispatch"] is True + rehydrated_mapped_payload = modal_executor_module._rehydrate_proxy_payload( + mapped_payload, + unique_id="7", + ) + assert rehydrated_mapped_payload["signal_parallel_local_dispatch"] is True assert { boundary_input["proxy_input_name"] for boundary_input in mapped_payload["boundary_inputs"]