657 lines
22 KiB
Python
657 lines
22 KiB
Python
"""Dynamic Modal proxy nodes for ComfyUI execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import threading
|
|
from collections import OrderedDict
|
|
from collections.abc import AsyncIterator, Mapping, Sequence
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
from comfy_api.latest import _io as io
|
|
|
|
if __package__:
|
|
from .serialization import (
|
|
deserialize_node_outputs,
|
|
serialize_node_inputs,
|
|
split_mapped_value,
|
|
)
|
|
from .proxy_node_factory import (
|
|
MODAL_ARTIFACT_FINALIZER_NODE_ID,
|
|
MODAL_COMPONENT_COMPLETION_OUTPUT_NAME,
|
|
MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID,
|
|
MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID,
|
|
ProxyNodeFactoryHooks,
|
|
_build_proxy_node_class,
|
|
_normalized_output_metadata,
|
|
_output_spec,
|
|
_proxy_node_id,
|
|
_register_modal_node,
|
|
configure_proxy_node_factory_hooks,
|
|
ensure_modal_artifact_finalizer_registered,
|
|
ensure_modal_component_proxy_node_registered,
|
|
ensure_modal_local_bridge_materializer_registered,
|
|
ensure_modal_parallel_local_passthrough_registered,
|
|
ensure_modal_proxy_node_registered,
|
|
)
|
|
from .proxy_payloads import (
|
|
MODAL_PROMPT_ID_EXTRA_PNGINFO_KEY,
|
|
_boost_modal_map_input_warmup,
|
|
_normalize_prompt_id,
|
|
_normalize_proxy_kwargs,
|
|
_normalize_proxy_payload,
|
|
_normalize_scheduler_list_outputs,
|
|
_payload_is_local_cache_safe,
|
|
_pop_proxy_hidden_value,
|
|
_prompt_id_from_extra_pnginfo,
|
|
_rehydrate_proxy_payload,
|
|
_sanitize_cache_surface_payload,
|
|
_unwrap_proxy_singleton,
|
|
register_cache_friendly_proxy_payload,
|
|
register_modal_map_input_warmup_context,
|
|
registered_proxy_execution_payload,
|
|
update_registered_proxy_payload_fields,
|
|
)
|
|
from .remote_executor_router import (
|
|
ModalRemoteExecutorClient,
|
|
RemoteExecutorClient,
|
|
RemoteExecutorRouterClient,
|
|
get_remote_executor_client,
|
|
set_remote_executor_client_factory,
|
|
)
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from serialization import (
|
|
deserialize_node_outputs,
|
|
serialize_node_inputs,
|
|
split_mapped_value,
|
|
)
|
|
from proxy_node_factory import (
|
|
MODAL_ARTIFACT_FINALIZER_NODE_ID,
|
|
MODAL_COMPONENT_COMPLETION_OUTPUT_NAME,
|
|
MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID,
|
|
MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID,
|
|
ProxyNodeFactoryHooks,
|
|
_build_proxy_node_class,
|
|
_normalized_output_metadata,
|
|
_output_spec,
|
|
_proxy_node_id,
|
|
_register_modal_node,
|
|
configure_proxy_node_factory_hooks,
|
|
ensure_modal_artifact_finalizer_registered,
|
|
ensure_modal_component_proxy_node_registered,
|
|
ensure_modal_local_bridge_materializer_registered,
|
|
ensure_modal_parallel_local_passthrough_registered,
|
|
ensure_modal_proxy_node_registered,
|
|
)
|
|
from proxy_payloads import (
|
|
MODAL_PROMPT_ID_EXTRA_PNGINFO_KEY,
|
|
_boost_modal_map_input_warmup,
|
|
_normalize_prompt_id,
|
|
_normalize_proxy_kwargs,
|
|
_normalize_proxy_payload,
|
|
_normalize_scheduler_list_outputs,
|
|
_payload_is_local_cache_safe,
|
|
_pop_proxy_hidden_value,
|
|
_prompt_id_from_extra_pnginfo,
|
|
_rehydrate_proxy_payload,
|
|
_sanitize_cache_surface_payload,
|
|
_unwrap_proxy_singleton,
|
|
register_cache_friendly_proxy_payload,
|
|
register_modal_map_input_warmup_context,
|
|
registered_proxy_execution_payload,
|
|
update_registered_proxy_payload_fields,
|
|
)
|
|
from remote_executor_router import (
|
|
ModalRemoteExecutorClient,
|
|
RemoteExecutorClient,
|
|
RemoteExecutorRouterClient,
|
|
get_remote_executor_client,
|
|
set_remote_executor_client_factory,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
MODAL_MAP_INPUT_NODE_ID = "ModalMapInput"
|
|
MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS = 100
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
_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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _acquire_modal_workflow_execution_slot(
|
|
prompt_id: str | None,
|
|
component_id: Any,
|
|
) -> bool:
|
|
"""Try to reserve a remote execution slot for one prompt."""
|
|
if prompt_id is None:
|
|
return False
|
|
|
|
global _MODAL_WORKFLOW_ACTIVE_PROMPT_ID
|
|
global _MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS
|
|
|
|
with _MODAL_WORKFLOW_EXECUTION_GATE:
|
|
if (
|
|
_MODAL_WORKFLOW_ACTIVE_PROMPT_ID is not None
|
|
and _MODAL_WORKFLOW_ACTIVE_PROMPT_ID != prompt_id
|
|
):
|
|
return False
|
|
|
|
_MODAL_WORKFLOW_ACTIVE_PROMPT_ID = prompt_id
|
|
_MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS += 1
|
|
logger.debug(
|
|
"Acquired Modal workflow execution slot for prompt=%s component=%s active_remote_calls=%d.",
|
|
prompt_id,
|
|
component_id,
|
|
_MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS,
|
|
)
|
|
return True
|
|
|
|
|
|
def _release_modal_workflow_execution_slot(prompt_id: str | None, component_id: Any) -> None:
|
|
"""Release a prompt-scoped remote execution slot and unblock the next prompt if idle."""
|
|
if prompt_id is None:
|
|
return
|
|
|
|
global _MODAL_WORKFLOW_ACTIVE_PROMPT_ID
|
|
global _MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS
|
|
global _MODAL_WORKFLOW_ABANDONED_RELEASES_BY_PROMPT_ID
|
|
|
|
with _MODAL_WORKFLOW_EXECUTION_GATE:
|
|
if _MODAL_WORKFLOW_ACTIVE_PROMPT_ID != prompt_id:
|
|
abandoned_release_count = _MODAL_WORKFLOW_ABANDONED_RELEASES_BY_PROMPT_ID.get(
|
|
prompt_id,
|
|
0,
|
|
)
|
|
if abandoned_release_count > 0:
|
|
logger.debug(
|
|
"Ignoring Modal workflow gate release for abandoned prompt=%s component=%s.",
|
|
prompt_id,
|
|
component_id,
|
|
)
|
|
if abandoned_release_count == 1:
|
|
_MODAL_WORKFLOW_ABANDONED_RELEASES_BY_PROMPT_ID.pop(prompt_id, None)
|
|
else:
|
|
_MODAL_WORKFLOW_ABANDONED_RELEASES_BY_PROMPT_ID[prompt_id] = (
|
|
abandoned_release_count - 1
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"Ignoring Modal workflow gate release for prompt=%s component=%s because active prompt is %s.",
|
|
prompt_id,
|
|
component_id,
|
|
_MODAL_WORKFLOW_ACTIVE_PROMPT_ID,
|
|
)
|
|
return
|
|
|
|
_MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS = max(0, _MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS - 1)
|
|
logger.debug(
|
|
"Released Modal workflow execution slot for prompt=%s component=%s active_remote_calls=%d.",
|
|
prompt_id,
|
|
component_id,
|
|
_MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS,
|
|
)
|
|
if _MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS == 0:
|
|
_MODAL_WORKFLOW_ACTIVE_PROMPT_ID = None
|
|
_MODAL_WORKFLOW_ABANDONED_RELEASES_BY_PROMPT_ID.pop(prompt_id, None)
|
|
_MODAL_WORKFLOW_EXECUTION_GATE.notify_all()
|
|
|
|
|
|
def abandon_modal_workflow_execution_prompt(prompt_id: str | None, reason: str) -> None:
|
|
"""Release a prompt-wide Modal gate when ComfyUI has cancelled that prompt locally."""
|
|
normalized_prompt_id = _normalize_prompt_id(prompt_id)
|
|
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
|
|
|
|
with _MODAL_WORKFLOW_EXECUTION_GATE:
|
|
if _MODAL_WORKFLOW_ACTIVE_PROMPT_ID != normalized_prompt_id:
|
|
return
|
|
|
|
logger.info(
|
|
"Abandoning Modal workflow execution slot for prompt=%s with %d active remote call(s): %s.",
|
|
normalized_prompt_id,
|
|
_MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS,
|
|
reason,
|
|
)
|
|
_MODAL_WORKFLOW_ABANDONED_RELEASES_BY_PROMPT_ID[normalized_prompt_id] = (
|
|
_MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS
|
|
)
|
|
_MODAL_WORKFLOW_ACTIVE_PROMPT_ID = None
|
|
_MODAL_WORKFLOW_ACTIVE_REMOTE_CALLS = 0
|
|
_MODAL_WORKFLOW_EXECUTION_GATE.notify_all()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _modal_workflow_execution_slot(payload: Mapping[str, Any]) -> AsyncIterator[None]:
|
|
"""Async context that serializes remote work across different prompt ids."""
|
|
prompt_id = _normalize_prompt_id(payload.get("prompt_id"))
|
|
component_id = payload.get("component_id")
|
|
if prompt_id is None:
|
|
yield
|
|
return
|
|
|
|
acquired = False
|
|
waiting_logged = False
|
|
while not acquired:
|
|
acquired = await asyncio.to_thread(
|
|
_acquire_modal_workflow_execution_slot,
|
|
prompt_id,
|
|
component_id,
|
|
)
|
|
if not acquired:
|
|
if not waiting_logged:
|
|
logger.info(
|
|
"Waiting to start Modal component=%s for prompt=%s until active prompt=%s finishes remote execution.",
|
|
component_id,
|
|
prompt_id,
|
|
_MODAL_WORKFLOW_ACTIVE_PROMPT_ID,
|
|
)
|
|
waiting_logged = True
|
|
await asyncio.sleep(0.1)
|
|
try:
|
|
yield
|
|
finally:
|
|
if acquired:
|
|
await asyncio.to_thread(
|
|
_release_modal_workflow_execution_slot,
|
|
prompt_id,
|
|
component_id,
|
|
)
|
|
|
|
|
|
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],
|
|
kwargs: Mapping[str, Any],
|
|
) -> Sequence[Any]:
|
|
"""Execute one Modal payload through the client, adapting sync clients when needed."""
|
|
execute_payload_async = getattr(client, "execute_payload_async", None)
|
|
if callable(execute_payload_async):
|
|
result = execute_payload_async(payload, kwargs)
|
|
if inspect.isawaitable(result):
|
|
return await result
|
|
return result
|
|
|
|
execute_payload = getattr(client, "execute_payload", None)
|
|
if not callable(execute_payload):
|
|
raise TypeError("Remote executor client must define execute_payload or execute_payload_async.")
|
|
return await asyncio.to_thread(execute_payload, payload, kwargs)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ModalUniversalExecutor(io.ComfyNode):
|
|
"""Base debug node for Modal execution routing."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""Define a minimal dev-only schema for the base proxy node."""
|
|
return io.Schema(
|
|
node_id="ModalUniversalExecutor",
|
|
display_name="Modal Universal Executor",
|
|
category="Modal",
|
|
description=(
|
|
"Debug entrypoint for Modal-backed execution. Production rewrites use "
|
|
"signature-preserving dynamic proxy variants."
|
|
),
|
|
inputs=[io.AnyType.Input("original_node_data")],
|
|
outputs=[io.AnyType.Output(display_name="output")],
|
|
accept_all_inputs=True,
|
|
is_dev_only=True,
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, original_node_data: Any, **kwargs: Any) -> io.NodeOutput:
|
|
"""Execute the base debug proxy node through the configured client."""
|
|
if isinstance(original_node_data, str):
|
|
original_node_data = json.loads(original_node_data)
|
|
outputs = tuple(get_remote_executor_client().execute_node(original_node_data, kwargs))
|
|
return io.NodeOutput(*outputs)
|
|
|
|
|
|
class ModalMapInput(io.ComfyNode):
|
|
"""Queue-time marker node that turns one remote boundary input into mapped parallel work."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""Expose a simple any-to-any adapter node for mapped remote execution."""
|
|
return io.Schema(
|
|
node_id=MODAL_MAP_INPUT_NODE_ID,
|
|
display_name="Modal Map Input",
|
|
category="Modal",
|
|
description=(
|
|
"Pass-through marker for data-parallel Modal execution. "
|
|
"When used inside a remote-marked component, list inputs and batched tensors "
|
|
"can fan out across multiple Modal executions and reassemble automatically."
|
|
),
|
|
inputs=[io.AnyType.Input("value")],
|
|
outputs=[io.AnyType.Output(display_name="value")],
|
|
hidden=[io.Hidden.unique_id],
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, value: Any, unique_id: str | None = None) -> io.NodeOutput:
|
|
"""Pass the input value through unchanged at runtime."""
|
|
_boost_modal_map_input_warmup(
|
|
_normalize_prompt_id(unique_id),
|
|
value,
|
|
)
|
|
return io.NodeOutput(value)
|
|
|
|
|
|
class ModalArtifactFinalizer(io.ComfyNode):
|
|
"""Internal output sink that anchors remote execution and artifact materialization."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""Accept one completion token from every rewritten terminal Modal component."""
|
|
completion_template = io.Autogrow.TemplatePrefix(
|
|
input=io.Boolean.Input("completion", force_input=True),
|
|
prefix="component_",
|
|
min=1,
|
|
max=MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS,
|
|
)
|
|
return io.Schema(
|
|
node_id=MODAL_ARTIFACT_FINALIZER_NODE_ID,
|
|
display_name="Modal Artifact Finalizer",
|
|
category="Modal",
|
|
description=(
|
|
"Internal output sink inserted into rewritten prompts so Modal components "
|
|
"run and materialize their remote output artifacts locally."
|
|
),
|
|
inputs=[
|
|
io.Autogrow.Input(
|
|
"components",
|
|
template=completion_template,
|
|
)
|
|
],
|
|
outputs=[],
|
|
is_output_node=True,
|
|
is_dev_only=True,
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, components: io.Autogrow.Type) -> io.NodeOutput:
|
|
"""Confirm every upstream Modal component completed before the prompt finishes."""
|
|
incomplete_component_names = [
|
|
str(component_name)
|
|
for component_name, completed in components.items()
|
|
if completed is not True
|
|
]
|
|
if incomplete_component_names:
|
|
raise RuntimeError(
|
|
"Modal artifact finalization received incomplete component tokens: "
|
|
f"{incomplete_component_names}."
|
|
)
|
|
logger.info(
|
|
"Modal artifact finalization completed after %d remote component(s).",
|
|
len(components),
|
|
)
|
|
return io.NodeOutput()
|
|
|
|
|
|
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 remote dispatch metadata."""
|
|
return io.Schema(
|
|
node_id=MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID,
|
|
display_name="Modal Parallel Local Passthrough",
|
|
category="Modal",
|
|
description=(
|
|
"Internal async gate that releases an independent local tap once "
|
|
"downstream Modal execution has started."
|
|
),
|
|
inputs=[
|
|
io.AnyType.Input("value"),
|
|
io.AnyType.Input("dispatch_context"),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="value")],
|
|
is_dev_only=True,
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
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)
|
|
|
|
|
|
class ModalLocalBridgeMaterializer(io.ComfyNode):
|
|
"""Download a durable remote bridge value for an independent local branch."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""Accept one durable bridge reference and return its local value."""
|
|
return io.Schema(
|
|
node_id=MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID,
|
|
display_name="Modal Local Bridge Materializer",
|
|
category="Modal",
|
|
description=(
|
|
"Internal async node that materializes a durable Modal bridge for "
|
|
"local-only work while remote execution continues."
|
|
),
|
|
inputs=[io.AnyType.Input("bridge_ref")],
|
|
outputs=[io.AnyType.Output(display_name="value")],
|
|
is_dev_only=True,
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(cls, bridge_ref: Mapping[str, Any]) -> io.NodeOutput:
|
|
"""Materialize the bridge without blocking ComfyUI's async scheduler."""
|
|
from .remote.modal_app import materialize_remote_session_bridge_ref_locally
|
|
|
|
value = await asyncio.to_thread(
|
|
materialize_remote_session_bridge_ref_locally,
|
|
bridge_ref,
|
|
)
|
|
return io.NodeOutput(value)
|
|
|
|
|
|
configure_proxy_node_factory_hooks(
|
|
ProxyNodeFactoryHooks(
|
|
modal_workflow_execution_slot=_modal_workflow_execution_slot,
|
|
signal_parallel_local_dispatch=_signal_parallel_local_dispatch,
|
|
execute_payload_async=_execute_payload_async,
|
|
modal_universal_executor=ModalUniversalExecutor,
|
|
modal_artifact_finalizer=ModalArtifactFinalizer,
|
|
modal_parallel_local_passthrough=ModalParallelLocalPassthrough,
|
|
modal_local_bridge_materializer=ModalLocalBridgeMaterializer,
|
|
)
|
|
)
|