"""Direct Vast worker launch settings, readiness, and watchdog state updates.""" from __future__ import annotations import json import logging import math import os import shlex import time from dataclasses import dataclass from pathlib import PurePosixPath from typing import Any, Callable, Mapping if __package__: from .remote.vast_watchdog import VastWatchdogSnapshot from .vast_leases import VastLeaseRecord from .vast_models import VastInstanceLaunchSpec, VastResourceProfile from .vast_ssh import VastSshError, VastSshRunner, VastSshVolumeBackend else: # pragma: no cover - direct debugging imports. from remote.vast_watchdog import VastWatchdogSnapshot from vast_leases import VastLeaseRecord from vast_models import VastInstanceLaunchSpec, VastResourceProfile from vast_ssh import VastSshError, VastSshRunner, VastSshVolumeBackend logger = logging.getLogger(__name__) VAST_IMAGE_ENV = "COMFY_MODAL_VAST_IMAGE" DEFAULT_VAST_REMOTE_REPO_ROOT = PurePosixPath("/opt/comfy-remote/repo") DEFAULT_VAST_REMOTE_COMFYUI_ROOT = PurePosixPath("/opt/comfy-remote/ComfyUI") DEFAULT_VAST_REMOTE_STORAGE_ROOT = PurePosixPath("/storage") DEFAULT_VAST_WATCHDOG_STATE_PATH = PurePosixPath( "/storage/comfy-vast-watchdog.json" ) class VastRuntimeFingerprintDriftError(VastSshError): """Report that a published Vast worker does not match the local runtime.""" def __init__( self, *, expected_fingerprint: str, actual_fingerprint: str, protocol_version: int, ) -> None: """Retain the compared identities for automated image replacement.""" self.expected_fingerprint = expected_fingerprint self.actual_fingerprint = actual_fingerprint self.protocol_version = protocol_version super().__init__( "Vast worker image source fingerprint drift: expected " f"{expected_fingerprint[:12]}, found " f"{actual_fingerprint[:12] or 'empty'}, protocol {protocol_version}." ) def _ssh_permission_repair_lines() -> tuple[str, ...]: """Return idempotent repairs for the SSH key path injected by Vast.""" return ( "chown root:root /root", "chmod go-w /root", "mkdir -p /root/.ssh", "chown root:root /root/.ssh", "chmod 700 /root/.ssh", "if [ -f /root/.ssh/authorized_keys ]; then", " chown root:root /root/.ssh/authorized_keys", " chmod 600 /root/.ssh/authorized_keys", "fi", ) @dataclass(frozen=True) class VastRuntimeConfiguration: """Describe the digest-pinned worker image and direct runtime layout.""" image: str runtime_fingerprint: str remote_repo_root: PurePosixPath = DEFAULT_VAST_REMOTE_REPO_ROOT remote_comfyui_root: PurePosixPath = DEFAULT_VAST_REMOTE_COMFYUI_ROOT remote_storage_root: PurePosixPath = DEFAULT_VAST_REMOTE_STORAGE_ROOT watchdog_state_path: PurePosixPath = DEFAULT_VAST_WATCHDOG_STATE_PATH startup_timeout_seconds: float = 900.0 readiness_poll_seconds: float = 1.0 readiness_backoff_multiplier: float = 1.5 def __post_init__(self) -> None: """Validate launch and fingerprint values.""" if not self.image.strip() or any( character in self.image for character in ("\x00", "\n", "\r") ): raise ValueError("Vast runtime image must be a non-empty single-line value.") if len(self.runtime_fingerprint) != 64: raise ValueError("Vast runtime fingerprint must be a SHA-256 hex digest.") for field_name, path in ( ("remote_repo_root", self.remote_repo_root), ("remote_comfyui_root", self.remote_comfyui_root), ("remote_storage_root", self.remote_storage_root), ("watchdog_state_path", self.watchdog_state_path), ): if not path.is_absolute() or path == PurePosixPath("/"): raise ValueError(f"{field_name} must be an absolute non-root path.") if any( not math.isfinite(value) or value <= 0 for value in ( self.startup_timeout_seconds, self.readiness_poll_seconds, ) ): raise ValueError("Vast runtime readiness timing must be positive.") if ( not math.isfinite(self.readiness_backoff_multiplier) or self.readiness_backoff_multiplier < 1.0 ): raise ValueError( "Vast readiness backoff multiplier must be finite and at least 1." ) @classmethod def from_environment( cls, runtime_fingerprint: str, *, environment: Mapping[str, str] | None = None, ) -> "VastRuntimeConfiguration": """Resolve the required worker image without inventing an unpublished tag.""" source = os.environ if environment is None else environment image = str(source.get(VAST_IMAGE_ENV) or "").strip() if not image: raise RuntimeError( f"Set {VAST_IMAGE_ENV} to a published ComfyUI-Modal Vast worker image " "before renting an instance." ) return cls(image=image, runtime_fingerprint=runtime_fingerprint) def launch_spec( self, profile: VastResourceProfile, label: str, ) -> VastInstanceLaunchSpec: """Return one direct-SSH launch with worker and watchdog supervision.""" environment = { "COMFY_MODAL_REMOTE_WORKER": "1", "COMFY_MODAL_LLM_EXECUTION_TARGET": "vast", "COMFYUI_ROOT": str(self.remote_comfyui_root), "COMFY_MODAL_COMFYUI_ROOT": str(self.remote_comfyui_root), "COMFY_MODAL_LOCAL_STORAGE_ROOT": str(self.remote_storage_root), "COMFY_MODAL_REMOTE_STORAGE_ROOT": str(self.remote_storage_root), "COMFY_MODAL_EXECUTION_MODE": "local", "PYTHONPATH": f"{self.remote_repo_root}:{self.remote_comfyui_root}", } exported_lines = "\n".join( f"{name}={value}" for name, value in sorted(environment.items()) ) supervisor_command = shlex.join( ( "python", "-m", "remote.vast_supervisor", "start", "--watchdog-state", str(self.watchdog_state_path), ) ) onstart = "\n".join( ( "set -e", *_ssh_permission_repair_lines(), "mkdir -p /run/comfy-remote /storage/logs", "cat >> /etc/environment <<'COMFY_VAST_ENV'", exported_lines, "COMFY_VAST_ENV", supervisor_command, ) ) return VastInstanceLaunchSpec( image=self.image, disk_gb=profile.allocated_disk_gb, label=label, onstart=onstart, environment=environment, rental_type=profile.rental_type, ) @dataclass class VastRuntimeManager: """Ensure the direct worker process matches the selected runtime image.""" runner: VastSshRunner configuration: VastRuntimeConfiguration fallback_runner: VastSshRunner | None = None fallback_selected: Callable[[], None] | None = None instance_validator: Callable[[], None] | None = None status_callback: Callable[[str], None] | None = None clock: Callable[[], float] = time.time monotonic: Callable[[], float] = time.monotonic sleep: Callable[[float], None] = time.sleep def ensure_worker(self) -> dict[str, Any]: """Wait for the supervised worker socket and exact runtime fingerprint.""" started_at = self.monotonic() deadline = started_at + self.configuration.startup_timeout_seconds last_error: str | None = None last_logged_error: str | None = None attempt = 0 retry_delay = self.configuration.readiness_poll_seconds while self.monotonic() < deadline: attempt += 1 ssh_error: str | None = None if self.instance_validator is not None: self.instance_validator() try: info = self.runtime_info() except VastSshError as exc: last_error = str(exc) ssh_error = last_error except ValueError as exc: last_error = str(exc) else: actual_fingerprint = str(info.get("runtime_fingerprint") or "") if actual_fingerprint != self.configuration.runtime_fingerprint: actual_protocol = int(info.get("protocol_version") or 0) raise VastRuntimeFingerprintDriftError( expected_fingerprint=self.configuration.runtime_fingerprint, actual_fingerprint=actual_fingerprint, protocol_version=actual_protocol, ) if bool(info.get("worker_socket_ready")): logger.info( "Vast worker became ready after %d probe attempts (%.1fs).", attempt, self.monotonic() - started_at, ) return info try: self._start_supervisor() except VastSshError as exc: last_error = str(exc) ssh_error = last_error else: last_error = "worker socket is not ready" remaining_seconds = max(0.0, deadline - self.monotonic()) if remaining_seconds <= 0: break actual_delay = min(retry_delay, remaining_seconds) if last_error != last_logged_error: logger.warning( "Vast worker readiness probe attempt=%d failed; retrying in " "%.2fs: %s", attempt, actual_delay, last_error, ) last_logged_error = last_error if ssh_error is not None: self._emit_status( f"Vast SSH attempt {attempt} failed: {ssh_error}; retrying in " f"{actual_delay:.2f}s" ) self.sleep(actual_delay) retry_delay *= self.configuration.readiness_backoff_multiplier raise TimeoutError( "Vast worker did not become ready within " f"{self.configuration.startup_timeout_seconds:.0f}s: " f"{last_error or 'no diagnostics'}." ) def runtime_info(self) -> dict[str, Any]: """Return validated JSON from the direct worker image.""" try: return self._runtime_info(self.runner) except VastSshError as primary_error: fallback_runner = self.fallback_runner if fallback_runner is None: raise try: payload = self._runtime_info(fallback_runner) except VastSshError as fallback_error: raise VastSshError( "Vast direct and proxy SSH endpoints are both unavailable. " f"Direct: {primary_error} Proxy: {fallback_error}" ) from fallback_error self.runner = fallback_runner self.fallback_runner = None if self.fallback_selected is not None: self.fallback_selected() self._emit_status( "Vast direct SSH failed; continuing through the Vast SSH proxy" ) logger.warning( "Vast direct SSH failed; continuing through the Vast SSH proxy." ) return payload def _emit_status(self, message: str) -> None: """Forward one readiness update to the workflow UI when available.""" if self.status_callback is not None: self.status_callback(message) @staticmethod def _runtime_info(runner: VastSshRunner) -> dict[str, Any]: """Return validated runtime metadata through one exact SSH endpoint.""" result = runner.run( ("python", "-m", "remote.ssh_worker", "runtime-info"), check=True, transport_attempts=1, ) try: payload = json.loads(result.stdout_text) except json.JSONDecodeError as exc: raise ValueError("Vast worker returned invalid runtime-info JSON.") from exc if not isinstance(payload, dict): raise ValueError("Vast worker runtime-info must be a JSON object.") return payload def restart_worker(self) -> None: """Restart only the managed worker while preserving watchdog and storage.""" self.runner.run( ( "python", "-m", "remote.vast_supervisor", "restart-worker", "--watchdog-state", str(self.configuration.watchdog_state_path), ), timeout_seconds=60.0, ) self.ensure_worker() def update_watchdog(self, lease: VastLeaseRecord) -> None: """Atomically publish controller activity for the in-instance fail-safe.""" if not math.isfinite(lease.idle_deadline_epoch): raise ValueError("Vast lease idle deadline must be finite.") snapshot = VastWatchdogSnapshot( instance_id=lease.instance_id, owner_label=lease.label, idle_deadline_epoch=lease.idle_deadline_epoch, active_invocations=lease.active_invocations, updated_at_epoch=self.clock(), ) backend = VastSshVolumeBackend( self.runner, storage_root=self.configuration.remote_storage_root, ) relative_path = self.configuration.watchdog_state_path.relative_to( self.configuration.remote_storage_root ) backend.put_bytes(snapshot.to_json_bytes(), relative_path.as_posix()) def _start_supervisor(self) -> None: """Idempotently start both direct managed processes.""" self.runner.run( ( "python", "-m", "remote.vast_supervisor", "start", "--watchdog-state", str(self.configuration.watchdog_state_path), ), timeout_seconds=60.0, ) __all__ = [ "DEFAULT_VAST_REMOTE_COMFYUI_ROOT", "DEFAULT_VAST_REMOTE_REPO_ROOT", "DEFAULT_VAST_REMOTE_STORAGE_ROOT", "DEFAULT_VAST_WATCHDOG_STATE_PATH", "VAST_IMAGE_ENV", "VastRuntimeConfiguration", "VastRuntimeFingerprintDriftError", "VastRuntimeManager", ]