nodes/ttulttul-ComfyUI-Modal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Run local previews alongside remote continuations

Browse Source
This commit is contained in:
Ken Simpson
2026-08-19 20:48:18 -07:00
parent c92af6a8c2
commit 093b2bdeec
5 changed files with 437 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -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.
api_intercept.py
+113 -56
View File
@@ -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": [
{
docs/LEARNINGS.md
+2 -1
View File
@@ -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.
modal_executor_node.py
+152 -33
View File
@@ -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)
tests/test_modal_executor_node.py
+169 -19
View File
@@ -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"]