"""Provider-neutral execution planning, capacity acquisition, and assignment metadata.""" from __future__ import annotations import hashlib import json import logging import math import threading import time from collections import defaultdict from concurrent.futures import ThreadPoolExecutor from dataclasses import replace from pathlib import Path from typing import Any, Callable, Iterator, Mapping, Sequence if __package__: from .execution_environments import ( ComponentResourceRequirements, CostAwareEnvironmentScheduler, EnvironmentCapabilities, EnvironmentHealth, EnvironmentSchedulingState, ExecutionAssignment, ExecutionPolicy, ExecutionProvider, GpuCapability, NoCompatibleExecutionEnvironmentError, WorkflowExecutionPreferences, ) from .execution_history import ExecutionHistory from .llm_profiles import get_llm_profile from .llm_resolver import resolve_model_profile from .modal_hardware import ( _MODAL_GPU_COST_USD_PER_SECOND, _MODAL_GPU_VRAM_GB, _capabilities_hardware_payload, _modal_hardware_payload, _vast_hardware_payload, ) from .r2_cache import R2CacheClient, R2CacheError, R2StorageUsage from .r2_credentials import R2CredentialError, R2CredentialStore from .remote_configuration_nodes import compile_remote_configuration_set from .remote_configurations import ( ModalRemoteConfiguration, RemoteConfiguration, RemoteConfigurationSet, R2StorageBackingConfiguration, SshRemoteConfiguration, VastRemoteConfiguration, ) from .remote_graph_analysis import _component_execution_stages, _is_link from .remote_hosts import RemoteHostRegistry, SshHostConfig from .remote_plan_types import ( ComponentExecutionPlan, ComponentMemoryEstimate, ModalPromptValidationError, RemoteComponentPlan, ) from .settings import ( ModalSyncSettings, discover_comfyui_user_directory, settings_for_modal_gpu, ) from .ssh_docker import SshDockerController, SshDockerVolumeBackend from .ssh_runtime import SshRuntimeManager from .sync_engine import ( MODEL_FILE_EXTENSIONS, ModalAssetSyncEngine, ModalVolumeBackend, SyncCancelledError, resolve_model_path, ) from .vast_config_node import extract_vast_profiles from .vast_service import VastProfileQuote, VastSearchRequirements, VastService else: # pragma: no cover - flat import inside the Modal container. from execution_environments import ( ComponentResourceRequirements, CostAwareEnvironmentScheduler, EnvironmentCapabilities, EnvironmentHealth, EnvironmentSchedulingState, ExecutionAssignment, ExecutionPolicy, ExecutionProvider, GpuCapability, NoCompatibleExecutionEnvironmentError, WorkflowExecutionPreferences, ) from execution_history import ExecutionHistory from llm_profiles import get_llm_profile from llm_resolver import resolve_model_profile from modal_hardware import ( _MODAL_GPU_COST_USD_PER_SECOND, _MODAL_GPU_VRAM_GB, _capabilities_hardware_payload, _modal_hardware_payload, _vast_hardware_payload, ) from r2_cache import R2CacheClient, R2CacheError, R2StorageUsage from r2_credentials import R2CredentialError, R2CredentialStore from remote_configuration_nodes import compile_remote_configuration_set from remote_configurations import ( ModalRemoteConfiguration, RemoteConfiguration, RemoteConfigurationSet, R2StorageBackingConfiguration, SshRemoteConfiguration, VastRemoteConfiguration, ) from remote_graph_analysis import _component_execution_stages, _is_link from remote_hosts import RemoteHostRegistry, SshHostConfig from remote_plan_types import ( ComponentExecutionPlan, ComponentMemoryEstimate, ModalPromptValidationError, RemoteComponentPlan, ) from settings import ( ModalSyncSettings, discover_comfyui_user_directory, settings_for_modal_gpu, ) from ssh_docker import SshDockerController, SshDockerVolumeBackend from ssh_runtime import SshRuntimeManager from sync_engine import ( MODEL_FILE_EXTENSIONS, ModalAssetSyncEngine, ModalVolumeBackend, SyncCancelledError, resolve_model_path, ) from vast_config_node import extract_vast_profiles from vast_service import VastProfileQuote, VastSearchRequirements, 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 ] _R2_STORAGE_USAGE_CACHE_SECONDS = 5 * 60 _R2_STORAGE_USAGE_CACHE_LOCK = threading.Lock() _R2_STORAGE_USAGE_CACHE: dict[ tuple[str, str, str], tuple[float, R2StorageUsage] ] = {} _MODEL_VRAM_WEIGHT_MULTIPLIER = 1.20 _MODEL_VRAM_HEADROOM_BYTES = 4 * 1024**3 _MODEL_RAM_HEADROOM_BYTES = 4 * 1024**3 _ADDITIVE_MODEL_CLASS_TOKENS = ( "adapter", "controlnet", "ipadapter", "lora", ) def _modal_environment_state(settings: ModalSyncSettings) -> EnvironmentSchedulingState: """Return the scheduler-facing state of the selected Modal GPU target.""" modal_gpu = settings.modal_gpu vram_bytes = int(_MODAL_GPU_VRAM_GB.get(modal_gpu, 0.0) * 1024**3) capabilities = EnvironmentCapabilities( architecture="x86_64", operating_system="linux", cpu_count=1, total_ram_bytes=max(vram_bytes, 1), available_ram_bytes=None, available_disk_bytes=None, docker_version="modal-managed", docker_rootless=False, nvidia_container_runtime=True, gpus=( GpuCapability( uuid=f"modal-{modal_gpu.lower()}", name=modal_gpu, total_vram_bytes=vram_bytes, ), ), ) return EnvironmentSchedulingState( environment_id=f"modal:{modal_gpu}", provider=ExecutionProvider.MODAL, enabled=True, health=EnvironmentHealth.READY, cost_usd_per_second=_MODAL_GPU_COST_USD_PER_SECOND.get(modal_gpu), capabilities=capabilities, maximum_workers=settings.max_containers or 1, ) def _modal_configuration_environment_state( configuration: ModalRemoteConfiguration, settings: ModalSyncSettings, ) -> EnvironmentSchedulingState: """Return one scheduler state for a workflow-declared Modal capacity pool.""" selected_settings = settings_for_modal_gpu(settings, configuration.gpu_type) base_state = _modal_environment_state(selected_settings) return replace( base_state, environment_id=( f"modal:{configuration.configuration_id}:{configuration.gpu_type}" ), configuration_id=configuration.configuration_id, display_name=configuration.display_name, maximum_workers=configuration.instance_count, ) def _ssh_host_registry(settings: ModalSyncSettings) -> RemoteHostRegistry | None: """Return the persistent SSH host registry when a user directory exists.""" user_directory = discover_comfyui_user_directory(settings) if user_directory is None: return None return RemoteHostRegistry.for_user_directory(user_directory) def _configured_ssh_hosts(settings: ModalSyncSettings) -> tuple[SshHostConfig, ...]: """Return configured SSH hosts without probing during queue submission.""" registry = _ssh_host_registry(settings) if registry is None: return () return registry.load().hosts def _schedulable_ssh_hosts(settings: ModalSyncSettings) -> tuple[SshHostConfig, ...]: """Probe every enabled host immediately before cost-aware scheduling.""" hosts = _configured_ssh_hosts(settings) registry = _ssh_host_registry(settings) if registry is None: return hosts probe_hosts = [host for host in hosts if host.enabled and not host.draining] if not probe_hosts: return hosts with ThreadPoolExecutor(max_workers=min(8, len(probe_hosts))) as executor: refreshed = { host.environment_id: future.result() for host, future in ( (host, executor.submit(_refresh_ssh_host, host, registry)) for host in probe_hosts ) } return tuple(refreshed.get(host.environment_id, host) for host in hosts) def _refresh_ssh_host( host: SshHostConfig, registry: RemoteHostRegistry, ) -> SshHostConfig: """Probe and persist one host without aborting other candidates.""" try: capabilities = SshDockerController(host).probe_capabilities() return registry.update_probe_result( host.environment_id, capabilities=capabilities, health=EnvironmentHealth.READY, last_error=None, ) except (OSError, RuntimeError, ValueError) as exc: logger.warning( "Remote environment probe failed environment=%s: %s", host.environment_id, exc, ) try: return registry.update_probe_result( host.environment_id, capabilities=host.capabilities, health=EnvironmentHealth.UNAVAILABLE, last_error=str(exc), ) except (KeyError, OSError, ValueError): return replace( host, health=EnvironmentHealth.UNAVAILABLE, last_error=str(exc), ) def _prompt_llm_model_references(prompt: Mapping[str, Any]) -> tuple[str, ...]: """Return every fixed Modal LLM model reference in a prompt.""" references: set[str] = set() for prompt_node in prompt.values(): if not isinstance(prompt_node, Mapping): continue if str(prompt_node.get("class_type") or "") != "ModalLLM": continue inputs = prompt_node.get("inputs") if not isinstance(inputs, Mapping): continue model_reference = inputs.get("model_profile") if isinstance(model_reference, str) and model_reference.strip(): references.add(model_reference.strip()) return tuple(sorted(references)) def _resolve_prompt_llm_profiles( prompt: Mapping[str, Any], settings: ModalSyncSettings, ) -> dict[str, Any]: """Resolve LLM metadata before environment admission and cost ranking.""" storage_root = Path( getattr(settings, "local_storage_root", "/tmp/comfyui-modal-sync-storage") ) profiles: dict[str, Any] = {} for model_reference in _prompt_llm_model_references(prompt): try: profile = get_llm_profile(model_reference, storage_root=storage_root) except ValueError as profile_error: if model_reference.startswith("hf-"): raise ModalPromptValidationError(str(profile_error)) from profile_error try: profile = resolve_model_profile(model_reference, storage_root).profile except ValueError as resolution_error: raise ModalPromptValidationError( str(resolution_error) ) from resolution_error profiles[model_reference] = profile logger.info( "Resolved planner LLM profile model=%s profile=%s weights_gib=%.2f " "estimated_vram_gib=%.2f.", model_reference, profile.profile_id, profile.artifact_bytes / 1024**3, profile.estimated_vram_gb, ) return profiles def _component_profile_memory_estimate( component: RemoteComponentPlan, prompt: Mapping[str, Any], preferences: WorkflowExecutionPreferences, resolved_profiles: Mapping[str, Any], ) -> ComponentMemoryEstimate: """Estimate a component's RAM and VRAM floors from resolved LLM profiles.""" minimum_vram_bytes = preferences.minimum_vram_bytes profiles: dict[str, Any] = {} for node_id in component.node_ids: prompt_node = prompt.get(node_id) if not isinstance(prompt_node, Mapping): continue if str(prompt_node.get("class_type") or "") != "ModalLLM": continue inputs = prompt_node.get("inputs") if not isinstance(inputs, Mapping): continue model_profile = inputs.get("model_profile") if not isinstance(model_profile, str) or not model_profile.strip(): continue profile = resolved_profiles.get(model_profile.strip()) if profile is None: continue profiles[profile.profile_id] = profile minimum_vram_bytes = max( minimum_vram_bytes, int(max(0.0, profile.estimated_vram_gb) * 1024**3), ) artifact_sizes = [int(profile.artifact_bytes) for profile in profiles.values()] return ComponentMemoryEstimate( minimum_vram_bytes=minimum_vram_bytes, minimum_ram_bytes=(sum(artifact_sizes) + _MODEL_RAM_HEADROOM_BYTES) if artifact_sizes else 0, model_asset_count=len(artifact_sizes), largest_model_bytes=max(artifact_sizes, default=0), ) def _component_required_provider( component: RemoteComponentPlan, prompt: Mapping[str, Any], resolved_profiles: Mapping[str, Any], ) -> ExecutionProvider | None: """Require SSH when a component uses an SSH-only resident backend.""" for node_id in component.node_ids: prompt_node = prompt.get(node_id) if not isinstance(prompt_node, Mapping): continue if str(prompt_node.get("class_type") or "") != "ModalLLM": continue inputs = prompt_node.get("inputs") if not isinstance(inputs, Mapping): continue model_reference = inputs.get("model_profile") if not isinstance(model_reference, str): continue profile = resolved_profiles.get(model_reference.strip()) if ( profile is not None and getattr(profile, "backend", "") == "llama_cpp_server" ): return ExecutionProvider.SSH_DOCKER return None def _iter_prompt_string_values(value: object) -> Iterator[str]: """Yield nested prompt string values while ignoring graph links and scalars.""" if isinstance(value, str): yield value return if isinstance(value, Mapping): for child in value.values(): yield from _iter_prompt_string_values(child) return if isinstance(value, (list, tuple)) and not _is_link(value): for child in value: yield from _iter_prompt_string_values(child) def _is_additive_model_node(class_type: str) -> bool: """Return whether a loader contributes weights alongside a primary model.""" normalized = class_type.casefold() return any(token in normalized for token in _ADDITIVE_MODEL_CLASS_TOKENS) def _component_model_asset_sizes( component: RemoteComponentPlan, prompt: Mapping[str, Any], settings: ModalSyncSettings, ) -> tuple[tuple[int, ...], tuple[int, ...]]: """Return unique primary and additive model sizes referenced by a component.""" asset_roles: dict[Path, bool] = {} for node_id in component.node_ids: prompt_node = prompt.get(node_id) if not isinstance(prompt_node, Mapping): continue inputs = prompt_node.get("inputs") if not isinstance(inputs, Mapping): continue additive = _is_additive_model_node(str(prompt_node.get("class_type") or "")) for value in _iter_prompt_string_values(inputs): resolved_path = resolve_model_path( value, comfyui_root=getattr(settings, "comfyui_root", None), extensions=MODEL_FILE_EXTENSIONS, ) if resolved_path is None: continue asset_roles[resolved_path] = ( asset_roles.get(resolved_path, True) and additive ) primary_sizes: list[int] = [] additive_sizes: list[int] = [] for asset_path, additive in asset_roles.items(): try: size_bytes = asset_path.stat().st_size except OSError as exc: logger.warning( "Unable to inspect model size for placement path=%s: %s", asset_path, exc, ) continue if size_bytes <= 0: continue target = additive_sizes if additive else primary_sizes target.append(size_bytes) return tuple(primary_sizes), tuple(additive_sizes) def _component_memory_estimate( component: RemoteComponentPlan, prompt: Mapping[str, Any], preferences: WorkflowExecutionPreferences, settings: ModalSyncSettings, resolved_llm_profiles: Mapping[str, Any] | None = None, ) -> ComponentMemoryEstimate: """Infer conservative RAM and VRAM floors from resident model weight sizes.""" profile_estimate = _component_profile_memory_estimate( component, prompt, preferences, resolved_llm_profiles if resolved_llm_profiles is not None else _resolve_prompt_llm_profiles(prompt, settings), ) primary_sizes, additive_sizes = _component_model_asset_sizes( component, prompt, settings, ) all_sizes = primary_sizes + additive_sizes if not all_sizes: return profile_estimate primary_peak_bytes = max(primary_sizes or all_sizes) additive_bytes = sum(additive_sizes) if primary_sizes else 0 resident_weight_bytes = primary_peak_bytes + additive_bytes model_vram_bytes = ( math.ceil(resident_weight_bytes * _MODEL_VRAM_WEIGHT_MULTIPLIER) + _MODEL_VRAM_HEADROOM_BYTES ) return ComponentMemoryEstimate( minimum_vram_bytes=max( profile_estimate.minimum_vram_bytes, model_vram_bytes, ), minimum_ram_bytes=max( profile_estimate.minimum_ram_bytes, resident_weight_bytes + _MODEL_RAM_HEADROOM_BYTES, ), model_asset_count=len(all_sizes) + profile_estimate.model_asset_count, largest_model_bytes=max( max(all_sizes), profile_estimate.largest_model_bytes, ), ) def _component_execution_signature( component: RemoteComponentPlan, prompt: Mapping[str, Any], ) -> str: """Return a stable workload signature for runtime-history estimates.""" cost_shaping_input_names = frozenset( { "batch_size", "duration", "frames", "height", "max_new_tokens", "model_profile", "num_frames", "steps", "width", } ) nodes: list[dict[str, Any]] = [] for node_id in sorted(component.node_ids): prompt_node = prompt.get(node_id) if not isinstance(prompt_node, Mapping): continue inputs = prompt_node.get("inputs") input_mapping = inputs if isinstance(inputs, Mapping) else {} nodes.append( { "class_type": str(prompt_node.get("class_type") or ""), "inputs": { name: input_mapping[name] for name in sorted(cost_shaping_input_names & set(input_mapping)) if not _is_link(input_mapping[name]) }, } ) serialized = json.dumps(nodes, separators=(",", ":"), sort_keys=True).encode( "utf-8" ) return hashlib.sha256(serialized).hexdigest() def _execution_history(settings: ModalSyncSettings) -> ExecutionHistory | None: """Return local runtime history when the ComfyUI user directory is known.""" user_directory = discover_comfyui_user_directory(settings) if user_directory is None: return None return ExecutionHistory.for_user_directory(user_directory) def _maximum_capacity_state( environment: EnvironmentSchedulingState, ) -> EnvironmentSchedulingState: """Return an optimistic state used only to identify reclaimable SSH capacity.""" capabilities = environment.capabilities if environment.provider is not ExecutionProvider.SSH_DOCKER or capabilities is None: return environment return replace( environment, capabilities=replace( capabilities, available_ram_bytes=capabilities.total_ram_bytes, gpus=tuple( replace(gpu, free_vram_bytes=gpu.total_vram_bytes) for gpu in capabilities.gpus ), ), ) def _reprobe_reclaimed_ssh_host( host: SshHostConfig, settings: ModalSyncSettings, ) -> SshHostConfig: """Probe a reclaimed host and persist its new free-memory measurements.""" registry = _ssh_host_registry(settings) if registry is not None: return _refresh_ssh_host(host, registry) try: capabilities = SshDockerController(host).probe_capabilities() except (OSError, RuntimeError, ValueError) as exc: logger.warning( "Remote environment re-probe failed after worker reclaim " "environment=%s: %s", host.environment_id, exc, ) return replace( host, health=EnvironmentHealth.UNAVAILABLE, last_error=str(exc), ) return replace( host, capabilities=capabilities, health=EnvironmentHealth.READY, last_error=None, ) def _probe_workflow_ssh_configuration( configuration: SshRemoteConfiguration, ) -> SshHostConfig: """Probe one workflow-declared SSH host without mutating the global registry.""" host = configuration.host try: capabilities = SshDockerController(host).probe_capabilities() except (OSError, RuntimeError, ValueError) as exc: logger.warning( "Workflow SSH environment probe failed configuration=%s environment=%s: %s", configuration.display_name, host.environment_id, exc, ) return replace( host, capabilities=None, health=EnvironmentHealth.UNAVAILABLE, last_error=str(exc), ) return replace( host, capabilities=capabilities, health=EnvironmentHealth.READY, last_error=None, ) def _optional_scheduler_choice( scheduler: CostAwareEnvironmentScheduler, environments: list[EnvironmentSchedulingState], requirements: ComponentResourceRequirements, ) -> tuple[ ExecutionAssignment | None, NoCompatibleExecutionEnvironmentError | None, ]: """Return a scheduler choice and preserve a compatibility failure for fallback.""" try: return scheduler.choose(environments, requirements), None except NoCompatibleExecutionEnvironmentError as exc: return None, exc def _require_scheduler_choice( assignment: ExecutionAssignment | None, error: NoCompatibleExecutionEnvironmentError | None, ) -> ExecutionAssignment: """Return an existing scheduler choice or raise its compatibility detail.""" if assignment is not None: return assignment if error is not None: raise error raise RuntimeError("Scheduler failed without an incompatibility reason.") def _remove_idle_ssh_workers_for_reclaim(host: SshHostConfig) -> tuple[str, ...]: """Recycle idle managed workers without disrupting fallback placement.""" try: return SshDockerController(host).remove_idle_managed_workers() except (OSError, RuntimeError, ValueError) as exc: logger.warning( "Could not reclaim managed SSH worker capacity environment=%s: %s", host.environment_id, exc, ) return () def _reclaim_improves_assignment( optimistic: ExecutionAssignment, actual: ExecutionAssignment | None, requirements: ComponentResourceRequirements, ) -> bool: """Return whether reclaim unlocks required or materially better placement.""" if actual is None: return True if optimistic.environment_id == actual.environment_id: return False preferred = requirements.preferred_environment_ids if optimistic.environment_id in preferred: if actual.environment_id not in preferred: return True return preferred.index(optimistic.environment_id) < preferred.index( actual.environment_id ) if optimistic.predicted_cost_usd is None: return False if actual.predicted_cost_usd is None: return True return optimistic.predicted_cost_usd < actual.predicted_cost_usd def _choose_with_idle_ssh_worker_reclaim( *, scheduler: CostAwareEnvironmentScheduler, environments: list[EnvironmentSchedulingState], requirements: ComponentResourceRequirements, ssh_hosts_by_id: dict[str, SshHostConfig], settings: ModalSyncSettings, ) -> ExecutionAssignment: """Choose an environment, recycling a cheaper idle SSH worker when necessary.""" actual_assignment, actual_error = _optional_scheduler_choice( scheduler, environments, requirements, ) optimistic_assignment, _ = _optional_scheduler_choice( scheduler, [_maximum_capacity_state(environment) for environment in environments], requirements, ) if optimistic_assignment is None: return _require_scheduler_choice(actual_assignment, actual_error) if ( optimistic_assignment.provider is not ExecutionProvider.SSH_DOCKER or not _reclaim_improves_assignment( optimistic_assignment, actual_assignment, requirements, ) ): return _require_scheduler_choice(actual_assignment, actual_error) host = ssh_hosts_by_id.get(optimistic_assignment.environment_id) if host is None: return _require_scheduler_choice(actual_assignment, actual_error) removed_workers = _remove_idle_ssh_workers_for_reclaim(host) if not removed_workers: return _require_scheduler_choice(actual_assignment, actual_error) logger.info( "Recycled idle managed SSH worker(s) environment=%s containers=%s so the " "planner can re-evaluate full host capacity.", host.environment_id, removed_workers, ) refreshed_host = _reprobe_reclaimed_ssh_host(host, settings) ssh_hosts_by_id[host.environment_id] = refreshed_host refreshed_environments = [ ( refreshed_host.scheduling_state() if environment.environment_id == host.environment_id else environment ) for environment in environments ] return scheduler.choose(refreshed_environments, requirements) def _plan_component_execution_assignments( *, components: list[RemoteComponentPlan], prompt: Mapping[str, Any], workflow: Mapping[str, Any] | None, settings: ModalSyncSettings, status_callback: SetupStatusCallback | None = None, resolved_llm_profiles: Mapping[str, Any] | None = None, ) -> dict[str, ExecutionAssignment]: """Assign components across Modal, SSH Docker, and workflow-declared Vast pools.""" preferences = WorkflowExecutionPreferences.from_workflow(workflow) active_llm_profiles = ( dict(resolved_llm_profiles) if resolved_llm_profiles is not None else _resolve_prompt_llm_profiles(prompt, settings) ) if preferences.policy is ExecutionPolicy.MODAL: incompatible_component_ids = [ component.representative_node_id for component in components if _component_required_provider( component, prompt, active_llm_profiles, ) is ExecutionProvider.SSH_DOCKER ] if incompatible_component_ids: raise ModalPromptValidationError( "Modal-only execution cannot run SSH-only component(s) " f"{incompatible_component_ids}. Select Automatic or Self-hosted " "execution for the workflow." ) modal_state = _modal_environment_state(settings) return { component.representative_node_id: ExecutionAssignment( environment_id=modal_state.environment_id, provider=ExecutionProvider.MODAL, predicted_cost_usd=None, predicted_completion_seconds=0.0, reasons=("workflow policy requires Modal",), ) for component in components } try: vast_profiles = extract_vast_profiles(prompt) except ValueError as exc: raise ModalPromptValidationError(str(exc)) from exc vast_service: VastService | None = None if preferences.policy in {ExecutionPolicy.VAST, ExecutionPolicy.AUTOMATIC}: if preferences.policy is ExecutionPolicy.VAST and not vast_profiles: raise ModalPromptValidationError( "Vast.ai-only execution requires at least one disconnected " "Vast.ai Lease Configuration node in the workflow." ) if vast_profiles: try: vast_service = VastService.from_environment( settings, repo_root=Path(__file__).resolve().parent, ) except (OSError, RuntimeError, ValueError) as exc: if preferences.policy is ExecutionPolicy.VAST: raise ModalPromptValidationError(str(exc)) from exc logger.warning("Skipping Vast.ai automatic placement: %s", exc) ssh_hosts = ( _schedulable_ssh_hosts(settings) if preferences.policy in {ExecutionPolicy.SELF_HOSTED, ExecutionPolicy.AUTOMATIC} else () ) ssh_hosts_by_id = {host.environment_id: host for host in ssh_hosts} modal_state = ( _modal_environment_state(settings) if preferences.policy is ExecutionPolicy.AUTOMATIC else None ) if not ssh_hosts and modal_state is None and vast_service is None: raise ModalPromptValidationError( "This workflow requests remote execution, but no compatible provider is configured." ) scheduler = CostAwareEnvironmentScheduler() history = _execution_history(settings) component_signatures: dict[str, str] = {} component_memory_estimates: dict[str, ComponentMemoryEstimate] = {} component_required_providers: dict[str, ExecutionProvider | None] = {} for component in components: component_id = component.representative_node_id component_signatures[component_id] = _component_execution_signature( component, prompt, ) memory_estimate = _component_memory_estimate( component, prompt, preferences, settings, active_llm_profiles, ) component_memory_estimates[component_id] = memory_estimate component_required_providers[component_id] = _component_required_provider( component, prompt, active_llm_profiles, ) if memory_estimate.model_asset_count: logger.info( "Estimated remote component=%s memory floor vram_gib=%.2f ram_gib=%.2f " "from model_assets=%d largest_model_gib=%.2f.", component_id, memory_estimate.minimum_vram_bytes / 1024**3, memory_estimate.minimum_ram_bytes / 1024**3, memory_estimate.model_asset_count, memory_estimate.largest_model_bytes / 1024**3, ) if vast_service is not None: try: vast_service.prefetch_offers_sync( vast_profiles, tuple( VastSearchRequirements( minimum_vram_bytes=memory_estimate.minimum_vram_bytes, minimum_ram_bytes=memory_estimate.minimum_ram_bytes, ) for memory_estimate in component_memory_estimates.values() ), ) except (OSError, RuntimeError, ValueError) as exc: if preferences.policy is ExecutionPolicy.VAST: raise ModalPromptValidationError(str(exc)) from exc logger.warning( "Unable to prefetch Vast.ai marketplace offers: %s", exc, ) assignments: dict[str, ExecutionAssignment] = {} for component in components: component_id = component.representative_node_id environments = [ host.scheduling_state() for host in ssh_hosts_by_id.values() ] if modal_state is not None: environments.append(modal_state) component_signature = component_signatures[component_id] memory_estimate = component_memory_estimates[component_id] vast_quote: VastProfileQuote | None = None if vast_service is not None: try: vast_quote = vast_service.quote_best_profile_sync( vast_profiles, minimum_vram_bytes=memory_estimate.minimum_vram_bytes, minimum_ram_bytes=memory_estimate.minimum_ram_bytes, predicted_execution_seconds=60.0, ) except (OSError, RuntimeError, ValueError) as exc: if preferences.policy is ExecutionPolicy.VAST: raise ModalPromptValidationError(str(exc)) from exc logger.warning( "No Vast.ai automatic candidate for component=%s: %s", component.representative_node_id, exc, ) else: environments.append(vast_service.scheduling_state(vast_quote)) environment_ids = [ environment.environment_id for environment in environments ] historical_estimates = ( history.estimates(component_signature, environment_ids) if history is not None else {} ) component_required_provider = component_required_providers[component_id] if preferences.policy is ExecutionPolicy.VAST: if component_required_provider is ExecutionProvider.SSH_DOCKER: raise ModalPromptValidationError( "Vast.ai-only execution cannot run an SSH-only component. " "Select Automatic or Self-hosted execution for this workflow." ) component_required_provider = ExecutionProvider.VAST requirements = ComponentResourceRequirements( minimum_vram_bytes=memory_estimate.minimum_vram_bytes, minimum_ram_bytes=memory_estimate.minimum_ram_bytes, gpu_required=True, architecture="x86_64", estimated_execution_seconds=60.0, estimated_execution_seconds_by_environment={ environment_id: estimate.execution_seconds for environment_id, estimate in historical_estimates.items() }, preferred_environment_ids=preferences.preferred_environment_ids, required_provider=component_required_provider, ) try: assignment = _choose_with_idle_ssh_worker_reclaim( scheduler=scheduler, environments=environments, requirements=requirements, ssh_hosts_by_id=ssh_hosts_by_id, settings=settings, ) except NoCompatibleExecutionEnvironmentError as exc: raise ModalPromptValidationError(str(exc)) from exc if assignment.provider is ExecutionProvider.VAST: if vast_service is None or vast_quote is None: raise ModalPromptValidationError( "Vast.ai was selected without a current marketplace quote." ) try: if status_callback is None: lease = vast_service.acquire_sync(vast_quote) else: lease = vast_service.acquire_sync( vast_quote, status_callback=lambda message: status_callback( message, None, None, ), ) except (OSError, RuntimeError, ValueError) as exc: raise ModalPromptValidationError( f"Unable to acquire Vast.ai capacity: {exc}" ) from exc assignment = replace( assignment, environment_id=lease.environment_id, predicted_cost_usd=vast_quote.predicted_incremental_cost_usd, reasons=assignment.reasons + ( f"Vast profile {vast_quote.profile.profile_name}", f"idle retention {lease.idle_retention_seconds / 3600:.1f}h", ), ) assignments[component.representative_node_id] = assignment logger.info( "Assigned remote component=%s provider=%s environment=%s predicted_cost_usd=%s reasons=%s.", component.representative_node_id, assignment.provider.value, assignment.environment_id, assignment.predicted_cost_usd, assignment.reasons, ) return assignments def _plan_component_execution( *, components: list[RemoteComponentPlan], prompt: Mapping[str, Any], workflow: Mapping[str, Any] | None, settings: ModalSyncSettings, status_callback: SetupStatusCallback | None = None, environment_status_callback: EnvironmentSetupStatusCallback | None = None, plan_callback: ExecutionPlanStatusCallback | None = None, occupied_environment_ids: frozenset[str] = frozenset(), ) -> ComponentExecutionPlan: """Plan through connected configurations or synthesize the legacy behavior.""" try: configuration_set = compile_remote_configuration_set(prompt) except (TypeError, ValueError) as exc: raise ModalPromptValidationError(str(exc)) from exc resolved_llm_profiles = _resolve_prompt_llm_profiles(prompt, settings) if configuration_set is None: return ComponentExecutionPlan( assignments=_plan_component_execution_assignments( components=components, prompt=prompt, workflow=workflow, settings=settings, status_callback=status_callback, resolved_llm_profiles=resolved_llm_profiles, ), resolved_llm_profiles=resolved_llm_profiles, ) return _plan_configured_component_execution( components=components, prompt=prompt, workflow=workflow, settings=settings, configuration_set=configuration_set, status_callback=status_callback, environment_status_callback=environment_status_callback, plan_callback=plan_callback, occupied_environment_ids=occupied_environment_ids, resolved_llm_profiles=resolved_llm_profiles, ) def _plan_configured_component_execution( *, components: list[RemoteComponentPlan], prompt: Mapping[str, Any], workflow: Mapping[str, Any] | None, settings: ModalSyncSettings, configuration_set: RemoteConfigurationSet, status_callback: SetupStatusCallback | None = None, environment_status_callback: EnvironmentSetupStatusCallback | None = None, plan_callback: ExecutionPlanStatusCallback | None = None, occupied_environment_ids: frozenset[str] = frozenset(), resolved_llm_profiles: Mapping[str, Any] | None = None, ) -> ComponentExecutionPlan: """Resolve and prepare a capacity-aware plan from connected configurations.""" configurations_by_id = { configuration.configuration_id: configuration for configuration in configuration_set.capacity_configurations } ssh_hosts_by_id = _probe_configured_ssh_hosts(configuration_set) vast_service, vast_unavailable_reason = _configured_vast_service( configuration_set, settings, ) requirements_by_component = _configured_component_requirements( components=components, prompt=prompt, workflow=workflow, settings=settings, resolved_llm_profiles=resolved_llm_profiles, ) _prefetch_configured_vast_offers( configuration_set=configuration_set, requirements_by_component=requirements_by_component, vast_service=vast_service, ) environments_by_component, vast_quotes = _configured_candidate_environments( configuration_set=configuration_set, requirements_by_component=requirements_by_component, settings=settings, ssh_hosts_by_id=ssh_hosts_by_id, vast_service=vast_service, vast_unavailable_reason=vast_unavailable_reason, occupied_environment_ids=occupied_environment_ids, ) if _reclaim_idle_configured_ssh_capacity( configuration_set=configuration_set, requirements_by_component=requirements_by_component, ssh_hosts_by_id=ssh_hosts_by_id, occupied_environment_ids=occupied_environment_ids, ): environments_by_component, vast_quotes = _configured_candidate_environments( configuration_set=configuration_set, requirements_by_component=requirements_by_component, settings=settings, ssh_hosts_by_id=ssh_hosts_by_id, vast_service=vast_service, vast_unavailable_reason=vast_unavailable_reason, occupied_environment_ids=occupied_environment_ids, ) _apply_historical_execution_estimates( components=components, environments_by_component=environments_by_component, requirements_by_component=requirements_by_component, prompt=prompt, settings=settings, ) execution_stages = _component_execution_stages( dict(prompt), { component.representative_node_id: set(component.node_ids) for component in components }, ) try: assignments = CostAwareEnvironmentScheduler().plan( execution_stages=execution_stages, environments_by_component=environments_by_component, requirements_by_component=requirements_by_component, ) except NoCompatibleExecutionEnvironmentError as exc: raise ModalPromptValidationError(str(exc)) from exc for component_id, assignment in tuple(assignments.items()): if assignment.environment_id not in occupied_environment_ids: continue assignments[component_id] = replace( assignment, reasons=assignment.reasons + ("capacity available after the earlier workflow completes",), ) safe_configurations = _safe_remote_configuration_payload(configuration_set) if plan_callback is not None: plan_callback( _planned_execution_assignments_payload( assignments, components, configurations_by_id=configurations_by_id, ssh_hosts_by_id=ssh_hosts_by_id, vast_quotes=vast_quotes, ), safe_configurations, ) vast_leases = _prepare_selected_vast_capacity( assignments=assignments, configuration_set=configuration_set, requirements_by_component=requirements_by_component, vast_quotes=vast_quotes, vast_service=vast_service, status_callback=status_callback, environment_status_callback=environment_status_callback, ) if plan_callback is not None and vast_leases: plan_callback( _planned_execution_assignments_payload( assignments, components, configurations_by_id=configurations_by_id, ssh_hosts_by_id=ssh_hosts_by_id, vast_quotes=vast_quotes, vast_leases_by_environment=vast_leases, ), safe_configurations, ) return ComponentExecutionPlan( assignments=assignments, configuration_set=configuration_set, configurations_by_id=configurations_by_id, safe_configurations=safe_configurations, ssh_hosts_by_id=ssh_hosts_by_id, vast_service=vast_service, vast_leases_by_environment=vast_leases, resolved_llm_profiles=dict(resolved_llm_profiles or {}), ) def _safe_remote_configuration_payload( configuration_set: RemoteConfigurationSet, ) -> list[dict[str, Any]]: """Return safe configuration metadata enriched with best-effort storage usage.""" payload = configuration_set.to_safe_list() payload_by_id = { str(configuration.get("configuration_id") or ""): configuration for configuration in payload } for storage in configuration_set.storage_configurations: if not isinstance(storage, R2StorageBackingConfiguration): continue safe_storage = payload_by_id.get(storage.configuration_id) if safe_storage is None: continue usage, error_code = _cached_r2_storage_usage(storage) if error_code is not None: safe_storage["credential_error_code"] = error_code if usage is None: continue safe_storage["storage_usage_bytes"] = usage.size_bytes safe_storage["storage_object_count"] = usage.object_count return payload def _cached_r2_storage_usage( storage: R2StorageBackingConfiguration, ) -> tuple[R2StorageUsage | None, str | None]: """Return cached bucket usage and an optional credential recovery code.""" cache_key = (storage.account_id, storage.bucket, storage.jurisdiction) now = time.monotonic() with _R2_STORAGE_USAGE_CACHE_LOCK: cached = _R2_STORAGE_USAGE_CACHE.get(cache_key) if cached is not None and now - cached[0] < _R2_STORAGE_USAGE_CACHE_SECONDS: return cached[1], None try: usage = _refresh_r2_storage_usage(storage) except R2CredentialError as exc: logger.warning( "Unable to read R2 storage usage for configuration=%s bucket=%s: %s", storage.configuration_id, storage.bucket, exc, ) return None, exc.code except (R2CacheError, RuntimeError, ValueError) as exc: logger.warning( "Unable to read R2 storage usage for configuration=%s bucket=%s: %s", storage.configuration_id, storage.bucket, exc, ) return None, None return usage, None def _refresh_r2_storage_usage( storage: R2StorageBackingConfiguration, ) -> R2StorageUsage: """Query current R2 bucket usage and replace the short-lived cache entry.""" configuration = R2CredentialStore().cache_configuration(storage) usage = R2CacheClient(configuration).storage_usage() cache_key = (storage.account_id, storage.bucket, storage.jurisdiction) with _R2_STORAGE_USAGE_CACHE_LOCK: _R2_STORAGE_USAGE_CACHE[cache_key] = (time.monotonic(), usage) return usage def _r2_storage_from_usage_payload( payload: Mapping[str, Any], ) -> R2StorageBackingConfiguration: """Build a validated R2 reference from a same-origin usage refresh request.""" configuration_id = str( payload.get("configuration_id") or "r2-storage-refresh" ).strip() return R2StorageBackingConfiguration( configuration_id=configuration_id, display_name=str(payload.get("display_name") or "R2 storage").strip(), account_id=str(payload.get("account_id") or "").strip(), bucket=str(payload.get("bucket") or "").strip(), credential_id=str(payload.get("credential_id") or "").strip(), jurisdiction=str(payload.get("jurisdiction") or "default") .strip() .casefold(), key_prefix=str( payload.get("key_prefix") or "comfy-modal-cache/v1/blobs/sha256" ).strip(), write_back_mode=str(payload.get("write_back_mode") or "async") .strip() .casefold(), ) def _planned_execution_assignments_payload( assignments: Mapping[str, ExecutionAssignment], components: Sequence[RemoteComponentPlan], *, configurations_by_id: Mapping[str, Any] | None = None, ssh_hosts_by_id: Mapping[str, SshHostConfig] | None = None, vast_quotes: Mapping[tuple[str, str], VastProfileQuote] | None = None, vast_leases_by_environment: Mapping[str, Any] | None = None, ) -> dict[str, dict[str, Any]]: """Serialize scheduler choices before provider capacity is acquired.""" component_nodes = { component.representative_node_id: list(component.node_ids) for component in components } return { component_id: { "provider": assignment.provider.value, "environment_id": assignment.environment_id, "configuration_id": assignment.configuration_id, "node_ids": component_nodes.get(component_id, [component_id]), "predicted_cost_usd": assignment.predicted_cost_usd, "predicted_completion_seconds": assignment.predicted_completion_seconds, "worker_index": assignment.capacity_slot_index, "reasons": list(assignment.reasons), "hardware": _assignment_hardware_payload( component_id=component_id, assignment=assignment, configurations_by_id=configurations_by_id or {}, ssh_hosts_by_id=ssh_hosts_by_id or {}, vast_quotes=vast_quotes or {}, vast_leases_by_environment=vast_leases_by_environment or {}, ), } for component_id, assignment in sorted(assignments.items()) } def _configuration_field(configuration: Any, field_name: str) -> Any: """Read one field from a configuration model or safe mapping.""" if isinstance(configuration, Mapping): return configuration.get(field_name) return getattr(configuration, field_name, None) def _configuration_host(configuration: Any) -> SshHostConfig | None: """Return a workflow SSH host from a model or safe configuration mapping.""" host = _configuration_field(configuration, "host") if isinstance(host, SshHostConfig): return host if isinstance(host, Mapping): try: return SshHostConfig.from_dict(host) except (TypeError, ValueError): return None return None def _assignment_hardware_payload( *, component_id: str, assignment: ExecutionAssignment, configurations_by_id: Mapping[str, Any], ssh_hosts_by_id: Mapping[str, SshHostConfig], vast_quotes: Mapping[tuple[str, str], VastProfileQuote], vast_leases_by_environment: Mapping[str, Any], ) -> dict[str, Any] | None: """Return the best hardware identity known for one planned assignment.""" configuration_id = str(assignment.configuration_id or "") configuration = configurations_by_id.get(configuration_id) if assignment.provider is ExecutionProvider.MODAL: gpu_type = str( _configuration_field(configuration, "gpu_type") or assignment.environment_id.rsplit(":", 1)[-1] ) return _modal_hardware_payload(gpu_type) if assignment.provider is ExecutionProvider.SSH_DOCKER: host = ssh_hosts_by_id.get(configuration_id) or _configuration_host( configuration ) return _capabilities_hardware_payload( host.capabilities if host is not None else None ) lease = vast_leases_by_environment.get(assignment.environment_id) if lease is not None: return _vast_hardware_payload(lease) quote = vast_quotes.get((component_id, configuration_id)) if quote is None: return None resource = getattr(quote, "existing_lease", None) or getattr( quote, "offer", None, ) return _vast_hardware_payload(resource) if resource is not None else None def _probe_configured_ssh_hosts( configuration_set: RemoteConfigurationSet, ) -> dict[str, SshHostConfig]: """Probe all workflow-declared SSH hosts concurrently.""" ssh_configurations = [ configuration for configuration in configuration_set.configurations if isinstance(configuration, SshRemoteConfiguration) ] if not ssh_configurations: return {} with ThreadPoolExecutor(max_workers=min(8, len(ssh_configurations))) as executor: probed_hosts = { configuration.configuration_id: future.result() for configuration, future in ( ( configuration, executor.submit( _probe_workflow_ssh_configuration, configuration, ), ) for configuration in ssh_configurations ) } return probed_hosts def _configured_vast_service( configuration_set: RemoteConfigurationSet, settings: ModalSyncSettings, ) -> tuple[VastService | None, str | None]: """Construct one Vast service only when the connected graph requests it.""" vast_configured = any( isinstance(configuration, VastRemoteConfiguration) for configuration in configuration_set.capacity_configurations ) if not vast_configured: return None, None try: return ( VastService.from_environment( settings, repo_root=Path(__file__).resolve().parent, ), None, ) except (OSError, RuntimeError, ValueError) as exc: has_fallback_provider = any( not isinstance(configuration, VastRemoteConfiguration) for configuration in configuration_set.capacity_configurations ) if not has_fallback_provider: raise ModalPromptValidationError(str(exc)) from exc logger.warning( "Skipping unavailable workflow Vast.ai capacity because another " "configured provider can be planned: %s", exc, ) return None, str(exc) def _configured_component_requirements( *, components: list[RemoteComponentPlan], prompt: Mapping[str, Any], workflow: Mapping[str, Any] | None, settings: ModalSyncSettings, resolved_llm_profiles: Mapping[str, Any] | None = None, ) -> dict[str, ComponentResourceRequirements]: """Build provider-neutral requirements for every remote component.""" preferences = WorkflowExecutionPreferences.from_workflow(workflow) active_llm_profiles = ( dict(resolved_llm_profiles) if resolved_llm_profiles is not None else _resolve_prompt_llm_profiles(prompt, settings) ) requirements: dict[str, ComponentResourceRequirements] = {} for component in components: memory_estimate = _component_memory_estimate( component, prompt, preferences, settings, active_llm_profiles, ) component_id = component.representative_node_id requirements[component_id] = ComponentResourceRequirements( minimum_vram_bytes=memory_estimate.minimum_vram_bytes, minimum_ram_bytes=memory_estimate.minimum_ram_bytes, gpu_required=True, architecture="x86_64", estimated_execution_seconds=60.0, preferred_environment_ids=preferences.preferred_environment_ids, required_provider=_component_required_provider( component, prompt, active_llm_profiles, ), ) return requirements def _configured_candidate_environments( *, configuration_set: RemoteConfigurationSet, requirements_by_component: Mapping[str, ComponentResourceRequirements], settings: ModalSyncSettings, ssh_hosts_by_id: Mapping[str, SshHostConfig], vast_service: VastService | None, vast_unavailable_reason: str | None, occupied_environment_ids: frozenset[str] = frozenset(), ) -> tuple[ dict[str, list[EnvironmentSchedulingState]], dict[tuple[str, str], VastProfileQuote], ]: """Resolve candidate pool states for every component without acquiring capacity.""" candidates: dict[str, list[EnvironmentSchedulingState]] = { component_id: [] for component_id in requirements_by_component } vast_quotes: dict[tuple[str, str], VastProfileQuote] = {} for component_id, requirements in requirements_by_component.items(): for configuration in configuration_set.capacity_configurations: state, quote = _configured_candidate_environment( configuration=configuration, requirements=requirements, settings=settings, ssh_hosts_by_id=ssh_hosts_by_id, vast_service=vast_service, vast_unavailable_reason=vast_unavailable_reason, occupied_environment_ids=occupied_environment_ids, ) if state is not None: candidates[component_id].append(state) if quote is not None: vast_quotes[(component_id, configuration.configuration_id)] = quote return candidates, vast_quotes def _reclaim_idle_configured_ssh_capacity( *, configuration_set: RemoteConfigurationSet, requirements_by_component: Mapping[str, ComponentResourceRequirements], ssh_hosts_by_id: dict[str, SshHostConfig], occupied_environment_ids: frozenset[str], ) -> bool: """Recycle idle configured workers when only their cached memory blocks work.""" scheduler = CostAwareEnvironmentScheduler() reclaimed = False for configuration in configuration_set.capacity_configurations: if not isinstance(configuration, SshRemoteConfiguration): continue host = ssh_hosts_by_id[configuration.configuration_id] if host.environment_id in occupied_environment_ids: continue state = replace( host.scheduling_state(), configuration_id=configuration.configuration_id, display_name=configuration.display_name, ) optimistic_state = _maximum_capacity_state(state) reclaimable = any( _optional_scheduler_choice(scheduler, [state], requirements)[0] is None and _optional_scheduler_choice( scheduler, [optimistic_state], requirements, )[0] is not None for requirements in requirements_by_component.values() ) if not reclaimable: continue removed_workers = _remove_idle_ssh_workers_for_reclaim(host) if not removed_workers: continue refreshed_host = _probe_workflow_ssh_configuration(configuration) ssh_hosts_by_id[configuration.configuration_id] = refreshed_host reclaimed = True logger.info( "Recycled idle workflow-configured SSH worker(s) environment=%s " "containers=%s and refreshed available memory.", host.environment_id, removed_workers, ) return reclaimed def _prefetch_configured_vast_offers( *, configuration_set: RemoteConfigurationSet, requirements_by_component: Mapping[str, ComponentResourceRequirements], vast_service: VastService | None, ) -> None: """Warm distinct Vast searches in parallel before deterministic planning.""" if vast_service is None: return profiles = [ configuration.profile for configuration in configuration_set.configurations if isinstance(configuration, VastRemoteConfiguration) ] prefetch = getattr(vast_service, "prefetch_offers_sync", None) if not callable(prefetch): return try: prefetch( profiles, tuple( VastSearchRequirements( minimum_vram_bytes=requirements.minimum_vram_bytes, minimum_ram_bytes=requirements.minimum_ram_bytes, ) for requirements in requirements_by_component.values() ), ) except (OSError, RuntimeError, ValueError) as exc: logger.warning( "Unable to prefetch configured Vast.ai marketplace offers: %s", exc, ) def _configured_candidate_environment( *, configuration: RemoteConfiguration, requirements: ComponentResourceRequirements, settings: ModalSyncSettings, ssh_hosts_by_id: Mapping[str, SshHostConfig], vast_service: VastService | None, vast_unavailable_reason: str | None, occupied_environment_ids: frozenset[str] = frozenset(), ) -> tuple[EnvironmentSchedulingState | None, VastProfileQuote | None]: """Resolve one configuration as a candidate for one component.""" if isinstance(configuration, ModalRemoteConfiguration): return _modal_configuration_environment_state(configuration, settings), None if isinstance(configuration, SshRemoteConfiguration): host = ssh_hosts_by_id[configuration.configuration_id] state = replace( host.scheduling_state(), configuration_id=configuration.configuration_id, display_name=configuration.display_name, ) if state.environment_id in occupied_environment_ids: state = _maximum_capacity_state(state) logger.info( "Planning queued work against reclaimable SSH capacity " "environment=%s because an earlier workflow owns that host.", state.environment_id, ) return state, None if not isinstance(configuration, VastRemoteConfiguration): raise TypeError( f"Unsupported remote configuration type {type(configuration).__name__}." ) if vast_service is None: return EnvironmentSchedulingState( environment_id=f"vast:{configuration.configuration_id}", provider=ExecutionProvider.VAST, enabled=True, health=EnvironmentHealth.UNAVAILABLE, cost_usd_per_second=None, capabilities=None, configuration_id=configuration.configuration_id, display_name=configuration.display_name, unavailable_reason=( vast_unavailable_reason or "Vast.ai controller is unavailable." ), maximum_workers=configuration.capacity_limit, ), None try: quote = vast_service.quote_best_profile_sync( [configuration.profile], minimum_vram_bytes=requirements.minimum_vram_bytes, minimum_ram_bytes=requirements.minimum_ram_bytes, predicted_execution_seconds=requirements.estimated_execution_seconds, ) except (OSError, RuntimeError, ValueError) as exc: logger.warning( "Vast configuration has no candidate configuration=%s: %s", configuration.display_name, exc, ) return EnvironmentSchedulingState( environment_id=f"vast:{configuration.configuration_id}", provider=ExecutionProvider.VAST, enabled=True, health=EnvironmentHealth.UNAVAILABLE, cost_usd_per_second=None, capabilities=None, configuration_id=configuration.configuration_id, display_name=configuration.display_name, unavailable_reason=str(exc), maximum_workers=configuration.capacity_limit, ), None state = vast_service.scheduling_state(quote) return replace( state, environment_id=f"vast:{configuration.configuration_id}", configuration_id=configuration.configuration_id, display_name=configuration.display_name, maximum_workers=configuration.capacity_limit, ), quote def _apply_historical_execution_estimates( *, components: list[RemoteComponentPlan], environments_by_component: Mapping[str, list[EnvironmentSchedulingState]], requirements_by_component: dict[str, ComponentResourceRequirements], prompt: Mapping[str, Any], settings: ModalSyncSettings, ) -> None: """Add environment-specific historical timing estimates in place.""" history = _execution_history(settings) if history is None: return for component in components: component_id = component.representative_node_id environment_ids = [ environment.environment_id for environment in environments_by_component.get(component_id, []) ] estimates = history.estimates( _component_execution_signature(component, prompt), environment_ids, ) requirements_by_component[component_id] = replace( requirements_by_component[component_id], estimated_execution_seconds_by_environment={ environment_id: estimate.execution_seconds for environment_id, estimate in estimates.items() }, ) def _prepare_selected_vast_capacity( *, assignments: dict[str, ExecutionAssignment], configuration_set: RemoteConfigurationSet, requirements_by_component: Mapping[str, ComponentResourceRequirements], vast_quotes: Mapping[tuple[str, str], VastProfileQuote], vast_service: VastService | None, status_callback: SetupStatusCallback | None = None, environment_status_callback: EnvironmentSetupStatusCallback | None = None, ) -> dict[str, Any]: """Acquire only Vast slots selected by the completed provider-neutral plan.""" selected_slots: dict[tuple[str, int], list[str]] = defaultdict(list) for component_id, assignment in assignments.items(): if assignment.provider is not ExecutionProvider.VAST: continue if assignment.configuration_id is None: raise ModalPromptValidationError( "Vast assignment is missing its configuration identity." ) selected_slots[ (assignment.configuration_id, assignment.capacity_slot_index) ].append(component_id) if not selected_slots: return {} if vast_service is None: raise ModalPromptValidationError( "Vast capacity was selected without an initialized Vast service." ) configurations_by_id = { configuration.configuration_id: configuration for configuration in configuration_set.configurations if isinstance(configuration, VastRemoteConfiguration) } leases_by_environment: dict[str, Any] = {} ordered_slots = sorted(selected_slots.items()) total_slots = len(ordered_slots) for slot_number, ((configuration_id, slot_index), component_ids) in enumerate( ordered_slots, start=1, ): configuration = configurations_by_id[configuration_id] planned_environment_id = assignments[component_ids[0]].environment_id quote = _quote_selected_vast_slot( configuration=configuration, component_ids=component_ids, requirements_by_component=requirements_by_component, vast_quotes=vast_quotes, vast_service=vast_service, ) try: _emit_vast_capacity_status( environment_id=planned_environment_id, message=f"Acquiring Vast.ai capacity {slot_number} of {total_slots}", current=slot_number - 1, total=total_slots, status_callback=status_callback, environment_status_callback=environment_status_callback, ) if status_callback is None and environment_status_callback is None: lease = vast_service.acquire_sync(quote, slot=slot_index) else: lease = vast_service.acquire_sync( quote, slot=slot_index, status_callback=( lambda message, environment_id=planned_environment_id, current=slot_number - 1: _emit_vast_capacity_status( environment_id=environment_id, message=message, current=current, total=total_slots, status_callback=status_callback, environment_status_callback=environment_status_callback, ) ), ) except SyncCancelledError: raise except (OSError, RuntimeError, ValueError) as exc: raise ModalPromptValidationError( f"Unable to acquire Vast.ai capacity for configuration " f"{configuration.display_name!r} slot {slot_index}: {exc}" ) from exc leases_by_environment[lease.environment_id] = lease cost_share = quote.predicted_incremental_cost_usd / len(component_ids) for component_id in component_ids: assignment = assignments[component_id] assignments[component_id] = replace( assignment, environment_id=lease.environment_id, predicted_cost_usd=cost_share, reasons=assignment.reasons + ( f"Vast profile {configuration.display_name}", f"Vast capacity slot {slot_index}", ), ) _emit_vast_capacity_status( environment_id=lease.environment_id, message=f"Vast.ai capacity {slot_number} of {total_slots} is ready", environment_message="Vast.ai worker ready; preparing remote assets next", current=slot_number, total=total_slots, status_callback=status_callback, environment_status_callback=environment_status_callback, ) return leases_by_environment def _emit_vast_capacity_status( *, environment_id: str, message: str, environment_message: str | None = None, current: int, total: int, status_callback: SetupStatusCallback | None, environment_status_callback: EnvironmentSetupStatusCallback | None, ) -> None: """Publish one Vast setup update at prompt and environment scopes.""" if status_callback is not None: status_callback(message, current, total) if environment_status_callback is not None: environment_status_callback( environment_id, environment_message or message, None, None, ) def _quote_selected_vast_slot( *, configuration: VastRemoteConfiguration, component_ids: list[str], requirements_by_component: Mapping[str, ComponentResourceRequirements], vast_quotes: Mapping[tuple[str, str], VastProfileQuote], vast_service: VastService, ) -> VastProfileQuote: """Quote aggregate slot requirements immediately before acquisition.""" minimum_vram_bytes = max( requirements_by_component[component_id].minimum_vram_bytes for component_id in component_ids ) minimum_ram_bytes = max( requirements_by_component[component_id].minimum_ram_bytes for component_id in component_ids ) predicted_execution_seconds = sum( requirements_by_component[component_id].estimated_execution_seconds for component_id in component_ids ) if len(component_ids) == 1: existing_quote = vast_quotes.get( (component_ids[0], configuration.configuration_id) ) if existing_quote is not None: return existing_quote return vast_service.quote_best_profile_sync( [configuration.profile], minimum_vram_bytes=minimum_vram_bytes, minimum_ram_bytes=minimum_ram_bytes, predicted_execution_seconds=predicted_execution_seconds, ) def _ssh_sync_engine( *, host: SshHostConfig, settings: ModalSyncSettings, r2_cache: R2CacheClient | None = None, ) -> ModalAssetSyncEngine: """Build a content-addressed sync engine for one SSH Docker host.""" ssh_settings = replace( settings, execution_mode="ssh_docker", volume_name=host.resolved_storage_volume_name, local_storage_root=( settings.local_storage_root / "ssh" / host.environment_id ).resolve(), remote_storage_root="/storage", ) resolved_r2_cache = r2_cache or R2CacheClient.from_environment() controller = SshDockerController(host) runtime_manager = SshRuntimeManager( controller=controller, repo_root=Path(__file__).resolve().parent, settings=settings, ) materializer_spec = runtime_manager.runtime_spec() def prepare_materializer_image() -> None: """Make the SSH runtime image available before an R2 helper starts.""" runtime_manager.ensure_image(materializer_spec) volume = SshDockerVolumeBackend( controller, host.resolved_storage_volume_name, materializer_image=materializer_spec.image_tag, materializer_image_preparer=( prepare_materializer_image if resolved_r2_cache is not None else None ), ) return ModalAssetSyncEngine( volume=volume, settings=ssh_settings, r2_cache=resolved_r2_cache, ) def _workflow_r2_cache( configuration_set: RemoteConfigurationSet | None, ) -> R2CacheClient | None: """Resolve the connected R2 backing through its opaque OS-keyring reference.""" if configuration_set is None: return None r2_configurations = [ configuration for configuration in configuration_set.storage_configurations if isinstance(configuration, R2StorageBackingConfiguration) ] if not r2_configurations: return None storage = r2_configurations[0] try: configuration = R2CredentialStore().cache_configuration(storage) except (RuntimeError, ValueError) as exc: raise ModalPromptValidationError(str(exc)) from exc return R2CacheClient(configuration) def _stamp_execution_assignment( payload: dict[str, Any], assignment: ExecutionAssignment, worker_index: int = 0, execution_history_signature: str | None = None, execution_location: str | None = None, provider_metadata: Mapping[str, Any] | None = None, ) -> None: """Attach provider placement to a payload and every nested proxy phase.""" payload["execution_provider"] = assignment.provider.value payload["execution_environment_id"] = assignment.environment_id payload["execution_worker_index"] = worker_index if execution_location: payload["execution_location"] = execution_location if execution_history_signature: payload["execution_history_signature"] = execution_history_signature if provider_metadata: payload.update(provider_metadata) split_payloads = payload.get("split_proxy_payloads") nested_payloads: list[dict[str, Any]] = [] if isinstance(split_payloads, dict): nested_payloads.extend( nested_payload for nested_payload in split_payloads.values() if isinstance(nested_payload, dict) ) elif isinstance(split_payloads, list): nested_payloads.extend( nested_payload for nested_payload in split_payloads if isinstance(nested_payload, dict) ) for nested_payload in nested_payloads: _stamp_execution_assignment( nested_payload, assignment, worker_index, execution_history_signature, execution_location, provider_metadata, ) def _ssh_hostname(ssh_target: str) -> str: """Return the host portion of one validated OpenSSH destination.""" host = ssh_target.rsplit("@", 1)[-1].strip() if host.startswith("[") and "]" in host: return host[1 : host.index("]")] return host def _execution_location_for_assignment( assignment: ExecutionAssignment, ssh_hosts_by_id: Mapping[str, SshHostConfig], vast_leases_by_environment: Mapping[str, Any] | None = None, ) -> str | None: """Return the runtime location label known before remote dispatch.""" if assignment.provider is ExecutionProvider.MODAL: return None if assignment.provider is ExecutionProvider.VAST: lease = (vast_leases_by_environment or {}).get(assignment.environment_id) if lease is not None: return str(lease.ssh_host or lease.gpu_name) return assignment.environment_id host = ssh_hosts_by_id.get(assignment.environment_id) return _ssh_hostname(host.ssh_target) if host is not None else assignment.environment_id def _vast_provider_metadata(lease: Any) -> dict[str, Any]: """Return credential-free lease details for execution and local status output.""" return { "vast_instance_id": lease.instance_id, "vast_profile_id": lease.profile_id, "vast_profile_name": lease.profile_name, "vast_gpu_name": lease.gpu_name, "vast_gpu_count": lease.gpu_count, "vast_gpu_ram_mb": lease.gpu_ram_mb, "vast_hourly_cost_usd": lease.hourly_cost_usd, "vast_idle_retention_seconds": lease.idle_retention_seconds, } def _configured_provider_metadata( *, execution_plan: ComponentExecutionPlan, assignment: ExecutionAssignment, vast_leases_by_environment: Mapping[str, Any], ) -> dict[str, Any] | None: """Return safe provider metadata needed to execute one prepared assignment.""" if assignment.provider is ExecutionProvider.VAST: return _vast_provider_metadata( vast_leases_by_environment[assignment.environment_id] ) configuration_id = assignment.configuration_id if configuration_id is None: return None configuration = execution_plan.configurations_by_id.get(configuration_id) if isinstance(configuration, ModalRemoteConfiguration): return { "remote_configuration_id": configuration.configuration_id, "remote_configuration_name": configuration.display_name, "modal_gpu": configuration.gpu_type, "modal_max_containers": configuration.instance_count, } if isinstance(configuration, SshRemoteConfiguration): host = execution_plan.ssh_hosts_by_id[configuration.configuration_id] portable_host = replace( host, health=EnvironmentHealth.UNKNOWN, last_error=None, ) return { "remote_configuration_id": configuration.configuration_id, "remote_configuration_name": configuration.display_name, "ssh_host_config": portable_host.to_dict(), } return None def _ensure_remote_sync_backend( settings: ModalSyncSettings, sync_engine: ModalAssetSyncEngine, ) -> None: """Fail before queueing when remote execution cannot write to Modal-visible storage.""" if settings.execution_mode == "local": return if isinstance(sync_engine.volume, ModalVolumeBackend): return raise ModalPromptValidationError( "Remote Modal execution requires asset sync to use the Modal volume backend, " f"but the active sync backend is {type(sync_engine.volume).__name__}. " "Restart ComfyUI with COMFY_MODAL_EXECUTION_MODE=remote and the Modal SDK available " "so synced assets and custom_nodes bundles are visible inside Modal workers." )