Add provider-aware workflow partitioning

This commit is contained in:
Ken Simpson
2026-08-22 20:10:22 -07:00
parent 3d067372c5
commit 9a92fe2ecb
13 changed files with 661 additions and 125 deletions
+3 -3
View File
@@ -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 user directory>/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_<hash>` proxy nodes.
+253 -86
View File
@@ -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,
+1
View File
@@ -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.
+5 -1
View File
@@ -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
)
+10
View File
@@ -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
+36 -4
View File
@@ -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))
+1
View File
@@ -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",
+234 -2
View File
@@ -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,
+63 -12
View File
@@ -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",
)
+16
View File
@@ -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,
+1
View File
@@ -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(
+6 -6
View File
@@ -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:
+32 -11
View File
@@ -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<string, string[]>, componentNodeIdsByMember: Map<string, string[]>, representativeNodeIdByMember: Map<string, string>, componentLabelByMember: Map<string, string>, laneNodeIdsByLane: Map<string, string> }}
* @returns {{ startedAt: number, modalGpu: string | null, executionLabel: string | null, remoteNodeIds: string[], componentsByRepresentative: Map<string, string[]>, componentNodeIdsByMember: Map<string, string[]>, representativeNodeIdByMember: Map<string, string>, componentLabelByMember: Map<string, string>, laneNodeIdsByLane: Map<string, string> }}
*/
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)) {