From 9a92fe2ecb138a76fca16f3fafe70e2770cef70e Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Sat, 22 Aug 2026 20:10:22 -0700 Subject: [PATCH] Add provider-aware workflow partitioning --- README.md | 6 +- api_intercept.py | 339 ++++++++++++++++++++++-------- docs/LEARNINGS.md | 1 + modal_llm_node.py | 6 +- modal_llm_runtime.py | 10 + remote/modal_app.py | 40 +++- ssh_runtime.py | 1 + tests/test_api_intercept.py | 236 ++++++++++++++++++++- tests/test_modal_executor_node.py | 75 +++++-- tests/test_modal_llm.py | 16 ++ tests/test_ssh_runtime.py | 1 + tests/test_web_modal_toggle.py | 12 +- web/modal_toggle.js | 43 +++- 13 files changed, 661 insertions(+), 125 deletions(-) diff --git a/README.md b/README.md index 7aa6c6f..ca55833 100644 --- a/README.md +++ b/README.md @@ -110,7 +110,7 @@ Workflow policy is stored under `workflow.extra.remote_execution` and travels wi - **Modal only** preserves existing behavior and is the default for old workflows. - **Self-hosted only** requires a ready, compatible SSH host and fails before queueing when none is available. -- **Automatic** considers ready SSH hosts and the selected Modal GPU, rejects hard incompatibilities such as insufficient VRAM, then prefers known lowest predicted cost with deterministic tie-breaking. +- **Automatic** considers ready SSH hosts and the selected Modal GPU, rejects hard incompatibilities such as insufficient VRAM or a provider-specific backend, then prefers known lowest predicted cost with deterministic tie-breaking. Provider-constrained nodes form their own acyclic components, and cross-provider outputs are materialized through ComfyUI instead of being hidden in one provider's session storage. - **Automatically place eligible nodes** makes the planner select the executable prompt without requiring each node toggle. Leave it off to keep manual component boundaries. Before comparing cost, the planner resolves model files referenced by each remote component and derives conservative memory floors. It budgets the largest primary model that may be resident, adds LoRA/adapter/ControlNet weights, applies 20% weight overhead, and reserves another 4 GiB for activations and runtime state. RAM receives the resident weights plus 4 GiB. Hosts are checked against currently free VRAM and available RAM when the capability probe reports them, falling back to total capacity only when live availability is unknown. An undersized zero-cost host is therefore rejected before any model upload; automatic policy can fall back to a compatible Modal GPU, while self-hosted-only policy returns the per-host memory shortfall. @@ -119,7 +119,7 @@ Images are immutable and named from the local node-pack source, selected ComfyUI Controller SSH sessions disable forwards inherited from the selected alias, so existing development tunnels cannot collide with worker-management connections. Image builds explicitly load their result into the remote daemon, including when its selected BuildKit builder uses the isolated `docker-container` driver. The worker image includes ComfyUI's top-level Python modules together with its selected package trees, so headless startup can import `nodes`, `execution`, `folder_paths`, and related runtime entry modules. The worker entrypoint keeps Modal-only package initialization lazy, allowing the same source tree to run as a top-level OCI application while ordinary ComfyUI imports retain package-relative behavior. Derived worker images disable health checks inherited from provider base images because readiness is established by the private worker socket rather than a base image's HTTP process. They also expose `/app` through `LD_LIBRARY_PATH`, and the GGUF adapter defensively exposes the selected server binary's directory to its child process, so shared llama.cpp libraries copied beside `llama-server` remain loadable. Custom-node requirements are installed into the image, except packages owned by the immutable Python/CUDA runtime—such as NumPy, OpenCV, Pillow, PyTorch, Triton, and vLLM—whose custom-node constraints are ignored before the final supported pins are applied. A custom node that is genuinely incompatible with those supported versions will fail explicitly during import rather than silently replacing the worker stack. -Enabled hosts are re-probed concurrently when their saved capability data is missing or older than five minutes. An unreachable host is excluded from new automatic placements, allowing Modal to remain the fallback; self-hosted-only workflows fail with the per-host rejection reason. A broken SSH/Docker relay removes and recreates its assigned worker slot and retries the stable invocation once, while remote application errors are returned without blind retry. +Enabled hosts are re-probed concurrently immediately before every placement decision, so recently saved free-VRAM and available-RAM values cannot outlive another workload that began moments later. An unreachable host is excluded from new automatic placements, allowing Modal to remain the fallback; self-hosted-only workflows fail with the per-host rejection reason. A broken SSH/Docker relay removes and recreates its assigned worker slot and retries the stable invocation once, while remote application errors are returned without blind retry. Successful components write a bounded 20-sample timing window to `/comfyui-modal/execution-history.sqlite3`. The planner uses the per-environment median for the same workload signature—node types plus cost-shaping inputs such as dimensions, steps, frames, batch size, model profile, and token limit—so a faster but more expensive-per-second GPU can win when its predicted total cost is lower. Until observations exist, planning uses a conservative common 60-second estimate. @@ -291,7 +291,7 @@ When a prompt is queued: 1. The frontend sends the prompt and `extra_pnginfo.workflow` metadata to `POST /modal/queue_prompt`. 2. The backend resolves marked workflow nodes onto queued prompt node ids, including nested subgraph ids such as `195:27`. Reusable definitions under `workflow.definitions.subgraphs` are expanded through each matching subgraph instance before markers are mapped. -3. Remote-marked nodes are partitioned into cost-aware components. Direct remote edges carrying large tensor or media values such as `LATENT`, `IMAGE`, `MASK`, `SIGMAS`, `AUDIO`, or `VIDEO` stay inside one component; inexpensive scalar edges may remain component boundaries. +3. Remote-marked nodes are partitioned into cost-aware, provider-aware components. Direct remote edges carrying large tensor or media values such as `LATENT`, `IMAGE`, `MASK`, `SIGMAS`, `AUDIO`, or `VIDEO` normally stay inside one component; inexpensive scalar edges may remain component boundaries. A backend that requires SSH cuts a provider boundary even across an expensive edge, and the partitioner refuses any later merge that would make the coarse component graph cyclic. 4. Components also expand across non-transportable inputs such as `MODEL`, `CLIP`, `VAE`, or `CONDITIONING`. 5. A terminal `SaveVideo` fed by a remote `VIDEO` is also placed remotely as an artifact sink. Its encoded file is downloaded into local `output/`; raw frames and audio are not exported merely to re-encode the same video locally. 6. Each component is replaced with one or more generated `ModalUniversalExecutor_` proxy nodes. diff --git a/api_intercept.py b/api_intercept.py index be74f3b..957b8bd 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -15,7 +15,7 @@ from collections import defaultdict, deque from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Any, Iterator, Mapping +from typing import Any, Callable, Iterator, Mapping from aiohttp import web @@ -87,7 +87,6 @@ _LOCAL_ONLY_REMOTE_CLASS_TYPES = frozenset( MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, } ) -_SSH_CAPABILITY_REFRESH_SECONDS = 5 * 60 _TRANSPORTABLE_OUTPUT_TYPES = frozenset( { "*", @@ -125,6 +124,8 @@ _SPECULATIVE_PREWARM_PAYLOAD_FIELDS = frozenset( "component_id", "prompt_id", "modal_gpu", + "execution_provider", + "execution_environment_id", "remote_worker_affinity_group", "remote_local_gap_pool", "subgraph_prompt", @@ -408,37 +409,26 @@ def _configured_ssh_hosts(settings: ModalSyncSettings) -> tuple[SshHostConfig, . def _schedulable_ssh_hosts(settings: ModalSyncSettings) -> tuple[SshHostConfig, ...]: - """Refresh stale enabled hosts concurrently before cost-aware scheduling.""" + """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 - now = time.time() - stale_hosts = [host for host in hosts if _ssh_host_probe_is_stale(host, now)] - if not stale_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(stale_hosts))) as executor: + 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 stale_hosts + for host in probe_hosts ) } return tuple(refreshed.get(host.environment_id, host) for host in hosts) -def _ssh_host_probe_is_stale(host: SshHostConfig, now: float) -> bool: - """Return whether an enabled host needs a fresh capability/health probe.""" - if not host.enabled or host.draining: - return False - capabilities = host.capabilities - if capabilities is None or capabilities.probed_at_epoch is None: - return True - return now - capabilities.probed_at_epoch > _SSH_CAPABILITY_REFRESH_SECONDS - - def _refresh_ssh_host( host: SshHostConfig, registry: RemoteHostRegistry, @@ -473,21 +463,6 @@ def _refresh_ssh_host( ) -def _component_minimum_vram_bytes( - component: RemoteComponentPlan, - prompt: Mapping[str, Any], - preferences: WorkflowExecutionPreferences, - settings: ModalSyncSettings, -) -> int: - """Return the conservative GPU VRAM floor for one remote component.""" - return _component_memory_estimate( - component, - prompt, - preferences, - settings, - ).minimum_vram_bytes - - 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() @@ -608,6 +583,30 @@ def _component_required_provider( return None +def _prompt_node_required_provider( + prompt_node: Mapping[str, Any], +) -> ExecutionProvider | None: + """Return the provider required by one curated node backend, when known.""" + if str(prompt_node.get("class_type") or "") != "ModalLLM": + return None + inputs = prompt_node.get("inputs") + if not isinstance(inputs, Mapping): + return None + model_reference = inputs.get("model_profile") + if not isinstance(model_reference, str) or not model_reference.strip(): + return None + + from .llm_profiles import get_llm_profile + + try: + profile = get_llm_profile(model_reference.strip()) + except ValueError: + return None + if 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): @@ -781,7 +780,24 @@ def _plan_component_execution_assignments( ) -> dict[str, ExecutionAssignment]: """Assign each component to Modal or one compatible self-hosted environment.""" preferences = WorkflowExecutionPreferences.from_workflow(workflow) + resolved_llm_profiles = _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, + resolved_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( @@ -806,7 +822,6 @@ def _plan_component_execution_assignments( scheduler = CostAwareEnvironmentScheduler() history = _execution_history(settings) environment_ids = [environment.environment_id for environment in environments] - resolved_llm_profiles = _resolve_prompt_llm_profiles(prompt, settings) assignments: dict[str, ExecutionAssignment] = {} for component in components: component_signature = _component_execution_signature(component, prompt) @@ -2315,6 +2330,53 @@ def _remote_output_is_list( return bool(output_is_list[output_index]) +def _remote_group_dependency_edges( + prompt: Mapping[str, Any], + remote_node_ids: set[str], + group_for_node: Callable[[str], str], +) -> dict[str, set[str]]: + """Return current coarse dependency edges for a remote-node grouping.""" + dependency_edges: dict[str, set[str]] = defaultdict(set) + for downstream_node_id in remote_node_ids: + downstream_node = prompt.get(downstream_node_id) + if not isinstance(downstream_node, Mapping): + continue + downstream_group = group_for_node(downstream_node_id) + for input_value in (downstream_node.get("inputs") or {}).values(): + if not _is_link(input_value): + continue + upstream_node_id = str(input_value[0]) + if upstream_node_id not in remote_node_ids: + continue + upstream_group = group_for_node(upstream_node_id) + if upstream_group != downstream_group: + dependency_edges[upstream_group].add(downstream_group) + return dependency_edges + + +def _has_alternate_group_path( + dependency_edges: Mapping[str, set[str]], + start_group: str, + target_group: str, +) -> bool: + """Return whether a non-direct path connects two coarse graph groups.""" + pending = [ + downstream_group + for downstream_group in dependency_edges.get(start_group, set()) + if downstream_group != target_group + ] + visited: set[str] = set() + while pending: + current_group = pending.pop() + if current_group == target_group: + return True + if current_group in visited: + continue + visited.add(current_group) + pending.extend(dependency_edges.get(current_group, set())) + return False + + def _remote_component_partition_groups( prompt: dict[str, Any], remote_node_ids: set[str], @@ -2324,6 +2386,11 @@ def _remote_component_partition_groups( """Return component groups after merging remote nodes across costly boundaries.""" parent: dict[str, str] = {node_id: node_id for node_id in remote_node_ids} downstream_remote_node_ids_by_node_id: dict[str, set[str]] = defaultdict(set) + provider_partitioning_active = any( + _prompt_node_required_provider(prompt_node) is not None + for node_id in remote_node_ids + if isinstance((prompt_node := prompt.get(node_id)), Mapping) + ) def find(node_id: str) -> str: """Return the canonical union-find representative for one remote node.""" @@ -2336,15 +2403,49 @@ def _remote_component_partition_groups( node_id = next_node_id return root - def union(left_node_id: str, right_node_id: str) -> None: - """Merge two remote nodes into the same mandatory component.""" + def union( + left_node_id: str, + right_node_id: str, + *, + preserve_coarse_dag: bool | None = None, + ) -> bool: + """Merge two remote nodes unless doing so creates a coarse dependency cycle.""" left_root = find(left_node_id) right_root = find(right_node_id) if left_root == right_root: - return + return True + should_preserve_coarse_dag = ( + provider_partitioning_active + if preserve_coarse_dag is None + else preserve_coarse_dag + ) + if should_preserve_coarse_dag: + dependency_edges = _remote_group_dependency_edges( + prompt, + remote_node_ids, + find, + ) + creates_cycle = _has_alternate_group_path( + dependency_edges, + left_root, + right_root, + ) or _has_alternate_group_path( + dependency_edges, + right_root, + left_root, + ) + if creates_cycle: + logger.info( + "Splitting remote nodes %s and %s to preserve an acyclic " + "component graph.", + left_node_id, + right_node_id, + ) + return False canonical_root = min(left_root, right_root) merged_root = max(left_root, right_root) parent[merged_root] = canonical_root + return True for node_id in sorted(remote_node_ids): prompt_node = prompt.get(node_id) @@ -2366,7 +2467,26 @@ def _remote_component_partition_groups( and str(upstream_prompt_node.get("class_type")) == MODAL_MAP_INPUT_NODE_ID ): - union(node_id, upstream_node_id) + union(node_id, upstream_node_id, preserve_coarse_dag=False) + continue + downstream_provider = _prompt_node_required_provider(prompt_node) + upstream_provider = ( + _prompt_node_required_provider(upstream_prompt_node) + if isinstance(upstream_prompt_node, Mapping) + else None + ) + if downstream_provider != upstream_provider and ( + downstream_provider is not None or upstream_provider is not None + ): + logger.info( + "Splitting remote nodes %s -> %s across provider boundary %s -> %s.", + upstream_node_id, + node_id, + upstream_provider.value if upstream_provider is not None else "any", + downstream_provider.value + if downstream_provider is not None + else "any", + ) continue if _output_supports_parallel_local_materialization( prompt=prompt, @@ -2415,7 +2535,11 @@ def _remote_component_partition_groups( if current_node_id in visited_node_ids: continue visited_node_ids.add(current_node_id) - union(remote_node_id, current_node_id) + union( + remote_node_id, + current_node_id, + preserve_coarse_dag=False, + ) for downstream_node_id in sorted( downstream_remote_node_ids_by_node_id.get(current_node_id, set()) ): @@ -3601,8 +3725,9 @@ def _mark_remote_to_remote_session_boundaries( prompt: dict[str, Any], components: list[RemoteComponentPlan], nodes_module: Any, + assignments_by_component_id: Mapping[str, ExecutionAssignment], ) -> set[str]: - """Mark remote edges for bridges and annotate independent local materializers.""" + """Keep only same-environment remote edges in provider-local session storage.""" consumers = _build_consumer_map(prompt) component_id_by_node_id = { node_id: component.representative_node_id @@ -3623,6 +3748,31 @@ def _mark_remote_to_remote_session_boundaries( ] if not remote_consumers: continue + producer_assignment = assignments_by_component_id[ + component.representative_node_id + ] + consumer_component_ids = { + component_id_by_node_id[consumer.node_id] + for consumer in remote_consumers + } + cross_environment_consumers = sorted( + consumer_component_id + for consumer_component_id in consumer_component_ids + if assignments_by_component_id[consumer_component_id].environment_id + != producer_assignment.environment_id + ) + if cross_environment_consumers: + logger.info( + "Materializing remote boundary output through ComfyUI across " + "execution environments source=%s:%d producer=%s environment=%s " + "consumer_components=%s.", + boundary_output.source.node_id, + boundary_output.source.output_index, + component.representative_node_id, + producer_assignment.environment_id, + cross_environment_consumers, + ) + continue ( non_returning_local_consumers, returning_local_consumers, @@ -5330,21 +5480,49 @@ def _configure_local_gap_keepalive_payloads( keepalive_component_ids: set[str] = set() continuation_component_ids: set[str] = set() for local_node_id in sandwiched_local_node_ids: + upstream_component_ids: set[str] = set() + downstream_component_ids: set[str] = set() for component_id in remote_component_ids: if _component_ancestors_of_local_source( prompt=rewritten_prompt, source_node_id=local_node_id, component_node_id_set={component_id}, ): - keepalive_component_ids.add(component_id) + upstream_component_ids.add(component_id) if _component_ancestors_of_local_source( prompt=rewritten_prompt, source_node_id=component_id, component_node_id_set={local_node_id}, ): - continuation_component_ids.add(component_id) + downstream_component_ids.add(component_id) + + for upstream_component_id in upstream_component_ids: + upstream_payload = payloads_by_component_id.get(upstream_component_id) + if upstream_payload is None: + continue + upstream_environment_id = str( + upstream_payload.get("execution_environment_id") or "" + ) + if upstream_payload.get("execution_provider") != ExecutionProvider.MODAL.value: + continue + for downstream_component_id in downstream_component_ids: + downstream_payload = payloads_by_component_id.get( + downstream_component_id + ) + if ( + downstream_payload is None + or downstream_payload.get("execution_provider") + != ExecutionProvider.MODAL.value + or str(downstream_payload.get("execution_environment_id") or "") + != upstream_environment_id + ): + continue + keepalive_component_ids.add(upstream_component_id) + continuation_component_ids.add(downstream_component_id) for component_id, embedded_payload in payloads_by_component_id.items(): + if component_id not in keepalive_component_ids | continuation_component_ids: + continue payload_fields: dict[str, Any] = {"remote_local_gap_pool": True} if component_id in keepalive_component_ids: payload_fields["keepalive_after_remote_component"] = True @@ -5359,7 +5537,7 @@ def _configure_local_gap_keepalive_payloads( ) logger.info( "Configured local-gap Modal pool for components=%s keepalive_producers=%s continuations=%s.", - sorted(payloads_by_component_id), + sorted(keepalive_component_ids | continuation_component_ids), sorted(keepalive_component_ids), sorted(continuation_component_ids), ) @@ -5474,6 +5652,11 @@ def _configure_speculative_affinity_prewarm_payloads( configured_targets: dict[str, str] = {} for component_id, payload in payloads_by_component_id.items(): + if payload.get("execution_provider") != ExecutionProvider.MODAL.value: + continue + execution_environment_id = str( + payload.get("execution_environment_id") or "" + ) affinity_group = ( str(payload.get("remote_worker_affinity_group") or "comfy").strip().lower() ) @@ -5485,6 +5668,17 @@ def _configure_speculative_affinity_prewarm_payloads( descendant_id for descendant_id in descendant_distances if descendant_id in payloads_by_component_id + and payloads_by_component_id[descendant_id].get( + "execution_provider" + ) + == ExecutionProvider.MODAL.value + and str( + payloads_by_component_id[descendant_id].get( + "execution_environment_id" + ) + or "" + ) + == execution_environment_id and str( payloads_by_component_id[descendant_id].get( "remote_worker_affinity_group" @@ -5577,10 +5771,23 @@ def rewrite_prompt_for_modal( expanded_remote_node_ids, resolved_nodes_module, ) + validate_remote_component_transport_compatibility( + prompt=rewritten_prompt, + components=components, + nodes_module=resolved_nodes_module, + ) + + assignments_by_component_id = _plan_component_execution_assignments( + components=components, + prompt=rewritten_prompt, + workflow=workflow, + settings=resolved_settings, + ) session_component_ids = _mark_remote_to_remote_session_boundaries( rewritten_prompt, components, resolved_nodes_module, + assignments_by_component_id, ) prompt_id = (extra_data or {}).get("prompt_id") remote_sessions_by_component_id = { @@ -5594,49 +5801,9 @@ def rewrite_prompt_for_modal( } if remote_sessions_by_component_id: logger.info( - "Enabled Modal-backed remote references for components=%s.", + "Enabled environment-local remote references for components=%s.", sorted(remote_sessions_by_component_id), ) - validate_remote_component_transport_compatibility( - prompt=rewritten_prompt, - components=components, - nodes_module=resolved_nodes_module, - ) - - assignments_by_component_id = _plan_component_execution_assignments( - components=components, - prompt=rewritten_prompt, - workflow=workflow, - settings=resolved_settings, - ) - if session_component_ids: - preferences = WorkflowExecutionPreferences.from_workflow(workflow) - session_components = [ - component - for component in components - if component.representative_node_id in session_component_ids - ] - anchor_component = max( - session_components, - key=lambda component: _component_minimum_vram_bytes( - component, - rewritten_prompt, - preferences, - resolved_settings, - ), - ) - anchor_assignment = assignments_by_component_id[ - anchor_component.representative_node_id - ] - for component in session_components: - assignments_by_component_id[ - component.representative_node_id - ] = anchor_assignment - logger.info( - "Co-located remote-session components=%s on environment=%s.", - sorted(session_component_ids), - anchor_assignment.environment_id, - ) summary.execution_assignments_by_representative = dict(assignments_by_component_id) sync_engines_by_environment: dict[str, ModalAssetSyncEngine] = {} @@ -6784,7 +6951,7 @@ def setup_modal_queue_route( ) if "prompt" in json_data: - emit_setup_status("Preparing Modal workflow") + emit_setup_status("Preparing remote workflow") rewrite_started_at = time.perf_counter() rewritten_prompt, summary = await rewrite_prompt_for_modal_async( prompt=json_data["prompt"], @@ -6866,7 +7033,7 @@ def setup_modal_queue_route( node_ids=remote_node_ids, modal_gpu=request_settings.modal_gpu, component_node_ids_by_representative=summary.component_node_ids_by_representative, - status_message="Submitting Modal workflow", + status_message="Submitting remote workflow", ) response = await _queue_prompt_json( prompt_server, diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 9ac212c..97b9a88 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -475,3 +475,4 @@ - A Hugging Face GGUF repository can contain hundreds of gigabytes of mutually exclusive quant files and no Transformers tokenizer files. Curated GGUF staging must pin the chosen weight and any matching multimodal projector by exact filename, download those artifacts rather than `*.gguf`, and independently pin a tokenizer source without allowing its safetensors. The deployed vLLM 0.27.1 image has no GGUF loader, so SSH workers use the official digest-pinned CUDA 13 llama.cpp server as a private loopback backend; copying Python 3.13 from its slim build stage into the newer Ubuntu runtime preserves both the existing worker and llama.cpp's required glibc. Image requests use its OpenAI-compatible multimodal chat endpoint while text-only requests retain the token-reporting completion endpoint. - Replacing a provider image's entrypoint does not replace its Docker health check. The llama-server base probes HTTP port 8080, while the managed SSH worker initially serves only its Unix socket; explicitly emit `HEALTHCHECK NONE` in the derived image so an otherwise ready worker is not trapped in an unhealthy startup state. - Copying an executable from a provider image does not make its sibling shared libraries discoverable after the derived image overrides `LD_LIBRARY_PATH`. Preserve `/app` in the SSH image library path and also prefix the configured llama-server binary directory in the spawned process environment; validate with `ldd` or an actual server launch because image construction alone does not reveal the missing dynamic dependency. +- Provider-aware placement must happen before remote-session bridging. Cut explicit backend boundaries during component formation, preserve an acyclic coarse graph when later expensive-edge merges would wrap around that cut, and materialize cross-environment values through ComfyUI; otherwise session co-location silently overwrites the scheduler's provider decisions. Modal speculative prewarm and local-gap keepalive metadata must carry and match both provider and environment identity, because recursively embedded SSH payloads can otherwise stage self-hosted-only models on Modal. diff --git a/modal_llm_node.py b/modal_llm_node.py index bfe922d..a0de880 100644 --- a/modal_llm_node.py +++ b/modal_llm_node.py @@ -217,7 +217,11 @@ class ModalLLM(io.ComfyNode): ) -> io.NodeOutput: """Run one cancellation-aware resident inference request.""" is_remote = os.getenv("COMFY_MODAL_REMOTE_WORKER") == "1" - execution_target = "modal" if is_remote else "local_apple" + execution_target = ( + os.getenv("COMFY_MODAL_LLM_EXECUTION_TARGET", "modal") + if is_remote + else "local_apple" + ) inference_runner = ( run_modal_llm_inference if is_remote else run_local_llm_inference ) diff --git a/modal_llm_runtime.py b/modal_llm_runtime.py index ad7823a..bca908d 100644 --- a/modal_llm_runtime.py +++ b/modal_llm_runtime.py @@ -2588,6 +2588,10 @@ def get_resident_llm_manager() -> ResidentLLMManager: global _RESIDENT_MANAGER with _RESIDENT_MANAGER_LOCK: if _RESIDENT_MANAGER is None: + execution_target = os.getenv( + "COMFY_MODAL_LLM_EXECUTION_TARGET", + "modal", + ) _RESIDENT_MANAGER = ResidentLLMManager( storage_root=os.getenv( "COMFY_MODAL_REMOTE_STORAGE_ROOT", _DEFAULT_STORAGE_ROOT @@ -2600,6 +2604,12 @@ def get_resident_llm_manager() -> ResidentLLMManager: "COMFY_MODAL_LLM_MEMORY_RECOVERY_TIMEOUT_SECONDS", _DEFAULT_MEMORY_RECOVERY_TIMEOUT_SECONDS, ), + execution_target=execution_target, + device_name=( + "cuda (SSH Docker)" + if execution_target == "ssh_docker" + else "cuda" + ), ) return _RESIDENT_MANAGER diff --git a/remote/modal_app.py b/remote/modal_app.py index c9feb1d..95c1be8 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -3596,12 +3596,25 @@ def _emit_local_modal_status( prompt_server.send_sync("modal_status", payload, client_id) +def _remote_execution_destination(payload: Mapping[str, Any]) -> str: + """Return the user-facing destination for a provider-stamped payload.""" + environment_id = str(payload.get("execution_environment_id") or "").strip() + provider = str(payload.get("execution_provider") or "modal").strip() + return ( + environment_id + if provider != "modal" and environment_id + else "Modal" + ) + + def _emit_local_remote_dispatch_status(payload: dict[str, Any]) -> None: - """Tell the local UI a remote component was dispatched before Modal streams progress.""" + """Tell the local UI a remote component was dispatched before progress streams.""" _emit_local_remote_startup_status( payload, phase="starting", - status_message="Starting Modal component", + status_message=( + f"Starting remote component on {_remote_execution_destination(payload)}" + ), ) @@ -4774,7 +4787,10 @@ def _consume_remote_payload_stream( phase="finalizing", node_ids=node_ids, modal_gpu=modal_gpu, - status_message="Receiving Modal outputs", + status_message=( + "Receiving remote outputs from " + f"{_remote_execution_destination(payload)}" + ), ) continue _emit_local_modal_status( @@ -7719,6 +7735,19 @@ def _schedule_speculative_affinity_prewarm( target_payload = payload.get(_SPECULATIVE_PREWARM_TARGET_KEY) if not isinstance(target_payload, Mapping): return False + if ( + payload.get("execution_provider") != "modal" + or target_payload.get("execution_provider") != "modal" + or payload.get("execution_environment_id") + != target_payload.get("execution_environment_id") + ): + logger.debug( + "Skipping speculative Modal prewarm across execution environments " + "source=%s target=%s.", + payload.get("execution_environment_id"), + target_payload.get("execution_environment_id"), + ) + return False normalized_target_payload = copy.deepcopy(dict(target_payload)) settings = _settings_for_payload(normalized_target_payload) @@ -7972,7 +8001,10 @@ def _run_local_gap_keepalive( def _start_local_gap_keepalive(payload: Mapping[str, Any]) -> bool: """Retain one remote affinity slot while a downstream local gap executes.""" - if not bool(payload.get("keepalive_after_remote_component")): + if ( + payload.get("execution_provider") != "modal" + or not bool(payload.get("keepalive_after_remote_component")) + ): return False settings = _settings_for_payload(payload) duration_seconds = max(0.0, float(settings.local_gap_keepalive_seconds)) diff --git a/ssh_runtime.py b/ssh_runtime.py index e941c81..9c7190c 100644 --- a/ssh_runtime.py +++ b/ssh_runtime.py @@ -340,6 +340,7 @@ class SshRuntimeManager: f"COMFY_MODAL_LOCAL_STORAGE_ROOT={_REMOTE_STORAGE_ROOT} " f"COMFY_MODAL_REMOTE_STORAGE_ROOT={_REMOTE_STORAGE_ROOT} " "COMFY_MODAL_EXECUTION_MODE=local COMFY_MODAL_REMOTE_WORKER=1 " + "COMFY_MODAL_LLM_EXECUTION_TARGET=ssh_docker " f"COMFY_MODAL_RUNTIME_FINGERPRINT={spec.identity.fingerprint}" ), "WORKDIR /opt/comfy-remote/repo", diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index ee50131..25c6ba1 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -356,21 +356,35 @@ def test_remote_environment_routes_save_and_probe_hosts( assert registry.get_host("gpu-one").capabilities == capabilities -def test_scheduler_refreshes_stale_ssh_capabilities_before_placement( +def test_scheduler_refreshes_recent_ssh_capabilities_before_placement( api_intercept_module: Any, remote_hosts_module: Any, execution_environments_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: - """Automatic placement should not trust stale persisted host health forever.""" + """Automatic placement must not trust even recently persisted free VRAM.""" registry = remote_hosts_module.RemoteHostRegistry.for_user_directory(tmp_path) + previous_capabilities = execution_environments_module.EnvironmentCapabilities( + architecture="x86_64", + operating_system="linux", + cpu_count=16, + total_ram_bytes=64 * 1024**3, + available_ram_bytes=60 * 1024**3, + available_disk_bytes=1024**4, + docker_version="28.0.0", + docker_rootless=False, + nvidia_container_runtime=True, + probed_at_epoch=2_000_000_000.0, + ) registry.replace_hosts( [ remote_hosts_module.SshHostConfig( environment_id="freshened", display_name="Freshened", ssh_target="freshened", + capabilities=previous_capabilities, + health=execution_environments_module.EnvironmentHealth.READY, ) ] ) @@ -407,6 +421,170 @@ def test_scheduler_refreshes_stale_ssh_capabilities_before_placement( assert hosts[0].capabilities == capabilities +def test_remote_partition_preserves_dag_around_ssh_only_llm( + api_intercept_module: Any, +) -> None: + """Provider boundaries must not be undone by a coarse fanout cycle.""" + prompt = { + "1": {"class_type": "RemoteImage", "inputs": {}}, + "2": { + "class_type": "RemoteImageConsumer", + "inputs": {"image": ["1", 0]}, + }, + "3": { + "class_type": "ModalLLM", + "inputs": { + "image": ["2", 0], + "model_profile": "huihui-qwen3.8-27b-abliterated-q2-k-gguf", + }, + }, + "4": { + "class_type": "RemoteImageConsumer", + "inputs": {"image": ["1", 0], "prompt": ["3", 0]}, + }, + } + fake_nodes_module = SimpleNamespace( + NODE_CLASS_MAPPINGS={ + "RemoteImage": _FakeRemoteImageNode, + "RemoteImageConsumer": _FakeRemoteImageConsumerNode, + "ModalLLM": _FakeTextNode, + }, + NODE_DISPLAY_NAME_MAPPINGS={}, + ) + + component_groups = api_intercept_module._remote_component_partition_groups( + prompt, + set(prompt), + api_intercept_module._build_consumer_map(prompt), + fake_nodes_module, + ) + components = api_intercept_module._component_topological_order( + prompt, + component_groups, + ) + + assert components == [["1", "2"], ["3"], ["4"]] + assert api_intercept_module._component_execution_stages( + prompt, + component_groups, + ) == [["1"], ["3"], ["4"]] + + +def test_modal_only_policy_rejects_ssh_only_llm_backend( + api_intercept_module: Any, + tmp_path: Path, +) -> None: + """A provider-specific backend must not be dispatched to Modal by policy.""" + component = api_intercept_module.RemoteComponentPlan( + node_ids=["257"], + representative_node_id="257", + boundary_inputs=[], + boundary_outputs=[], + execute_node_ids=["257"], + contains_output_node=False, + ) + + with pytest.raises( + api_intercept_module.ModalPromptValidationError, + match="Modal-only execution cannot run SSH-only component", + ): + api_intercept_module._plan_component_execution_assignments( + components=[component], + prompt={ + "257": { + "class_type": "ModalLLM", + "inputs": { + "model_profile": ( + "huihui-qwen3.8-27b-abliterated-q2-k-gguf" + ) + }, + } + }, + workflow={"extra": {"remote_execution": {"policy": "modal"}}}, + settings=SimpleNamespace( + modal_gpu="H200", + max_containers=1, + local_storage_root=tmp_path, + ), + ) + + +def test_cross_provider_boundary_uses_transport_instead_of_remote_session( + api_intercept_module: Any, + execution_environments_module: Any, +) -> None: + """Session-backed references must never cross provider storage boundaries.""" + source = api_intercept_module.LinkedOutputRef("1", 0) + boundary_output = api_intercept_module.BoundaryOutputSpec( + proxy_output_name="remote_output_0", + source=source, + io_type="IMAGE", + is_list=False, + ) + producer = api_intercept_module.RemoteComponentPlan( + node_ids=["1"], + representative_node_id="1", + boundary_inputs=[], + boundary_outputs=[boundary_output], + execute_node_ids=["1"], + contains_output_node=False, + ) + consumer = api_intercept_module.RemoteComponentPlan( + node_ids=["2"], + representative_node_id="2", + boundary_inputs=[ + api_intercept_module.BoundaryInputSpec( + proxy_input_name="remote_input_0", + source=source, + io_type="IMAGE", + targets=[api_intercept_module.InputTarget("2", "image")], + ) + ], + boundary_outputs=[], + execute_node_ids=["2"], + contains_output_node=False, + ) + assignment_type = execution_environments_module.ExecutionAssignment + provider_type = execution_environments_module.ExecutionProvider + + session_component_ids = ( + api_intercept_module._mark_remote_to_remote_session_boundaries( + { + "1": {"class_type": "RemoteImage", "inputs": {}}, + "2": { + "class_type": "RemoteImageConsumer", + "inputs": {"image": ["1", 0]}, + }, + }, + [producer, consumer], + SimpleNamespace( + NODE_CLASS_MAPPINGS={ + "RemoteImage": _FakeRemoteImageNode, + "RemoteImageConsumer": _FakeRemoteImageConsumerNode, + } + ), + { + "1": assignment_type( + "modal:H200", + provider_type.MODAL, + None, + 0.0, + ), + "2": assignment_type( + "lambda", + provider_type.SSH_DOCKER, + 0.0, + 0.0, + ), + }, + ) + ) + + assert session_component_ids == set() + assert boundary_output.session_output is False + assert boundary_output.session_consumer_node_ids == [] + + def test_automatic_policy_assigns_component_to_lower_cost_ready_host( api_intercept_module: Any, remote_hosts_module: Any, @@ -3580,6 +3758,8 @@ def test_planner_attaches_next_distinct_affinity_as_speculative_prewarm_target( "component_id": "spec-a", "prompt_id": "prompt-spec", "modal_gpu": "RTX-PRO-6000", + "execution_provider": "modal", + "execution_environment_id": "modal:RTX-PRO-6000", "remote_worker_affinity_group": "llm", "subgraph_prompt": {"1": {"class_type": "ModalLLM", "inputs": {}}}, } @@ -3593,6 +3773,8 @@ def test_planner_attaches_next_distinct_affinity_as_speculative_prewarm_target( "component_id": "spec-b", "prompt_id": "prompt-spec", "modal_gpu": "RTX-PRO-6000", + "execution_provider": "modal", + "execution_environment_id": "modal:RTX-PRO-6000", "remote_worker_affinity_group": "comfy", "remote_local_gap_pool": True, "snapshot_profile_key": "loader-profile:abc", @@ -3617,6 +3799,8 @@ def test_planner_attaches_next_distinct_affinity_as_speculative_prewarm_target( "component_id": "spec-c", "prompt_id": "prompt-spec", "modal_gpu": "RTX-PRO-6000", + "execution_provider": "modal", + "execution_environment_id": "modal:RTX-PRO-6000", "remote_worker_affinity_group": "llm", "subgraph_prompt": {"3": {"class_type": "ModalLLM", "inputs": {}}}, }, @@ -3649,6 +3833,54 @@ def test_planner_attaches_next_distinct_affinity_as_speculative_prewarm_target( assert "speculative_remote_prewarm_target" not in third_payload +def test_planner_does_not_bridge_local_gap_keepalive_across_providers( + api_intercept_module: Any, +) -> None: + """A Modal producer must not retain a slot for an SSH continuation.""" + rewritten_prompt = { + "modal-producer": { + "class_type": "ModalProxy", + "inputs": { + "original_node_data": { + "component_id": "modal-producer", + "execution_provider": "modal", + "execution_environment_id": "modal:H200", + } + }, + }, + "local-gap": { + "class_type": "PreviewAny", + "inputs": {"source": ["modal-producer", 0]}, + }, + "ssh-consumer": { + "class_type": "ModalProxy", + "inputs": { + "source": ["local-gap", 0], + "original_node_data": { + "component_id": "ssh-consumer", + "execution_provider": "ssh_docker", + "execution_environment_id": "lambda", + }, + }, + }, + } + + api_intercept_module._configure_local_gap_keepalive_payloads( + rewritten_prompt=rewritten_prompt, + remote_component_ids=["modal-producer", "ssh-consumer"], + sandwiched_local_node_ids={"local-gap"}, + ) + + for component_id in ("modal-producer", "ssh-consumer"): + payload = api_intercept_module.registered_proxy_execution_payload( + component_id, + rewritten_prompt[component_id]["inputs"]["original_node_data"], + ) + assert "remote_local_gap_pool" not in payload + assert "keepalive_after_remote_component" not in payload + assert "stop_local_gap_keepalive_before_remote_component" not in payload + + def test_rewrite_keeps_unmapped_remote_siblings_without_local_reentry_together( api_intercept_module: Any, settings_module: Any, diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 368d206..c9fa447 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -1920,11 +1920,15 @@ def test_speculative_affinity_prewarm_is_distinct_and_deduplicated( payload = { "prompt_id": "prompt-spec", "component_id": "llm-component", + "execution_provider": "modal", + "execution_environment_id": "modal:RTX-PRO-6000", "remote_worker_affinity_group": "llm", "speculative_remote_prewarm_target": { "prompt_id": "prompt-spec", "component_id": "comfy-component", "modal_gpu": "RTX-PRO-6000", + "execution_provider": "modal", + "execution_environment_id": "modal:RTX-PRO-6000", "remote_worker_affinity_group": "comfy", "remote_local_gap_pool": True, "subgraph_prompt": { @@ -1962,6 +1966,49 @@ def test_speculative_affinity_prewarm_is_distinct_and_deduplicated( ] == ["UNETLoader"] +def test_speculative_affinity_prewarm_rejects_cross_provider_target( + remote_modal_app_module: Any, + monkeypatch: Any, +) -> None: + """An SSH continuation must never stage its payload on a Modal worker.""" + submitted_tasks: list[tuple[Any, tuple[Any, ...]]] = [] + + class FakeExecutor: + """Executor double that records unexpected warmup jobs.""" + + def submit(self, fn: Any, *args: Any) -> Future[Any]: + """Capture one scheduled job.""" + submitted_tasks.append((fn, args)) + return Future() + + monkeypatch.setattr( + remote_modal_app_module, + "_REMOTE_MODAL_WARMUP_EXECUTOR", + FakeExecutor(), + ) + payload = { + "prompt_id": "prompt-cross-provider", + "component_id": "modal-component", + "execution_provider": "modal", + "execution_environment_id": "modal:H200", + "speculative_remote_prewarm_target": { + "prompt_id": "prompt-cross-provider", + "component_id": "ssh-component", + "execution_provider": "ssh_docker", + "execution_environment_id": "lambda", + }, + } + + assert ( + remote_modal_app_module._schedule_speculative_affinity_prewarm( + payload, + reason="provider_boundary", + ) + is False + ) + assert submitted_tasks == [] + + def test_local_gap_keepalive_is_bounded_and_stopped_by_next_remote_component( remote_modal_app_module: Any, monkeypatch: Any, @@ -1988,15 +2035,19 @@ def test_local_gap_keepalive_is_bounded_and_stopped_by_next_remote_component( remote_modal_app_module.get_settings.cache_clear() with remote_modal_app_module._LOCAL_GAP_KEEPALIVES_LOCK: remote_modal_app_module._LOCAL_GAP_KEEPALIVES.clear() - producer_payload = { - "prompt_id": "prompt-gap", - "component_id": "component-a", - "keepalive_after_remote_component": True, - } - consumer_payload = { - "prompt_id": "prompt-gap", - "component_id": "component-b", - "stop_local_gap_keepalive_before_remote_component": True, + producer_payload = { + "prompt_id": "prompt-gap", + "component_id": "component-a", + "execution_provider": "modal", + "execution_environment_id": "modal:RTX-PRO-6000", + "keepalive_after_remote_component": True, + } + consumer_payload = { + "prompt_id": "prompt-gap", + "component_id": "component-b", + "execution_provider": "modal", + "execution_environment_id": "modal:RTX-PRO-6000", + "stop_local_gap_keepalive_before_remote_component": True, } try: @@ -6405,7 +6456,7 @@ def test_remote_modal_auto_deploys_missing_app_by_default( "llm_staged", "LLM staging complete (0.0 GiB downloaded); starting GPU worker", ), - ("starting", "Starting Modal component"), + ("starting", "Starting remote component on Modal"), ] @@ -8120,7 +8171,7 @@ def test_remote_modal_consumes_streamed_progress_and_result( "prompt_id": "prompt-1", "node_ids": ["7", "8"], "modal_gpu": "B300", - "status_message": "Receiving Modal outputs", + "status_message": "Receiving remote outputs from Modal", }, "client-1", ), @@ -8200,7 +8251,7 @@ def test_emit_local_remote_dispatch_status_marks_component_starting( "prompt_id": "prompt-1", "node_ids": ["7", "8"], "modal_gpu": "B300", - "status_message": "Starting Modal component", + "status_message": "Starting remote component on Modal", }, "client-1", ) diff --git a/tests/test_modal_llm.py b/tests/test_modal_llm.py index a1e53ed..097b447 100644 --- a/tests/test_modal_llm.py +++ b/tests/test_modal_llm.py @@ -2503,6 +2503,22 @@ def test_modal_llm_node_routes_between_local_and_remote_runtimes( ] +def test_resident_manager_reports_ssh_docker_execution_target( + modal_llm_runtime_module: Any, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """A self-hosted CUDA worker must not identify its inference as Modal.""" + monkeypatch.setenv("COMFY_MODAL_LLM_EXECUTION_TARGET", "ssh_docker") + monkeypatch.setenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", str(tmp_path)) + monkeypatch.setattr(modal_llm_runtime_module, "_RESIDENT_MANAGER", None) + + manager = modal_llm_runtime_module.get_resident_llm_manager() + + assert manager.execution_target == "ssh_docker" + assert manager.device_name == "cuda (SSH Docker)" + + def test_remote_dispatch_stages_llm_once_and_forces_volume_reload( remote_modal_app_module: Any, monkeypatch: pytest.MonkeyPatch, diff --git a/tests/test_ssh_runtime.py b/tests/test_ssh_runtime.py index d4c2f6b..b30cf5c 100644 --- a/tests/test_ssh_runtime.py +++ b/tests/test_ssh_runtime.py @@ -128,6 +128,7 @@ def test_worker_dockerfile_disables_inherited_base_image_healthcheck( assert "HEALTHCHECK NONE" in dockerfile assert dockerfile.index("HEALTHCHECK NONE") < dockerfile.index("ENTRYPOINT") assert "LD_LIBRARY_PATH=/app:${LD_LIBRARY_PATH}" in dockerfile + assert "COMFY_MODAL_LLM_EXECUTION_TARGET=ssh_docker" in dockerfile def test_worker_indices_are_distributed_across_discovered_gpus( diff --git a/tests/test_web_modal_toggle.py b/tests/test_web_modal_toggle.py index c801f77..a4f181b 100644 --- a/tests/test_web_modal_toggle.py +++ b/tests/test_web_modal_toggle.py @@ -37,10 +37,10 @@ def test_global_modal_status_badge_is_installed() -> None: assert 'class="modal-status-gpu" hidden' in source assert "gpuText.textContent = activeState.modalGpu ?? \"\";" in source assert "font-size: 10px;" in source - assert "Syncing graph with Modal" in source + assert "Preparing remote workflow" in source assert "Waiting for Modal startup" in source - assert "Receiving Modal outputs" in source - assert "Modal workflow running on" in source + assert "Receiving remote outputs" in source + assert "Remote workflow running on" in source assert "setGlobalStatusBatchProgress(promptId, value, maxValue)" in source assert "batchValue: state.batchValue ?? null," in source assert "batchMax: state.batchMax ?? null," in source @@ -58,8 +58,8 @@ def test_global_modal_status_badge_is_installed() -> None: assert "nodeCount: promptRemoteNodeCount(promptId, nodeCount)," in source assert "hasRemoteExecutionStarted: false" in source assert "Waiting for Modal container" in source - assert "Starting Modal component" in source - assert "Cancelling Modal workflow" in source + assert "Starting remote component" in source + assert "Cancelling remote workflow" in source assert "function promptActiveNodeIsLive(promptId)" in source @@ -282,7 +282,7 @@ def test_global_modal_status_badge_supports_setup_and_finalizing_details() -> No assert "if (detail.phase === STATE_FINALIZING) {" in source assert "dot.style.background = FINALIZING_BORDER_COLOR;" in source assert "READY_BORDER_COLOR" not in source - assert 'text.textContent = activeState.statusMessage ?? "Receiving Modal outputs";' in source + assert 'text.textContent = activeState.statusMessage ?? "Receiving remote outputs";' in source def test_global_modal_status_badge_polls_and_renders_active_containers() -> None: diff --git a/web/modal_toggle.js b/web/modal_toggle.js index dc1048c..7f8c1a5 100644 --- a/web/modal_toggle.js +++ b/web/modal_toggle.js @@ -671,6 +671,7 @@ function currentGlobalStatus() { statusCurrent: state.statusCurrent ?? null, statusTotal: state.statusTotal ?? null, modalGpu: state.modalGpu ?? null, + executionLabel: state.executionLabel ?? null, updatedAt: state.updatedAt, })); phases.sort((left, right) => right.updatedAt - left.updatedAt); @@ -1236,21 +1237,21 @@ function refreshGlobalStatusElement() { dot.style.background = SETUP_BORDER_COLOR; dot.style.boxShadow = "0 0 0 6px rgba(245, 158, 11, 0.18)"; dot.style.animation = "modal-status-pulse 1.1s ease-in-out infinite"; - text.textContent = activeState.statusMessage ?? "Syncing graph with Modal"; + text.textContent = activeState.statusMessage ?? "Preparing remote workflow"; } else if (activeState.phase === STATE_STARTING) { element.style.borderColor = "rgba(234, 179, 8, 0.58)"; element.style.background = "rgba(54, 45, 6, 0.94)"; dot.style.background = STARTING_BORDER_COLOR; dot.style.boxShadow = "0 0 0 6px rgba(234, 179, 8, 0.2)"; dot.style.animation = "modal-status-pulse 0.85s ease-in-out infinite"; - text.textContent = activeState.statusMessage ?? "Starting Modal component"; + text.textContent = activeState.statusMessage ?? "Starting remote component"; } else if (activeState.phase === STATE_WAITING) { element.style.borderColor = "rgba(245, 158, 11, 0.55)"; element.style.background = "rgba(61, 42, 9, 0.94)"; dot.style.background = SETUP_BORDER_COLOR; dot.style.boxShadow = "0 0 0 6px rgba(245, 158, 11, 0.18)"; dot.style.animation = "modal-status-pulse 1.1s ease-in-out infinite"; - text.textContent = activeState.statusMessage ?? "Waiting for Modal startup"; + text.textContent = activeState.statusMessage ?? "Waiting for remote startup"; } else if (activeState.phase === EXECUTION_PHASE) { element.style.borderColor = "rgba(34, 197, 94, 0.55)"; element.style.background = hasBatchProgress @@ -1261,37 +1262,40 @@ function refreshGlobalStatusElement() { dot.style.background = ACTIVE_BORDER_COLOR; dot.style.boxShadow = "0 0 0 6px rgba(34, 197, 94, 0.18)"; dot.style.animation = "modal-status-pulse 1.1s ease-in-out infinite"; + const executionLabel = activeState.executionLabel + ? ` via ${activeState.executionLabel}` + : ""; text.textContent = hasBatchProgress - ? `Modal workflow running on ${activeState.nodeCount} ${nodeLabel} · ${Math.round(batchValue)}/${Math.round(batchMax)}` - : `Modal workflow running on ${activeState.nodeCount} ${nodeLabel}`; + ? `Remote workflow running on ${activeState.nodeCount} ${nodeLabel}${executionLabel} · ${Math.round(batchValue)}/${Math.round(batchMax)}` + : `Remote workflow running on ${activeState.nodeCount} ${nodeLabel}${executionLabel}`; } else if (activeState.phase === STATE_FINALIZING) { element.style.borderColor = "rgba(59, 130, 246, 0.55)"; element.style.background = "rgba(15, 23, 42, 0.94)"; dot.style.background = FINALIZING_BORDER_COLOR; dot.style.boxShadow = "0 0 0 6px rgba(59, 130, 246, 0.18)"; dot.style.animation = "modal-status-pulse 1.1s ease-in-out infinite"; - text.textContent = activeState.statusMessage ?? "Receiving Modal outputs"; + text.textContent = activeState.statusMessage ?? "Receiving remote outputs"; } else if (activeState.phase === STATE_CANCELLING) { element.style.borderColor = "rgba(251, 113, 133, 0.58)"; element.style.background = "rgba(76, 5, 25, 0.94)"; dot.style.background = CANCELLING_BORDER_COLOR; dot.style.boxShadow = "0 0 0 6px rgba(251, 113, 133, 0.2)"; dot.style.animation = "modal-status-pulse 0.75s ease-in-out infinite"; - text.textContent = activeState.statusMessage ?? "Cancelling Modal workflow"; + text.textContent = activeState.statusMessage ?? "Cancelling remote workflow"; } else if (activeState.phase === STATE_ERROR) { element.style.borderColor = "rgba(239, 68, 68, 0.55)"; element.style.background = "rgba(69, 10, 10, 0.94)"; dot.style.background = ERROR_BORDER_COLOR; dot.style.boxShadow = "0 0 0 6px rgba(239, 68, 68, 0.18)"; dot.style.animation = "none"; - text.textContent = "Modal workflow failed"; + text.textContent = "Remote workflow failed"; } else { element.style.borderColor = "rgba(29, 155, 240, 0.5)"; element.style.background = "rgba(15, 23, 42, 0.94)"; dot.style.background = IDLE_BORDER_COLOR; dot.style.boxShadow = "0 0 0 6px rgba(29, 155, 240, 0.18)"; dot.style.animation = "none"; - text.textContent = "Modal workflow active"; + text.textContent = "Remote workflow active"; } dot.style.width = "10px"; @@ -1330,6 +1334,8 @@ function setGlobalStatusPhase(promptId, phase, nodeCount, details = null) { statusCurrent: details?.current ?? existingState?.statusCurrent ?? null, statusTotal: details?.total ?? existingState?.statusTotal ?? null, modalGpu: detailModalGpu ?? existingState?.modalGpu ?? promptState?.modalGpu ?? null, + executionLabel: + existingState?.executionLabel ?? promptState?.executionLabel ?? null, updatedAt: nowMs(), }); refreshGlobalStatusElement(); @@ -1361,6 +1367,10 @@ function setGlobalStatusBatchProgress(promptId, value, maxValue) { statusTotal: existingState?.statusTotal ?? null, modalGpu: existingState?.modalGpu ?? modalPromptStates.get(promptId)?.modalGpu ?? null, + executionLabel: + existingState?.executionLabel ?? + modalPromptStates.get(promptId)?.executionLabel ?? + null, updatedAt: nowMs(), }); refreshGlobalStatusElement(); @@ -1399,7 +1409,7 @@ function markPromptCancellationRequested(promptId) { endSyntheticExecutionUi(promptId); setPromptActiveNode(promptId, null); setGlobalStatusPhase(promptId, STATE_CANCELLING, remoteNodeIds.length || 1, { - message: "Cancelling Modal workflow", + message: "Cancelling remote workflow", }); if (remoteNodeIds.length > 0) { setNodesPhase(remoteNodeIds, STATE_CANCELLING, promptId); @@ -1745,7 +1755,7 @@ function refreshModalUiAfterVisibilityChange() { /** * Return the prompt metadata bucket, creating it if needed. * @param {string} promptId - * @returns {{ startedAt: number, modalGpu: string | null, remoteNodeIds: string[], componentsByRepresentative: Map, componentNodeIdsByMember: Map, representativeNodeIdByMember: Map, componentLabelByMember: Map, laneNodeIdsByLane: Map }} + * @returns {{ startedAt: number, modalGpu: string | null, executionLabel: string | null, remoteNodeIds: string[], componentsByRepresentative: Map, componentNodeIdsByMember: Map, representativeNodeIdByMember: Map, componentLabelByMember: Map, laneNodeIdsByLane: Map }} */ function ensurePromptState(promptId) { if (isPromptTerminal(promptId)) { @@ -1756,6 +1766,7 @@ function ensurePromptState(promptId) { modalPromptStates.set(promptId, { startedAt, modalGpu: null, + executionLabel: null, remoteNodeIds: [], componentsByRepresentative: new Map(), componentNodeIdsByMember: new Map(), @@ -4972,6 +4983,15 @@ function patchQueuePrompt() { const usesSshExecution = executionAssignments.some( (assignment) => assignment?.provider === "ssh_docker", ); + const executionEnvironmentLabels = Array.from( + new Set( + executionAssignments.map((assignment) => + assignment?.provider === "modal" + ? `Modal ${acceptedModalGpu}` + : String(assignment?.environment_id ?? "self-hosted"), + ), + ), + ); const resolvedRemoteNodeIds = (responsePayload.modal_remote_node_ids ?? []).map((nodeIdValue) => String(nodeIdValue), ); @@ -4985,6 +5005,7 @@ function patchQueuePrompt() { if (!promptState) { return responsePayload; } + promptState.executionLabel = executionEnvironmentLabels.join(" + ") || null; const acceptedRemoteNodeIds = promptState.remoteNodeIds.length > 0 ? promptState.remoteNodeIds : remoteNodeIds; if (!isPromptQueuedBehindActiveModal(promptId)) {