2182 lines
74 KiB
Python
2182 lines
74 KiB
Python
"""Remote Modal runtime and local execution fallback."""
|
|
|
|
import asyncio
|
|
import copy
|
|
from dataclasses import dataclass, field, replace
|
|
import hashlib
|
|
import importlib
|
|
import importlib.util
|
|
import inspect
|
|
from io import BytesIO
|
|
import json
|
|
import logging
|
|
import os
|
|
import queue
|
|
import shutil
|
|
import statistics
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from types import ModuleType
|
|
import uuid
|
|
import zipfile
|
|
from concurrent.futures import (
|
|
Future,
|
|
ThreadPoolExecutor,
|
|
TimeoutError as FutureTimeoutError,
|
|
)
|
|
from contextlib import contextmanager, nullcontext
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Iterable, Iterator, Mapping
|
|
|
|
from ..serialization import (
|
|
coerce_serialized_node_outputs,
|
|
deserialize_value,
|
|
deserialize_node_inputs,
|
|
deserialize_node_outputs,
|
|
is_mapped_output_value,
|
|
join_mapped_values_for_scheduler,
|
|
serialize_mapping,
|
|
serialize_node_inputs,
|
|
serialize_node_outputs,
|
|
serialize_value,
|
|
split_mapped_value,
|
|
unwrap_mapped_output_value,
|
|
)
|
|
from ..durable_state import (
|
|
DurableObjectRef,
|
|
DurableStateError,
|
|
FileDurableObjectStore,
|
|
read_modal_volume_file,
|
|
stable_remote_invocation_id,
|
|
)
|
|
from ..output_artifacts import (
|
|
RemoteExecutionResult,
|
|
materialize_remote_output_artifacts,
|
|
unpack_remote_execution_result,
|
|
)
|
|
from ..remote_protocol import (
|
|
BOUNDARY_INPUT_SIGNATURES_KEY as _BOUNDARY_INPUT_SIGNATURES_KEY,
|
|
PRIMITIVE_WIDGET_INPUT_TYPES as _PRIMITIVE_WIDGET_INPUT_TYPES,
|
|
)
|
|
from ..session_state import (
|
|
InMemoryRemoteSessionBridgeStore,
|
|
InMemoryRemoteSessionStore,
|
|
RemoteSessionBridgeRecord,
|
|
RemoteSessionBridgeRecoveryKind,
|
|
RemoteSessionBridgeRef,
|
|
RemoteSessionHandle,
|
|
RemoteSessionStateError,
|
|
RemoteSessionValueRef,
|
|
is_remote_session_bridge_ref_payload,
|
|
is_remote_session_handle_payload,
|
|
is_remote_session_value_ref_payload,
|
|
stable_session_bridge_key,
|
|
)
|
|
from ..runtime_environment import (
|
|
REMOTE_APP_PROTOCOL_VERSION,
|
|
RemoteRuntimeIdentity,
|
|
build_remote_runtime_identity,
|
|
)
|
|
from ..llm_profiles import (
|
|
get_llm_profile,
|
|
llm_model_reference_node_ids_from_payload,
|
|
llm_model_references_from_payload,
|
|
resolved_llm_profile_payloads,
|
|
rewrite_llm_model_references,
|
|
)
|
|
from ..llm_recovery import (
|
|
LLM_FORCE_VLLM_THROUGHPUT_PAYLOAD_KEY,
|
|
exhausted_recovery_used_vllm_throughput,
|
|
is_llm_memory_recovery_exhausted,
|
|
)
|
|
from ..staging_process import staging_no_progress_timeout_seconds
|
|
from ..settings import (
|
|
MODAL_GPU_TYPES,
|
|
ModalSyncSettings,
|
|
get_settings,
|
|
modal_deployment_app_name,
|
|
settings_for_modal_gpu,
|
|
)
|
|
from .modal_billing import (
|
|
ModalBillingStatusError,
|
|
ModalHourlyBillingStatus,
|
|
_completed_modal_billing_interval,
|
|
_fetch_modal_hourly_billing_synchronously,
|
|
_matching_modal_hourly_billing_rows,
|
|
_modal_billing_cost,
|
|
_modal_billing_row_value,
|
|
_modal_hourly_billing_status_from_rows,
|
|
_prune_modal_hourly_billing_cache,
|
|
get_hourly_modal_app_billing,
|
|
)
|
|
from .modal_container_logs import (
|
|
MODAL_GPU_ESTIMATED_USD_PER_SECOND,
|
|
MODAL_GPU_PRICING_EFFECTIVE_DATE,
|
|
ModalContainerStatus,
|
|
ModalContainerStatusError,
|
|
_coerce_modal_task_id,
|
|
_is_remote_container_log_stream_enabled,
|
|
_release_remote_container_log_stream,
|
|
_retain_remote_container_log_stream,
|
|
_stop_modal_task_synchronously,
|
|
list_active_modal_containers,
|
|
stop_managed_modal_container,
|
|
)
|
|
from .host_session_bridge import (
|
|
_RemoteSessionBridgeResolutionStats,
|
|
_build_remote_session_bridge_record,
|
|
_log_remote_session_resolution_summary,
|
|
_payload_remote_session_handle,
|
|
_rehydrate_remote_session_bridge_value,
|
|
_resolve_remote_session_inputs,
|
|
_restore_serialized_remote_session_bridge_value,
|
|
_store_remote_session_bridge_value,
|
|
materialize_remote_session_bridge_ref_locally,
|
|
)
|
|
from . import host_session_bridge as _host_session_bridge
|
|
from .local_execution import (
|
|
RemoteSubgraphExecutionError,
|
|
_execute_subgraph_prompt,
|
|
_is_link,
|
|
_iter_prompt_links,
|
|
_load_nodes_module,
|
|
_materialize_local_custom_node_assets,
|
|
_node_input_type_map,
|
|
_resolve_required_subgraph_nodes,
|
|
execute_node_locally,
|
|
execute_subgraph_locally,
|
|
)
|
|
from .local_ui_events import (
|
|
_emit_local_executed_output,
|
|
_emit_local_modal_progress,
|
|
_emit_local_modal_status,
|
|
_emit_local_preview_boundary_output,
|
|
_emit_local_preview_image,
|
|
_emit_local_remote_dispatch_status,
|
|
_emit_local_remote_startup_status,
|
|
_lookup_local_prompt_server,
|
|
_progress_stream_event_metadata,
|
|
_remote_execution_destination,
|
|
_remote_execution_identity,
|
|
_remote_prompt_ancestor_node_ids,
|
|
_should_forward_suppressed_stream_event,
|
|
_should_stream_remote_progress,
|
|
)
|
|
from .modal_deployment import (
|
|
_MODAL_APP_STOP_TIMEOUT_SECONDS,
|
|
_MODAL_CLOUD_MODULE_NAME,
|
|
_REMOTE_APP_PROTOCOL_VERSION,
|
|
ModalDeploymentHooks,
|
|
ModalRemoteAppOutOfDateError,
|
|
ModalRemoteInvocationError,
|
|
_auto_deploy_modal_app,
|
|
_call_modal_method,
|
|
_component_pool_slot_index,
|
|
_ensure_remote_engine_protocol_current,
|
|
_expected_remote_runtime_fingerprint,
|
|
_install_modal_cloud_exception_compatibility_module,
|
|
_is_missing_modal_deployment_error,
|
|
_load_modal_cloud_module,
|
|
_lookup_deployed_remote_engine,
|
|
_lookup_deployed_remote_engine_with_retry,
|
|
_lookup_protocol_current_remote_engine,
|
|
_mapped_lane_affinity_key,
|
|
_modal_auto_deploy_state,
|
|
_modal_cloud_settings_override,
|
|
_modal_deploy_cache_key,
|
|
_modal_environment_name,
|
|
_modal_lookup_error_types,
|
|
_modal_runtime_cache_key,
|
|
_remote_runtime_identity_for_settings,
|
|
_remote_worker_affinity_key,
|
|
_remote_worker_pool_affinity_key,
|
|
_replace_outdated_modal_app,
|
|
_runtime_fingerprint_from_payload,
|
|
_settings_for_payload,
|
|
_stop_modal_app_for_replacement,
|
|
_stop_modal_app_via_cli,
|
|
_stop_modal_app_via_sdk,
|
|
configure_modal_deployment_hooks,
|
|
)
|
|
from .modal_warmup import (
|
|
ModalWarmupHooks,
|
|
_await_matching_speculative_prewarm,
|
|
_await_prompt_warmup_slots,
|
|
_build_llm_prewarm_plans,
|
|
_build_loader_prewarm_plans,
|
|
_build_prompt_warmup_request,
|
|
_component_parallelism_metadata,
|
|
_ensure_prompt_warmup_state,
|
|
_invoke_modal_warmup_blocking,
|
|
_loader_prewarm_plan_signature,
|
|
_loader_snapshot_profile_key,
|
|
_prepare_snapshot_profile_fields,
|
|
_prompt_parallelism_target,
|
|
_prompt_warmup_head_start_seconds,
|
|
_record_snapshot_warmup_measurement,
|
|
_register_exact_component_parallelism,
|
|
_run_speculative_affinity_prewarm,
|
|
_schedule_post_deploy_runtime_seed,
|
|
_schedule_speculative_affinity_prewarm,
|
|
_select_gpu_snapshot_for_profile,
|
|
_speculative_warmup_identity,
|
|
_start_local_gap_keepalive,
|
|
_stop_local_gap_keepalive,
|
|
_track_prompt_warmup_future,
|
|
_warmup_prompt_id,
|
|
boost_mapped_component_warmup,
|
|
configure_modal_warmup_hooks,
|
|
ensure_remote_warm_capacity,
|
|
)
|
|
from .mapped_execution import (
|
|
MappedExecutionHooks,
|
|
_aggregate_mapped_outputs,
|
|
_annotate_implicit_batched_boundary_outputs,
|
|
_build_mapped_item_payload,
|
|
_build_phase_subgraph_payload,
|
|
_build_remote_session_cleanup_payload,
|
|
_build_static_mapped_payload,
|
|
_clear_local_mapped_lane_progress,
|
|
_emit_local_mapped_lane_progress_start,
|
|
_emit_local_mapped_progress,
|
|
_execute_mapped_subgraph_payload,
|
|
_implicit_batch_boundary_effective_io_type,
|
|
_implicit_batch_boundary_target_input_types,
|
|
_implicit_batch_input_is_list_target_node_ids,
|
|
_implicit_batch_preserving_targets,
|
|
_invoke_implicitly_mapped_subgraph_async,
|
|
_invoke_mapped_remote_engine_async,
|
|
_is_latent_like_mapping,
|
|
_is_mapped_boundary_output,
|
|
_list_is_latent_like_batch,
|
|
_log_detached_mapped_lane_result,
|
|
_mapped_execution_parallelism,
|
|
_mapped_phase_definition,
|
|
_mapped_progress_owner_component_id,
|
|
_merge_static_and_mapped_outputs,
|
|
_partition_implicit_batched_execute_nodes,
|
|
_pop_mapped_lane_node_id,
|
|
_remember_mapped_lane_node_id,
|
|
_shared_subgraph_payload_fields,
|
|
_split_batch_boundary_inputs,
|
|
_split_phase_outputs,
|
|
configure_mapped_execution_hooks,
|
|
)
|
|
from .modal_llm_profile_staging import (
|
|
_bounded_modal_stage_events,
|
|
_close_modal_stage_events,
|
|
_emit_local_llm_staging_progress,
|
|
_ensure_llm_profiles_staged,
|
|
_read_modal_stage_events,
|
|
_rewrite_staged_llm_kwargs_payload,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Remote durable-object handles.
|
|
_MODAL_INTERRUPT_DICTS_LOCK = threading.Lock()
|
|
_MODAL_INTERRUPT_DICTS: dict[tuple[str, str | None], Any] = {}
|
|
|
|
# Active invocation and local durable-session state.
|
|
_ACTIVE_REMOTE_INVOCATIONS_LOCK = threading.Lock()
|
|
_ACTIVE_REMOTE_INVOCATIONS_BY_PROMPT: dict[
|
|
str, dict[str, "_ActiveRemoteInvocation"]
|
|
] = {}
|
|
|
|
def _remote_modal_call_worker_count() -> int:
|
|
"""Return the number of local worker threads reserved for blocking Modal calls."""
|
|
return max(1, int(get_settings().max_inflight_calls))
|
|
|
|
|
|
# Blocking Modal call executor state.
|
|
_REMOTE_MODAL_CALL_EXECUTOR_LOCK = threading.Lock()
|
|
_REMOTE_MODAL_CALL_EXECUTOR: ThreadPoolExecutor | None = None
|
|
|
|
|
|
def _remote_modal_call_executor() -> ThreadPoolExecutor:
|
|
"""Return the lazily constructed executor for blocking Modal calls."""
|
|
global _REMOTE_MODAL_CALL_EXECUTOR
|
|
with _REMOTE_MODAL_CALL_EXECUTOR_LOCK:
|
|
if _REMOTE_MODAL_CALL_EXECUTOR is None:
|
|
_REMOTE_MODAL_CALL_EXECUTOR = ThreadPoolExecutor(
|
|
max_workers=_remote_modal_call_worker_count()
|
|
)
|
|
return _REMOTE_MODAL_CALL_EXECUTOR
|
|
|
|
|
|
@dataclass
|
|
class _ActiveRemoteInvocation:
|
|
"""Track one local proxy call that is currently waiting on remote Modal work."""
|
|
|
|
prompt_id: str
|
|
component_id: str
|
|
cancellation_event: threading.Event | None
|
|
interrupt_remote_call: Callable[[], Any] | None
|
|
|
|
|
|
|
|
|
|
try:
|
|
import modal # type: ignore
|
|
except ModuleNotFoundError: # pragma: no cover - exercised by local fallback tests.
|
|
modal = None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _close_remote_payload_stream(stream_events: Iterator[dict[str, Any]]) -> None:
|
|
"""Best-effort close one streamed Modal iterator after the terminal result arrives."""
|
|
close_callable = getattr(stream_events, "close", None)
|
|
if not callable(close_callable):
|
|
return
|
|
try:
|
|
close_result = close_callable()
|
|
except (AttributeError, RuntimeError, TypeError, ValueError) as exc:
|
|
logger.debug("Ignoring remote payload stream close failure: %s", exc)
|
|
return
|
|
if not asyncio.iscoroutine(close_result):
|
|
return
|
|
try:
|
|
asyncio.run(close_result)
|
|
except (RuntimeError, ValueError) as exc:
|
|
logger.debug("Ignoring async remote payload stream close failure: %s", exc)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _local_processing_interrupted() -> bool:
|
|
"""Return whether the current local ComfyUI execution was interrupted."""
|
|
try:
|
|
import comfy.model_management
|
|
except ModuleNotFoundError:
|
|
return False
|
|
|
|
return bool(comfy.model_management.processing_interrupted())
|
|
|
|
|
|
def _raise_local_interrupt() -> None:
|
|
"""Raise ComfyUI's native interruption exception for the current execution."""
|
|
import comfy.model_management
|
|
|
|
raise comfy.model_management.InterruptProcessingException()
|
|
|
|
|
|
def _exception_indicates_interruption(exc: BaseException) -> bool:
|
|
"""Return whether an exception represents cancellation or interrupted execution."""
|
|
if isinstance(exc, asyncio.CancelledError):
|
|
return True
|
|
message = str(exc).lower()
|
|
return "interrupt" in message or "cancel" in message
|
|
|
|
|
|
def _remote_interrupt_key(payload: dict[str, Any]) -> tuple[str, str]:
|
|
"""Return the prompt/component pair used to interrupt one remote execution."""
|
|
prompt_id = str(
|
|
payload.get("prompt_id") or payload.get("component_id") or "modal-subgraph"
|
|
)
|
|
component_id = str(payload.get("component_id") or "single-node")
|
|
return prompt_id, component_id
|
|
|
|
|
|
def _remote_interrupt_flag_key(prompt_id: str, component_id: str) -> str:
|
|
"""Return the shared Modal interrupt-store key for one payload execution."""
|
|
return f"{prompt_id}:{component_id}"
|
|
|
|
|
|
def _lookup_modal_interrupt_store() -> Any | None:
|
|
"""Return the shared Modal Dict used to signal remote cancellation requests."""
|
|
if modal is None or not hasattr(modal, "Dict"):
|
|
return None
|
|
|
|
settings = get_settings()
|
|
cache_key = (settings.interrupt_dict_name, _modal_environment_name())
|
|
with _MODAL_INTERRUPT_DICTS_LOCK:
|
|
cached_store = _MODAL_INTERRUPT_DICTS.get(cache_key)
|
|
if cached_store is not None:
|
|
return cached_store
|
|
|
|
interrupt_store = modal.Dict.from_name(
|
|
settings.interrupt_dict_name,
|
|
environment_name=cache_key[1],
|
|
create_if_missing=True,
|
|
)
|
|
with _MODAL_INTERRUPT_DICTS_LOCK:
|
|
_MODAL_INTERRUPT_DICTS[cache_key] = interrupt_store
|
|
return interrupt_store
|
|
|
|
|
|
def _remote_interrupt_flag_value() -> dict[str, float]:
|
|
"""Return the shared Modal interrupt-store value for one cancellation request."""
|
|
return {"requested_at": time.time()}
|
|
|
|
|
|
def _write_remote_interrupt_flag(
|
|
interrupt_store: Any, prompt_id: str, component_id: str
|
|
) -> None:
|
|
"""Write one remote cancellation request with the blocking Modal Dict API."""
|
|
interrupt_store.put(
|
|
_remote_interrupt_flag_key(prompt_id, component_id),
|
|
_remote_interrupt_flag_value(),
|
|
)
|
|
|
|
|
|
async def _write_remote_interrupt_flag_async(
|
|
interrupt_store: Any,
|
|
prompt_id: str,
|
|
component_id: str,
|
|
) -> None:
|
|
"""Write one remote cancellation request without blocking the async caller."""
|
|
put_method = getattr(interrupt_store, "put", None)
|
|
put_async = getattr(put_method, "aio", None)
|
|
if callable(put_async):
|
|
result = put_async(
|
|
_remote_interrupt_flag_key(prompt_id, component_id),
|
|
_remote_interrupt_flag_value(),
|
|
)
|
|
if inspect.isawaitable(result):
|
|
await result
|
|
return
|
|
|
|
await asyncio.to_thread(
|
|
_write_remote_interrupt_flag, interrupt_store, prompt_id, component_id
|
|
)
|
|
|
|
|
|
def _request_remote_interrupt(payload: dict[str, Any]) -> bool:
|
|
"""Write one remote cancellation request into the shared Modal interrupt store."""
|
|
_abandon_local_modal_workflow_gate(payload, "local interrupt requested")
|
|
interrupt_store = _lookup_modal_interrupt_store()
|
|
if interrupt_store is None:
|
|
return False
|
|
|
|
prompt_id, component_id = _remote_interrupt_key(payload)
|
|
_write_remote_interrupt_flag(interrupt_store, prompt_id, component_id)
|
|
logger.info(
|
|
"Propagated local interrupt to Modal prompt=%s component=%s through shared control state.",
|
|
prompt_id,
|
|
component_id,
|
|
)
|
|
return True
|
|
|
|
|
|
async def _request_remote_interrupt_async(payload: dict[str, Any]) -> bool:
|
|
"""Write one remote cancellation request into the shared Modal interrupt store asynchronously."""
|
|
_abandon_local_modal_workflow_gate(payload, "local interrupt requested")
|
|
interrupt_store = await asyncio.to_thread(_lookup_modal_interrupt_store)
|
|
if interrupt_store is None:
|
|
return False
|
|
|
|
prompt_id, component_id = _remote_interrupt_key(payload)
|
|
await _write_remote_interrupt_flag_async(interrupt_store, prompt_id, component_id)
|
|
logger.info(
|
|
"Propagated local interrupt to Modal prompt=%s component=%s through shared control state.",
|
|
prompt_id,
|
|
component_id,
|
|
)
|
|
return True
|
|
|
|
|
|
def active_remote_modal_prompt_ids() -> set[str]:
|
|
"""Return prompt ids that currently have local proxies waiting on Modal work."""
|
|
with _ACTIVE_REMOTE_INVOCATIONS_LOCK:
|
|
return set(_ACTIVE_REMOTE_INVOCATIONS_BY_PROMPT)
|
|
|
|
|
|
@contextmanager
|
|
def _registered_active_remote_invocation(
|
|
payload: dict[str, Any],
|
|
cancellation_event: threading.Event | None,
|
|
interrupt_remote_call: Callable[[], Any] | None,
|
|
) -> Iterator[None]:
|
|
"""Register one active Modal call so targeted ComfyUI interrupts can find it."""
|
|
prompt_id, component_id = _remote_interrupt_key(payload)
|
|
invocation = _ActiveRemoteInvocation(
|
|
prompt_id=prompt_id,
|
|
component_id=component_id,
|
|
cancellation_event=cancellation_event,
|
|
interrupt_remote_call=interrupt_remote_call,
|
|
)
|
|
with _ACTIVE_REMOTE_INVOCATIONS_LOCK:
|
|
prompt_invocations = _ACTIVE_REMOTE_INVOCATIONS_BY_PROMPT.setdefault(
|
|
prompt_id, {}
|
|
)
|
|
prompt_invocations[component_id] = invocation
|
|
logger.info(
|
|
"Registered active Modal invocation prompt=%s component=%s for targeted cancellation.",
|
|
prompt_id,
|
|
component_id,
|
|
)
|
|
try:
|
|
yield
|
|
finally:
|
|
with _ACTIVE_REMOTE_INVOCATIONS_LOCK:
|
|
prompt_invocations = _ACTIVE_REMOTE_INVOCATIONS_BY_PROMPT.get(prompt_id)
|
|
if prompt_invocations is not None:
|
|
prompt_invocations.pop(component_id, None)
|
|
if not prompt_invocations:
|
|
_ACTIVE_REMOTE_INVOCATIONS_BY_PROMPT.pop(prompt_id, None)
|
|
logger.info(
|
|
"Unregistered active Modal invocation prompt=%s component=%s.",
|
|
prompt_id,
|
|
component_id,
|
|
)
|
|
|
|
|
|
def request_remote_modal_prompt_interrupt(prompt_id: str) -> bool:
|
|
"""Request cancellation for every active Modal invocation belonging to one prompt."""
|
|
normalized_prompt_id = str(prompt_id)
|
|
_abandon_local_modal_workflow_gate(
|
|
{"prompt_id": normalized_prompt_id},
|
|
"prompt-level interrupt requested",
|
|
)
|
|
with _ACTIVE_REMOTE_INVOCATIONS_LOCK:
|
|
invocations = list(
|
|
_ACTIVE_REMOTE_INVOCATIONS_BY_PROMPT.get(normalized_prompt_id, {}).values()
|
|
)
|
|
if not invocations:
|
|
return False
|
|
|
|
logger.info(
|
|
"Requesting remote Modal cancellation for prompt=%s across %d active component(s).",
|
|
normalized_prompt_id,
|
|
len(invocations),
|
|
)
|
|
for invocation in invocations:
|
|
if invocation.cancellation_event is not None:
|
|
invocation.cancellation_event.set()
|
|
_propagate_remote_interrupt_request(
|
|
{
|
|
"prompt_id": invocation.prompt_id,
|
|
"component_id": invocation.component_id,
|
|
},
|
|
invocation.interrupt_remote_call,
|
|
)
|
|
return True
|
|
|
|
|
|
async def request_remote_modal_prompt_interrupt_async(prompt_id: str) -> bool:
|
|
"""Request cancellation for every active Modal invocation belonging to one prompt asynchronously."""
|
|
normalized_prompt_id = str(prompt_id)
|
|
_abandon_local_modal_workflow_gate(
|
|
{"prompt_id": normalized_prompt_id},
|
|
"prompt-level interrupt requested",
|
|
)
|
|
with _ACTIVE_REMOTE_INVOCATIONS_LOCK:
|
|
invocations = list(
|
|
_ACTIVE_REMOTE_INVOCATIONS_BY_PROMPT.get(normalized_prompt_id, {}).values()
|
|
)
|
|
if not invocations:
|
|
return False
|
|
|
|
logger.info(
|
|
"Requesting async remote Modal cancellation for prompt=%s across %d active component(s).",
|
|
normalized_prompt_id,
|
|
len(invocations),
|
|
)
|
|
for invocation in invocations:
|
|
if invocation.cancellation_event is not None:
|
|
invocation.cancellation_event.set()
|
|
await _request_remote_interrupt_async(
|
|
{"prompt_id": invocation.prompt_id, "component_id": invocation.component_id}
|
|
)
|
|
return True
|
|
|
|
|
|
def _sync_local_interrupt_to_cancellation_event(
|
|
payload: dict[str, Any],
|
|
cancellation_event: threading.Event | None,
|
|
) -> bool:
|
|
"""Mirror ComfyUI's interrupt flag into the current Modal cancellation event."""
|
|
if cancellation_event is not None and cancellation_event.is_set():
|
|
return True
|
|
if not _local_processing_interrupted():
|
|
return False
|
|
if cancellation_event is not None and not cancellation_event.is_set():
|
|
logger.info(
|
|
"Observed local interrupt while Modal component=%s was running; requesting remote cancellation.",
|
|
payload.get("component_id"),
|
|
)
|
|
cancellation_event.set()
|
|
_abandon_local_modal_workflow_gate(payload, "observed local interrupt")
|
|
return True
|
|
|
|
|
|
def _abandon_local_modal_workflow_gate(payload: dict[str, Any], reason: str) -> None:
|
|
"""Release the local prompt gate for a Modal prompt that ComfyUI has cancelled."""
|
|
prompt_id = payload.get("prompt_id")
|
|
if prompt_id is None:
|
|
return
|
|
|
|
try:
|
|
from ..modal_executor_node import abandon_modal_workflow_execution_prompt
|
|
except ImportError:
|
|
logger.debug(
|
|
"Unable to import Modal workflow gate helper while abandoning prompt."
|
|
)
|
|
return
|
|
|
|
abandon_modal_workflow_execution_prompt(str(prompt_id), reason)
|
|
|
|
|
|
def _propagate_remote_interrupt_request(
|
|
payload: dict[str, Any],
|
|
interrupt_remote_call: Callable[[], Any] | None,
|
|
) -> None:
|
|
"""Send one best-effort remote cancellation request for an active Modal payload."""
|
|
prompt_id, component_id = _remote_interrupt_key(payload)
|
|
if interrupt_remote_call is None:
|
|
logger.warning(
|
|
"Local interrupt requested for component=%s, but no remote interrupt method is available.",
|
|
component_id,
|
|
)
|
|
return
|
|
try:
|
|
interrupt_remote_call()
|
|
logger.info(
|
|
"Propagated local interrupt to Modal prompt=%s component=%s.",
|
|
prompt_id,
|
|
component_id,
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to propagate local interrupt to Modal prompt=%s component=%s.",
|
|
prompt_id,
|
|
component_id,
|
|
)
|
|
|
|
|
|
def _handle_modal_wait_cancellation(
|
|
payload: dict[str, Any],
|
|
cancellation_event: threading.Event,
|
|
*,
|
|
interrupt_sent: bool,
|
|
cancellation_started_at: float | None,
|
|
) -> tuple[bool, float | None]:
|
|
"""Propagate and bound local waiting after cancellation during any Modal wait phase."""
|
|
if not _sync_local_interrupt_to_cancellation_event(payload, cancellation_event):
|
|
return interrupt_sent, cancellation_started_at
|
|
|
|
if not interrupt_sent:
|
|
_request_remote_interrupt(payload)
|
|
return True, time.monotonic()
|
|
|
|
if cancellation_started_at is None:
|
|
return interrupt_sent, time.monotonic()
|
|
|
|
grace_seconds = max(0.0, get_settings().remote_cancel_grace_seconds)
|
|
if time.monotonic() - cancellation_started_at >= grace_seconds:
|
|
logger.info(
|
|
"Modal component=%s did not reach a cancellable remote call within %.3fs of local interrupt; releasing the local prompt while remote cancellation continues.",
|
|
payload.get("component_id"),
|
|
grace_seconds,
|
|
)
|
|
raise ModalRemoteInvocationError(
|
|
"Remote Modal call did not reach a cancellable remote phase after local interrupt propagation."
|
|
)
|
|
|
|
return interrupt_sent, cancellation_started_at
|
|
|
|
|
|
async def _handle_modal_wait_cancellation_async(
|
|
payload: dict[str, Any],
|
|
cancellation_event: threading.Event,
|
|
*,
|
|
interrupt_sent: bool,
|
|
cancellation_started_at: float | None,
|
|
) -> tuple[bool, float | None]:
|
|
"""Propagate and bound local waiting after cancellation during an async Modal wait phase."""
|
|
if not _sync_local_interrupt_to_cancellation_event(payload, cancellation_event):
|
|
return interrupt_sent, cancellation_started_at
|
|
|
|
if not interrupt_sent:
|
|
await _request_remote_interrupt_async(payload)
|
|
return True, time.monotonic()
|
|
|
|
if cancellation_started_at is None:
|
|
return interrupt_sent, time.monotonic()
|
|
|
|
grace_seconds = max(0.0, get_settings().remote_cancel_grace_seconds)
|
|
if time.monotonic() - cancellation_started_at >= grace_seconds:
|
|
logger.info(
|
|
"Modal component=%s did not reach a cancellable remote call within %.3fs of local interrupt; releasing the local prompt while remote cancellation continues.",
|
|
payload.get("component_id"),
|
|
grace_seconds,
|
|
)
|
|
raise ModalRemoteInvocationError(
|
|
"Remote Modal call did not reach a cancellable remote phase after local interrupt propagation."
|
|
)
|
|
|
|
return interrupt_sent, cancellation_started_at
|
|
|
|
|
|
def _invoke_remote_call_with_interrupts(
|
|
*,
|
|
payload: dict[str, Any],
|
|
invoke_remote_call: Callable[[], bytes],
|
|
interrupt_remote_call: Callable[[], Any] | None,
|
|
cancellation_event: threading.Event | None,
|
|
) -> bytes:
|
|
"""Run one blocking remote call while optionally propagating cancellation to Modal."""
|
|
result_queue: queue.Queue[tuple[str, Any]] = queue.Queue()
|
|
cancellation_started_at: float | None = None
|
|
|
|
def execute_remote_call() -> None:
|
|
"""Run the blocking Modal request in a worker thread."""
|
|
try:
|
|
result_queue.put(("result", invoke_remote_call()))
|
|
except BaseException as exc:
|
|
result_queue.put(("error", exc))
|
|
|
|
request_thread = threading.Thread(
|
|
target=execute_remote_call,
|
|
name=f"modal-request-{payload.get('component_id', 'payload')}",
|
|
daemon=True,
|
|
)
|
|
request_thread.start()
|
|
interrupt_sent = False
|
|
try:
|
|
with _registered_active_remote_invocation(
|
|
payload, cancellation_event, interrupt_remote_call
|
|
):
|
|
while True:
|
|
try:
|
|
result_kind, result_payload = result_queue.get(timeout=0.1)
|
|
except queue.Empty:
|
|
if _sync_local_interrupt_to_cancellation_event(
|
|
payload, cancellation_event
|
|
):
|
|
if not interrupt_sent:
|
|
_propagate_remote_interrupt_request(
|
|
payload, interrupt_remote_call
|
|
)
|
|
interrupt_sent = True
|
|
cancellation_started_at = time.monotonic()
|
|
elif cancellation_started_at is not None:
|
|
grace_seconds = max(
|
|
0.0, get_settings().remote_cancel_grace_seconds
|
|
)
|
|
if (
|
|
time.monotonic() - cancellation_started_at
|
|
>= grace_seconds
|
|
):
|
|
logger.info(
|
|
"Modal component=%s did not return within %.3fs of local interrupt propagation; releasing the local prompt while remote cancellation continues.",
|
|
payload.get("component_id"),
|
|
grace_seconds,
|
|
)
|
|
raise ModalRemoteInvocationError(
|
|
"Remote Modal call did not finish after local interrupt propagation."
|
|
)
|
|
continue
|
|
|
|
if result_kind == "result":
|
|
return bytes(result_payload)
|
|
raise result_payload
|
|
finally:
|
|
request_thread.join(
|
|
timeout=0.1
|
|
if cancellation_event is not None and cancellation_event.is_set()
|
|
else 1.0
|
|
)
|
|
|
|
|
|
def _consume_remote_payload_stream(
|
|
payload: dict[str, Any],
|
|
stream_events: Iterator[dict[str, Any]],
|
|
) -> bytes:
|
|
"""Forward remote progress events into the local UI and return the final payload bytes."""
|
|
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
|
|
)
|
|
node_ids = [str(node_id) for node_id in payload.get("component_node_ids", [])]
|
|
modal_gpu = (
|
|
str(payload["modal_gpu"]) if payload.get("modal_gpu") is not None else None
|
|
)
|
|
suppress_status_stream = bool(payload.get("suppress_status_stream"))
|
|
result_payload: bytes | bytearray | None = None
|
|
suppressed_progress_node_metadata: dict[str, dict[str, str | None]] = {}
|
|
active_remote_task_id: str | None = None
|
|
active_remote_log_task_id: str | None = None
|
|
should_close_stream = False
|
|
component_id = str(payload.get("component_id") or "payload")
|
|
invocation_id = str(payload.get("invocation_id") or "none")
|
|
stream_started_at = time.monotonic()
|
|
previous_event_at = stream_started_at
|
|
event_count = 0
|
|
progress_event_count = 0
|
|
speculative_prewarm_checked = False
|
|
logger.info(
|
|
"Starting local Modal stream consumption component=%s prompt_id=%s invocation_id=%s.",
|
|
component_id,
|
|
prompt_id or "none",
|
|
invocation_id,
|
|
)
|
|
|
|
try:
|
|
for stream_event in stream_events:
|
|
event_received_at = time.monotonic()
|
|
seconds_since_previous_event = event_received_at - previous_event_at
|
|
previous_event_at = event_received_at
|
|
event_count += 1
|
|
if not speculative_prewarm_checked:
|
|
speculative_prewarm_checked = True
|
|
_schedule_speculative_affinity_prewarm(
|
|
payload,
|
|
reason="current_remote_stream_started",
|
|
)
|
|
event_kind = str(stream_event.get("kind", ""))
|
|
if event_kind == "remote_logs":
|
|
task_id = _coerce_modal_task_id(stream_event.get("task_id"))
|
|
if task_id is not None and active_remote_task_id is None:
|
|
active_remote_task_id = task_id
|
|
if (
|
|
task_id is not None
|
|
and active_remote_log_task_id is None
|
|
and _is_remote_container_log_stream_enabled()
|
|
):
|
|
active_remote_log_task_id = _retain_remote_container_log_stream(
|
|
task_id
|
|
)
|
|
continue
|
|
if event_kind == "progress":
|
|
progress_event_count += 1
|
|
event_type = str(stream_event.get("event_type", ""))
|
|
if event_type == "node_progress":
|
|
progress_metadata = _progress_stream_event_metadata(stream_event)
|
|
filter_node_id = (
|
|
progress_metadata["filter_node_id"]
|
|
if progress_metadata is not None
|
|
else None
|
|
)
|
|
lane_id = (
|
|
str(stream_event["lane_id"])
|
|
if stream_event.get("lane_id") is not None
|
|
else (
|
|
str(payload["mapped_progress_lane_id"])
|
|
if payload.get("mapped_progress_lane_id") is not None
|
|
else None
|
|
)
|
|
)
|
|
aggregate_only = bool(stream_event.get("aggregate_only", False))
|
|
if (
|
|
suppress_status_stream
|
|
and lane_id is None
|
|
and not aggregate_only
|
|
and not _should_forward_suppressed_stream_event(
|
|
payload, filter_node_id
|
|
)
|
|
):
|
|
logger.debug(
|
|
"Suppressing streamed Modal node progress for component=%s node_id=%s real_node_id=%s because it does not belong to this mapped/static payload.",
|
|
payload.get("component_id"),
|
|
stream_event.get("node_id"),
|
|
stream_event.get("real_node_id"),
|
|
)
|
|
continue
|
|
reported_node_id = (
|
|
progress_metadata["node_id"]
|
|
if progress_metadata is not None
|
|
else None
|
|
)
|
|
if reported_node_id is not None:
|
|
display_node_id = (
|
|
progress_metadata["display_node_id"]
|
|
if progress_metadata is not None
|
|
else str(reported_node_id)
|
|
)
|
|
real_node_id = (
|
|
progress_metadata["real_node_id"]
|
|
if progress_metadata is not None
|
|
else None
|
|
)
|
|
progress_node_id = real_node_id or display_node_id
|
|
completed_ancestor_node_ids = _remote_prompt_ancestor_node_ids(
|
|
payload,
|
|
progress_node_id,
|
|
)
|
|
if lane_id is not None:
|
|
_remember_mapped_lane_node_id(
|
|
payload, lane_id, progress_node_id
|
|
)
|
|
elif (
|
|
suppress_status_stream
|
|
and not aggregate_only
|
|
and progress_metadata is not None
|
|
):
|
|
suppressed_progress_node_metadata[
|
|
str(progress_metadata["filter_node_id"])
|
|
] = {
|
|
"node_id": str(reported_node_id),
|
|
"display_node_id": display_node_id,
|
|
"real_node_id": real_node_id,
|
|
}
|
|
logger.debug(
|
|
"Forwarding streamed Modal node progress for component=%s node_id=%s real_node_id=%s value=%s max=%s lane_id=%s.",
|
|
payload.get("component_id"),
|
|
reported_node_id,
|
|
real_node_id,
|
|
stream_event.get("value"),
|
|
stream_event.get("max"),
|
|
lane_id,
|
|
)
|
|
progress_kwargs = {
|
|
"prompt_id": prompt_id,
|
|
"client_id": client_id,
|
|
"node_id": str(reported_node_id),
|
|
"value": float(stream_event.get("value", 0.0)),
|
|
"max_value": float(stream_event.get("max", 1.0)),
|
|
"display_node_id": display_node_id,
|
|
"real_node_id": real_node_id,
|
|
"lane_id": lane_id,
|
|
"clear": bool(stream_event.get("clear", False)),
|
|
"item_index": (
|
|
int(stream_event["item_index"])
|
|
if stream_event.get("item_index") is not None
|
|
else (
|
|
int(payload["map_item_index"])
|
|
if payload.get("map_item_index") is not None
|
|
else None
|
|
)
|
|
),
|
|
"aggregate_only": aggregate_only,
|
|
}
|
|
progress_kwargs.update(
|
|
_remote_execution_identity(payload, active_remote_task_id)
|
|
)
|
|
for string_field in ("stage", "message", "unit"):
|
|
if stream_event.get(string_field) is not None:
|
|
progress_kwargs[string_field] = str(
|
|
stream_event[string_field]
|
|
)
|
|
if stream_event.get("indeterminate") is not None:
|
|
progress_kwargs["indeterminate"] = bool(
|
|
stream_event["indeterminate"]
|
|
)
|
|
for numeric_field in (
|
|
"elapsed_seconds",
|
|
"time_to_first_token_seconds",
|
|
"tokens_per_second",
|
|
):
|
|
if stream_event.get(numeric_field) is not None:
|
|
progress_kwargs[numeric_field] = float(
|
|
stream_event[numeric_field]
|
|
)
|
|
if completed_ancestor_node_ids:
|
|
progress_kwargs[
|
|
"completed_ancestor_node_ids"
|
|
] = completed_ancestor_node_ids
|
|
_emit_local_modal_progress(**progress_kwargs)
|
|
continue
|
|
if event_type == "executed":
|
|
reported_node_id = stream_event.get("node_id")
|
|
if reported_node_id is not None:
|
|
if not _should_forward_suppressed_stream_event(
|
|
payload, reported_node_id
|
|
):
|
|
logger.debug(
|
|
"Suppressing streamed Modal executed output for component=%s node_id=%s because it does not belong to this mapped/static payload.",
|
|
payload.get("component_id"),
|
|
reported_node_id,
|
|
)
|
|
continue
|
|
logger.debug(
|
|
"Forwarding streamed Modal executed output for component=%s node_id=%s.",
|
|
payload.get("component_id"),
|
|
reported_node_id,
|
|
)
|
|
_emit_local_executed_output(
|
|
prompt_id=prompt_id,
|
|
client_id=client_id,
|
|
node_id=str(reported_node_id),
|
|
display_node_id=(
|
|
str(stream_event["display_node_id"])
|
|
if stream_event.get("display_node_id") is not None
|
|
else None
|
|
),
|
|
output_payload=deserialize_value(
|
|
stream_event.get("output")
|
|
),
|
|
)
|
|
continue
|
|
if event_type == "preview":
|
|
reported_node_id = stream_event.get("node_id")
|
|
image_bytes = deserialize_value(stream_event.get("image_bytes"))
|
|
if reported_node_id is not None and isinstance(image_bytes, bytes):
|
|
if not _should_forward_suppressed_stream_event(
|
|
payload, reported_node_id
|
|
):
|
|
logger.debug(
|
|
"Suppressing streamed Modal preview image for component=%s node_id=%s because it does not belong to this mapped/static payload.",
|
|
payload.get("component_id"),
|
|
reported_node_id,
|
|
)
|
|
continue
|
|
logger.debug(
|
|
"Forwarding streamed Modal preview image for component=%s node_id=%s.",
|
|
payload.get("component_id"),
|
|
reported_node_id,
|
|
)
|
|
_emit_local_preview_image(
|
|
prompt_id=prompt_id,
|
|
client_id=client_id,
|
|
node_id=str(reported_node_id),
|
|
display_node_id=(
|
|
str(stream_event["display_node_id"])
|
|
if stream_event.get("display_node_id") is not None
|
|
else None
|
|
),
|
|
parent_node_id=(
|
|
str(stream_event["parent_node_id"])
|
|
if stream_event.get("parent_node_id") is not None
|
|
else None
|
|
),
|
|
real_node_id=(
|
|
str(stream_event["real_node_id"])
|
|
if stream_event.get("real_node_id") is not None
|
|
else None
|
|
),
|
|
image_type=str(stream_event.get("image_type", "PNG")),
|
|
image_bytes=image_bytes,
|
|
max_size=(
|
|
int(stream_event["max_size"])
|
|
if stream_event.get("max_size") is not None
|
|
else None
|
|
),
|
|
)
|
|
continue
|
|
if event_type == "boundary_output":
|
|
preview_target_node_ids = [
|
|
str(node_id)
|
|
for node_id in stream_event.get("preview_target_node_ids", [])
|
|
if str(node_id)
|
|
]
|
|
if preview_target_node_ids:
|
|
logger.debug(
|
|
"Forwarding streamed Modal boundary output previews for component=%s source_node=%s targets=%s.",
|
|
payload.get("component_id"),
|
|
stream_event.get("node_id"),
|
|
preview_target_node_ids,
|
|
)
|
|
_emit_local_preview_boundary_output(
|
|
prompt_id=prompt_id,
|
|
client_id=client_id,
|
|
preview_target_node_ids=preview_target_node_ids,
|
|
image_value=deserialize_value(stream_event.get("value")),
|
|
)
|
|
continue
|
|
if event_type == "node_cached":
|
|
progress_metadata = _progress_stream_event_metadata(stream_event)
|
|
filter_node_id = (
|
|
progress_metadata["filter_node_id"]
|
|
if progress_metadata is not None
|
|
else None
|
|
)
|
|
if (
|
|
suppress_status_stream
|
|
and not _should_forward_suppressed_stream_event(
|
|
payload, filter_node_id
|
|
)
|
|
):
|
|
logger.debug(
|
|
"Suppressing streamed Modal cached-node marker for component=%s node_id=%s real_node_id=%s because it does not belong to this mapped/static payload.",
|
|
payload.get("component_id"),
|
|
stream_event.get("node_id"),
|
|
stream_event.get("real_node_id"),
|
|
)
|
|
continue
|
|
reported_node_id = (
|
|
progress_metadata["node_id"]
|
|
if progress_metadata is not None
|
|
else None
|
|
)
|
|
if reported_node_id is not None:
|
|
display_node_id = (
|
|
progress_metadata["display_node_id"]
|
|
if progress_metadata is not None
|
|
else str(reported_node_id)
|
|
)
|
|
real_node_id = (
|
|
progress_metadata["real_node_id"]
|
|
if progress_metadata is not None
|
|
else None
|
|
)
|
|
logger.debug(
|
|
"Forwarding streamed Modal cached-node marker for component=%s node_id=%s real_node_id=%s.",
|
|
payload.get("component_id"),
|
|
reported_node_id,
|
|
real_node_id,
|
|
)
|
|
_emit_local_modal_progress(
|
|
prompt_id=prompt_id,
|
|
client_id=client_id,
|
|
node_id=str(reported_node_id),
|
|
value=0.0,
|
|
max_value=1.0,
|
|
display_node_id=display_node_id,
|
|
real_node_id=real_node_id,
|
|
cached_hit=True,
|
|
**_remote_execution_identity(
|
|
payload, active_remote_task_id
|
|
),
|
|
)
|
|
continue
|
|
logger.info(
|
|
"Forwarding streamed Modal progress for component=%s phase=%s active_node_id=%s.",
|
|
payload.get("component_id"),
|
|
stream_event.get("phase"),
|
|
stream_event.get("active_node_id"),
|
|
)
|
|
if suppress_status_stream:
|
|
remote_phase = str(stream_event.get("phase", "executing"))
|
|
if remote_phase in {
|
|
"execution_success",
|
|
"execution_error",
|
|
"execution_interrupted",
|
|
}:
|
|
for (
|
|
progress_metadata
|
|
) in suppressed_progress_node_metadata.values():
|
|
_emit_local_modal_progress(
|
|
prompt_id=prompt_id,
|
|
client_id=client_id,
|
|
node_id=str(progress_metadata["node_id"]),
|
|
value=0.0,
|
|
max_value=1.0,
|
|
display_node_id=progress_metadata["display_node_id"],
|
|
real_node_id=progress_metadata["real_node_id"],
|
|
clear=True,
|
|
)
|
|
suppressed_progress_node_metadata.clear()
|
|
continue
|
|
remote_phase = str(stream_event.get("phase", "executing"))
|
|
if remote_phase == "execution_success":
|
|
_emit_local_modal_status(
|
|
prompt_id=prompt_id,
|
|
client_id=client_id,
|
|
phase="finalizing",
|
|
node_ids=node_ids,
|
|
modal_gpu=modal_gpu,
|
|
status_message=(
|
|
"Receiving remote outputs from "
|
|
f"{_remote_execution_destination(payload)}"
|
|
),
|
|
**_remote_execution_identity(payload, active_remote_task_id),
|
|
)
|
|
continue
|
|
_emit_local_modal_status(
|
|
prompt_id=prompt_id,
|
|
client_id=client_id,
|
|
phase=remote_phase,
|
|
node_ids=node_ids,
|
|
modal_gpu=modal_gpu,
|
|
active_node_id=(
|
|
str(stream_event["active_node_id"])
|
|
if stream_event.get("active_node_id") is not None
|
|
else None
|
|
),
|
|
completed_ancestor_node_ids=_remote_prompt_ancestor_node_ids(
|
|
payload,
|
|
str(stream_event["active_node_id"])
|
|
if stream_event.get("active_node_id") is not None
|
|
else None,
|
|
)
|
|
or None,
|
|
active_node_class_type=(
|
|
str(stream_event["active_node_class_type"])
|
|
if stream_event.get("active_node_class_type") is not None
|
|
else None
|
|
),
|
|
active_node_role=(
|
|
str(stream_event["active_node_role"])
|
|
if stream_event.get("active_node_role") is not None
|
|
else None
|
|
),
|
|
**_remote_execution_identity(payload, active_remote_task_id),
|
|
)
|
|
continue
|
|
if event_kind == "result":
|
|
candidate_outputs = stream_event.get("outputs")
|
|
candidate_bytes = (
|
|
len(candidate_outputs)
|
|
if isinstance(candidate_outputs, bytes | bytearray)
|
|
else -1
|
|
)
|
|
logger.info(
|
|
"Received streamed Modal result component=%s prompt_id=%s invocation_id=%s "
|
|
"result_bytes=%d stream_elapsed_seconds=%.3f seconds_since_previous_event=%.3f "
|
|
"event_count=%d progress_event_count=%d.",
|
|
component_id,
|
|
prompt_id or "none",
|
|
invocation_id,
|
|
candidate_bytes,
|
|
event_received_at - stream_started_at,
|
|
seconds_since_previous_event,
|
|
event_count,
|
|
progress_event_count,
|
|
)
|
|
try:
|
|
result_payload = coerce_serialized_node_outputs(candidate_outputs)
|
|
except TypeError as exc:
|
|
raise ModalRemoteInvocationError(
|
|
"Modal streamed payload result did not include transport-safe outputs."
|
|
) from exc
|
|
should_close_stream = True
|
|
break
|
|
logger.debug(
|
|
"Ignoring unexpected streamed Modal event kind=%s for component=%s.",
|
|
event_kind,
|
|
payload.get("component_id"),
|
|
)
|
|
finally:
|
|
if should_close_stream:
|
|
close_started_at = time.monotonic()
|
|
logger.info(
|
|
"Starting local Modal result stream close component=%s prompt_id=%s invocation_id=%s.",
|
|
component_id,
|
|
prompt_id or "none",
|
|
invocation_id,
|
|
)
|
|
_close_remote_payload_stream(stream_events)
|
|
logger.info(
|
|
"Finished local Modal result stream close in %.3fs component=%s prompt_id=%s "
|
|
"invocation_id=%s.",
|
|
time.monotonic() - close_started_at,
|
|
component_id,
|
|
prompt_id or "none",
|
|
invocation_id,
|
|
)
|
|
if active_remote_log_task_id is not None:
|
|
_release_remote_container_log_stream(active_remote_log_task_id)
|
|
|
|
if result_payload is None:
|
|
raise ModalRemoteInvocationError(
|
|
f"Modal streamed payload for component={payload.get('component_id')!r} did not yield a final result."
|
|
)
|
|
logger.info(
|
|
"Finished local Modal stream consumption in %.3fs component=%s prompt_id=%s invocation_id=%s "
|
|
"result_bytes=%d event_count=%d progress_event_count=%d.",
|
|
time.monotonic() - stream_started_at,
|
|
component_id,
|
|
prompt_id or "none",
|
|
invocation_id,
|
|
len(result_payload),
|
|
event_count,
|
|
progress_event_count,
|
|
)
|
|
return bytes(result_payload)
|
|
|
|
|
|
async def _invoke_bound_remote_engine_async(
|
|
remote_engine: Any,
|
|
payload: dict[str, Any],
|
|
kwargs_payload: bytes,
|
|
) -> bytes:
|
|
"""Invoke one pre-bound remote engine handle asynchronously with local interrupt mirroring."""
|
|
logger.info(
|
|
"Dispatching async Modal remote invocation via bound engine for component=%s payload_kind=%s.",
|
|
payload.get("component_id"),
|
|
payload.get("payload_kind"),
|
|
)
|
|
cancellation_event = threading.Event()
|
|
future = _remote_modal_call_executor().submit(
|
|
_invoke_remote_engine_payload_with_recovery,
|
|
remote_engine,
|
|
dict(payload),
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
wrapped_future = asyncio.wrap_future(future)
|
|
interrupt_sent = False
|
|
cancellation_started_at: float | None = None
|
|
try:
|
|
while True:
|
|
try:
|
|
response = await asyncio.wait_for(
|
|
asyncio.shield(wrapped_future), timeout=0.1
|
|
)
|
|
break
|
|
except asyncio.TimeoutError:
|
|
(
|
|
interrupt_sent,
|
|
cancellation_started_at,
|
|
) = await _handle_modal_wait_cancellation_async(
|
|
payload,
|
|
cancellation_event,
|
|
interrupt_sent=interrupt_sent,
|
|
cancellation_started_at=cancellation_started_at,
|
|
)
|
|
continue
|
|
except asyncio.CancelledError:
|
|
cancellation_event.set()
|
|
prompt_id = (
|
|
str(payload.get("prompt_id"))
|
|
if payload.get("prompt_id") is not None
|
|
else None
|
|
)
|
|
if prompt_id is not None:
|
|
await request_remote_modal_prompt_interrupt_async(prompt_id)
|
|
raise
|
|
except Exception:
|
|
if cancellation_event.is_set() or _local_processing_interrupted():
|
|
logger.info(
|
|
"Reraising async Modal failure as a local interrupt for component=%s after cancellation.",
|
|
payload.get("component_id"),
|
|
)
|
|
_raise_local_interrupt()
|
|
logger.exception(
|
|
"Async Modal remote invocation via bound engine failed for component=%s.",
|
|
payload.get("component_id"),
|
|
)
|
|
raise
|
|
if cancellation_event.is_set() or _local_processing_interrupted():
|
|
logger.info(
|
|
"Async bound-engine invocation for component=%s finished after interruption; raising local interrupt.",
|
|
payload.get("component_id"),
|
|
)
|
|
_raise_local_interrupt()
|
|
logger.info(
|
|
"Async Modal remote invocation via bound engine completed for component=%s.",
|
|
payload.get("component_id"),
|
|
)
|
|
return response
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _build_remote_interrupt_callback(
|
|
_remote_engine: Any, payload: dict[str, Any]
|
|
) -> Callable[[], Any] | None:
|
|
"""Return a callable that requests interruption for one active Modal payload."""
|
|
interrupt_store = _lookup_modal_interrupt_store()
|
|
if interrupt_store is None:
|
|
return None
|
|
return lambda: _request_remote_interrupt(payload)
|
|
|
|
|
|
def _invoke_remote_engine_payload(
|
|
remote_engine: Any,
|
|
payload: dict[str, Any],
|
|
kwargs_payload: bytes,
|
|
cancellation_event: threading.Event | None,
|
|
) -> bytes:
|
|
"""Invoke one prepared remote engine instance with optional progress streaming."""
|
|
_install_modal_cloud_exception_compatibility_module()
|
|
if cancellation_event is not None and cancellation_event.is_set():
|
|
_request_remote_interrupt(payload)
|
|
raise ModalRemoteInvocationError(
|
|
"Remote Modal payload dispatch was cancelled before the remote call started."
|
|
)
|
|
|
|
stream_method = getattr(remote_engine, "execute_payload_stream", None)
|
|
interrupt_remote_call = _build_remote_interrupt_callback(remote_engine, payload)
|
|
if _should_stream_remote_progress(payload) and hasattr(stream_method, "remote_gen"):
|
|
logger.info(
|
|
"Using streamed Modal progress path for component=%s via execute_payload_stream.remote_gen(...).",
|
|
payload.get("component_id"),
|
|
)
|
|
return _invoke_remote_call_with_interrupts(
|
|
payload=payload,
|
|
invoke_remote_call=lambda: _consume_remote_payload_stream(
|
|
payload,
|
|
stream_method.remote_gen(payload, kwargs_payload),
|
|
),
|
|
interrupt_remote_call=interrupt_remote_call,
|
|
cancellation_event=cancellation_event,
|
|
)
|
|
|
|
if _should_stream_remote_progress(payload):
|
|
logger.warning(
|
|
"Streamed Modal progress is unavailable for component=%s; falling back to execute_payload.remote(...).",
|
|
payload.get("component_id"),
|
|
)
|
|
return _invoke_remote_call_with_interrupts(
|
|
payload=payload,
|
|
invoke_remote_call=lambda: remote_engine.execute_payload.remote(
|
|
payload, kwargs_payload
|
|
),
|
|
interrupt_remote_call=interrupt_remote_call,
|
|
cancellation_event=cancellation_event,
|
|
)
|
|
|
|
|
|
def _fresh_llm_memory_recovery_affinity_key(payload: Mapping[str, Any]) -> str:
|
|
"""Return a one-use affinity identity that cannot select the dirty worker."""
|
|
return (
|
|
f"{_remote_worker_affinity_key(dict(payload))}:llm-memory-recovery:"
|
|
f"{uuid.uuid4().hex[:12]}"
|
|
)
|
|
|
|
|
|
def _retry_exhausted_llm_memory_on_fresh_worker(
|
|
*,
|
|
payload: dict[str, Any],
|
|
kwargs_payload: bytes,
|
|
cancellation_event: threading.Event | None,
|
|
error: BaseException,
|
|
) -> bytes:
|
|
"""Retry one exhausted post-eviction admission on a distinct Modal worker."""
|
|
recovery_payload = dict(payload)
|
|
if exhausted_recovery_used_vllm_throughput(error):
|
|
recovery_payload[LLM_FORCE_VLLM_THROUGHPUT_PAYLOAD_KEY] = True
|
|
recovery_affinity_key = _fresh_llm_memory_recovery_affinity_key(payload)
|
|
logger.warning(
|
|
"Modal LLM memory recovery timed out for component=%s; retrying "
|
|
"invocation_id=%s once on fresh worker_affinity=%s.",
|
|
payload.get("component_id"),
|
|
payload["invocation_id"],
|
|
recovery_affinity_key,
|
|
)
|
|
_emit_local_remote_startup_status(
|
|
recovery_payload,
|
|
phase="starting",
|
|
status_message="Retrying LLM on a fresh Modal worker",
|
|
)
|
|
recovered_remote_engine = _lookup_deployed_remote_engine(
|
|
recovery_payload,
|
|
affinity_key_override=recovery_affinity_key,
|
|
)
|
|
return _invoke_remote_engine_payload(
|
|
recovered_remote_engine,
|
|
recovery_payload,
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
|
|
|
|
def _invoke_remote_engine_payload_with_recovery(
|
|
remote_engine: Any,
|
|
payload: dict[str, Any],
|
|
kwargs_payload: bytes,
|
|
cancellation_event: threading.Event | None,
|
|
) -> bytes:
|
|
"""Recover missing deployments and exhausted LLM workers with one safe retry."""
|
|
payload = dict(payload)
|
|
settings = _settings_for_payload(payload)
|
|
if llm_model_references_from_payload(payload):
|
|
kwargs_payload = _rewrite_staged_llm_kwargs_payload(
|
|
kwargs_payload,
|
|
modal_deployment_app_name(settings),
|
|
)
|
|
payload.setdefault("capture_remote_outputs", True)
|
|
payload.setdefault(
|
|
"invocation_id",
|
|
stable_remote_invocation_id(payload, kwargs_payload),
|
|
)
|
|
lookup_error_types = _modal_lookup_error_types()
|
|
invocation_error_types = tuple(
|
|
dict.fromkeys((*lookup_error_types, RemoteSubgraphExecutionError, RuntimeError))
|
|
)
|
|
try:
|
|
response = _invoke_remote_engine_payload(
|
|
remote_engine,
|
|
payload,
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
except invocation_error_types as exc:
|
|
if is_llm_memory_recovery_exhausted(exc):
|
|
response = _retry_exhausted_llm_memory_on_fresh_worker(
|
|
payload=payload,
|
|
kwargs_payload=kwargs_payload,
|
|
cancellation_event=cancellation_event,
|
|
error=exc,
|
|
)
|
|
elif not isinstance(exc, lookup_error_types):
|
|
raise
|
|
elif not settings.auto_deploy or not _is_missing_modal_deployment_error(exc):
|
|
raise
|
|
else:
|
|
logger.warning(
|
|
"Modal payload invocation failed for component=%s because the deployed app was missing at call time: %s. Recreating the app and retrying.",
|
|
payload.get("component_id"),
|
|
exc,
|
|
)
|
|
recovered_remote_engine = _auto_deploy_modal_app(payload, exc)
|
|
_ensure_llm_profiles_staged(
|
|
payload,
|
|
modal_deployment_app_name(settings),
|
|
)
|
|
response = _invoke_remote_engine_payload(
|
|
recovered_remote_engine,
|
|
payload,
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
return _materialize_remote_execution_result(response, settings=settings)
|
|
|
|
|
|
def _local_comfy_output_directory(settings: ModalSyncSettings) -> Path:
|
|
"""Return the effective output directory of the local ComfyUI process."""
|
|
try:
|
|
import folder_paths
|
|
except ModuleNotFoundError as exc:
|
|
if exc.name != "folder_paths":
|
|
raise
|
|
else:
|
|
get_output_directory = getattr(folder_paths, "get_output_directory", None)
|
|
if callable(get_output_directory):
|
|
return Path(get_output_directory()).expanduser().resolve()
|
|
if settings.comfyui_root is None:
|
|
raise ModalRemoteInvocationError(
|
|
"Modal returned ComfyUI output files, but the local ComfyUI output directory "
|
|
"could not be resolved. Set COMFY_MODAL_COMFYUI_ROOT."
|
|
)
|
|
return (settings.comfyui_root / "output").resolve()
|
|
|
|
|
|
def _materialize_remote_execution_result(
|
|
response: bytes | bytearray,
|
|
*,
|
|
settings: ModalSyncSettings | None = None,
|
|
) -> bytes:
|
|
"""Download bundled remote files and return the ordinary serialized node outputs."""
|
|
result: RemoteExecutionResult = unpack_remote_execution_result(response)
|
|
if not result.artifacts:
|
|
return result.outputs
|
|
resolved_settings = settings or get_settings()
|
|
output_directory = _local_comfy_output_directory(resolved_settings)
|
|
logger.info(
|
|
"Materializing %d remote ComfyUI output file(s) into %s.",
|
|
len(result.artifacts),
|
|
output_directory,
|
|
)
|
|
materialized_paths = materialize_remote_output_artifacts(
|
|
result,
|
|
output_directory=output_directory,
|
|
app_name=modal_deployment_app_name(resolved_settings),
|
|
)
|
|
logger.info(
|
|
"Finished downloading %d remote ComfyUI output file(s): %s.",
|
|
len(materialized_paths),
|
|
[str(path) for path in materialized_paths],
|
|
)
|
|
return result.outputs
|
|
|
|
|
|
def _invoke_modal_payload_blocking(
|
|
payload: dict[str, Any],
|
|
kwargs_payload: bytes,
|
|
cancellation_event: threading.Event | None = None,
|
|
) -> bytes:
|
|
"""Invoke the Modal runtime from a worker thread using deployed or ephemeral app state."""
|
|
if modal is None:
|
|
raise ModalRemoteInvocationError("Modal SDK is unavailable.")
|
|
|
|
_await_matching_speculative_prewarm(payload, cancellation_event)
|
|
lookup_error_types = _modal_lookup_error_types()
|
|
settings = _settings_for_payload(payload)
|
|
deployment_app_name = modal_deployment_app_name(settings)
|
|
if lookup_error_types:
|
|
try:
|
|
remote_engine = _lookup_protocol_current_remote_engine(payload)
|
|
_ensure_llm_profiles_staged(payload, deployment_app_name)
|
|
logger.info(
|
|
"Using deployed Modal app %s for component %s.",
|
|
deployment_app_name,
|
|
payload.get("component_id"),
|
|
)
|
|
return _invoke_remote_engine_payload_with_recovery(
|
|
remote_engine,
|
|
payload,
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
except lookup_error_types as exc:
|
|
missing_deployment = _is_missing_modal_deployment_error(exc)
|
|
if settings.auto_deploy and missing_deployment:
|
|
remote_engine = _auto_deploy_modal_app(payload, exc)
|
|
try:
|
|
_ensure_llm_profiles_staged(payload, deployment_app_name)
|
|
logger.info(
|
|
"Using auto-deployed Modal app %s for component %s.",
|
|
deployment_app_name,
|
|
payload.get("component_id"),
|
|
)
|
|
return _invoke_remote_engine_payload_with_recovery(
|
|
remote_engine,
|
|
payload,
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
except lookup_error_types as retry_exc:
|
|
exc = retry_exc
|
|
missing_deployment = _is_missing_modal_deployment_error(exc)
|
|
if not missing_deployment:
|
|
raise
|
|
if not settings.allow_ephemeral_fallback:
|
|
raise ModalRemoteInvocationError(
|
|
"Remote execution requires a deployed Modal app or a successful first-run auto-deploy. "
|
|
f"Lookup failed for app={deployment_app_name!r} component={payload.get('component_id')!r}: {exc}. "
|
|
"Ensure Modal credentials are configured so the custom node can auto-deploy, "
|
|
"or set COMFY_MODAL_ALLOW_EPHEMERAL_FALLBACK=true to allow slow ephemeral app.run() fallback behavior."
|
|
) from exc
|
|
logger.warning(
|
|
"Deployed Modal app lookup failed for app=%s component=%s: %s. Falling back to ephemeral app.run(); this creates a temporary Modal app session, not a persistent deployment or endpoint.",
|
|
deployment_app_name,
|
|
payload.get("component_id"),
|
|
exc,
|
|
)
|
|
else:
|
|
remote_engine = _lookup_protocol_current_remote_engine(payload)
|
|
_ensure_llm_profiles_staged(payload, deployment_app_name)
|
|
logger.info(
|
|
"Using deployed Modal app %s for component %s.",
|
|
deployment_app_name,
|
|
payload.get("component_id"),
|
|
)
|
|
return _invoke_remote_engine_payload_with_recovery(
|
|
remote_engine,
|
|
payload,
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
|
|
with _modal_cloud_settings_override(settings):
|
|
cloud_module = _load_modal_cloud_module()
|
|
cloud_app = getattr(cloud_module, "app", None)
|
|
cloud_remote_engine = getattr(cloud_module, "RemoteEngine", None)
|
|
if cloud_app is None or cloud_remote_engine is None:
|
|
raise ModalRemoteInvocationError(
|
|
"Stable Modal cloud entry module did not expose app and RemoteEngine."
|
|
)
|
|
logger.info(
|
|
"Starting ephemeral Modal app.run() for component %s. This does not create a persistent deployed app or web endpoint.",
|
|
payload.get("component_id"),
|
|
)
|
|
run_context = cloud_app.run() if hasattr(cloud_app, "run") else nullcontext()
|
|
with run_context:
|
|
remote_engine = cloud_remote_engine()
|
|
result = _invoke_remote_engine_payload_with_recovery(
|
|
remote_engine,
|
|
payload,
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
logger.info(
|
|
"Ephemeral Modal app.run() invocation completed for component %s.",
|
|
payload.get("component_id"),
|
|
)
|
|
return result
|
|
|
|
|
|
def invoke_remote_engine(
|
|
payload: dict[str, Any],
|
|
kwargs_payload: bytes,
|
|
*,
|
|
allow_implicit_mapping: bool = True,
|
|
) -> bytes:
|
|
"""Invoke Modal when configured, or fall back to local in-process execution."""
|
|
execution_mode = get_settings().execution_mode
|
|
if payload.get("payload_kind") == "mapped_subgraph":
|
|
if execution_mode == "local" or modal is None:
|
|
hydrated_inputs = deserialize_node_inputs(kwargs_payload)
|
|
return serialize_node_outputs(
|
|
_execute_mapped_subgraph_payload(payload, hydrated_inputs)
|
|
)
|
|
if allow_implicit_mapping and payload.get("payload_kind") == "subgraph":
|
|
hydrated_inputs = deserialize_node_inputs(kwargs_payload)
|
|
if _split_batch_boundary_inputs(payload, hydrated_inputs) is not None:
|
|
return asyncio.run(
|
|
_invoke_implicitly_mapped_subgraph_async(payload, kwargs_payload)
|
|
)
|
|
if execution_mode == "local" or modal is None:
|
|
if payload.get("payload_kind") == "subgraph":
|
|
return execute_subgraph_locally(payload, kwargs_payload)
|
|
return execute_node_locally(payload, kwargs_payload)
|
|
|
|
logger.info(
|
|
"Dispatching Modal remote invocation for component=%s payload_kind=%s.",
|
|
payload.get("component_id"),
|
|
payload.get("payload_kind"),
|
|
)
|
|
_emit_local_remote_dispatch_status(payload)
|
|
cancellation_event = threading.Event()
|
|
future = _remote_modal_call_executor().submit(
|
|
_invoke_modal_payload_blocking,
|
|
dict(payload),
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
interrupt_sent = False
|
|
cancellation_started_at: float | None = None
|
|
try:
|
|
while True:
|
|
try:
|
|
response = future.result(timeout=0.1)
|
|
break
|
|
except FutureTimeoutError:
|
|
(
|
|
interrupt_sent,
|
|
cancellation_started_at,
|
|
) = _handle_modal_wait_cancellation(
|
|
payload,
|
|
cancellation_event,
|
|
interrupt_sent=interrupt_sent,
|
|
cancellation_started_at=cancellation_started_at,
|
|
)
|
|
continue
|
|
except Exception:
|
|
if cancellation_event.is_set() or _local_processing_interrupted():
|
|
logger.info(
|
|
"Reraising Modal failure as a local interrupt for component=%s after cancellation.",
|
|
payload.get("component_id"),
|
|
)
|
|
_raise_local_interrupt()
|
|
logger.exception(
|
|
"Modal remote invocation failed for component=%s.",
|
|
payload.get("component_id"),
|
|
)
|
|
raise
|
|
if cancellation_event.is_set() or _local_processing_interrupted():
|
|
logger.info(
|
|
"Remote invocation for component=%s finished after interruption; raising local interrupt.",
|
|
payload.get("component_id"),
|
|
)
|
|
_raise_local_interrupt()
|
|
logger.info(
|
|
"Modal remote invocation completed for component=%s.",
|
|
payload.get("component_id"),
|
|
)
|
|
return response
|
|
|
|
|
|
async def invoke_remote_engine_async(
|
|
payload: dict[str, Any],
|
|
kwargs_payload: bytes,
|
|
*,
|
|
allow_implicit_mapping: bool = True,
|
|
) -> bytes:
|
|
"""Invoke Modal asynchronously so multiple proxy nodes can wait on remote work in parallel."""
|
|
execution_mode = get_settings().execution_mode
|
|
if (
|
|
execution_mode != "local"
|
|
and modal is not None
|
|
and bool(payload.get("stop_local_gap_keepalive_before_remote_component"))
|
|
):
|
|
_stop_local_gap_keepalive(payload, reason="next_remote_component_started")
|
|
if payload.get("payload_kind") == "mapped_subgraph":
|
|
if execution_mode == "local" or modal is None:
|
|
return await _invoke_mapped_remote_engine_async(payload, kwargs_payload)
|
|
if allow_implicit_mapping and payload.get("payload_kind") == "subgraph":
|
|
hydrated_inputs = deserialize_node_inputs(kwargs_payload)
|
|
if _split_batch_boundary_inputs(payload, hydrated_inputs) is not None:
|
|
if execution_mode != "local" and modal is not None:
|
|
await asyncio.to_thread(
|
|
_lookup_protocol_current_remote_engine,
|
|
payload,
|
|
)
|
|
return await _invoke_implicitly_mapped_subgraph_async(
|
|
payload, kwargs_payload
|
|
)
|
|
if execution_mode == "local" or modal is None:
|
|
return await asyncio.to_thread(
|
|
invoke_remote_engine,
|
|
payload,
|
|
kwargs_payload,
|
|
allow_implicit_mapping=allow_implicit_mapping,
|
|
)
|
|
|
|
logger.info(
|
|
"Dispatching async Modal remote invocation for component=%s payload_kind=%s.",
|
|
payload.get("component_id"),
|
|
payload.get("payload_kind"),
|
|
)
|
|
_emit_local_remote_dispatch_status(payload)
|
|
cancellation_event = threading.Event()
|
|
future = _remote_modal_call_executor().submit(
|
|
_invoke_modal_payload_blocking,
|
|
dict(payload),
|
|
kwargs_payload,
|
|
cancellation_event,
|
|
)
|
|
wrapped_future = asyncio.wrap_future(future)
|
|
interrupt_sent = False
|
|
cancellation_started_at: float | None = None
|
|
try:
|
|
while True:
|
|
try:
|
|
response = await asyncio.wait_for(
|
|
asyncio.shield(wrapped_future), timeout=0.1
|
|
)
|
|
break
|
|
except asyncio.TimeoutError:
|
|
(
|
|
interrupt_sent,
|
|
cancellation_started_at,
|
|
) = await _handle_modal_wait_cancellation_async(
|
|
payload,
|
|
cancellation_event,
|
|
interrupt_sent=interrupt_sent,
|
|
cancellation_started_at=cancellation_started_at,
|
|
)
|
|
continue
|
|
except asyncio.CancelledError:
|
|
cancellation_event.set()
|
|
prompt_id = (
|
|
str(payload.get("prompt_id"))
|
|
if payload.get("prompt_id") is not None
|
|
else None
|
|
)
|
|
if prompt_id is not None:
|
|
await request_remote_modal_prompt_interrupt_async(prompt_id)
|
|
raise
|
|
except Exception:
|
|
if cancellation_event.is_set() or _local_processing_interrupted():
|
|
logger.info(
|
|
"Reraising async Modal failure as a local interrupt for component=%s after cancellation.",
|
|
payload.get("component_id"),
|
|
)
|
|
_raise_local_interrupt()
|
|
logger.exception(
|
|
"Async Modal remote invocation failed for component=%s.",
|
|
payload.get("component_id"),
|
|
)
|
|
raise
|
|
if cancellation_event.is_set() or _local_processing_interrupted():
|
|
logger.info(
|
|
"Async remote invocation for component=%s finished after interruption; raising local interrupt.",
|
|
payload.get("component_id"),
|
|
)
|
|
_raise_local_interrupt()
|
|
logger.info(
|
|
"Async Modal remote invocation completed for component=%s.",
|
|
payload.get("component_id"),
|
|
)
|
|
_start_local_gap_keepalive(payload)
|
|
return response
|
|
|
|
|
|
configure_mapped_execution_hooks(
|
|
MappedExecutionHooks(
|
|
invoke_remote_engine_async=lambda *args, **kwargs: invoke_remote_engine_async(
|
|
*args, **kwargs
|
|
),
|
|
invoke_bound_remote_engine_async=lambda remote_engine, payload, kwargs_payload: _invoke_bound_remote_engine_async(
|
|
remote_engine, payload, kwargs_payload
|
|
),
|
|
local_processing_interrupted=lambda: _local_processing_interrupted(),
|
|
raise_local_interrupt=lambda: _raise_local_interrupt(),
|
|
exception_indicates_interruption=lambda exc: _exception_indicates_interruption(
|
|
exc
|
|
),
|
|
)
|
|
)
|
|
|
|
|
|
configure_modal_warmup_hooks(
|
|
ModalWarmupHooks(
|
|
ensure_llm_profiles_staged=lambda payload, app_name: _ensure_llm_profiles_staged(
|
|
payload, app_name
|
|
),
|
|
mapped_execution_parallelism=lambda item_count: _mapped_execution_parallelism(
|
|
item_count
|
|
),
|
|
)
|
|
)
|
|
|
|
|
|
configure_modal_deployment_hooks(
|
|
ModalDeploymentHooks(
|
|
ensure_llm_profiles_staged=lambda payload, app_name: _ensure_llm_profiles_staged(
|
|
payload, app_name
|
|
),
|
|
schedule_post_deploy_runtime_seed=lambda payload: _schedule_post_deploy_runtime_seed(
|
|
payload
|
|
),
|
|
await_matching_speculative_prewarm=lambda payload, cancellation_event: _await_matching_speculative_prewarm(
|
|
payload, cancellation_event
|
|
),
|
|
prepare_snapshot_profile_fields=lambda payload: _prepare_snapshot_profile_fields(
|
|
payload
|
|
),
|
|
select_gpu_snapshot_for_profile=lambda payload, profile_key: _select_gpu_snapshot_for_profile(
|
|
payload, profile_key
|
|
),
|
|
prompt_parallelism_target=lambda payload: _prompt_parallelism_target(payload),
|
|
)
|
|
)
|