Harden remote LLM staging lifecycle

README.md

@@ -291,7 +291,7 @@ Local snapshots are stored under `<ComfyUI models>/modal_llm`; set `COMFY_MODAL_
 
 Loaded models are cached in a per-target LRU and reused across executions. Before a cold load, the manager asks ComfyUI to release idle managed models, evicts older resident LLMs if needed, and enforces a free-memory reserve (`local_reserve_free_memory_gb` for macOS unified memory, `reserve_free_vram_gb` for remote GPUs). On Modal, vLLM-backed profiles default to eager mode for a container's first workflow and are promoted to throughput mode — backed by a persistent compilation cache — when the container serves a second one; pin either behavior with `COMFY_MODAL_LLM_VLLM_EXECUTION_MODE`.
 
-Model staging is deliberately kept off billed GPU time: weights are downloaded by a CPU-only stager (Modal) or inside the persistent worker (SSH) before the GPU allocation begins, and the planner applies each model's VRAM/RAM floors before placement so a too-small GPU cannot win cost ranking for a checkpoint it cannot load.
+Model staging is deliberately kept off billed GPU time: weights are downloaded by a CPU-only stager (Modal) or inside the persistent worker (SSH) before the GPU allocation begins, and the planner applies each model's VRAM/RAM floors before placement so a too-small GPU cannot win cost ranking for a checkpoint it cannot load. The planner's already validated, credential-free model profile travels with the execution payload, so the worker reports **Using planner-resolved metadata** instead of repeating Hugging Face inspection. Before downloading, the stager checks that the missing artifacts fit while retaining a free-disk reserve. Snapshot leases carry a host, PID, process-start identity, and heartbeat; a dead same-host owner is reclaimed immediately, while a live download cannot lose its lease merely because it runs for a long time. Cancelling a Vast or SSH staging phase targets that exact owner process without recycling the whole worker, and every provider aborts a staging stream that produces no progress within the configured deadline.
 
 ### Modal Endpoint Chat
 
@@ -467,6 +467,10 @@ New workflow SSH host settings (destinations, cost, worker limits, VRAM reserve,
 | `COMFY_MODAL_LLM_RESERVE_FREE_GB` | `24.0` | Default minimum free VRAM retained for ComfyUI-managed image and video models; the node can override per request. |
 | `COMFY_MODAL_LLM_VLLM_EXECUTION_MODE` | `auto` | `auto` uses eager for a container's first workflow and promotes to throughput on its second; `eager` and `throughput` pin one mode. |
 | `COMFY_MODAL_LLM_MEMORY_RECOVERY_TIMEOUT_SECONDS` | `15.0` | Maximum post-eviction wait for CUDA free memory before retiring the worker and retrying once on a fresh worker. |
+| `COMFY_MODAL_LLM_MIN_FREE_DISK_GB` | `8.0` | Free storage that must remain after staging the model bytes not already present in a resumable snapshot. |
+| `COMFY_MODAL_LLM_STAGE_LEASE_TIMEOUT_SECONDS` | `7200` | Maximum wait for a foreign or legacy snapshot lease; dead local process owners are reclaimed immediately and live leases heartbeat. |
+| `COMFY_MODAL_LLM_STAGE_LEASE_HEARTBEAT_STALE_SECONDS` | `300` | Age after which a structured lease with a missing heartbeat is reclaimable, including after its former worker container disappears. Must exceed two heartbeat intervals. |
+| `COMFY_MODAL_LLM_STAGE_NO_PROGRESS_TIMEOUT_SECONDS` | `600` | Maximum silence between model-staging events before the controller terminates the transport and requests provider-side cancellation. |
 | `COMFY_MODAL_STREAM_EVENT_QUEUE_MAXSIZE` | `256` | Maximum buffered remote progress/result envelopes; stale progress is coalesced when full. |
 | `COMFY_MODAL_ENABLE_PROACTIVE_WARMUP` | `true` | Start background warmup from runtime parallelism signals and planner lookahead. |
 | `COMFY_MODAL_ENABLE_LOADER_PREWARM` | `true` | During warmup, execute synthetic loader prompts for root literal model-loader nodes. |

api_intercept.py

@@ -347,6 +347,7 @@ class ComponentExecutionPlan:
     ssh_hosts_by_id: dict[str, SshHostConfig] = field(default_factory=dict)
     vast_service: VastService | None = None
     vast_leases_by_environment: dict[str, Any] = field(default_factory=dict)
+    resolved_llm_profiles: dict[str, Any] = field(default_factory=dict)
 
 
 @dataclass(frozen=True)
@@ -1185,10 +1186,15 @@ def _plan_component_execution_assignments(
     workflow: Mapping[str, Any] | None,
     settings: ModalSyncSettings,
     status_callback: SetupStatusCallback | None = None,
+    resolved_llm_profiles: Mapping[str, Any] | None = None,
 ) -> dict[str, ExecutionAssignment]:
     """Assign components across Modal, SSH Docker, and workflow-declared Vast pools."""
     preferences = WorkflowExecutionPreferences.from_workflow(workflow)
-    resolved_llm_profiles = _resolve_prompt_llm_profiles(prompt, settings)
+    active_llm_profiles = (
+        dict(resolved_llm_profiles)
+        if resolved_llm_profiles is not None
+        else _resolve_prompt_llm_profiles(prompt, settings)
+    )
     if preferences.policy is ExecutionPolicy.MODAL:
         incompatible_component_ids = [
             component.representative_node_id
@@ -1196,7 +1202,7 @@ def _plan_component_execution_assignments(
             if _component_required_provider(
                 component,
                 prompt,
-                resolved_llm_profiles,
+                active_llm_profiles,
             )
             is ExecutionProvider.SSH_DOCKER
         ]
@@ -1272,13 +1278,13 @@ def _plan_component_execution_assignments(
             prompt,
             preferences,
             settings,
-            resolved_llm_profiles,
+            active_llm_profiles,
         )
         component_memory_estimates[component_id] = memory_estimate
         component_required_providers[component_id] = _component_required_provider(
             component,
             prompt,
-            resolved_llm_profiles,
+            active_llm_profiles,
         )
         if memory_estimate.model_asset_count:
             logger.info(
@@ -1436,6 +1442,7 @@ def _plan_component_execution(
         configuration_set = compile_remote_configuration_set(prompt)
     except (TypeError, ValueError) as exc:
         raise ModalPromptValidationError(str(exc)) from exc
+    resolved_llm_profiles = _resolve_prompt_llm_profiles(prompt, settings)
     if configuration_set is None:
         return ComponentExecutionPlan(
             assignments=_plan_component_execution_assignments(
@@ -1444,7 +1451,9 @@ def _plan_component_execution(
                 workflow=workflow,
                 settings=settings,
                 status_callback=status_callback,
-            )
+                resolved_llm_profiles=resolved_llm_profiles,
+            ),
+            resolved_llm_profiles=resolved_llm_profiles,
         )
     return _plan_configured_component_execution(
         components=components,
@@ -1456,6 +1465,7 @@ def _plan_component_execution(
         environment_status_callback=environment_status_callback,
         plan_callback=plan_callback,
         occupied_environment_ids=occupied_environment_ids,
+        resolved_llm_profiles=resolved_llm_profiles,
     )
 
 
@@ -1470,6 +1480,7 @@ def _plan_configured_component_execution(
     environment_status_callback: EnvironmentSetupStatusCallback | None = None,
     plan_callback: ExecutionPlanStatusCallback | None = None,
     occupied_environment_ids: frozenset[str] = frozenset(),
+    resolved_llm_profiles: Mapping[str, Any] | None = None,
 ) -> ComponentExecutionPlan:
     """Resolve and prepare a capacity-aware plan from connected configurations."""
     configurations_by_id = {
@@ -1486,6 +1497,7 @@ def _plan_configured_component_execution(
         prompt=prompt,
         workflow=workflow,
         settings=settings,
+        resolved_llm_profiles=resolved_llm_profiles,
     )
     _prefetch_configured_vast_offers(
         configuration_set=configuration_set,
@@ -1590,6 +1602,7 @@ def _plan_configured_component_execution(
         ssh_hosts_by_id=ssh_hosts_by_id,
         vast_service=vast_service,
         vast_leases_by_environment=vast_leases,
+        resolved_llm_profiles=dict(resolved_llm_profiles or {}),
     )
 
 
@@ -1840,10 +1853,15 @@ def _configured_component_requirements(
     prompt: Mapping[str, Any],
     workflow: Mapping[str, Any] | None,
     settings: ModalSyncSettings,
+    resolved_llm_profiles: Mapping[str, Any] | None = None,
 ) -> dict[str, ComponentResourceRequirements]:
     """Build provider-neutral requirements for every remote component."""
     preferences = WorkflowExecutionPreferences.from_workflow(workflow)
-    resolved_llm_profiles = _resolve_prompt_llm_profiles(prompt, settings)
+    active_llm_profiles = (
+        dict(resolved_llm_profiles)
+        if resolved_llm_profiles is not None
+        else _resolve_prompt_llm_profiles(prompt, settings)
+    )
     requirements: dict[str, ComponentResourceRequirements] = {}
     for component in components:
         memory_estimate = _component_memory_estimate(
@@ -1851,7 +1869,7 @@ def _configured_component_requirements(
             prompt,
             preferences,
             settings,
-            resolved_llm_profiles,
+            active_llm_profiles,
         )
         component_id = component.representative_node_id
         requirements[component_id] = ComponentResourceRequirements(
@@ -1864,7 +1882,7 @@ def _configured_component_requirements(
             required_provider=_component_required_provider(
                 component,
                 prompt,
-                resolved_llm_profiles,
+                active_llm_profiles,
             ),
         )
     return requirements
@@ -3065,6 +3083,68 @@ def _attach_snapshot_profile_key(
     return payload
 
 
+def _resolved_llm_profile_entry(
+    profile: Any,
+    settings: ModalSyncSettings,
+) -> dict[str, Any]:
+    """Serialize planner metadata plus its persisted security-scan result."""
+    from .llm_profiles import generated_profile_manifest_path
+
+    security_scan_complete = True
+    if getattr(profile, "source", "curated") == "generated":
+        manifest_path = generated_profile_manifest_path(
+            settings.local_storage_root,
+            profile.profile_id,
+        )
+        try:
+            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+        except (FileNotFoundError, OSError, json.JSONDecodeError) as exc:
+            raise ModalPromptValidationError(
+                f"Planner-resolved LLM manifest {manifest_path} is unavailable."
+            ) from exc
+        if not isinstance(manifest, Mapping):
+            raise ModalPromptValidationError(
+                f"Planner-resolved LLM manifest {manifest_path} is invalid."
+            )
+        security_scan_complete = bool(
+            manifest.get("security_scan_complete", False)
+        )
+    return {
+        "profile": profile.to_mapping(),
+        "security_scan_complete": security_scan_complete,
+    }
+
+
+def _attach_resolved_llm_profiles(
+    payload: dict[str, Any],
+    resolved_profiles: Mapping[str, Any],
+    settings: ModalSyncSettings,
+) -> None:
+    """Attach only the planner-resolved profiles used by each payload subtree."""
+    from .llm_profiles import llm_model_references_from_payload
+
+    if not resolved_profiles:
+        return
+    references = llm_model_references_from_payload(payload)
+    entries = {
+        reference: _resolved_llm_profile_entry(resolved_profiles[reference], settings)
+        for reference in references
+        if reference in resolved_profiles
+    }
+    if entries:
+        payload["resolved_llm_profiles"] = entries
+    split_payloads = payload.get("split_proxy_payloads")
+    if isinstance(split_payloads, Mapping):
+        children = split_payloads.values()
+    elif isinstance(split_payloads, list):
+        children = split_payloads
+    else:
+        children = ()
+    for child in children:
+        if isinstance(child, dict):
+            _attach_resolved_llm_profiles(child, resolved_profiles, settings)
+
+
 def _boundary_source_signature(
     prompt: dict[str, Any],
     source: LinkedOutputRef,
@@ -8087,6 +8167,11 @@ def rewrite_prompt_for_modal(
                 component.representative_node_id
             ),
         )
+        _attach_resolved_llm_profiles(
+            payload,
+            execution_plan.resolved_llm_profiles,
+            resolved_settings,
+        )
         _stamp_execution_assignment(
             payload,
             assignment,

comfyui_modal_sync_cloud.py

@@ -7824,12 +7824,15 @@ if modal is not None:  # pragma: no branch - remote entrypoint configuration.
         def _stage_profiles(
             self,
             model_references: list[str],
+            resolved_profiles: Mapping[str, Any] | None = None,
             progress_callback: Callable[[dict[str, Any]], None] | None = None,
         ) -> list[dict[str, Any]]:
             """Resolve and stage profiles while optionally publishing progress."""
             staged_profiles = resolve_and_stage_model_references(
                 model_references,
                 settings.remote_storage_root,
+                resolved_profiles=resolved_profiles,
+                owner_id=f"modal:{os.getpid()}:{time.time_ns()}",
                 progress_callback=(
                     lambda progress: progress_callback(
                         {
@@ -7865,14 +7868,19 @@ if modal is not None:  # pragma: no branch - remote entrypoint configuration.
             return results
 
         @modal.method()
-        def stage_profiles(self, model_references: list[str]) -> list[dict[str, Any]]:
+        def stage_profiles(
+            self,
+            model_references: list[str],
+            resolved_profiles: Mapping[str, Any] | None = None,
+        ) -> list[dict[str, Any]]:
             """Resolve model references, stage snapshots, and return metadata."""
-            return self._stage_profiles(model_references)
+            return self._stage_profiles(model_references, resolved_profiles)
 
         @modal.method()
         def stage_profiles_stream(
             self,
             model_references: list[str],
+            resolved_profiles: Mapping[str, Any] | None = None,
         ) -> Iterator[dict[str, Any]]:
             """Stream CPU staging progress and finish with resolved profile data."""
             progress_events: queue.Queue[dict[str, Any]] = queue.Queue()
@@ -7885,6 +7893,7 @@ if modal is not None:  # pragma: no branch - remote entrypoint configuration.
                     results.append(
                         self._stage_profiles(
                             model_references,
+                            resolved_profiles,
                             progress_events.put,
                         )
                     )

docs/LEARNINGS.md

@@ -585,3 +585,7 @@
 - Fingerprinted worker tags need an ownership-scoped garbage-collection policy: a busy development host accumulated 23 obsolete 10–13 GiB `comfy-remote` images and filled a 1.8 TiB filesystem even though Docker reported most image storage as reclaimable. Before rebuilding, remove only superseded `comfy-remote` and `comfy-remote-deps` tags carrying this extension's labels, without forcing deletion; Docker then protects any image still referenced by a running or stopped container while unrelated repositories remain untouched.
 - Component-level staging progress cannot use the component representative as its UI node: representatives are selected for graph execution and may be an unrelated dependency such as `VAELoader`. Carry the fixed model reference through every staging event, recover the matching `ModalLLM` node IDs from the nested payload, and emit to those nodes. Also report the pre-download state machine explicitly; a single `Inspecting` label otherwise conceals snapshot checks, lease contention, and Hugging Face's pre-transfer preparation.
 - A useful execution plan is target-centric rather than component-centric. Attach a credential-free hardware snapshot to every assignment while provider data is still available, group components beneath each concrete environment in the Configurator, and distinguish known HBM/VRAM and system RAM from dynamic provider-managed capacity instead of inventing precision.
+- File age is not process liveness. A multi-hour model download can legitimately outlive a stale-lock threshold, so a snapshot lease must publish host, PID, process-start identity, and a random ownership token, heartbeat while work is live, reclaim a provably dead same-host owner immediately, and remove the lease only while its token still matches.
+- Stopping an SSH client does not prove the remote process stopped. Give every staging invocation a controller-issued owner ID, publish its exact remote PID and start identity, validate both the owner record and `/proc/<pid>/cmdline` before signalling, and keep cancellation separate from worker-container recycling. Terminate the local transport too, retry briefly for the owner-publication race, and apply a bounded no-progress deadline to SSH, Vast, and Modal staging streams.
+- Queue-time model inspection should be authoritative for the queued execution. Serialize the validated, credential-free profile and its security-scan result into only the payload subtrees that use it; remote stagers can persist the generated manifest and begin snapshot checks without another mutable Hub metadata request. Bump the worker protocol when adding this payload/CLI contract so an older image cannot silently reject it.
+- Model admission needs storage as well as GPU and RAM capacity. Under the per-snapshot lease, subtract bytes already present for a resumable snapshot from the profile's artifact size, require that remainder plus a configurable post-download reserve to fit on the actual mounted filesystem, and fail before opening network transfers.

llm_profiles.py

@@ -2,14 +2,15 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, replace
-from functools import lru_cache
+import base64
 import hashlib
 import json
 import logging
-from pathlib import Path
 import re
-from typing import Any, Mapping
+from dataclasses import dataclass, replace
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Mapping, Sequence
 
 if __package__:
     from .llm_compatibility import (
@@ -505,6 +506,60 @@ def rewrite_llm_model_references(
     visit(payload)
 
 
+def resolved_llm_profile_payloads(
+    payload: Mapping[str, Any],
+    model_references: Sequence[str] | None = None,
+) -> dict[str, dict[str, Any]]:
+    """Return validated planner-resolved metadata attached to a remote payload."""
+    raw_profiles = payload.get("resolved_llm_profiles")
+    if raw_profiles is None:
+        return {}
+    if not isinstance(raw_profiles, Mapping):
+        raise ValueError("resolved_llm_profiles must be an object.")
+    selected_references = (
+        set(model_references)
+        if model_references is not None
+        else {str(reference) for reference in raw_profiles}
+    )
+    profiles: dict[str, dict[str, Any]] = {}
+    for raw_reference, raw_entry in raw_profiles.items():
+        reference = str(raw_reference).strip()
+        if reference not in selected_references:
+            continue
+        if not isinstance(raw_entry, Mapping):
+            raise ValueError(
+                f"Resolved LLM profile entry for {reference!r} must be an object."
+            )
+        raw_profile = raw_entry.get("profile", raw_entry)
+        if not isinstance(raw_profile, Mapping):
+            raise ValueError(
+                f"Resolved LLM profile entry for {reference!r} has no profile object."
+            )
+        security_scan_complete = raw_entry.get("security_scan_complete", False)
+        if not isinstance(security_scan_complete, bool):
+            raise ValueError(
+                f"Resolved LLM profile entry for {reference!r} has an invalid "
+                "security scan state."
+            )
+        LLMModelProfile.from_mapping(raw_profile)
+        profiles[reference] = dict(raw_entry)
+    return profiles
+
+
+def encoded_resolved_llm_profile_payloads(
+    payload: Mapping[str, Any],
+    model_references: Sequence[str],
+) -> str | None:
+    """Encode selected planner metadata for the direct worker command line."""
+    profiles = resolved_llm_profile_payloads(payload, model_references)
+    if not profiles:
+        return None
+    encoded = base64.urlsafe_b64encode(
+        json.dumps(profiles, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    )
+    return encoded.decode("ascii")
+
+
 def llm_profile_ids_from_payload(payload: Mapping[str, Any]) -> tuple[str, ...]:
     """Find curated LLM profiles referenced anywhere in a remote payload."""
     profile_ids: set[str] = set()
@@ -551,4 +606,6 @@ __all__ = [
     "load_llm_profiles",
     "profile_for_execution_target",
     "rewrite_llm_model_references",
+    "resolved_llm_profile_payloads",
+    "encoded_resolved_llm_profile_payloads",
 ]

llm_staging.py

@@ -4,12 +4,17 @@ from __future__ import annotations
 
 import json
 import logging
+import math
 import os
+import shutil
+import socket
+import threading
 import time
 from contextlib import contextmanager
 from dataclasses import asdict, dataclass, replace
 from pathlib import Path
-from typing import Any, Callable, Iterator
+from typing import Any, Callable, Iterator, Mapping
+from uuid import uuid4
 
 if __package__:
     from .llm_profiles import LLMModelProfile, get_llm_profile, load_llm_profiles
@@ -21,6 +26,9 @@ logger = logging.getLogger(__name__)
 _COMPLETE_MARKER_FILENAME = ".comfy-modal-llm-complete.json"
 _DEFAULT_LEASE_TIMEOUT_SECONDS = 7200.0
 _LEASE_POLL_SECONDS = 2.0
+_LEASE_HEARTBEAT_SECONDS = 30.0
+_DEFAULT_LEASE_HEARTBEAT_STALE_SECONDS = 300.0
+_DEFAULT_MINIMUM_FREE_DISK_BYTES = 8 * 1024**3
 _SNAPSHOT_ALLOW_PATTERNS = (
     "*.safetensors",
     "*.safetensors.index.json",
@@ -280,71 +288,326 @@ def _legacy_snapshot_matches(
     )
 
 
+def _process_start_identity(pid: int) -> str | None:
+    """Return the Linux process start tick used to reject PID reuse."""
+    try:
+        stat_record = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
+    except (FileNotFoundError, OSError):
+        return None
+    fields_after_command = stat_record.rpartition(") ")[2].split()
+    return fields_after_command[19] if len(fields_after_command) > 19 else None
+
+
+def _lease_owner_payload(owner_id: str) -> dict[str, Any]:
+    """Return an owner record that can be validated on this worker."""
+    pid = os.getpid()
+    return {
+        "version": 1,
+        "owner_id": owner_id,
+        "host_id": socket.gethostname(),
+        "pid": pid,
+        "process_start": _process_start_identity(pid),
+        "token": uuid4().hex,
+        "acquired_at": time.time(),
+    }
+
+
+def _read_lease_owner(lease_path: Path) -> dict[str, Any] | None:
+    """Read a snapshot lease owner, tolerating legacy text records."""
+    try:
+        raw_value = lease_path.read_text(encoding="utf-8")
+    except (FileNotFoundError, OSError):
+        return None
+    try:
+        payload = json.loads(raw_value)
+    except json.JSONDecodeError:
+        return None
+    return dict(payload) if isinstance(payload, Mapping) else None
+
+
+def _local_lease_owner_is_alive(owner: Mapping[str, Any]) -> bool | None:
+    """Return local owner liveness, or None for another worker host."""
+    if str(owner.get("host_id") or "") != socket.gethostname():
+        return None
+    try:
+        pid = int(owner.get("pid"))
+    except (TypeError, ValueError):
+        return False
+    if pid <= 0:
+        return False
+    expected_start = str(owner.get("process_start") or "")
+    actual_start = _process_start_identity(pid)
+    if actual_start is None:
+        if expected_start:
+            return False
+        try:
+            os.kill(pid, 0)
+        except ProcessLookupError:
+            return False
+        except PermissionError:
+            return True
+        return True
+    return not expected_start or actual_start == expected_start
+
+
+def _lease_record_matches(lease_path: Path, owner: Mapping[str, Any]) -> bool:
+    """Return whether a path still contains the observed owner token."""
+    current = _read_lease_owner(lease_path)
+    return bool(
+        current
+        and current.get("token")
+        and current.get("token") == owner.get("token")
+    )
+
+
+def _remove_owned_lease(lease_path: Path, owner: Mapping[str, Any]) -> None:
+    """Remove a lock only while its ownership token still matches."""
+    if not _lease_record_matches(lease_path, owner):
+        return
+    try:
+        lease_path.unlink()
+    except FileNotFoundError:
+        pass
+
+
+def _heartbeat_snapshot_lease(
+    lease_path: Path,
+    owner: Mapping[str, Any],
+    stop_event: threading.Event,
+) -> None:
+    """Refresh an active lease so another container never steals live work."""
+    while not stop_event.wait(_LEASE_HEARTBEAT_SECONDS):
+        if not _lease_record_matches(lease_path, owner):
+            return
+        try:
+            lease_path.touch()
+        except FileNotFoundError:
+            return
+        except OSError as exc:
+            logger.warning(
+                "Unable to heartbeat LLM staging lease %s: %s",
+                lease_path,
+                exc,
+            )
+            return
+
+
+def _snapshot_lease_wait_timeout_seconds() -> float:
+    """Return the configured maximum wait to acquire one snapshot lease."""
+    timeout_seconds = float(
+        os.getenv("COMFY_MODAL_LLM_STAGE_LEASE_TIMEOUT_SECONDS", "7200")
+    )
+    if not math.isfinite(timeout_seconds) or timeout_seconds <= 0:
+        raise ValueError(
+            "COMFY_MODAL_LLM_STAGE_LEASE_TIMEOUT_SECONDS must be positive."
+        )
+    return timeout_seconds
+
+
+def _snapshot_lease_stale_seconds(
+    existing_owner: Mapping[str, Any] | None,
+) -> float:
+    """Return the age that proves a heartbeating or legacy lease abandoned."""
+    if existing_owner is None or not existing_owner.get("token"):
+        return _DEFAULT_LEASE_TIMEOUT_SECONDS
+    raw_value = os.getenv(
+        "COMFY_MODAL_LLM_STAGE_LEASE_HEARTBEAT_STALE_SECONDS",
+        str(_DEFAULT_LEASE_HEARTBEAT_STALE_SECONDS),
+    )
+    stale_seconds = float(raw_value)
+    if (
+        not math.isfinite(stale_seconds)
+        or stale_seconds <= _LEASE_HEARTBEAT_SECONDS * 2
+    ):
+        raise ValueError(
+            "COMFY_MODAL_LLM_STAGE_LEASE_HEARTBEAT_STALE_SECONDS must exceed "
+            f"{_LEASE_HEARTBEAT_SECONDS * 2:.0f} seconds."
+        )
+    return stale_seconds
+
+
+def _remove_expired_snapshot_lease(
+    lease_path: Path,
+    existing_owner: Mapping[str, Any] | None,
+) -> None:
+    """Remove one expired lease without deleting a replacement owner."""
+    logger.warning(
+        "Removing expired LLM staging lease %s owner=%s.",
+        lease_path,
+        (existing_owner or {}).get("owner_id"),
+    )
+    if existing_owner is not None:
+        _remove_owned_lease(lease_path, existing_owner)
+        return
+    try:
+        lease_path.unlink()
+    except FileNotFoundError:
+        pass
+
+
+def _report_snapshot_lease_wait(
+    progress_callback: StagingProgressCallback | None,
+    model_label: str,
+) -> None:
+    """Publish the first wait state for a contended snapshot."""
+    if progress_callback is None:
+        return
+    progress_callback(
+        LLMStagingProgress(
+            stage="waiting_for_download",
+            message=f"Waiting for another download of {model_label} to finish",
+            indeterminate=True,
+        )
+    )
+
+
+def _wait_for_existing_snapshot_lease(
+    lease_path: Path,
+    *,
+    started_at: float,
+    timeout_seconds: float,
+    progress_callback: StagingProgressCallback | None,
+    model_label: str,
+    report_wait: bool,
+) -> bool:
+    """Wait once or reclaim an abandoned existing snapshot lease."""
+    existing_owner = _read_lease_owner(lease_path)
+    if (
+        existing_owner is not None
+        and _local_lease_owner_is_alive(existing_owner) is False
+    ):
+        logger.warning(
+            "Removing abandoned LLM staging lease %s owner=%s.",
+            lease_path,
+            existing_owner.get("owner_id"),
+        )
+        _remove_owned_lease(lease_path, existing_owner)
+        return False
+    try:
+        lease_age = time.time() - lease_path.stat().st_mtime
+    except FileNotFoundError:
+        return False
+    if lease_age >= _snapshot_lease_stale_seconds(existing_owner):
+        _remove_expired_snapshot_lease(lease_path, existing_owner)
+        return False
+    if time.monotonic() - started_at >= timeout_seconds:
+        raise TimeoutError(
+            f"Timed out waiting {timeout_seconds:.0f}s for model staging lease "
+            f"{lease_path}."
+        )
+    if report_wait:
+        _report_snapshot_lease_wait(progress_callback, model_label)
+    time.sleep(_LEASE_POLL_SECONDS)
+    return True
+
+
+def _acquire_snapshot_lease(
+    lease_path: Path,
+    owner: Mapping[str, Any],
+    *,
+    progress_callback: StagingProgressCallback | None,
+    model_label: str,
+) -> None:
+    """Acquire one exclusive snapshot lease, waiting or reclaiming as needed."""
+    timeout_seconds = _snapshot_lease_wait_timeout_seconds()
+    started_at = time.monotonic()
+    waiting_reported = False
+    while True:
+        try:
+            descriptor = os.open(lease_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
+        except FileExistsError:
+            waited = _wait_for_existing_snapshot_lease(
+                lease_path,
+                started_at=started_at,
+                timeout_seconds=timeout_seconds,
+                progress_callback=progress_callback,
+                model_label=model_label,
+                report_wait=not waiting_reported,
+            )
+            waiting_reported = waiting_reported or waited
+            continue
+        with os.fdopen(descriptor, "w", encoding="utf-8") as lease_file:
+            json.dump(owner, lease_file, sort_keys=True, separators=(",", ":"))
+        return
+
+
 @contextmanager
 def _snapshot_lease(
     snapshot_path: Path,
     *,
     progress_callback: StagingProgressCallback | None = None,
     model_label: str,
+    owner_id: str,
 ) -> Iterator[None]:
-    """Serialize a snapshot download and report waits for another downloader."""
+    """Serialize downloads with liveness-checked, heartbeating ownership."""
     snapshot_path.parent.mkdir(parents=True, exist_ok=True)
     lease_path = snapshot_path.parent / f".{snapshot_path.name}.download.lock"
-    timeout_seconds = float(
-        os.getenv("COMFY_MODAL_LLM_STAGE_LEASE_TIMEOUT_SECONDS", "7200")
+    owner = _lease_owner_payload(owner_id)
+    _acquire_snapshot_lease(
+        lease_path,
+        owner,
+        progress_callback=progress_callback,
+        model_label=model_label,
     )
-    if timeout_seconds <= 0:
-        raise ValueError(
-            "COMFY_MODAL_LLM_STAGE_LEASE_TIMEOUT_SECONDS must be positive."
-        )
-    started_at = time.monotonic()
-    acquired = False
-    waiting_reported = False
-    while not acquired:
-        try:
-            descriptor = os.open(lease_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
-        except FileExistsError:
-            if progress_callback is not None and not waiting_reported:
-                progress_callback(
-                    LLMStagingProgress(
-                        stage="waiting_for_download",
-                        message=(
-                            f"Waiting for another download of {model_label} to finish"
-                        ),
-                        indeterminate=True,
-                    )
-                )
-                waiting_reported = True
-            try:
-                lease_age = time.time() - lease_path.stat().st_mtime
-            except FileNotFoundError:
-                continue
-            if lease_age >= _DEFAULT_LEASE_TIMEOUT_SECONDS:
-                logger.warning(
-                    "Removing abandoned Modal LLM staging lease %s.", lease_path
-                )
-                try:
-                    lease_path.unlink()
-                except FileNotFoundError:
-                    pass
-                continue
-            if time.monotonic() - started_at >= timeout_seconds:
-                raise TimeoutError(
-                    f"Timed out waiting {timeout_seconds:.0f}s for model staging lease "
-                    f"{lease_path}."
-                )
-            time.sleep(_LEASE_POLL_SECONDS)
-        else:
-            with os.fdopen(descriptor, "w", encoding="utf-8") as lease_file:
-                lease_file.write(f"pid={os.getpid()} acquired_at={time.time()}\n")
-            acquired = True
+    heartbeat_stop = threading.Event()
+    heartbeat_thread = threading.Thread(
+        target=_heartbeat_snapshot_lease,
+        args=(lease_path, owner, heartbeat_stop),
+        name="llm-snapshot-lease-heartbeat",
+        daemon=True,
+    )
+    heartbeat_thread.start()
     try:
         yield
     finally:
+        heartbeat_stop.set()
+        heartbeat_thread.join(timeout=1.0)
+        _remove_owned_lease(lease_path, owner)
+
+
+def _snapshot_existing_bytes(snapshot_path: Path) -> int:
+    """Return bytes already present for a resumable model snapshot."""
+    total = 0
+    if not snapshot_path.exists():
+        return total
+    for candidate in snapshot_path.rglob("*"):
         try:
-            lease_path.unlink()
-        except FileNotFoundError:
-            pass
+            if candidate.is_file() and not candidate.name.endswith(".lock"):
+                total += candidate.stat().st_size
+        except (FileNotFoundError, OSError):
+            continue
+    return total
+
+
+def _minimum_free_disk_bytes() -> int:
+    """Return the configured free-space reserve retained after staging."""
+    raw_value = os.getenv("COMFY_MODAL_LLM_MIN_FREE_DISK_GB")
+    if raw_value is None:
+        return _DEFAULT_MINIMUM_FREE_DISK_BYTES
+    value_gb = float(raw_value)
+    if not math.isfinite(value_gb) or value_gb < 0:
+        raise ValueError("COMFY_MODAL_LLM_MIN_FREE_DISK_GB must not be negative.")
+    return int(value_gb * 1024**3)
+
+
+def _preflight_snapshot_capacity(
+    snapshot_path: Path,
+    profile: LLMModelProfile,
+) -> tuple[int, int, int]:
+    """Require room for missing artifacts plus a post-download safety reserve."""
+    snapshot_path.parent.mkdir(parents=True, exist_ok=True)
+    free_bytes = shutil.disk_usage(snapshot_path.parent).free
+    existing_bytes = _snapshot_existing_bytes(snapshot_path)
+    remaining_bytes = max(0, profile.artifact_bytes - existing_bytes)
+    required_bytes = remaining_bytes + _minimum_free_disk_bytes()
+    if free_bytes < required_bytes:
+        raise RuntimeError(
+            "Insufficient disk space to stage LLM profile "
+            f"{profile.profile_id!r}: {free_bytes / 1024**3:.2f} GiB free, "
+            f"{remaining_bytes / 1024**3:.2f} GiB still required, and "
+            f"{_minimum_free_disk_bytes() / 1024**3:.2f} GiB must remain free."
+        )
+    return free_bytes, remaining_bytes, required_bytes
 
 
 def stage_model_profile(
@@ -355,6 +618,7 @@ def stage_model_profile(
     snapshot_download: Callable[..., str] | None = None,
     progress_callback: StagingProgressCallback | None = None,
     model_reference: str | None = None,
+    owner_id: str | None = None,
 ) -> StagedModelSnapshot:
     """Download one pinned curated or generated snapshot on a CPU worker."""
     profile = profile or get_llm_profile(profile_id, storage_root=storage_root)
@@ -408,6 +672,7 @@ def stage_model_profile(
         snapshot_path,
         progress_callback=report,
         model_label=model_label,
+        owner_id=owner_id or f"pid-{os.getpid()}",
     ):
         if is_model_snapshot_staged(storage_root, profile):
             report(
@@ -429,6 +694,23 @@ def stage_model_profile(
                 downloaded=False,
                 elapsed_seconds=time.perf_counter() - started_at,
             )
+        free_bytes, remaining_bytes, _required_bytes = _preflight_snapshot_capacity(
+            snapshot_path,
+            profile,
+        )
+        report(
+            LLMStagingProgress(
+                stage="disk_check",
+                message=(
+                    f"Storage ready for {model_label}: "
+                    f"{free_bytes / 1024**3:.1f} GiB free, "
+                    f"{remaining_bytes / 1024**3:.1f} GiB remaining"
+                ),
+                value=free_bytes,
+                maximum=free_bytes,
+                unit="bytes",
+            )
+        )
         if snapshot_download is None:
             from huggingface_hub import (
                 snapshot_download as huggingface_snapshot_download,
@@ -513,11 +795,148 @@ def stage_model_profile(
     return result
 
 
+def _persist_supplied_profile(
+    model_reference: str,
+    storage_root: str | Path,
+    supplied: Mapping[str, Any],
+) -> tuple[LLMModelProfile, str | None, bool, bool]:
+    """Validate planner metadata and persist a generated manifest remotely."""
+    raw_profile = supplied.get("profile", supplied)
+    if not isinstance(raw_profile, Mapping):
+        raise ValueError(
+            f"Resolved LLM profile for {model_reference!r} must contain an object."
+        )
+    profile = LLMModelProfile.from_mapping(raw_profile)
+    _validate_supplied_profile_reference(model_reference, profile)
+    if profile.source != "generated":
+        return profile, None, False, True
+    if __package__:
+        from .llm_profiles import generated_profile_manifest_path
+    else:  # pragma: no cover - stable cloud entrypoint imports top-level modules.
+        from llm_profiles import generated_profile_manifest_path
+
+    manifest_path = generated_profile_manifest_path(storage_root, profile.profile_id)
+    manifest_payload = _supplied_profile_manifest_payload(
+        model_reference,
+        profile,
+        supplied,
+    )
+    existing_scan = _existing_supplied_profile_scan(
+        manifest_path,
+        raw_profile,
+    )
+    if existing_scan is not None:
+        return profile, str(manifest_path), False, existing_scan
+    manifest_path.parent.mkdir(parents=True, exist_ok=True)
+    temporary_path = manifest_path.with_suffix(f".{uuid4().hex}.tmp")
+    temporary_path.write_text(
+        json.dumps(manifest_payload, sort_keys=True, separators=(",", ":")),
+        encoding="utf-8",
+    )
+    os.replace(temporary_path, manifest_path)
+    return profile, str(manifest_path), True, bool(
+        manifest_payload["security_scan_complete"]
+    )
+
+
+def _supplied_profile_manifest_payload(
+    model_reference: str,
+    profile: LLMModelProfile,
+    supplied: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Build the generated manifest persisted from trusted planner metadata."""
+    security_scan_complete = supplied.get("security_scan_complete", False)
+    if not isinstance(security_scan_complete, bool):
+        raise ValueError(
+            f"Resolved LLM profile for {model_reference!r} has an invalid "
+            "security scan state."
+        )
+    return {
+        "schema_version": profile.schema_version,
+        "compatibility_policy_version": profile.compatibility_policy_version,
+        "requested_reference": model_reference,
+        "resolved_at_unix": time.time(),
+        "security_scan_complete": security_scan_complete,
+        "profile": profile.to_mapping(),
+    }
+
+
+def _existing_supplied_profile_scan(
+    manifest_path: Path,
+    raw_profile: Mapping[str, Any],
+) -> bool | None:
+    """Return an existing matching manifest's scan state, or None if absent."""
+    if manifest_path.is_file():
+        try:
+            existing = json.loads(manifest_path.read_text(encoding="utf-8"))
+        except (json.JSONDecodeError, OSError) as exc:
+            raise RuntimeError(
+                f"Resolved LLM manifest {manifest_path} is unreadable: {exc}"
+            ) from exc
+        if not isinstance(existing, Mapping) or existing.get("profile") != raw_profile:
+            raise RuntimeError(
+                f"Resolved LLM profile collision at {manifest_path}."
+            )
+        return bool(existing.get("security_scan_complete", False))
+    return None
+
+
+def _validate_supplied_profile_reference(
+    model_reference: str,
+    profile: LLMModelProfile,
+) -> None:
+    """Reject planner metadata that does not describe its requested reference."""
+    curated_profile = load_llm_profiles().get(model_reference)
+    if curated_profile is not None:
+        if profile != curated_profile:
+            raise ValueError(
+                f"Resolved metadata for curated profile {model_reference!r} does "
+                "not match the checked-in profile."
+            )
+        return
+    if model_reference.startswith("hf-"):
+        if profile.profile_id != model_reference:
+            raise ValueError(
+                f"Resolved metadata profile id {profile.profile_id!r} does not "
+                f"match requested id {model_reference!r}."
+            )
+        return
+    if profile.source != "generated":
+        raise ValueError(
+            f"Resolved metadata for Hugging Face reference {model_reference!r} "
+            "must contain a generated profile."
+        )
+    if __package__:
+        from .llm_resolver import HuggingFaceModelReference
+    else:  # pragma: no cover - stable cloud entrypoint imports top-level modules.
+        from llm_resolver import HuggingFaceModelReference
+
+    parsed = HuggingFaceModelReference.parse(model_reference)
+    if profile.repository != parsed.repository:
+        raise ValueError(
+            f"Resolved metadata repository {profile.repository!r} does not match "
+            f"requested repository {parsed.repository!r}."
+        )
+    requested_revision = (parsed.requested_revision or "").lower()
+    if len(requested_revision) == 40 and requested_revision != profile.revision:
+        raise ValueError(
+            f"Resolved metadata revision {profile.revision!r} does not match "
+            f"requested revision {requested_revision!r}."
+        )
+
+
 def _resolve_profile_for_staging(
     model_reference: str,
     storage_root: str | Path,
+    supplied_profile: Mapping[str, Any] | None = None,
 ) -> tuple[LLMModelProfile, str | None, bool, bool]:
     """Resolve one curated, generated, or Hugging Face model reference."""
+    if supplied_profile is not None:
+        return _persist_supplied_profile(
+            model_reference,
+            storage_root,
+            supplied_profile,
+        )
     curated_profiles = load_llm_profiles()
     if model_reference in curated_profiles:
         return curated_profiles[model_reference], None, False, True
@@ -547,10 +966,13 @@ def resolve_and_stage_model_references(
     storage_root: str | Path,
     *,
     progress_callback: StagingProgressCallback | None = None,
+    resolved_profiles: Mapping[str, Mapping[str, Any]] | None = None,
+    owner_id: str | None = None,
 ) -> list[ResolvedStagedModelProfile]:
     """Resolve and stage model references on any CPU-backed remote worker."""
     results: list[ResolvedStagedModelProfile] = []
     for model_reference in model_references:
+        supplied_profile = (resolved_profiles or {}).get(model_reference)
         if progress_callback is not None:
             is_hugging_face_reference = (
                 "/" in model_reference and not model_reference.startswith("hf-")
@@ -558,12 +980,16 @@ def resolve_and_stage_model_references(
             progress_callback(
                 LLMStagingProgress(
                     stage=(
-                        "metadata"
+                        "resolved_metadata"
+                        if supplied_profile is not None
+                        else "metadata"
                         if is_hugging_face_reference
                         else "profile_resolution"
                     ),
                     message=(
-                        f"Inspecting Hugging Face metadata for {model_reference}"
+                        f"Using planner-resolved metadata for {model_reference}"
+                        if supplied_profile is not None
+                        else f"Inspecting Hugging Face metadata for {model_reference}"
                         if is_hugging_face_reference
                         else f"Resolving model profile {model_reference}"
                     ),
@@ -577,7 +1003,15 @@ def resolve_and_stage_model_references(
             manifest_path,
             manifest_created,
             scan_complete,
-        ) = _resolve_profile_for_staging(model_reference, storage_root)
+        ) = (
+            _resolve_profile_for_staging(
+                model_reference,
+                storage_root,
+                supplied_profile,
+            )
+            if supplied_profile is not None
+            else _resolve_profile_for_staging(model_reference, storage_root)
+        )
         resolve_elapsed_seconds = time.perf_counter() - resolve_started_at
         staged = stage_model_profile(
             profile.profile_id,
@@ -585,6 +1019,7 @@ def resolve_and_stage_model_references(
             profile=profile,
             progress_callback=progress_callback,
             model_reference=model_reference,
+            owner_id=owner_id,
         )
         results.append(
             ResolvedStagedModelProfile(

remote/modal_app.py

@@ -85,6 +85,7 @@ from ..llm_profiles import (
     get_llm_profile,
     llm_model_reference_node_ids_from_payload,
     llm_model_references_from_payload,
+    resolved_llm_profile_payloads,
     rewrite_llm_model_references,
 )
 from ..llm_recovery import (
@@ -92,6 +93,7 @@ from ..llm_recovery import (
     exhausted_recovery_used_vllm_throughput,
     is_llm_memory_recovery_exhausted,
 )
+from ..staging_process import staging_no_progress_timeout_seconds
 from ..settings import (
     MODAL_GPU_TYPES,
     ModalSyncSettings,
@@ -128,6 +130,16 @@ _LOCAL_GAP_KEEPALIVES_LOCK = threading.Lock()
 _LOCAL_GAP_KEEPALIVES: dict[tuple[str, int], "_LocalGapKeepaliveState"] = {}
 _SNAPSHOT_PROFILE_RECORDS_LOCK = threading.Lock()
 _SNAPSHOT_PROFILE_RECORDS: dict[str, dict[str, Any]] = {}
+_MODAL_STAGE_STREAM_END = object()
+
+
+@dataclass(frozen=True)
+class _ModalStageStreamFailure:
+    """Carry an arbitrary Modal stream exception across a reader thread."""
+
+    error: Exception
+
+
 _PRIMITIVE_WIDGET_INPUT_TYPES = frozenset({"INT", "FLOAT", "BOOLEAN", "STRING"})
 _ROOT_LOADER_PREWARM_CLASS_TYPES = frozenset(
     {
@@ -6288,6 +6300,65 @@ def _emit_local_llm_staging_progress(
         )
 
 
+def _read_modal_stage_events(
+    stage_events: Iterable[Any],
+    output: queue.Queue[Any],
+) -> None:
+    """Read a blocking Modal generator while exposing controller timeouts."""
+    try:
+        for event in stage_events:
+            output.put(event)
+    except Exception as exc:
+        output.put(_ModalStageStreamFailure(exc))
+    finally:
+        output.put(_MODAL_STAGE_STREAM_END)
+
+
+def _close_modal_stage_events(stage_events: Iterable[Any]) -> None:
+    """Ask Modal to cancel a stage generator that stopped reporting progress."""
+    close = getattr(stage_events, "close", None)
+    if not callable(close):
+        return
+    try:
+        close()
+    except (OSError, RuntimeError, ValueError) as exc:
+        logger.warning("Unable to close stalled Modal model staging stream: %s", exc)
+
+
+def _bounded_modal_stage_events(stage_events: Iterable[Any]) -> Iterator[Any]:
+    """Yield Modal staging events with a bounded interval between updates."""
+    try:
+        timeout_seconds = staging_no_progress_timeout_seconds()
+    except ValueError:
+        _close_modal_stage_events(stage_events)
+        raise
+    output: queue.Queue[Any] = queue.Queue()
+    reader = threading.Thread(
+        target=_read_modal_stage_events,
+        args=(stage_events, output),
+        name="modal-llm-stage-stream",
+        daemon=True,
+    )
+    reader.start()
+    try:
+        while True:
+            try:
+                item = output.get(timeout=timeout_seconds)
+            except queue.Empty as exc:
+                _close_modal_stage_events(stage_events)
+                raise ModalRemoteInvocationError(
+                    "Modal model staging produced no progress for "
+                    f"{timeout_seconds:.0f} seconds; the staging call was cancelled."
+                ) from exc
+            if item is _MODAL_STAGE_STREAM_END:
+                return
+            if isinstance(item, _ModalStageStreamFailure):
+                raise item.error
+            yield item
+    finally:
+        reader.join(timeout=1.0)
+
+
 def _ensure_llm_profiles_staged(
     payload: dict[str, Any],
     deployment_app_name: str,
@@ -6307,6 +6378,10 @@ def _ensure_llm_profiles_staged(
             if (deployment_app_name, reference) not in _STAGED_LLM_PROFILE_RESULTS
         ]
         if missing_model_references:
+            resolved_profiles = resolved_llm_profile_payloads(
+                payload,
+                missing_model_references,
+            )
             _emit_local_remote_startup_status(
                 payload,
                 phase="llm_staging",
@@ -6327,7 +6402,12 @@ def _ensure_llm_profiles_staged(
             stage_stream = getattr(stager, "stage_profiles_stream", None)
             remote_generator = getattr(stage_stream, "remote_gen", None)
             if callable(remote_generator):
-                for stage_event in remote_generator(missing_model_references):
+                stage_events = (
+                    remote_generator(missing_model_references, resolved_profiles)
+                    if resolved_profiles
+                    else remote_generator(missing_model_references)
+                )
+                for stage_event in _bounded_modal_stage_events(stage_events):
                     if not isinstance(stage_event, Mapping):
                         continue
                     if stage_event.get("kind") == "result":
@@ -6339,8 +6419,13 @@ def _ensure_llm_profiles_staged(
                         continue
                     _emit_local_llm_staging_progress(payload, stage_event)
             else:
-                stage_results = stager.stage_profiles.remote(
-                    missing_model_references
+                stage_results = (
+                    stager.stage_profiles.remote(
+                        missing_model_references,
+                        resolved_profiles,
+                    )
+                    if resolved_profiles
+                    else stager.stage_profiles.remote(missing_model_references)
                 )
             confirmed_references: set[str] = set()
             for stage_result in stage_results:

remote/ssh_worker.py

@@ -3,15 +3,21 @@
 from __future__ import annotations
 
 import argparse
+import base64
 import json
 import logging
 import os
+import re
+import signal
 import socket
 import socketserver
 import sys
 import threading
+import time
+from contextlib import contextmanager
 from pathlib import Path
 from typing import Any, BinaryIO, Iterator, Mapping, Sequence
+from uuid import uuid4
 
 try:
     from ..remote_protocol import (
@@ -40,6 +46,8 @@ logger = logging.getLogger(__name__)
 
 DEFAULT_WORKER_SOCKET_PATH = Path("/run/comfy-remote/worker.sock")
 DEFAULT_STORAGE_ROOT = Path("/storage")
+_STAGING_OWNER_ID_PATTERN = re.compile(r"^[A-Za-z0-9._:-]{1,200}$")
+_STAGING_TERMINATE_GRACE_SECONDS = 5.0
 
 
 class SshWorkerError(RuntimeError):
@@ -249,6 +257,9 @@ def runtime_info() -> dict[str, Any]:
 def stage_profiles(
     model_references: list[str],
     storage_root: Path = DEFAULT_STORAGE_ROOT,
+    *,
+    resolved_profiles: Mapping[str, Mapping[str, Any]] | None = None,
+    owner_id: str | None = None,
 ) -> list[dict[str, Any]]:
     """Resolve and stage immutable LLM profiles in the worker's storage volume."""
     try:
@@ -280,6 +291,8 @@ def stage_profiles(
         model_references,
         storage_root,
         progress_callback=emit_progress,
+        resolved_profiles=resolved_profiles,
+        owner_id=owner_id,
     )
     result_payload = [result.to_dict() for result in results]
     print(
@@ -293,6 +306,175 @@ def stage_profiles(
     return result_payload
 
 
+def _validated_staging_owner_id(owner_id: str) -> str:
+    """Return a filesystem-safe controller-issued staging owner identity."""
+    normalized = owner_id.strip()
+    if not _STAGING_OWNER_ID_PATTERN.fullmatch(normalized):
+        raise ValueError(f"Invalid staging owner id {owner_id!r}.")
+    return normalized
+
+
+def _staging_owner_path(storage_root: Path, owner_id: str) -> Path:
+    """Return the isolated process record for one staging invocation."""
+    return (
+        storage_root.resolve()
+        / "llm_staging"
+        / "owners"
+        / f"{_validated_staging_owner_id(owner_id)}.json"
+    )
+
+
+def _linux_process_start(pid: int) -> str | None:
+    """Return one Linux process start tick for PID-reuse validation."""
+    try:
+        stat_record = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
+    except (FileNotFoundError, OSError):
+        return None
+    fields_after_command = stat_record.rpartition(") ")[2].split()
+    return fields_after_command[19] if len(fields_after_command) > 19 else None
+
+
+def _owner_record_matches(path: Path, expected: Mapping[str, Any]) -> bool:
+    """Return whether a process record still belongs to the expected owner."""
+    try:
+        current = json.loads(path.read_text(encoding="utf-8"))
+    except (FileNotFoundError, json.JSONDecodeError, OSError):
+        return False
+    return bool(
+        isinstance(current, Mapping)
+        and current.get("pid") == expected.get("pid")
+        and current.get("process_start") == expected.get("process_start")
+    )
+
+
+def _remove_owned_snapshot_leases(
+    storage_root: Path,
+    owner: Mapping[str, Any],
+) -> None:
+    """Remove snapshot leases that still name one terminated staging owner."""
+    model_root = storage_root.resolve() / "llm_models"
+    if not model_root.is_dir():
+        return
+    for lease_path in model_root.rglob(".*.download.lock"):
+        try:
+            lease_owner = json.loads(lease_path.read_text(encoding="utf-8"))
+        except (FileNotFoundError, json.JSONDecodeError, OSError):
+            continue
+        if not isinstance(lease_owner, Mapping):
+            continue
+        if any(
+            lease_owner.get(field) != owner.get(field)
+            for field in ("owner_id", "pid", "process_start")
+        ):
+            continue
+        try:
+            lease_path.unlink()
+        except FileNotFoundError:
+            pass
+
+
+@contextmanager
+def staging_process_owner(storage_root: Path, owner_id: str) -> Iterator[None]:
+    """Publish an exact process record that a controller can safely terminate."""
+    owner_path = _staging_owner_path(storage_root, owner_id)
+    owner_path.parent.mkdir(parents=True, exist_ok=True)
+    pid = os.getpid()
+    owner_record = {
+        "version": 1,
+        "owner_id": _validated_staging_owner_id(owner_id),
+        "pid": pid,
+        "process_start": _linux_process_start(pid),
+        "created_at": time.time(),
+    }
+    temporary_path = owner_path.with_suffix(f".{uuid4().hex}.tmp")
+    try:
+        temporary_path.write_text(
+            json.dumps(owner_record, sort_keys=True, separators=(",", ":")),
+            encoding="utf-8",
+        )
+        os.chmod(temporary_path, 0o600)
+        os.replace(temporary_path, owner_path)
+    finally:
+        temporary_path.unlink(missing_ok=True)
+    try:
+        yield
+    finally:
+        if _owner_record_matches(owner_path, owner_record):
+            try:
+                owner_path.unlink()
+            except FileNotFoundError:
+                pass
+
+
+def cancel_staging_process(storage_root: Path, owner_id: str) -> bool:
+    """Terminate only the live process matching one controller owner record."""
+    owner_path = _staging_owner_path(storage_root, owner_id)
+    try:
+        owner = json.loads(owner_path.read_text(encoding="utf-8"))
+    except FileNotFoundError:
+        return False
+    except (json.JSONDecodeError, OSError) as exc:
+        raise SshWorkerError(f"Staging owner record is unreadable: {exc}") from exc
+    if not isinstance(owner, Mapping):
+        raise SshWorkerError("Staging owner record must be an object.")
+    try:
+        pid = int(owner.get("pid"))
+    except (TypeError, ValueError) as exc:
+        raise SshWorkerError("Staging owner record has an invalid PID.") from exc
+    if pid <= 0:
+        raise SshWorkerError("Staging owner record has an invalid PID.")
+    expected_start = str(owner.get("process_start") or "")
+    if _linux_process_start(pid) != expected_start:
+        _remove_owned_snapshot_leases(storage_root, owner)
+        owner_path.unlink(missing_ok=True)
+        return False
+    try:
+        command_line = Path(f"/proc/{pid}/cmdline").read_bytes().split(b"\0")
+    except (FileNotFoundError, OSError):
+        owner_path.unlink(missing_ok=True)
+        return False
+    encoded_owner = owner_id.encode("utf-8")
+    if (
+        b"stage-profiles" not in command_line
+        or b"--owner-id" not in command_line
+        or encoded_owner not in command_line
+    ):
+        raise SshWorkerError(
+            f"Refusing to terminate PID {pid}; its command does not match owner "
+            f"{owner_id!r}."
+        )
+    os.kill(pid, signal.SIGTERM)
+    deadline = time.monotonic() + _STAGING_TERMINATE_GRACE_SECONDS
+    while time.monotonic() < deadline:
+        if _linux_process_start(pid) != expected_start:
+            _remove_owned_snapshot_leases(storage_root, owner)
+            owner_path.unlink(missing_ok=True)
+            return True
+        time.sleep(0.1)
+    os.kill(pid, signal.SIGKILL)
+    _remove_owned_snapshot_leases(storage_root, owner)
+    owner_path.unlink(missing_ok=True)
+    return True
+
+
+def _decode_resolved_profiles(encoded_payload: str | None) -> dict[str, dict[str, Any]]:
+    """Decode planner-resolved profile metadata from a credential-free CLI value."""
+    if not encoded_payload:
+        return {}
+    try:
+        decoded = base64.urlsafe_b64decode(encoded_payload.encode("ascii"))
+        payload = json.loads(decoded.decode("utf-8"))
+    except (UnicodeDecodeError, ValueError, json.JSONDecodeError) as exc:
+        raise ValueError("Resolved LLM profile payload is invalid.") from exc
+    if not isinstance(payload, Mapping):
+        raise ValueError("Resolved LLM profile payload must be an object.")
+    return {
+        str(reference): dict(profile)
+        for reference, profile in payload.items()
+        if isinstance(reference, str) and isinstance(profile, Mapping)
+    }
+
+
 def _copy_request_frames(source: BinaryIO, destination: socket.socket) -> None:
     """Copy exactly one request command and its optional input frame."""
     first_frame = read_frame(source)
@@ -341,7 +523,13 @@ def _argument_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
         "command",
-        choices=("serve", "client", "runtime-info", "stage-profiles"),
+        choices=(
+            "serve",
+            "client",
+            "runtime-info",
+            "stage-profiles",
+            "cancel-staging",
+        ),
     )
     parser.add_argument(
         "--socket",
@@ -355,6 +543,16 @@ def _argument_parser() -> argparse.ArgumentParser:
         default=[],
         help="Curated profile or Hugging Face model reference to stage.",
     )
+    parser.add_argument(
+        "--owner-id",
+        default="",
+        help="Controller-issued identity used for targeted staging cancellation.",
+    )
+    parser.add_argument(
+        "--resolved-profiles",
+        default="",
+        help="URL-safe base64 encoded planner-resolved profile metadata.",
+    )
     parser.add_argument(
         "--storage-root",
         type=Path,
@@ -376,7 +574,34 @@ def main(argv: Sequence[str] | None = None) -> int:
     if arguments.command == "stage-profiles":
         if not arguments.model_reference:
             raise ValueError("stage-profiles requires at least one --model-reference.")
-        stage_profiles(arguments.model_reference, arguments.storage_root)
+        resolved_profiles = _decode_resolved_profiles(arguments.resolved_profiles)
+        if arguments.owner_id:
+            with staging_process_owner(arguments.storage_root, arguments.owner_id):
+                stage_profiles(
+                    arguments.model_reference,
+                    arguments.storage_root,
+                    resolved_profiles=resolved_profiles,
+                    owner_id=arguments.owner_id,
+                )
+        else:
+            stage_profiles(
+                arguments.model_reference,
+                arguments.storage_root,
+                resolved_profiles=resolved_profiles,
+            )
+        return 0
+    if arguments.command == "cancel-staging":
+        if not arguments.owner_id:
+            raise ValueError("cancel-staging requires --owner-id.")
+        print(
+            json.dumps(
+                {"cancelled": cancel_staging_process(
+                    arguments.storage_root,
+                    arguments.owner_id,
+                )},
+                sort_keys=True,
+            )
+        )
         return 0
     print(json.dumps(runtime_info(), sort_keys=True))
     return 0

runtime_environment.py

@@ -21,7 +21,7 @@ else:  # pragma: no cover - the stable cloud entrypoint imports this module top-
 
 logger = logging.getLogger(__name__)
 
-REMOTE_APP_PROTOCOL_VERSION = 9
+REMOTE_APP_PROTOCOL_VERSION = 10
 REMOTE_RUNTIME_DEPENDENCY_LAYER_VERSION = 1
 REMOTE_PYTHON_VERSION = "3.13"
 REMOTE_MODAL_SDK_SPEC = "modal==1.4.2"

ssh_executor.py

@@ -6,14 +6,17 @@ import asyncio
 import logging
 import subprocess
 import threading
+import time
 from collections.abc import Iterator, Mapping, Sequence
 from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Any
+from uuid import uuid4
 
 if __package__:
     from .durable_state import stable_remote_invocation_id
     from .llm_profiles import (
+        encoded_resolved_llm_profile_payloads,
         llm_model_references_from_payload,
         rewrite_llm_model_references,
     )
@@ -28,11 +31,17 @@ if __package__:
     )
     from .serialization import deserialize_node_outputs, serialize_node_inputs
     from .settings import ModalSyncSettings, get_settings
+    from .staging_process import (
+        RemoteStagingProcessError,
+        consume_staging_process,
+        terminate_staging_transport,
+    )
     from .ssh_docker import SshDockerController, SshDockerError
     from .ssh_runtime import SshRuntimeManager, SshRuntimeSpec
 else:  # pragma: no cover - top-level remote imports.
     from durable_state import stable_remote_invocation_id
     from llm_profiles import (
+        encoded_resolved_llm_profile_payloads,
         llm_model_references_from_payload,
         rewrite_llm_model_references,
     )
@@ -47,6 +56,11 @@ else:  # pragma: no cover - top-level remote imports.
     )
     from serialization import deserialize_node_outputs, serialize_node_inputs
     from settings import ModalSyncSettings, get_settings
+    from staging_process import (
+        RemoteStagingProcessError,
+        consume_staging_process,
+        terminate_staging_transport,
+    )
     from ssh_docker import SshDockerController, SshDockerError
     from ssh_runtime import SshRuntimeManager, SshRuntimeSpec
 
@@ -56,7 +70,7 @@ _STAGED_SSH_PROFILE_RESULTS: dict[tuple[str, str], dict[str, Any]] = {}
 _ACTIVE_SSH_STAGERS_LOCK = threading.Lock()
 _ACTIVE_SSH_STAGERS: dict[
     str,
-    tuple[SshRuntimeManager, SshRuntimeSpec],
+    tuple[SshRuntimeManager, SshRuntimeSpec, str, subprocess.Popen[bytes]],
 ] = {}
 
 
@@ -119,17 +133,23 @@ class SshDockerExecutorClient:
         with _ACTIVE_SSH_STAGERS_LOCK:
             active_stager = _ACTIVE_SSH_STAGERS.get(invocation_id)
         if active_stager is not None:
-            active_manager, active_spec = active_stager
-            stopped = active_manager.stop_worker(active_spec.worker_index)
+            active_manager, active_spec, owner_id, process = active_stager
+            terminate_staging_transport(process)
+            remote_cancelled = self._cancel_remote_stager(
+                active_manager,
+                active_spec,
+                owner_id,
+                wait_for_owner_seconds=5.0,
+            )
             logger.info(
-                "Stopped SSH worker to cancel active model staging invocation=%s "
-                "environment=%s worker_index=%d stopped=%s.",
+                "Cancelled SSH model staging invocation=%s "
+                "environment=%s worker_index=%d remote_process_found=%s.",
                 invocation_id,
                 active_manager.controller.host.environment_id,
                 active_spec.worker_index,
-                stopped,
+                remote_cancelled,
             )
-            return stopped
+            return True
         manager, spec = self._runtime(payload)
         request = encode_json_frame(
             RemoteFrameKind.CANCEL,
@@ -296,18 +316,39 @@ class SshDockerExecutorClient:
             "remote.ssh_worker",
             "stage-profiles",
         ]
+        invocation_id = str(payload["invocation_id"])
+        owner_id = f"ssh:{invocation_id}:{uuid4().hex}"
+        arguments.extend(("--owner-id", owner_id))
         for model_reference in model_references:
             arguments.extend(("--model-reference", model_reference))
+        resolved_profiles = encoded_resolved_llm_profile_payloads(
+            payload,
+            model_references,
+        )
+        if resolved_profiles is not None:
+            arguments.extend(("--resolved-profiles", resolved_profiles))
         process = manager.controller.docker_popen(tuple(arguments))
-        invocation_id = str(payload["invocation_id"])
         with _ACTIVE_SSH_STAGERS_LOCK:
-            _ACTIVE_SSH_STAGERS[invocation_id] = (manager, spec)
+            _ACTIVE_SSH_STAGERS[invocation_id] = (
+                manager,
+                spec,
+                owner_id,
+                process,
+            )
         try:
-            results = self._consume_stager_output(
+            results = consume_staging_process(
                 process,
                 payload,
                 _emit_local_llm_staging_progress,
+                provider_label="SSH",
+                abort_remote=lambda: self._cancel_remote_stager(
+                    manager,
+                    spec,
+                    owner_id,
+                ),
             )
+        except RemoteStagingProcessError as exc:
+            raise SshRemoteInvocationError(str(exc)) from exc
         finally:
             with _ACTIVE_SSH_STAGERS_LOCK:
                 _ACTIVE_SSH_STAGERS.pop(invocation_id, None)
@@ -325,44 +366,47 @@ class SshDockerExecutorClient:
         )
         return results
 
-    def _consume_stager_output(
-        self,
-        process: subprocess.Popen[bytes],
-        payload: dict[str, Any],
-        progress_callback: Any,
-    ) -> list[dict[str, Any]]:
-        """Consume JSON-line progress and the terminal SSH staging result."""
-        if process.stdout is None or process.stderr is None:
-            process.kill()
-            raise SshRemoteInvocationError("SSH model stager did not expose output streams.")
-        stderr_chunks: list[bytes] = []
-        stderr_thread = threading.Thread(
-            target=_collect_bounded_stream,
-            args=(process.stderr, stderr_chunks),
-            daemon=True,
-        )
-        stderr_thread.start()
-        results: list[dict[str, Any]] | None = None
-        for raw_line in process.stdout:
-            try:
-                event = decode_json_payload(raw_line)
-            except (RemoteProtocolError, UnicodeDecodeError) as exc:
-                process.terminate()
-                raise SshRemoteInvocationError(
-                    "SSH model stager returned invalid progress JSON."
-                ) from exc
-            if event.get("kind") == "progress":
-                progress_callback(payload, event)
-            elif event.get("kind") == "result" and isinstance(event.get("results"), list):
-                results = [dict(result) for result in event["results"] if isinstance(result, Mapping)]
-        returncode = process.wait()
-        stderr_thread.join(timeout=1.0)
-        if returncode != 0 or results is None:
-            diagnostics = b"".join(stderr_chunks).decode("utf-8", errors="replace").strip()
-            raise SshRemoteInvocationError(
-                f"SSH model staging failed: {diagnostics or 'no result returned'}"
+    @staticmethod
+    def _cancel_remote_stager(
+        manager: SshRuntimeManager,
+        spec: SshRuntimeSpec,
+        owner_id: str,
+        *,
+        wait_for_owner_seconds: float = 0.0,
+    ) -> bool:
+        """Terminate one exact owner-tagged stager without recycling its worker."""
+        deadline = time.monotonic() + max(0.0, wait_for_owner_seconds)
+        while True:
+            result = manager.controller.docker(
+                (
+                    "exec",
+                    spec.container_name,
+                    "python",
+                    "-m",
+                    "remote.ssh_worker",
+                    "cancel-staging",
+                    "--owner-id",
+                    owner_id,
+                ),
+                timeout_seconds=15.0,
+                check=False,
             )
-        return results
+            if result.returncode != 0:
+                logger.warning(
+                    "SSH remote stager cancellation failed owner=%s error=%s.",
+                    owner_id,
+                    result.stderr_text.strip() or result.stdout_text.strip(),
+                )
+                return False
+            try:
+                response = decode_json_payload(result.stdout)
+            except (RemoteProtocolError, UnicodeDecodeError):
+                return False
+            if response.get("cancelled"):
+                return True
+            if time.monotonic() >= deadline:
+                return False
+            time.sleep(0.25)
 
     def _cache_staged_profiles(
         self,

staging_process.py

@@ -0,0 +1,244 @@
+"""Bounded controller-side consumption for remote model-staging processes."""
+
+from __future__ import annotations
+
+import logging
+import math
+import os
+import queue
+import subprocess
+import threading
+import time
+from collections.abc import Callable, Mapping
+from typing import Any
+
+if __package__:
+    from .remote_protocol import RemoteProtocolError, decode_json_payload
+else:  # pragma: no cover - direct debugging imports.
+    from remote_protocol import RemoteProtocolError, decode_json_payload
+
+_DEFAULT_NO_PROGRESS_TIMEOUT_SECONDS = 600.0
+_STREAM_END = object()
+logger = logging.getLogger(__name__)
+
+
+class RemoteStagingProcessError(RuntimeError):
+    """Raised when a remote staging subprocess fails or stops reporting progress."""
+
+
+def staging_no_progress_timeout_seconds() -> float:
+    """Return the bounded interval allowed without a staging JSON event."""
+    raw_value = os.getenv("COMFY_MODAL_LLM_STAGE_NO_PROGRESS_TIMEOUT_SECONDS")
+    timeout = (
+        _DEFAULT_NO_PROGRESS_TIMEOUT_SECONDS
+        if raw_value is None
+        else float(raw_value)
+    )
+    if not math.isfinite(timeout) or timeout <= 0:
+        raise ValueError(
+            "COMFY_MODAL_LLM_STAGE_NO_PROGRESS_TIMEOUT_SECONDS must be positive."
+        )
+    return timeout
+
+
+def _collect_bounded_stream(stream: Any, chunks: list[bytes]) -> None:
+    """Drain diagnostics while retaining at most the most recent MiB."""
+    while chunk := stream.read(65536):
+        chunks.append(chunk)
+        if sum(len(value) for value in chunks) > 1024 * 1024:
+            del chunks[:-8]
+
+
+def _read_progress_lines(stream: Any, output: queue.Queue[Any]) -> None:
+    """Read blocking process output without blocking timeout enforcement."""
+    try:
+        for raw_line in stream:
+            output.put(raw_line)
+    finally:
+        output.put(_STREAM_END)
+
+
+def _stop_process(process: subprocess.Popen[bytes]) -> None:
+    """Boundedly terminate one local transport process."""
+    if process.poll() is not None:
+        return
+    process.terminate()
+    try:
+        process.wait(timeout=5.0)
+    except subprocess.TimeoutExpired:
+        process.kill()
+        process.wait(timeout=5.0)
+
+
+def terminate_staging_transport(process: subprocess.Popen[bytes]) -> None:
+    """Terminate a controller-side SSH transport for one active stager."""
+    _stop_process(process)
+
+
+def _abort_remote_safely(abort_remote: Callable[[], None]) -> None:
+    """Attempt remote cleanup without hiding the primary staging failure."""
+    try:
+        abort_remote()
+    except (OSError, RuntimeError, ValueError) as exc:
+        logger.warning("Unable to clean up the remote model stager: %s", exc)
+
+
+def _next_progress_line(
+    progress_lines: queue.Queue[Any],
+    deadline: float,
+    timeout_seconds: float,
+    provider_label: str,
+) -> Any:
+    """Return the next line or raise the provider-specific silence failure."""
+    remaining = deadline - time.monotonic()
+    if remaining <= 0:
+        raise RemoteStagingProcessError(
+            f"{provider_label} model staging produced no progress for "
+            f"{timeout_seconds:.0f} seconds."
+        )
+    try:
+        return progress_lines.get(timeout=remaining)
+    except queue.Empty as exc:
+        raise RemoteStagingProcessError(
+            f"{provider_label} model staging produced no progress for "
+            f"{timeout_seconds:.0f} seconds."
+        ) from exc
+
+
+def _consume_progress_lines(
+    progress_lines: queue.Queue[Any],
+    payload: dict[str, Any],
+    progress_callback: Callable[[dict[str, Any], Mapping[str, Any]], None],
+    *,
+    provider_label: str,
+    timeout_seconds: float,
+) -> list[dict[str, Any]] | None:
+    """Consume staging JSON until stdout closes, resetting timeout per event."""
+    results: list[dict[str, Any]] | None = None
+    deadline = time.monotonic() + timeout_seconds
+    while True:
+        raw_line = _next_progress_line(
+            progress_lines,
+            deadline,
+            timeout_seconds,
+            provider_label,
+        )
+        if raw_line is _STREAM_END:
+            return results
+        deadline = time.monotonic() + timeout_seconds
+        try:
+            event = decode_json_payload(raw_line)
+        except (RemoteProtocolError, UnicodeDecodeError) as exc:
+            raise RemoteStagingProcessError(
+                f"{provider_label} model stager returned invalid progress JSON."
+            ) from exc
+        if event.get("kind") == "progress":
+            progress_callback(payload, event)
+        elif event.get("kind") == "result" and isinstance(
+            event.get("results"), list
+        ):
+            results = [
+                dict(result)
+                for result in event["results"]
+                if isinstance(result, Mapping)
+            ]
+
+
+def _validated_staging_result(
+    process: subprocess.Popen[bytes],
+    results: list[dict[str, Any]] | None,
+    stderr_chunks: list[bytes],
+    provider_label: str,
+) -> list[dict[str, Any]]:
+    """Wait for a clean process exit and require one terminal result event."""
+    try:
+        returncode = process.wait(timeout=5.0)
+    except subprocess.TimeoutExpired as exc:
+        raise RemoteStagingProcessError(
+            f"{provider_label} model stager kept running after closing output."
+        ) from exc
+    if returncode == 0 and results is not None:
+        return results
+    diagnostics = b"".join(stderr_chunks).decode("utf-8", errors="replace").strip()
+    raise RemoteStagingProcessError(
+        f"{provider_label} model staging failed: "
+        f"{diagnostics or f'exit status {returncode}'}."
+    )
+
+
+def consume_staging_process(
+    process: subprocess.Popen[bytes],
+    payload: dict[str, Any],
+    progress_callback: Callable[[dict[str, Any], Mapping[str, Any]], None],
+    *,
+    provider_label: str,
+    abort_remote: Callable[[], None],
+    timeout_seconds: float | None = None,
+) -> list[dict[str, Any]]:
+    """Consume JSON events and abort local plus remote work after inactivity."""
+    if process.stdout is None or process.stderr is None:
+        _stop_process(process)
+        _abort_remote_safely(abort_remote)
+        raise RemoteStagingProcessError(
+            f"{provider_label} model stager did not expose output streams."
+        )
+    try:
+        resolved_timeout = (
+            staging_no_progress_timeout_seconds()
+            if timeout_seconds is None
+            else timeout_seconds
+        )
+        if not math.isfinite(resolved_timeout) or resolved_timeout <= 0:
+            raise ValueError("Staging no-progress timeout must be positive.")
+    except ValueError:
+        _stop_process(process)
+        _abort_remote_safely(abort_remote)
+        raise
+    stderr_chunks: list[bytes] = []
+    stderr_thread = threading.Thread(
+        target=_collect_bounded_stream,
+        args=(process.stderr, stderr_chunks),
+        name=f"{provider_label.lower()}-stager-stderr",
+        daemon=True,
+    )
+    progress_lines: queue.Queue[Any] = queue.Queue()
+    stdout_thread = threading.Thread(
+        target=_read_progress_lines,
+        args=(process.stdout, progress_lines),
+        name=f"{provider_label.lower()}-stager-stdout",
+        daemon=True,
+    )
+    stderr_thread.start()
+    stdout_thread.start()
+    completed = False
+    try:
+        results = _consume_progress_lines(
+            progress_lines,
+            payload,
+            progress_callback,
+            provider_label=provider_label,
+            timeout_seconds=resolved_timeout,
+        )
+        stderr_thread.join(timeout=1.0)
+        validated_results = _validated_staging_result(
+            process,
+            results,
+            stderr_chunks,
+            provider_label,
+        )
+        completed = True
+        return validated_results
+    finally:
+        if not completed:
+            _stop_process(process)
+            _abort_remote_safely(abort_remote)
+        stdout_thread.join(timeout=1.0)
+        stderr_thread.join(timeout=1.0)
+
+
+__all__ = [
+    "RemoteStagingProcessError",
+    "consume_staging_process",
+    "staging_no_progress_timeout_seconds",
+    "terminate_staging_transport",
+]

tests/conftest.py

@@ -335,6 +335,12 @@ def llm_staging_module(extension_package: object) -> object:
     return importlib.import_module(f"{PACKAGE_NAME}.llm_staging")
 
 
+@pytest.fixture(scope="session")
+def staging_process_module(extension_package: object) -> object:
+    """Return the bounded remote staging process controller module."""
+    return importlib.import_module(f"{PACKAGE_NAME}.staging_process")
+
+
 @pytest.fixture(scope="session")
 def modal_llm_runtime_module(extension_package: object) -> object:
     """Return the resident Modal LLM runtime module."""

tests/test_api_intercept.py

@@ -4400,6 +4400,34 @@ def test_snapshot_profile_stamping_excludes_llm_phase_from_comfy_profile(
     assert phases[2]["snapshot_profile_key"] == snapshot_profile_key
 
 
+def test_planner_resolved_llm_profile_is_attached_to_matching_payload(
+    api_intercept_module: Any,
+    llm_profiles_module: Any,
+    tmp_path: Path,
+) -> None:
+    """The execution payload should carry metadata already resolved by planning."""
+    profile = llm_profiles_module.get_llm_profile("smolvlm2-2.2b-instruct")
+    payload = {
+        "component_id": "llm",
+        "subgraph_prompt": {
+            "llm": {
+                "class_type": "ModalLLM",
+                "inputs": {"model_profile": profile.profile_id},
+            }
+        },
+    }
+
+    api_intercept_module._attach_resolved_llm_profiles(
+        payload,
+        {profile.profile_id: profile},
+        SimpleNamespace(local_storage_root=tmp_path),
+    )
+
+    entry = payload["resolved_llm_profiles"][profile.profile_id]
+    assert entry["profile"] == profile.to_mapping()
+    assert entry["security_scan_complete"] is True
+
+
 def test_planner_attaches_next_distinct_affinity_as_speculative_prewarm_target(
     api_intercept_module: Any,
 ) -> None:

tests/test_modal_llm.py

@@ -7,6 +7,7 @@ import base64
 import json
 import logging
 import sys
+import threading
 from dataclasses import dataclass, replace
 from fractions import Fraction
 from io import BytesIO, StringIO
@@ -823,6 +824,7 @@ def test_cpu_stager_writes_completion_marker_and_reuses_snapshot(
     ]
     assert [event.stage for event in progress] == [
         "snapshot_check",
+        "disk_check",
         "download_preparing",
         "download",
         "download",
@@ -951,6 +953,201 @@ def test_provider_neutral_stager_resolves_and_stages_model_reference(
     assert progress[0].model_reference == "owner/model"
 
 
+def test_provider_neutral_stager_uses_planner_resolved_profile(
+    llm_staging_module: Any,
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    """Remote staging must not repeat Hugging Face inspection done by the planner."""
+    digest = "c" * 64
+    raw_profile = {
+        "id": f"hf-{digest}",
+        "display_name": "Owner Model",
+        "repository": "owner/model",
+        "revision": "9" * 40,
+        "dtype": "bfloat16",
+        "modalities": ["text"],
+        "estimated_vram_gb": 12,
+        "max_context_tokens": 4096,
+        "max_images": 1,
+        "max_video_frames": 1,
+        "max_file_bytes": 1024,
+        "max_file_characters": 1024,
+        "allow_mixed_image_video": False,
+        "trust_remote_code": False,
+        "schema_version": 2,
+        "source": "generated",
+        "profile_digest": digest,
+        "backend": "transformers",
+        "artifact_bytes": 1024,
+    }
+    staged = llm_staging_module.StagedModelSnapshot(
+        profile_id=raw_profile["id"],
+        repository="owner/model",
+        revision="9" * 40,
+        path=str(tmp_path / "snapshot"),
+        downloaded=False,
+        elapsed_seconds=0.0,
+    )
+    monkeypatch.setattr(
+        llm_staging_module,
+        "stage_model_profile",
+        lambda *_args, **_kwargs: staged,
+    )
+    progress: list[Any] = []
+
+    results = llm_staging_module.resolve_and_stage_model_references(
+        ["owner/model"],
+        tmp_path,
+        resolved_profiles={
+            "owner/model": {
+                "profile": raw_profile,
+                "security_scan_complete": False,
+            }
+        },
+        progress_callback=progress.append,
+        owner_id="vast:invocation:owner",
+    )
+
+    assert results[0].profile_id == raw_profile["id"]
+    assert results[0].security_scan_complete is False
+    assert results[0].manifest_created is True
+    assert progress[0].stage == "resolved_metadata"
+    assert "planner-resolved" in progress[0].message
+
+
+def test_snapshot_lease_reclaims_dead_local_owner(
+    llm_staging_module: Any,
+    tmp_path: Path,
+) -> None:
+    """A dead process record should never force a two-hour stale-lock wait."""
+    snapshot_path = tmp_path / "llm_models" / "profile" / ("1" * 40)
+    snapshot_path.parent.mkdir(parents=True)
+    lease_path = snapshot_path.parent / f".{snapshot_path.name}.download.lock"
+    lease_path.write_text(
+        json.dumps(
+            {
+                "owner_id": "orphan",
+                "host_id": llm_staging_module.socket.gethostname(),
+                "pid": 2_000_000_000,
+                "process_start": "missing",
+                "token": "old-token",
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    with llm_staging_module._snapshot_lease(
+        snapshot_path,
+        model_label="test model",
+        owner_id="replacement",
+    ):
+        owner = json.loads(lease_path.read_text(encoding="utf-8"))
+        assert owner["owner_id"] == "replacement"
+
+    assert not lease_path.exists()
+
+
+def test_snapshot_lease_recognizes_live_owner_without_procfs(
+    llm_staging_module: Any,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Portable PID probing should keep a live macOS staging owner valid."""
+    monkeypatch.setattr(
+        llm_staging_module,
+        "_process_start_identity",
+        lambda _pid: None,
+    )
+
+    alive = llm_staging_module._local_lease_owner_is_alive(
+        {
+            "host_id": llm_staging_module.socket.gethostname(),
+            "pid": llm_staging_module.os.getpid(),
+            "process_start": None,
+        }
+    )
+
+    assert alive is True
+
+
+def test_snapshot_lease_reclaims_missing_foreign_heartbeat(
+    llm_staging_module: Any,
+    tmp_path: Path,
+) -> None:
+    """A vanished prior container should not leave a structured lease for hours."""
+    snapshot_path = tmp_path / "llm_models" / "profile" / ("2" * 40)
+    snapshot_path.parent.mkdir(parents=True)
+    lease_path = snapshot_path.parent / f".{snapshot_path.name}.download.lock"
+    lease_path.write_text(
+        json.dumps(
+            {
+                "owner_id": "old-container",
+                "host_id": "different-container",
+                "pid": 123,
+                "process_start": "456",
+                "token": "old-token",
+            }
+        ),
+        encoding="utf-8",
+    )
+    stale_time = (
+        llm_staging_module.time.time()
+        - llm_staging_module._DEFAULT_LEASE_HEARTBEAT_STALE_SECONDS
+        - 1
+    )
+    llm_staging_module.os.utime(lease_path, (stale_time, stale_time))
+
+    with llm_staging_module._snapshot_lease(
+        snapshot_path,
+        model_label="test model",
+        owner_id="replacement",
+    ):
+        owner = json.loads(lease_path.read_text(encoding="utf-8"))
+        assert owner["owner_id"] == "replacement"
+
+    assert not lease_path.exists()
+
+
+def test_snapshot_disk_preflight_rejects_insufficient_capacity(
+    llm_staging_module: Any,
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    """Staging should fail before download when artifacts plus reserve do not fit."""
+    profile = llm_staging_module.get_llm_profile("smolvlm2-2.2b-instruct")
+    monkeypatch.setattr(
+        llm_staging_module.shutil,
+        "disk_usage",
+        lambda _path: SimpleNamespace(free=profile.artifact_bytes),
+    )
+
+    with pytest.raises(RuntimeError, match="Insufficient disk space"):
+        llm_staging_module._preflight_snapshot_capacity(
+            tmp_path / "snapshot",
+            profile,
+        )
+
+
+def test_direct_stager_payload_encodes_planner_profiles(
+    llm_profiles_module: Any,
+) -> None:
+    """SSH and Vast CLI staging should receive the validated planner envelope."""
+    profile = llm_profiles_module.get_llm_profile("smolvlm2-2.2b-instruct")
+    entry = {
+        "profile": profile.to_mapping(),
+        "security_scan_complete": True,
+    }
+
+    encoded = llm_profiles_module.encoded_resolved_llm_profile_payloads(
+        {"resolved_llm_profiles": {profile.profile_id: entry}},
+        [profile.profile_id],
+    )
+
+    assert encoded is not None
+    decoded = json.loads(base64.urlsafe_b64decode(encoded).decode("utf-8"))
+    assert decoded == {profile.profile_id: entry}
+
+
 def test_weight_snapshot_is_shared_across_runtime_profiles(
     llm_staging_module: Any,
     tmp_path: Path,
@@ -2751,6 +2948,43 @@ def test_remote_dispatch_streams_cpu_llm_staging_progress(
     )
 
 
+def test_modal_staging_stream_has_bounded_no_progress_wait(
+    remote_modal_app_module: Any,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A silent Modal generator should be closed instead of hanging forever."""
+
+    class SilentStream:
+        """Block until the controller asks this fake remote stream to close."""
+
+        def __init__(self) -> None:
+            self.closed = threading.Event()
+
+        def __iter__(self) -> Any:
+            """Wait without producing a staging event."""
+            self.closed.wait(timeout=5)
+            return
+            yield  # pragma: no cover - makes this method a generator.
+
+        def close(self) -> None:
+            """Model successful cancellation of the remote generator."""
+            self.closed.set()
+
+    monkeypatch.setenv(
+        "COMFY_MODAL_LLM_STAGE_NO_PROGRESS_TIMEOUT_SECONDS",
+        "0.05",
+    )
+    stream = SilentStream()
+
+    with pytest.raises(
+        remote_modal_app_module.ModalRemoteInvocationError,
+        match="produced no progress",
+    ):
+        list(remote_modal_app_module._bounded_modal_stage_events(stream))
+
+    assert stream.closed.is_set()
+
+
 def test_llm_staging_progress_targets_the_matching_llm_node(
     remote_modal_app_module: Any,
     monkeypatch: pytest.MonkeyPatch,

tests/test_runtime_environment.py

@@ -64,7 +64,7 @@ def test_remote_environment_is_fully_pinned(runtime_environment_module: Any) ->
     torch_packages = runtime_environment_module.remote_torch_packages()
 
     assert runtime_environment_module.DEFAULT_MODAL_GPU == "RTX-PRO-6000"
-    assert runtime_environment_module.REMOTE_APP_PROTOCOL_VERSION == 9
+    assert runtime_environment_module.REMOTE_APP_PROTOCOL_VERSION == 10
     assert runtime_environment_module.REMOTE_PYTHON_VERSION == "3.13"
     assert apt_packages == ("build-essential", "libgl1", "libglib2.0-0")
     assert all("==" in requirement for requirement in runtime_packages)

tests/test_ssh_executor.py

@@ -273,17 +273,17 @@ def test_ssh_dispatch_stages_and_rewrites_hugging_face_model_reference(
     assert stage_calls == [[requested_model]]
 
 
-def test_ssh_cancel_stops_worker_during_model_staging(
+def test_ssh_cancel_targets_owner_during_model_staging(
     ssh_executor_module: Any,
+    monkeypatch: Any,
     tmp_path: Path,
 ) -> None:
-    """Cancellation must stop a worker whose CPU stager is still downloading."""
-    stopped: list[int] = []
+    """Cancellation must target one stager without stopping its worker."""
     manager = SimpleNamespace(
         controller=SimpleNamespace(host=SimpleNamespace(environment_id="lambda")),
-        stop_worker=lambda worker_index: stopped.append(worker_index) or True,
     )
     spec = SimpleNamespace(worker_index=3)
+    process = SimpleNamespace()
     invocation_id = "RIV_staging"
     client = ssh_executor_module.SshDockerExecutorClient(
         registry=SimpleNamespace(),
@@ -291,7 +291,24 @@ def test_ssh_cancel_stops_worker_during_model_staging(
         settings=SimpleNamespace(comfyui_root=tmp_path, app_name="app"),
     )
     with ssh_executor_module._ACTIVE_SSH_STAGERS_LOCK:
-        ssh_executor_module._ACTIVE_SSH_STAGERS[invocation_id] = (manager, spec)
+        ssh_executor_module._ACTIVE_SSH_STAGERS[invocation_id] = (
+            manager,
+            spec,
+            "ssh:RIV_staging:owner",
+            process,
+        )
+    targeted: list[tuple[Any, Any, str]] = []
+    terminated: list[Any] = []
+    monkeypatch.setattr(
+        ssh_executor_module,
+        "terminate_staging_transport",
+        terminated.append,
+    )
+    client._cancel_remote_stager = (  # type: ignore[method-assign]
+        lambda active_manager, active_spec, owner_id, **_kwargs: (
+            targeted.append((active_manager, active_spec, owner_id)) or True
+        )
+    )
     try:
         cancelled = client.cancel({}, invocation_id)
     finally:
@@ -299,4 +316,5 @@ def test_ssh_cancel_stops_worker_during_model_staging(
             ssh_executor_module._ACTIVE_SSH_STAGERS.pop(invocation_id, None)
 
     assert cancelled is True
-    assert stopped == [3]
+    assert terminated == [process]
+    assert targeted == [(manager, spec, "ssh:RIV_staging:owner")]

tests/test_ssh_worker.py

@@ -93,6 +93,68 @@ def test_worker_stage_profiles_streams_progress_and_result(
     assert events[1] == {"kind": "result", "results": results}
 
 
+def test_worker_cancels_only_matching_owned_stager(
+    ssh_worker_module: Any,
+    monkeypatch: Any,
+    tmp_path: Path,
+) -> None:
+    """The cancellation command should validate identity before killing a PID."""
+    owner_id = "vast:RIV_test:owner"
+    pid = 12345
+    owner_path = ssh_worker_module._staging_owner_path(tmp_path, owner_id)
+    owner_path.parent.mkdir(parents=True)
+    owner_path.write_text(
+        json.dumps(
+            {
+                "owner_id": owner_id,
+                "pid": pid,
+                "process_start": "start-tick",
+            }
+        ),
+        encoding="utf-8",
+    )
+    lease_path = tmp_path / "llm_models" / "repo" / ".revision.download.lock"
+    lease_path.parent.mkdir(parents=True)
+    lease_path.write_text(
+        json.dumps(
+            {
+                "owner_id": owner_id,
+                "pid": pid,
+                "process_start": "start-tick",
+                "token": "lease-token",
+            }
+        ),
+        encoding="utf-8",
+    )
+    starts = iter(("start-tick", None))
+    killed: list[tuple[int, int]] = []
+    monkeypatch.setattr(
+        ssh_worker_module,
+        "_linux_process_start",
+        lambda _pid: next(starts),
+    )
+    monkeypatch.setattr(
+        ssh_worker_module.Path,
+        "read_bytes",
+        lambda _path: (
+            b"python\0-m\0remote.ssh_worker\0stage-profiles\0--owner-id\0"
+            + owner_id.encode("utf-8")
+            + b"\0"
+        ),
+    )
+    monkeypatch.setattr(
+        ssh_worker_module.os,
+        "kill",
+        lambda target_pid, sent_signal: killed.append((target_pid, sent_signal)),
+    )
+
+    assert ssh_worker_module.cancel_staging_process(tmp_path, owner_id) is True
+
+    assert not owner_path.exists()
+    assert not lease_path.exists()
+    assert killed == [(pid, ssh_worker_module.signal.SIGTERM)]
+
+
 def test_worker_request_relay_preserves_framed_binary_payload(
     ssh_worker_module: Any,
     remote_protocol_module: Any,

tests/test_staging_process.py

@@ -0,0 +1,90 @@
+"""Tests for bounded controller-side remote staging processes."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from typing import Any
+
+import pytest
+
+
+def test_staging_process_aborts_remote_after_no_progress(
+    staging_process_module: Any,
+) -> None:
+    """A silent remote stager must be terminated and explicitly cancelled."""
+    process = subprocess.Popen(
+        [sys.executable, "-c", "import time; time.sleep(60)"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    aborted: list[bool] = []
+
+    with pytest.raises(
+        staging_process_module.RemoteStagingProcessError,
+        match="produced no progress",
+    ):
+        staging_process_module.consume_staging_process(
+            process,
+            {},
+            lambda _payload, _event: None,
+            provider_label="Test",
+            abort_remote=lambda: aborted.append(True),
+            timeout_seconds=0.05,
+        )
+
+    assert process.poll() is not None
+    assert aborted == [True]
+
+
+def test_staging_process_returns_result_after_progress(
+    staging_process_module: Any,
+) -> None:
+    """Valid progress and terminal result envelopes should pass through intact."""
+    script = (
+        "print('{\"kind\":\"progress\",\"stage\":\"download\"}');"
+        "print('{\"kind\":\"result\",\"results\":[{\"profile_id\":\"p\"}]}')"
+    )
+    process = subprocess.Popen(
+        [sys.executable, "-c", script],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    observed: list[dict[str, Any]] = []
+
+    results = staging_process_module.consume_staging_process(
+        process,
+        {"component_id": "llm"},
+        lambda _payload, event: observed.append(dict(event)),
+        provider_label="Test",
+        abort_remote=lambda: pytest.fail("successful staging must not abort"),
+        timeout_seconds=1.0,
+    )
+
+    assert observed[0]["stage"] == "download"
+    assert results == [{"profile_id": "p"}]
+
+
+def test_invalid_staging_timeout_still_cleans_up_process(
+    staging_process_module: Any,
+) -> None:
+    """A bad timeout setting must not strand the subprocess it was meant to bound."""
+    process = subprocess.Popen(
+        [sys.executable, "-c", "import time; time.sleep(60)"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    aborted: list[bool] = []
+
+    with pytest.raises(ValueError, match="must be positive"):
+        staging_process_module.consume_staging_process(
+            process,
+            {},
+            lambda _payload, _event: None,
+            provider_label="Test",
+            abort_remote=lambda: aborted.append(True),
+            timeout_seconds=0,
+        )
+
+    assert process.poll() is not None
+    assert aborted == [True]

tests/test_vast_executor.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
@@ -80,3 +81,47 @@ def test_runtime_rejects_unregistered_or_draining_lease(
         match="not present in the local registry",
     ):
         client._runtime({"vast_instance_id": 9001})
+
+
+def test_vast_cancel_targets_owned_stager_before_worker_request(
+    vast_executor_module: Any,
+    monkeypatch: Any,
+    tmp_path: Path,
+) -> None:
+    """A staging cancellation must terminate its transport and exact remote owner."""
+    client = vast_executor_module.VastExecutorClient(
+        registry=SimpleNamespace(),
+        activity_manager=SimpleNamespace(),
+        runtime_configuration=SimpleNamespace(),
+        user_directory=tmp_path,
+    )
+    invocation_id = "RIV_staging"
+    runner = SimpleNamespace()
+    process = SimpleNamespace()
+    terminated: list[Any] = []
+    targeted: list[tuple[Any, str]] = []
+    monkeypatch.setattr(
+        vast_executor_module,
+        "terminate_staging_transport",
+        terminated.append,
+    )
+    client._cancel_remote_stager = (  # type: ignore[method-assign]
+        lambda active_runner, owner_id, **_kwargs: (
+            targeted.append((active_runner, owner_id)) or True
+        )
+    )
+    with vast_executor_module._ACTIVE_VAST_STAGERS_LOCK:
+        vast_executor_module._ACTIVE_VAST_STAGERS[invocation_id] = (
+            runner,
+            "vast:RIV_staging:owner",
+            process,
+        )
+    try:
+        cancelled = client.cancel({}, invocation_id)
+    finally:
+        with vast_executor_module._ACTIVE_VAST_STAGERS_LOCK:
+            vast_executor_module._ACTIVE_VAST_STAGERS.pop(invocation_id, None)
+
+    assert cancelled is True
+    assert terminated == [process]
+    assert targeted == [(runner, "vast:RIV_staging:owner")]

vast_executor.py

@@ -6,14 +6,17 @@ import asyncio
 import logging
 import subprocess
 import threading
+import time
 from collections.abc import Iterator, Mapping, Sequence
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable, Protocol
+from uuid import uuid4
 
 if __package__:
     from .durable_state import stable_remote_invocation_id
     from .llm_profiles import (
+        encoded_resolved_llm_profile_payloads,
         llm_model_references_from_payload,
         rewrite_llm_model_references,
     )
@@ -27,6 +30,11 @@ if __package__:
     )
     from .serialization import deserialize_node_outputs, serialize_node_inputs
     from .settings import ModalSyncSettings, get_settings
+    from .staging_process import (
+        RemoteStagingProcessError,
+        consume_staging_process,
+        terminate_staging_transport,
+    )
     from .vast_leases import VastLeaseRecord, VastLeaseRegistry
     from .vast_runtime import VastRuntimeConfiguration, VastRuntimeManager
     from .vast_ssh import (
@@ -37,6 +45,7 @@ if __package__:
 else:  # pragma: no cover - direct debugging imports.
     from durable_state import stable_remote_invocation_id
     from llm_profiles import (
+        encoded_resolved_llm_profile_payloads,
         llm_model_references_from_payload,
         rewrite_llm_model_references,
     )
@@ -50,6 +59,11 @@ else:  # pragma: no cover - direct debugging imports.
     )
     from serialization import deserialize_node_outputs, serialize_node_inputs
     from settings import ModalSyncSettings, get_settings
+    from staging_process import (
+        RemoteStagingProcessError,
+        consume_staging_process,
+        terminate_staging_transport,
+    )
     from vast_leases import VastLeaseRecord, VastLeaseRegistry
     from vast_runtime import VastRuntimeConfiguration, VastRuntimeManager
     from vast_ssh import VastSshError, VastSshRunner, vast_connection_from_lease
@@ -57,6 +71,11 @@ else:  # pragma: no cover - direct debugging imports.
 logger = logging.getLogger(__name__)
 _STAGED_VAST_PROFILES_LOCK = threading.Lock()
 _STAGED_VAST_PROFILE_RESULTS: dict[tuple[int, str], dict[str, Any]] = {}
+_ACTIVE_VAST_STAGERS_LOCK = threading.Lock()
+_ACTIVE_VAST_STAGERS: dict[
+    str,
+    tuple[VastSshRunner, str, subprocess.Popen[bytes]],
+] = {}
 
 
 class VastRemoteInvocationError(RuntimeError):
@@ -128,6 +147,23 @@ class VastExecutorClient:
 
     def cancel(self, payload: Mapping[str, Any], invocation_id: str) -> bool:
         """Signal one active direct worker invocation."""
+        with _ACTIVE_VAST_STAGERS_LOCK:
+            active_stager = _ACTIVE_VAST_STAGERS.get(invocation_id)
+        if active_stager is not None:
+            runner, owner_id, process = active_stager
+            terminate_staging_transport(process)
+            remote_cancelled = self._cancel_remote_stager(
+                runner,
+                owner_id,
+                wait_for_owner_seconds=5.0,
+            )
+            logger.info(
+                "Cancelled Vast model staging invocation=%s "
+                "remote_process_found=%s.",
+                invocation_id,
+                remote_cancelled,
+            )
+            return True
         _lease, runner, runtime = self._runtime(payload)
         runtime.ensure_worker()
         request = encode_json_frame(
@@ -305,15 +341,43 @@ class VastExecutorClient:
             phase="llm_staging",
             status_message="Preparing LLM model snapshots on the Vast.ai instance",
         )
-        arguments = ["python", "-m", "remote.ssh_worker", "stage-profiles"]
+        invocation_id = str(payload["invocation_id"])
+        owner_id = f"vast:{invocation_id}:{uuid4().hex}"
+        arguments = [
+            "python",
+            "-m",
+            "remote.ssh_worker",
+            "stage-profiles",
+            "--owner-id",
+            owner_id,
+        ]
         for reference in model_references:
             arguments.extend(("--model-reference", reference))
-        process = runner.popen(tuple(arguments))
-        results = _consume_stager_output(
-            process,
+        resolved_profiles = encoded_resolved_llm_profile_payloads(
             payload,
-            _emit_local_llm_staging_progress,
+            model_references,
         )
+        if resolved_profiles is not None:
+            arguments.extend(("--resolved-profiles", resolved_profiles))
+        process = runner.popen(tuple(arguments))
+        with _ACTIVE_VAST_STAGERS_LOCK:
+            _ACTIVE_VAST_STAGERS[invocation_id] = (runner, owner_id, process)
+        try:
+            results = consume_staging_process(
+                process,
+                payload,
+                _emit_local_llm_staging_progress,
+                provider_label="Vast",
+                abort_remote=lambda: self._cancel_remote_stager(
+                    runner,
+                    owner_id,
+                ),
+            )
+        except RemoteStagingProcessError as exc:
+            raise VastRemoteInvocationError(str(exc)) from exc
+        finally:
+            with _ACTIVE_VAST_STAGERS_LOCK:
+                _ACTIVE_VAST_STAGERS.pop(invocation_id, None)
         _emit_local_remote_startup_status(
             payload,
             phase="llm_staged",
@@ -321,6 +385,46 @@ class VastExecutorClient:
         )
         return results
 
+    @staticmethod
+    def _cancel_remote_stager(
+        runner: VastSshRunner,
+        owner_id: str,
+        *,
+        wait_for_owner_seconds: float = 0.0,
+    ) -> bool:
+        """Terminate one exact owner-tagged stager on a Vast instance."""
+        deadline = time.monotonic() + max(0.0, wait_for_owner_seconds)
+        while True:
+            result = runner.run(
+                (
+                    "python",
+                    "-m",
+                    "remote.ssh_worker",
+                    "cancel-staging",
+                    "--owner-id",
+                    owner_id,
+                ),
+                timeout_seconds=15.0,
+                check=False,
+                transport_attempts=1,
+            )
+            if result.returncode != 0:
+                logger.warning(
+                    "Vast remote stager cancellation failed owner=%s error=%s.",
+                    owner_id,
+                    result.stderr_text.strip() or result.stdout_text.strip(),
+                )
+                return False
+            try:
+                response = decode_json_payload(result.stdout)
+            except (RemoteProtocolError, UnicodeDecodeError):
+                return False
+            if response.get("cancelled"):
+                return True
+            if time.monotonic() >= deadline:
+                return False
+            time.sleep(0.25)
+
     def _cache_staged_profiles(
         self,
         instance_id: int,
@@ -512,54 +616,6 @@ def _payload_retention_seconds(payload: Mapping[str, Any]) -> float:
     return retention
 
 
-def _consume_stager_output(
-    process: subprocess.Popen[bytes],
-    payload: dict[str, Any],
-    progress_callback: Callable[[dict[str, Any], Mapping[str, Any]], None],
-) -> list[dict[str, Any]]:
-    """Consume JSON-line progress and terminal direct staging results."""
-    if process.stdout is None or process.stderr is None:
-        process.kill()
-        raise VastRemoteInvocationError(
-            "Vast model stager did not expose output streams."
-        )
-    stderr_chunks: list[bytes] = []
-    stderr_thread = threading.Thread(
-        target=_collect_bounded_stream,
-        args=(process.stderr, stderr_chunks),
-        daemon=True,
-    )
-    stderr_thread.start()
-    results: list[dict[str, Any]] | None = None
-    for raw_line in process.stdout:
-        try:
-            event = decode_json_payload(raw_line)
-        except (RemoteProtocolError, UnicodeDecodeError) as exc:
-            process.terminate()
-            raise VastRemoteInvocationError(
-                "Vast model stager returned invalid progress JSON."
-            ) from exc
-        if event.get("kind") == "progress":
-            progress_callback(payload, event)
-        elif event.get("kind") == "result" and isinstance(event.get("results"), list):
-            results = [
-                dict(result)
-                for result in event["results"]
-                if isinstance(result, Mapping)
-            ]
-    returncode = process.wait()
-    stderr_thread.join(timeout=1.0)
-    if returncode != 0 or results is None:
-        diagnostics = (
-            b"".join(stderr_chunks).decode("utf-8", errors="replace").strip()
-        )
-        raise VastRemoteInvocationError(
-            "Vast model staging failed: "
-            f"{diagnostics or f'exit status {returncode}'}."
-        )
-    return results
-
-
 def _collect_bounded_stream(stream: Any, chunks: list[bytes]) -> None:
     """Drain diagnostics while retaining at most the recent MiB."""
     while chunk := stream.read(65536):