"""Remote executor client for direct workers inside managed Vast.ai leases.""" from __future__ import annotations 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, ) from .remote_protocol import ( RemoteFrameKind, RemoteProtocolError, decode_json_payload, encode_frame, encode_json_frame, read_frame, ) 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, ) 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, ) from remote_protocol import ( RemoteFrameKind, RemoteProtocolError, decode_json_payload, encode_frame, encode_json_frame, read_frame, ) 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 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): """Raised when a Vast worker rejects one application invocation.""" class VastRemoteTransportError(VastRemoteInvocationError): """Raised when the direct SSH relay ends without an application result.""" class VastLeaseActivityManager(Protocol): """Track activity and retention for one managed Vast lease.""" def begin_activity(self, instance_id: int) -> VastLeaseRecord: """Mark a lease busy.""" def finish_activity( self, instance_id: int, *, idle_retention_seconds: float, error: str | None = None, ) -> VastLeaseRecord: """Release activity and reset the idle deadline.""" @dataclass class VastExecutorClient: """Execute serialized ComfyUI components on one direct Vast worker.""" registry: VastLeaseRegistry activity_manager: VastLeaseActivityManager runtime_configuration: VastRuntimeConfiguration user_directory: Path settings: ModalSyncSettings | None = None identity_file: Path | None = None runner_factory: Callable[[VastLeaseRecord], VastSshRunner] | None = None def execute_payload( self, payload: Mapping[str, Any], kwargs: Mapping[str, Any], ) -> Sequence[Any]: """Execute one payload synchronously and deserialize its outputs.""" inputs_payload = serialize_node_inputs(kwargs) prepared_payload = self._prepare_payload(payload, inputs_payload) response = self._execute_with_activity(prepared_payload, inputs_payload) return deserialize_node_outputs(response) async def execute_payload_async( self, payload: Mapping[str, Any], kwargs: Mapping[str, Any], ) -> Sequence[Any]: """Execute without blocking ComfyUI's event loop.""" inputs_payload = serialize_node_inputs(kwargs) prepared_payload = self._prepare_payload(payload, inputs_payload) invocation_id = str(prepared_payload["invocation_id"]) try: response = await asyncio.to_thread( self._execute_with_activity, prepared_payload, inputs_payload, ) except asyncio.CancelledError: await asyncio.to_thread(self.cancel, prepared_payload, invocation_id) raise return deserialize_node_outputs(response) 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( RemoteFrameKind.CANCEL, {"invocation_id": invocation_id}, ) process = runner.popen(self._relay_arguments()) try: stdout, stderr = process.communicate(request, timeout=30.0) except subprocess.TimeoutExpired: process.kill() process.communicate() logger.warning("Vast cancellation timed out invocation=%s.", invocation_id) return False if process.returncode != 0: logger.warning( "Vast cancellation failed invocation=%s error=%s.", invocation_id, stderr.decode("utf-8", errors="replace").strip(), ) return False frame = read_frame(_BytesReader(stdout)) if frame is None or frame[0] is not RemoteFrameKind.ACKNOWLEDGEMENT: return False return bool(decode_json_payload(frame[1]).get("cancelled")) def _prepare_payload( self, payload: Mapping[str, Any], inputs_payload: bytes, ) -> dict[str, Any]: """Attach provider and stable invocation identity.""" prepared = dict(payload) prepared["execution_provider"] = "vast" prepared.setdefault( "invocation_id", stable_remote_invocation_id(prepared, inputs_payload), ) prepared.setdefault("capture_remote_outputs", True) return prepared def _execute_with_activity( self, payload: dict[str, Any], inputs_payload: bytes, ) -> bytes: """Keep registry and in-instance watchdog state around one invocation.""" lease, _runner, runtime = self._runtime(payload) active_lease = self.activity_manager.begin_activity(lease.instance_id) runtime.update_watchdog(active_lease) terminal_error: str | None = None try: return self._consume_stream(payload, inputs_payload) except (RuntimeError, OSError, ValueError) as exc: terminal_error = str(exc) raise finally: retention_seconds = _payload_retention_seconds(payload) try: finished = self.activity_manager.finish_activity( lease.instance_id, idle_retention_seconds=retention_seconds, error=terminal_error, ) runtime.update_watchdog(finished) except (KeyError, OSError, RuntimeError, ValueError) as state_error: logger.error( "Unable to publish terminal Vast lease activity instance=%d: %s", lease.instance_id, state_error, ) def _consume_stream( self, payload: dict[str, Any], inputs_payload: bytes, ) -> bytes: """Consume direct worker frames through the shared Modal-compatible UI relay.""" from .remote.modal_app import ( _consume_remote_payload_stream, _materialize_remote_execution_result, ) response: bytes | None = None for attempt in range(1, 3): lease, runner, runtime = self._runtime(payload) try: runtime.ensure_worker() self._ensure_llm_profiles_staged(lease, runner, payload) response = _consume_remote_payload_stream( payload, self._invoke_stream(runner, payload, inputs_payload), ) break except VastRemoteTransportError as exc: transport_error: BaseException = exc except VastRemoteInvocationError: raise except ( BrokenPipeError, ConnectionResetError, OSError, RemoteProtocolError, VastSshError, subprocess.SubprocessError, ) as exc: transport_error = exc if attempt >= 2: raise VastRemoteTransportError( "Vast worker transport failed after one recovery attempt: " f"{transport_error}" ) from transport_error logger.warning( "Recovering Vast worker transport instance=%d error=%s.", lease.instance_id, transport_error, ) runtime.restart_worker() if response is None: raise VastRemoteTransportError("Vast worker returned no execution result.") return _materialize_remote_execution_result( response, settings=self.settings or get_settings(), ) def _ensure_llm_profiles_staged( self, lease: VastLeaseRecord, runner: VastSshRunner, payload: dict[str, Any], ) -> None: """Resolve and stage LLM profiles on the selected instance storage.""" model_references = llm_model_references_from_payload(payload) if not model_references: return with _STAGED_VAST_PROFILES_LOCK: missing_references = [ reference for reference in model_references if (lease.instance_id, reference) not in _STAGED_VAST_PROFILE_RESULTS ] if missing_references: results = self._run_profile_stager(runner, payload, missing_references) self._cache_staged_profiles( lease.instance_id, missing_references, results, ) resolved_results = { reference: _STAGED_VAST_PROFILE_RESULTS[(lease.instance_id, reference)] for reference in model_references } rewrite_llm_model_references( payload, { reference: str(result["profile_id"]) for reference, result in resolved_results.items() }, ) def _run_profile_stager( self, runner: VastSshRunner, payload: dict[str, Any], model_references: list[str], ) -> list[dict[str, Any]]: """Run the shared CPU stager directly in the Vast container.""" from .remote.modal_app import ( _emit_local_llm_staging_progress, _emit_local_remote_startup_status, ) _emit_local_remote_startup_status( payload, phase="llm_staging", status_message="Preparing LLM model snapshots on the Vast.ai instance", ) 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)) resolved_profiles = encoded_resolved_llm_profile_payloads( payload, 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", status_message="Vast.ai LLM staging complete", ) 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, requested_references: list[str], results: list[dict[str, Any]], ) -> None: """Cache stager aliases and require every requested model result.""" confirmed: set[str] = set() by_requested = { str(result.get("requested_reference") or ""): result for result in results } for requested_reference in requested_references: result = by_requested.get(requested_reference) if result is None: continue profile_id = str(result.get("profile_id") or "").strip() if not profile_id: continue _STAGED_VAST_PROFILE_RESULTS[(instance_id, requested_reference)] = result _STAGED_VAST_PROFILE_RESULTS[(instance_id, profile_id)] = result confirmed.add(requested_reference) missing = set(requested_references) - confirmed if missing: raise VastRemoteInvocationError( f"Vast model stager did not confirm models {sorted(missing)}." ) def _invoke_stream( self, runner: VastSshRunner, payload: dict[str, Any], inputs_payload: bytes, ) -> Iterator[dict[str, Any]]: """Start one direct worker relay and yield shared event mappings.""" request = b"".join( ( encode_json_frame( RemoteFrameKind.REQUEST, { "invocation_id": str(payload["invocation_id"]), "payload": payload, }, ), encode_frame(RemoteFrameKind.INPUTS, inputs_payload), ) ) process = runner.popen(self._relay_arguments()) if process.stdin is None or process.stdout is None or process.stderr is None: process.kill() raise VastRemoteTransportError( "Vast worker relay did not expose binary streams." ) stderr_chunks: list[bytes] = [] stderr_thread = threading.Thread( target=_collect_bounded_stream, args=(process.stderr, stderr_chunks), daemon=True, ) stderr_thread.start() process.stdin.write(request) process.stdin.close() terminal_received = False try: while True: frame = read_frame(process.stdout) if frame is None: break kind, frame_payload = frame if kind is RemoteFrameKind.PROGRESS: yield decode_json_payload(frame_payload) continue if kind is RemoteFrameKind.RESULT: terminal_received = True yield {"kind": "result", "outputs": frame_payload} break if kind is RemoteFrameKind.ERROR: terminal_received = True error = decode_json_payload(frame_payload) raise VastRemoteInvocationError( f"Vast worker {error.get('error_type', 'Error')}: " f"{error.get('message', 'remote execution failed')}" ) raise RemoteProtocolError( f"Unexpected Vast worker response frame {kind.name}." ) finally: if not terminal_received and process.poll() is None: process.terminate() try: returncode = process.wait(timeout=30.0) except subprocess.TimeoutExpired: process.kill() returncode = process.wait(timeout=5.0) stderr_thread.join(timeout=1.0) if returncode != 0 and terminal_received: logger.warning( "Vast worker relay exited after terminal frame status=%d stderr=%s.", returncode, b"".join(stderr_chunks) .decode("utf-8", errors="replace") .strip(), ) if not terminal_received: diagnostics = ( b"".join(stderr_chunks).decode("utf-8", errors="replace").strip() ) raise VastRemoteTransportError( "Vast worker stream ended without a result: " f"{diagnostics or 'no diagnostics'}." ) def _runtime( self, payload: Mapping[str, Any], ) -> tuple[VastLeaseRecord, VastSshRunner, VastRuntimeManager]: """Resolve one assigned lease and direct runtime manager.""" instance_id = _payload_instance_id(payload) lease = next( ( candidate for candidate in self.registry.load().leases if candidate.instance_id == instance_id ), None, ) if lease is None: raise VastRemoteInvocationError( f"Assigned Vast lease {instance_id} is not present in the local registry." ) if lease.draining: raise VastRemoteInvocationError( f"Assigned Vast lease {instance_id} is draining." ) runner = ( self.runner_factory(lease) if self.runner_factory is not None else VastSshRunner( vast_connection_from_lease( ssh_host=lease.ssh_host, ssh_port=lease.ssh_port, user_directory=self.user_directory, identity_file=self.identity_file, ) ) ) runtime = VastRuntimeManager( runner=runner, configuration=self.runtime_configuration, ) return lease, runner, runtime @staticmethod def _relay_arguments() -> tuple[str, ...]: """Return the direct worker binary relay command.""" return ("python", "-m", "remote.ssh_worker", "client") def _payload_instance_id(payload: Mapping[str, Any]) -> int: """Return the concrete Vast instance assigned by queue-time planning.""" raw_instance_id = payload.get("vast_instance_id") if raw_instance_id is None: environment_id = str(payload.get("execution_environment_id") or "") raw_instance_id = environment_id.rsplit(":", maxsplit=1)[-1] try: instance_id = int(raw_instance_id) except (TypeError, ValueError) as exc: raise VastRemoteInvocationError( "Vast execution payload is missing a concrete instance identity." ) from exc if instance_id <= 0: raise VastRemoteInvocationError( "Vast execution payload contains an invalid instance identity." ) return instance_id def _payload_retention_seconds(payload: Mapping[str, Any]) -> float: """Return the validated per-profile idle retention stamped at queue time.""" try: retention = float(payload.get("vast_idle_retention_seconds", 24 * 3600)) except (TypeError, ValueError) as exc: raise VastRemoteInvocationError( "Vast execution payload contains invalid idle retention." ) from exc if retention < 0: raise VastRemoteInvocationError( "Vast execution payload idle retention must not be negative." ) return retention def _collect_bounded_stream(stream: Any, chunks: list[bytes]) -> None: """Drain diagnostics while retaining at most the recent MiB.""" while chunk := stream.read(65536): chunks.append(chunk) if sum(len(value) for value in chunks) > 1024 * 1024: del chunks[:-8] class _BytesReader: """Minimal immutable byte cursor for one-shot protocol responses.""" def __init__(self, payload: bytes) -> None: """Initialize the cursor.""" self._payload = payload self._offset = 0 def read(self, length: int = -1) -> bytes: """Read up to ``length`` bytes.""" if length < 0: length = len(self._payload) - self._offset end = min(len(self._payload), self._offset + length) result = self._payload[self._offset : end] self._offset = end return result __all__ = [ "VastExecutorClient", "VastLeaseActivityManager", "VastRemoteInvocationError", "VastRemoteTransportError", ]