"""Queue-time remote prompt analysis, asset preparation, and proxy rewriting.""" from __future__ import annotations import asyncio import copy from collections import defaultdict from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field import logging from pathlib import Path import threading import time from typing import Any, Callable, Mapping, Sequence import uuid if __package__: from .component_planning import ( _build_component_plans, _implicitly_mapped_boundary_output_sources, _mark_payload_scheduler_list_outputs, _mark_remote_to_remote_session_boundaries, _workflow_visible_remote_node_ids, validate_remote_component_transport_compatibility, ) from .execution_assignment_runtime import ( _configured_provider_metadata, _ensure_remote_sync_backend, _execution_location_for_assignment, _ssh_sync_engine, _stamp_execution_assignment, _workflow_r2_cache, ) from .execution_environments import ExecutionAssignment, ExecutionProvider from .execution_resource_estimates import _component_execution_signature from .execution_scheduling import ( _configured_ssh_hosts, _plan_component_execution, ) from .prompt_affinity_planning import ( _configure_local_gap_keepalive_payloads, _configure_speculative_affinity_prewarm_payloads, _parallelize_non_returning_local_branches, ) from .prompt_diagnostics import _log_modal_rewritten_prompt_diagnostics from .prompt_payload_metadata import _attach_resolved_llm_profiles from .prompt_payload_building import _build_component_payload from .prompt_rewrite import ( _attach_modal_artifact_finalizer, _component_uploaded_volume_paths, _deduplicate_synced_assets, _rewrite_component_into_proxy, _sync_component_prompt_inputs, ) from .r2_cache import R2CacheError from .remote_graph_analysis import ( _component_dependency_graph, _component_execution_stages, _estimated_stage_parallelism, _expand_remote_node_ids_for_non_transportable_inputs, _expand_remote_node_ids_for_terminal_video_sinks, _sandwiched_local_node_ids, requested_remote_node_ids, ) from .remote_plan_types import ( ComponentExecutionPlan, ModalPromptValidationError, RemoteComponentPlan, RewriteSummary, _EnvironmentAssetPreparationResult, ) from .session_state import RemoteSessionHandle from .settings import ModalSyncSettings, get_settings from .sync_engine import ( ModalAssetSyncEngine, SyncCancelledError, SyncedAsset, ) from .vast_service import VastService else: # pragma: no cover - flat import inside the Modal container. from component_planning import ( _build_component_plans, _implicitly_mapped_boundary_output_sources, _mark_payload_scheduler_list_outputs, _mark_remote_to_remote_session_boundaries, _workflow_visible_remote_node_ids, validate_remote_component_transport_compatibility, ) from execution_assignment_runtime import ( _configured_provider_metadata, _ensure_remote_sync_backend, _execution_location_for_assignment, _ssh_sync_engine, _stamp_execution_assignment, _workflow_r2_cache, ) from execution_environments import ExecutionAssignment, ExecutionProvider from execution_resource_estimates import _component_execution_signature from execution_scheduling import ( _configured_ssh_hosts, _plan_component_execution, ) from prompt_affinity_planning import ( _configure_local_gap_keepalive_payloads, _configure_speculative_affinity_prewarm_payloads, _parallelize_non_returning_local_branches, ) from prompt_diagnostics import _log_modal_rewritten_prompt_diagnostics from prompt_payload_metadata import _attach_resolved_llm_profiles from prompt_payload_building import _build_component_payload from prompt_rewrite import ( _attach_modal_artifact_finalizer, _component_uploaded_volume_paths, _deduplicate_synced_assets, _rewrite_component_into_proxy, _sync_component_prompt_inputs, ) from r2_cache import R2CacheError from remote_graph_analysis import ( _component_dependency_graph, _component_execution_stages, _estimated_stage_parallelism, _expand_remote_node_ids_for_non_transportable_inputs, _expand_remote_node_ids_for_terminal_video_sinks, _sandwiched_local_node_ids, requested_remote_node_ids, ) from remote_plan_types import ( ComponentExecutionPlan, ModalPromptValidationError, RemoteComponentPlan, RewriteSummary, _EnvironmentAssetPreparationResult, ) from session_state import RemoteSessionHandle from settings import ModalSyncSettings, get_settings from sync_engine import ( ModalAssetSyncEngine, SyncCancelledError, SyncedAsset, ) from vast_service import VastService logger = logging.getLogger(__name__) SetupStatusCallback = Callable[[str, int | None, int | None], None] EnvironmentSetupStatusCallback = Callable[ [str, str, int | None, int | None], None ] ExecutionPlanStatusCallback = Callable[ [dict[str, dict[str, Any]], list[dict[str, Any]]], None ] _REMOTE_ASSET_PREPARATION_MAX_WORKERS = 8 def _get_nodes_module() -> Any: """Import the ComfyUI nodes module lazily.""" import nodes return nodes def _environment_setup_status_callback( environment_id: str, status_callback: SetupStatusCallback | None, environment_status_callback: EnvironmentSetupStatusCallback | None, ) -> SetupStatusCallback | None: """Return a setup callback that updates prompt-wide and environment UI.""" if status_callback is None and environment_status_callback is None: return None def emit_status( message: str, current: int | None, total: int | None, ) -> None: """Forward one setup update with its concrete execution environment.""" if status_callback is not None: status_callback(message, current, total) if environment_status_callback is not None: environment_status_callback(environment_id, message, current, total) return emit_status def _prepare_environment_assets( *, environment_id: str, components: Sequence[RemoteComponentPlan], sync_engine: ModalAssetSyncEngine, engine_lock: threading.Lock, rewritten_prompt: dict[str, Any], sync_custom_nodes: bool, completion_message: str, status_callback: SetupStatusCallback | None, environment_status_callback: EnvironmentSetupStatusCallback | None, ) -> _EnvironmentAssetPreparationResult: """Prepare one environment's custom nodes and prompt assets in order.""" started_at = time.perf_counter() logger.info( "Starting remote asset preparation environment=%s components=%d.", environment_id, len(components), ) environment_callback = _environment_setup_status_callback( environment_id, status_callback, environment_status_callback, ) if environment_callback is not None: environment_callback("Preparing remote assets", None, None) with engine_lock: sync_engine.preflight_r2_access(status_callback=environment_callback) custom_nodes_bundle = ( sync_engine.sync_custom_nodes_directory( status_callback=environment_callback, ) if sync_custom_nodes else None ) component_prompts, assets_by_component_id = ( _sync_environment_prompt_assets( components=components, sync_engine=sync_engine, rewritten_prompt=rewritten_prompt, status_callback=environment_callback, ) ) if environment_callback is not None: environment_callback(completion_message, None, None) logger.info( "Finished remote asset preparation environment=%s components=%d assets=%d elapsed_seconds=%.3f.", environment_id, len(components), sum(len(assets) for assets in assets_by_component_id.values()), time.perf_counter() - started_at, ) return _EnvironmentAssetPreparationResult( environment_id=environment_id, custom_nodes_bundle=custom_nodes_bundle, component_prompts=component_prompts, assets_by_component_id=assets_by_component_id, ) def _sync_environment_prompt_assets( *, components: Sequence[RemoteComponentPlan], sync_engine: ModalAssetSyncEngine, rewritten_prompt: dict[str, Any], status_callback: SetupStatusCallback | None, ) -> tuple[dict[str, dict[str, Any]], dict[str, list[SyncedAsset]]]: """Sync every component assigned to one environment with one request cache.""" request_cache = sync_engine.create_request_asset_cache( rewritten_prompt[node_id].get("inputs", {}) for component in components for node_id in component.node_ids ) component_prompts: dict[str, dict[str, Any]] = {} assets_by_component_id: dict[str, list[SyncedAsset]] = {} for component in components: component_prompt, synced_assets = _sync_component_prompt_inputs( component=component, rewritten_prompt=rewritten_prompt, sync_engine=sync_engine, request_cache=request_cache, status_callback=status_callback, ) component_id = component.representative_node_id component_prompts[component_id] = component_prompt assets_by_component_id[component_id] = list(synced_assets) return component_prompts, assets_by_component_id def _prepare_remote_environment_assets( *, components: Sequence[RemoteComponentPlan], assignments_by_component_id: Mapping[str, ExecutionAssignment], sync_engines_by_environment: Mapping[str, ModalAssetSyncEngine], rewritten_prompt: dict[str, Any], sync_custom_nodes: bool, status_callback: SetupStatusCallback | None, environment_status_callback: EnvironmentSetupStatusCallback | None, ) -> dict[str, _EnvironmentAssetPreparationResult]: """Prepare distinct remote environments concurrently in a bounded pool.""" components_by_environment: dict[str, list[RemoteComponentPlan]] = { environment_id: [] for environment_id in sync_engines_by_environment } providers_by_environment: dict[str, ExecutionProvider] = {} for component in components: assignment = assignments_by_component_id[component.representative_node_id] components_by_environment[assignment.environment_id].append(component) providers_by_environment[assignment.environment_id] = assignment.provider locks_by_engine: dict[int, threading.Lock] = {} environment_ids = list(components_by_environment) if not environment_ids: return {} max_workers = min(_REMOTE_ASSET_PREPARATION_MAX_WORKERS, len(environment_ids)) logger.info( "Preparing remote assets across %d environments with %d workers.", len(environment_ids), max_workers, ) with ThreadPoolExecutor( max_workers=max_workers, thread_name_prefix="comfy-remote-assets", ) as executor: futures_by_environment = { environment_id: executor.submit( _prepare_environment_assets, environment_id=environment_id, components=components_by_environment[environment_id], sync_engine=sync_engines_by_environment[environment_id], engine_lock=locks_by_engine.setdefault( id(sync_engines_by_environment[environment_id]), threading.Lock(), ), rewritten_prompt=rewritten_prompt, sync_custom_nodes=sync_custom_nodes, completion_message=_asset_preparation_completion_message( providers_by_environment[environment_id] ), status_callback=status_callback, environment_status_callback=environment_status_callback, ) for environment_id in environment_ids } return { environment_id: futures_by_environment[environment_id].result() for environment_id in environment_ids } def _asset_preparation_completion_message(provider: ExecutionProvider) -> str: """Return the truthful state after assets, but not runtime startup, finish.""" if provider is ExecutionProvider.SSH_DOCKER: return "Remote assets prepared; SSH runtime starts on dispatch" if provider is ExecutionProvider.MODAL: return "Remote assets prepared; Modal runtime starts on dispatch" return "Ready for remote execution" @dataclass class _PromptRewriteState: """Hold mutable queue-time state across prompt rewrite stages.""" prompt: dict[str, Any] workflow: dict[str, Any] | None rewritten_prompt: dict[str, Any] settings: ModalSyncSettings nodes_module: Any sync_engine: ModalAssetSyncEngine extra_data: dict[str, Any] | None summary: RewriteSummary components: list[RemoteComponentPlan] = field(default_factory=list) execution_plan: ComponentExecutionPlan | None = None assignments: dict[str, ExecutionAssignment] = field(default_factory=dict) remote_sessions: dict[str, dict[str, Any]] = field(default_factory=dict) vast_service: VastService | None = None vast_leases: dict[str, Any] = field(default_factory=dict) workflow_r2_cache: Any | None = None ssh_hosts: dict[str, Any] = field(default_factory=dict) sync_engines: dict[str, ModalAssetSyncEngine] = field(default_factory=dict) synced_component_prompts: dict[str, dict[str, Any]] = field( default_factory=dict ) synced_assets_by_component: dict[str, list[SyncedAsset]] = field( default_factory=dict ) mapped_proxy_component_ids: set[str] = field(default_factory=set) def _analyze_remote_components( state: _PromptRewriteState, remote_node_ids: set[str], ) -> None: """Expand remote selection, build components, and validate transport.""" expanded_node_ids, _ = _expand_remote_node_ids_for_non_transportable_inputs( prompt=state.rewritten_prompt, remote_node_ids=remote_node_ids, nodes_module=state.nodes_module, ) expanded_node_ids = _expand_remote_node_ids_for_terminal_video_sinks( prompt=state.rewritten_prompt, remote_node_ids=expanded_node_ids, nodes_module=state.nodes_module, ) state.summary.remote_node_ids = sorted(expanded_node_ids) state.summary.sandwiched_local_node_ids = sorted( _sandwiched_local_node_ids(state.rewritten_prompt, expanded_node_ids) ) if state.summary.sandwiched_local_node_ids: logger.warning( "Detected local nodes sandwiched between remote graph regions; " "these nodes may force additional remote phases: %s", state.summary.sandwiched_local_node_ids, ) state.components = _build_component_plans( state.rewritten_prompt, expanded_node_ids, state.nodes_module, ) validate_remote_component_transport_compatibility( prompt=state.rewritten_prompt, components=state.components, nodes_module=state.nodes_module, ) def _remote_sessions_for_components( state: _PromptRewriteState, session_component_ids: set[str], ) -> dict[str, dict[str, Any]]: """Create prompt-scoped sessions for components with remote-to-remote links.""" prompt_id = (state.extra_data or {}).get("prompt_id") sessions = { component.representative_node_id: RemoteSessionHandle( session_id=uuid.uuid4().hex, prompt_id=(str(prompt_id) if prompt_id is not None else None), owner_component_id=component.representative_node_id, ).to_payload() for component in state.components if component.representative_node_id in session_component_ids } if sessions: logger.info( "Enabled environment-local remote references for components=%s.", sorted(sessions), ) return sessions def _plan_remote_component_execution( state: _PromptRewriteState, *, status_callback: SetupStatusCallback | None, environment_status_callback: EnvironmentSetupStatusCallback | None, plan_callback: ExecutionPlanStatusCallback | None, occupied_environment_ids: frozenset[str], ) -> set[str]: """Plan providers, record locations, and mark session boundaries.""" execution_plan = _plan_component_execution( components=state.components, prompt=state.rewritten_prompt, workflow=state.workflow, settings=state.settings, status_callback=status_callback, environment_status_callback=environment_status_callback, plan_callback=plan_callback, occupied_environment_ids=occupied_environment_ids, ) state.execution_plan = execution_plan state.assignments = execution_plan.assignments if execution_plan.configuration_set is not None: state.summary.remote_configurations = list(execution_plan.safe_configurations) for assignment in state.assignments.values(): execution_location = _execution_location_for_assignment( assignment, execution_plan.ssh_hosts_by_id, execution_plan.vast_leases_by_environment, ) if execution_location: state.summary.execution_locations_by_environment[ assignment.environment_id ] = execution_location session_component_ids = _mark_remote_to_remote_session_boundaries( state.rewritten_prompt, state.components, state.nodes_module, state.assignments, ) state.remote_sessions = _remote_sessions_for_components( state, session_component_ids ) state.summary.execution_assignments_by_representative = dict(state.assignments) return session_component_ids def _resolve_vast_service(state: _PromptRewriteState) -> None: """Resolve the planner-provided or registry-backed Vast controller state.""" assert state.execution_plan is not None state.vast_service = state.execution_plan.vast_service state.vast_leases = dict(state.execution_plan.vast_leases_by_environment) if state.vast_service is not None or not any( assignment.provider is ExecutionProvider.VAST for assignment in state.assignments.values() ): return try: runtime_fingerprints = { lease.runtime_fingerprint for lease in state.vast_leases.values() } worker_images = { lease.worker_image for lease in state.vast_leases.values() if lease.worker_image is not None } runtime_fingerprint = ( next(iter(runtime_fingerprints)) if len(runtime_fingerprints) == 1 else None ) worker_image = next(iter(worker_images)) if len(worker_images) == 1 else None state.vast_service = VastService.from_environment( state.settings, repo_root=Path(__file__).resolve().parent, runtime_fingerprint=runtime_fingerprint, worker_image=worker_image, ) state.vast_leases = { assignment.environment_id: state.vast_service.lease_for_environment_id( assignment.environment_id ) for assignment in state.assignments.values() if assignment.provider is ExecutionProvider.VAST } except (KeyError, OSError, RuntimeError, ValueError) as exc: raise ModalPromptValidationError( f"Unable to resolve the acquired Vast.ai lease: {exc}" ) from exc def _resolve_execution_resources(state: _PromptRewriteState) -> None: """Resolve Vast, R2, and SSH resources needed after placement.""" assert state.execution_plan is not None _resolve_vast_service(state) state.workflow_r2_cache = ( _workflow_r2_cache(state.execution_plan.configuration_set) if any( assignment.provider is not ExecutionProvider.MODAL for assignment in state.assignments.values() ) else None ) if state.workflow_r2_cache is not None and state.vast_service is not None: state.vast_service.r2_cache = state.workflow_r2_cache state.ssh_hosts = dict(state.execution_plan.ssh_hosts_by_id) if not state.ssh_hosts: state.ssh_hosts = { host.environment_id: host for host in _configured_ssh_hosts(state.settings) } def _assign_component_worker_indices( state: _PromptRewriteState, session_component_ids: set[str], ) -> None: """Assign deterministic worker slots while reserving session-affine slot zero.""" assert state.execution_plan is not None session_environment_ids = { state.assignments[component_id].environment_id for component_id in session_component_ids } worker_counts: dict[str, int] = defaultdict(int) for component in state.components: component_id = component.representative_node_id assignment = state.assignments[component_id] if assignment.provider in {ExecutionProvider.MODAL, ExecutionProvider.VAST}: worker_index = 0 elif state.execution_plan.configuration_set is not None: worker_index = assignment.capacity_slot_index logger.info( "Assigned workflow-configured remote component=%s environment=%s worker_index=%d.", component_id, assignment.environment_id, worker_index, ) else: host = state.ssh_hosts.get(assignment.environment_id) maximum_workers = max(1, host.maximum_workers if host is not None else 1) if component_id in session_component_ids: worker_index = 0 else: first_worker = ( 1 if assignment.environment_id in session_environment_ids and maximum_workers > 1 else 0 ) available_workers = max(1, maximum_workers - first_worker) worker_index = first_worker + ( worker_counts[assignment.environment_id] % available_workers ) worker_counts[assignment.environment_id] += 1 logger.info( "Assigned remote component=%s environment=%s worker_index=%d.", component_id, assignment.environment_id, worker_index, ) state.summary.execution_worker_indices_by_representative[component_id] = ( worker_index ) def _sync_engine_for_assignment( state: _PromptRewriteState, assignment: ExecutionAssignment, cancellation_check: Callable[[], bool] | None, ) -> ModalAssetSyncEngine: """Return the provider-specific sync engine for one assignment.""" if assignment.provider is ExecutionProvider.MODAL: _ensure_remote_sync_backend(state.settings, state.sync_engine) return state.sync_engine if assignment.provider is ExecutionProvider.VAST: if state.vast_service is None: raise ModalPromptValidationError( "Vast.ai assignment has no active controller service." ) lease = state.vast_leases[assignment.environment_id] sync_engine = state.vast_service.sync_engine(lease) sync_engine.cancellation_check = cancellation_check return sync_engine host = state.ssh_hosts.get(assignment.environment_id) if host is None: raise ModalPromptValidationError( f"Assigned SSH execution environment {assignment.environment_id!r} is no longer configured." ) return _ssh_sync_engine( host=host, settings=state.settings, r2_cache=state.workflow_r2_cache, ) def _build_environment_sync_engines( state: _PromptRewriteState, cancellation_check: Callable[[], bool] | None, ) -> None: """Build one asset synchronization engine per selected environment.""" for assignment in state.assignments.values(): if assignment.environment_id in state.sync_engines: continue state.sync_engines[assignment.environment_id] = _sync_engine_for_assignment( state, assignment, cancellation_check ) def _record_environment_preparations( state: _PromptRewriteState, preparations: dict[str, _EnvironmentAssetPreparationResult], ) -> None: """Record custom-node bundles, synced component prompts, and unique assets.""" if state.settings.sync_custom_nodes: state.summary.custom_nodes_bundles_by_environment = { environment_id: preparation.custom_nodes_bundle for environment_id, preparation in preparations.items() } modal_bundle = next( ( state.summary.custom_nodes_bundles_by_environment[ assignment.environment_id ] for assignment in state.assignments.values() if assignment.provider is ExecutionProvider.MODAL ), None, ) state.summary.custom_nodes_bundle = modal_bundle or next( iter(state.summary.custom_nodes_bundles_by_environment.values()), None ) else: logger.info( "Skipping custom_nodes bundle sync because sync is disabled for execution_mode=%s.", state.settings.execution_mode, ) for component in state.components: component_id = component.representative_node_id assignment = state.assignments[component_id] preparation = preparations[assignment.environment_id] state.synced_component_prompts[component_id] = ( preparation.component_prompts[component_id] ) assets = preparation.assets_by_component_id[component_id] state.synced_assets_by_component[component_id] = assets state.summary.synced_assets.extend(assets) state.summary.synced_assets = _deduplicate_synced_assets( state.summary.synced_assets ) def _prepare_rewrite_assets( state: _PromptRewriteState, *, status_callback: SetupStatusCallback | None, environment_status_callback: EnvironmentSetupStatusCallback | None, ) -> None: """Prepare all environment assets and record their component-scoped results.""" if status_callback is not None: status_callback("Preparing assets for remote execution", None, None) try: preparations = _prepare_remote_environment_assets( components=state.components, assignments_by_component_id=state.assignments, sync_engines_by_environment=state.sync_engines, rewritten_prompt=state.rewritten_prompt, sync_custom_nodes=state.settings.sync_custom_nodes, status_callback=status_callback, environment_status_callback=environment_status_callback, ) except R2CacheError as exc: raise ModalPromptValidationError(str(exc)) from exc _record_environment_preparations(state, preparations) def _resolve_modal_volume_reload(state: _PromptRewriteState) -> None: """Compute request-wide Modal volume reload state after asset preparation.""" modal_environment_ids = { assignment.environment_id for assignment in state.assignments.values() if assignment.provider is ExecutionProvider.MODAL } modal_component_ids = { component.representative_node_id for component in state.components if state.assignments[component.representative_node_id].provider is ExecutionProvider.MODAL } requires_reload = any( asset.uploaded for component_id, assets in state.synced_assets_by_component.items() if component_id in modal_component_ids for asset in assets ) or any( bundle is not None and bundle.uploaded for environment_id, bundle in ( state.summary.custom_nodes_bundles_by_environment.items() ) if environment_id in modal_environment_ids ) state.summary.requires_volume_reload = requires_reload state.summary.volume_reload_marker = uuid.uuid4().hex if requires_reload else None state.summary.uploaded_volume_paths = [ asset.remote_path for asset in state.summary.synced_assets if asset.uploaded ] logger.info( "Resolved request-wide Modal volume reload requirement: requires_volume_reload=%s volume_reload_marker=%s synced_assets=%d custom_nodes_uploaded=%s", requires_reload, state.summary.volume_reload_marker, len(state.summary.synced_assets), bool( state.summary.custom_nodes_bundle is not None and state.summary.custom_nodes_bundle.uploaded ), ) def _build_stamped_component_payload( state: _PromptRewriteState, component: RemoteComponentPlan, ) -> tuple[dict[str, Any], set[Any]]: """Build one component payload and stamp its provider assignment metadata.""" assert state.execution_plan is not None component_id = component.representative_node_id assignment = state.assignments[component_id] custom_nodes_bundle = state.summary.custom_nodes_bundles_by_environment.get( assignment.environment_id ) uploaded_paths = _component_uploaded_volume_paths( component_prompt=state.synced_component_prompts[component_id], synced_assets=state.synced_assets_by_component[component_id], custom_nodes_bundle=custom_nodes_bundle, ) payload = _build_component_payload( component=component, component_prompt=state.synced_component_prompts[component_id], signature_prompt=state.prompt, extra_data=state.extra_data, settings=state.settings, requires_volume_reload=( assignment.provider is ExecutionProvider.MODAL and bool(uploaded_paths) ), volume_reload_marker=state.summary.volume_reload_marker, custom_nodes_bundle=custom_nodes_bundle, uploaded_volume_paths=uploaded_paths, terminate_container_on_error=state.settings.terminate_container_on_error, nodes_module=state.nodes_module, remote_session=state.remote_sessions.get(component_id), ) _attach_resolved_llm_profiles( payload, state.execution_plan.resolved_llm_profiles, state.settings ) _stamp_execution_assignment( payload, assignment, state.summary.execution_worker_indices_by_representative.get(component_id, 0), _component_execution_signature(component, state.prompt), _execution_location_for_assignment( assignment, state.ssh_hosts, state.vast_leases ), _configured_provider_metadata( execution_plan=state.execution_plan, assignment=assignment, vast_leases_by_environment=state.vast_leases, ), ) implicitly_mapped_sources = _implicitly_mapped_boundary_output_sources( component=component, original_prompt=state.prompt, rewritten_prompt=state.rewritten_prompt, nodes_module=state.nodes_module, ) _mark_payload_scheduler_list_outputs(payload, implicitly_mapped_sources) return payload, implicitly_mapped_sources def _record_hybrid_component_summary( state: _PromptRewriteState, component: RemoteComponentPlan, proxy_node_ids: list[str], ) -> None: """Record node mappings for paired static and mapped proxies.""" static_proxy_node_id, mapped_proxy_node_id = proxy_node_ids state.summary.remote_component_ids.extend(proxy_node_ids) for proxy_node_id, component_node_ids in ( (static_proxy_node_id, component.static_node_ids), (mapped_proxy_node_id, component.mapped_node_ids), ): visible_node_ids = _workflow_visible_remote_node_ids(component_node_ids) state.summary.component_node_ids_by_representative[proxy_node_id] = ( visible_node_ids ) for node_id in visible_node_ids: state.summary.rewritten_node_id_map[node_id] = proxy_node_id state.mapped_proxy_component_ids.add(mapped_proxy_node_id) def _record_ordered_component_summary( state: _PromptRewriteState, component: RemoteComponentPlan, proxy_node_ids: list[str], split_payloads: list[Any], ) -> None: """Record node mappings for an ordered proxy sequence.""" state.summary.remote_component_ids.extend(proxy_node_ids) mapped_node_ids = set(component.mapped_node_ids) for phase_payload in split_payloads: proxy_node_id = str(phase_payload["component_id"]) component_node_ids = [ str(node_id) for node_id in phase_payload["component_node_ids"] ] visible_node_ids = _workflow_visible_remote_node_ids(component_node_ids) state.summary.component_node_ids_by_representative[proxy_node_id] = ( visible_node_ids ) for node_id in visible_node_ids: state.summary.rewritten_node_id_map[node_id] = proxy_node_id if mapped_node_ids.intersection(component_node_ids): state.mapped_proxy_component_ids.add(proxy_node_id) def _record_single_component_summary( state: _PromptRewriteState, component: RemoteComponentPlan, proxy_node_id: str, implicitly_mapped_sources: set[Any], ) -> None: """Record node mappings for one ordinary component proxy.""" state.summary.remote_component_ids.append(proxy_node_id) visible_node_ids = _workflow_visible_remote_node_ids(component.node_ids) state.summary.component_node_ids_by_representative[proxy_node_id] = ( visible_node_ids ) for node_id in visible_node_ids: state.summary.rewritten_node_id_map[node_id] = proxy_node_id if component.mapped_boundary_input_name or implicitly_mapped_sources: state.mapped_proxy_component_ids.add(proxy_node_id) def _record_component_proxy_summary( state: _PromptRewriteState, component: RemoteComponentPlan, payload: dict[str, Any], proxy_node_ids: list[str], implicitly_mapped_sources: set[Any], ) -> None: """Record assignment and workflow-node metadata for rewritten proxies.""" assignment = state.assignments[component.representative_node_id] worker_index = state.summary.execution_worker_indices_by_representative.get( component.representative_node_id, 0 ) for proxy_node_id in proxy_node_ids: state.summary.execution_assignments_by_representative[proxy_node_id] = ( assignment ) state.summary.execution_worker_indices_by_representative[proxy_node_id] = ( worker_index ) split_payloads = payload.get("split_proxy_payloads") if isinstance(split_payloads, dict): _record_hybrid_component_summary(state, component, proxy_node_ids) elif isinstance(split_payloads, list): _record_ordered_component_summary( state, component, proxy_node_ids, split_payloads ) else: _record_single_component_summary( state, component, proxy_node_ids[0], implicitly_mapped_sources ) def _rewrite_remote_components(state: _PromptRewriteState) -> None: """Build, stamp, and rewrite every planned remote component.""" for component in state.components: logger.info( "Rewriting remote component %s covering nodes %s.", component.representative_node_id, component.node_ids, ) payload, implicitly_mapped_sources = _build_stamped_component_payload( state, component ) proxy_node_ids = _rewrite_component_into_proxy( component=component, rewritten_prompt=state.rewritten_prompt, payload=payload, nodes_module=state.nodes_module, ) _record_component_proxy_summary( state, component, payload, proxy_node_ids, implicitly_mapped_sources, ) def _parallelism_upper_bound( state: _PromptRewriteState, estimated_parallelism: int, ) -> int: """Return the configured or runtime upper bound for remote parallelism.""" assert state.execution_plan is not None if state.execution_plan.configuration_set is not None: configured_capacity = sum( configuration.capacity_limit for configuration in ( state.execution_plan.configuration_set.capacity_configurations ) ) return min(estimated_parallelism, configured_capacity) if state.settings.max_containers is not None: return min(estimated_parallelism, state.settings.max_containers) return estimated_parallelism def _finalize_rewritten_prompt(state: _PromptRewriteState) -> None: """Attach affinity/finalizer nodes and calculate dependency summary metadata.""" _configure_local_gap_keepalive_payloads( rewritten_prompt=state.rewritten_prompt, remote_component_ids=state.summary.remote_component_ids, sandwiched_local_node_ids=set(state.summary.sandwiched_local_node_ids), ) state.summary.parallel_local_branch_node_ids = ( _parallelize_non_returning_local_branches( rewritten_prompt=state.rewritten_prompt, remote_component_ids=state.summary.remote_component_ids, nodes_module=state.nodes_module, ) ) state.summary.artifact_finalizer_node_id = _attach_modal_artifact_finalizer( rewritten_prompt=state.rewritten_prompt, remote_component_ids=state.summary.remote_component_ids, nodes_module=state.nodes_module, ) proxy_groups = { component_id: {component_id} for component_id in state.summary.remote_component_ids } _, dependency_edges, _ = _component_dependency_graph( state.rewritten_prompt, proxy_groups ) execution_stages = _component_execution_stages( state.rewritten_prompt, proxy_groups ) _configure_speculative_affinity_prewarm_payloads( rewritten_prompt=state.rewritten_prompt, execution_stages=execution_stages, ) state.summary.component_dependency_ids_by_representative = { component_id: sorted( upstream_id for upstream_id, downstream_ids in dependency_edges.items() if component_id in downstream_ids ) for component_id in sorted(proxy_groups) } state.summary.component_execution_stages = [ list(stage) for stage in execution_stages ] state.summary.mapped_component_ids = sorted(state.mapped_proxy_component_ids) state.summary.estimated_max_parallel_requests = _estimated_stage_parallelism( execution_stages, state.mapped_proxy_component_ids, mapped_component_weight=1, ) state.summary.max_parallel_requests_upper_bound = _parallelism_upper_bound( state, state.summary.estimated_max_parallel_requests ) def _log_rewrite_completion(state: _PromptRewriteState) -> None: """Log final parallelism and diagnostics for one rewritten prompt.""" logger.info( "Estimated remote parallelism after proxy rewrite: known_max_parallel_requests=%d max_parallel_requests_upper_bound=%s mapped_components=%s execution_stages=%s", state.summary.estimated_max_parallel_requests, state.summary.max_parallel_requests_upper_bound, state.summary.mapped_component_ids, state.summary.component_execution_stages, ) _log_modal_rewritten_prompt_diagnostics( prompt_id=( str(state.extra_data.get("prompt_id")) if isinstance(state.extra_data, Mapping) and state.extra_data.get("prompt_id") else None ), prompt=state.rewritten_prompt, summary=state.summary, reason="post_rewrite", ) def rewrite_prompt_for_modal( prompt: dict[str, Any], workflow: dict[str, Any] | None, sync_engine: ModalAssetSyncEngine | None = None, settings: ModalSyncSettings | None = None, nodes_module: Any | None = None, extra_data: dict[str, Any] | None = None, status_callback: SetupStatusCallback | None = None, environment_status_callback: EnvironmentSetupStatusCallback | None = None, plan_callback: ExecutionPlanStatusCallback | None = None, cancellation_check: Callable[[], bool] | None = None, occupied_environment_ids: frozenset[str] = frozenset(), ) -> tuple[dict[str, Any], RewriteSummary]: """Rewrite connected remote components into provider-backed proxy nodes.""" if cancellation_check is not None and cancellation_check(): raise SyncCancelledError("Remote workflow preparation was cancelled.") resolved_settings = settings or get_settings() remote_node_ids = requested_remote_node_ids( prompt=prompt, workflow=workflow, settings=resolved_settings, ) summary = RewriteSummary(remote_node_ids=sorted(remote_node_ids)) logger.info( "Found %d workflow nodes marked for remote execution.", len(remote_node_ids) ) if not remote_node_ids: return copy.deepcopy(prompt), summary state = _PromptRewriteState( prompt=prompt, workflow=workflow, rewritten_prompt=copy.deepcopy(prompt), settings=resolved_settings, nodes_module=nodes_module or _get_nodes_module(), sync_engine=sync_engine or ModalAssetSyncEngine.from_environment(resolved_settings), extra_data=extra_data, summary=summary, ) _analyze_remote_components(state, remote_node_ids) session_component_ids = _plan_remote_component_execution( state, status_callback=status_callback, environment_status_callback=environment_status_callback, plan_callback=plan_callback, occupied_environment_ids=occupied_environment_ids, ) _resolve_execution_resources(state) _assign_component_worker_indices(state, session_component_ids) _build_environment_sync_engines(state, cancellation_check) _prepare_rewrite_assets( state, status_callback=status_callback, environment_status_callback=environment_status_callback, ) _resolve_modal_volume_reload(state) _rewrite_remote_components(state) _finalize_rewritten_prompt(state) _log_rewrite_completion(state) return state.rewritten_prompt, state.summary async def rewrite_prompt_for_modal_async( prompt: dict[str, Any], workflow: dict[str, Any] | None, sync_engine: ModalAssetSyncEngine | None = None, settings: ModalSyncSettings | None = None, nodes_module: Any | None = None, extra_data: dict[str, Any] | None = None, status_callback: SetupStatusCallback | None = None, environment_status_callback: EnvironmentSetupStatusCallback | None = None, plan_callback: ExecutionPlanStatusCallback | None = None, cancellation_check: Callable[[], bool] | None = None, occupied_environment_ids: frozenset[str] = frozenset(), ) -> tuple[dict[str, Any], RewriteSummary]: """Prepare one Modal prompt without blocking ComfyUI's event loop.""" return await asyncio.to_thread( rewrite_prompt_for_modal, prompt=prompt, workflow=workflow, sync_engine=sync_engine, settings=settings, nodes_module=nodes_module, extra_data=extra_data, status_callback=status_callback, environment_status_callback=environment_status_callback, plan_callback=plan_callback, cancellation_check=cancellation_check, occupied_environment_ids=occupied_environment_ids, )