diff --git a/README.md b/README.md index f6a2a26..4d8699d 100644 --- a/README.md +++ b/README.md @@ -291,7 +291,7 @@ Local snapshots are stored under `/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. | diff --git a/api_intercept.py b/api_intercept.py index 1ceb9f4..83aa1a6 100644 --- a/api_intercept.py +++ b/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, diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index db0e5a6..6485929 100644 --- a/comfyui_modal_sync_cloud.py +++ b/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, ) ) diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 92b7619..976128f 100644 --- a/docs/LEARNINGS.md +++ b/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//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. diff --git a/llm_profiles.py b/llm_profiles.py index f96b925..37b6ba4 100644 --- a/llm_profiles.py +++ b/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", ] diff --git a/llm_staging.py b/llm_staging.py index 0c75968..80c6163 100644 --- a/llm_staging.py +++ b/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( diff --git a/remote/modal_app.py b/remote/modal_app.py index 575b8ab..8578a1c 100644 --- a/remote/modal_app.py +++ b/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: diff --git a/remote/ssh_worker.py b/remote/ssh_worker.py index ed9fa46..c59d087 100644 --- a/remote/ssh_worker.py +++ b/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 diff --git a/runtime_environment.py b/runtime_environment.py index 90a0364..224ffdc 100644 --- a/runtime_environment.py +++ b/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" diff --git a/ssh_executor.py b/ssh_executor.py index 742191b..828cd05 100644 --- a/ssh_executor.py +++ b/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, diff --git a/staging_process.py b/staging_process.py new file mode 100644 index 0000000..f9351da --- /dev/null +++ b/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", +] diff --git a/tests/conftest.py b/tests/conftest.py index ec8da6a..40e2768 100644 --- a/tests/conftest.py +++ b/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.""" diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index 2ef9f70..78d1e4d 100644 --- a/tests/test_api_intercept.py +++ b/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: diff --git a/tests/test_modal_llm.py b/tests/test_modal_llm.py index 933d71e..47dcd8b 100644 --- a/tests/test_modal_llm.py +++ b/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, diff --git a/tests/test_runtime_environment.py b/tests/test_runtime_environment.py index 70f2777..624dfd7 100644 --- a/tests/test_runtime_environment.py +++ b/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) diff --git a/tests/test_ssh_executor.py b/tests/test_ssh_executor.py index f18c8c3..f68efaf 100644 --- a/tests/test_ssh_executor.py +++ b/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")] diff --git a/tests/test_ssh_worker.py b/tests/test_ssh_worker.py index a4d37fe..74881bc 100644 --- a/tests/test_ssh_worker.py +++ b/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, diff --git a/tests/test_staging_process.py b/tests/test_staging_process.py new file mode 100644 index 0000000..01113a9 --- /dev/null +++ b/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] diff --git a/tests/test_vast_executor.py b/tests/test_vast_executor.py index 3991737..94ccdbc 100644 --- a/tests/test_vast_executor.py +++ b/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")] diff --git a/vast_executor.py b/vast_executor.py index 6c70a4c..598da04 100644 --- a/vast_executor.py +++ b/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):