567 lines
22 KiB
Python
567 lines
22 KiB
Python
"""Remote executor client backed by warm Docker workers reached over SSH."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import subprocess
|
|
import threading
|
|
from collections.abc import Iterator, Mapping, Sequence
|
|
from dataclasses import dataclass, replace
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
if __package__:
|
|
from .durable_state import stable_remote_invocation_id
|
|
from .llm_profiles import (
|
|
llm_model_references_from_payload,
|
|
rewrite_llm_model_references,
|
|
)
|
|
from .remote_hosts import RemoteHostRegistry, SshHostConfig
|
|
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 .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 (
|
|
llm_model_references_from_payload,
|
|
rewrite_llm_model_references,
|
|
)
|
|
from remote_hosts import RemoteHostRegistry, SshHostConfig
|
|
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 ssh_docker import SshDockerController, SshDockerError
|
|
from ssh_runtime import SshRuntimeManager, SshRuntimeSpec
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_STAGED_SSH_PROFILES_LOCK = threading.Lock()
|
|
_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],
|
|
] = {}
|
|
|
|
|
|
class SshRemoteInvocationError(RuntimeError):
|
|
"""Raised when an SSH worker rejects or loses an invocation."""
|
|
|
|
|
|
class SshRemoteTransportError(SshRemoteInvocationError):
|
|
"""Raised when the SSH relay ends without a remote application error."""
|
|
|
|
|
|
@dataclass
|
|
class SshDockerExecutorClient:
|
|
"""Execute serialized ComfyUI components on one configured SSH host."""
|
|
|
|
registry: RemoteHostRegistry | None
|
|
repo_root: Path
|
|
settings: ModalSyncSettings | 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)
|
|
result = self._consume_stream(
|
|
prepared_payload,
|
|
inputs_payload,
|
|
)
|
|
return deserialize_node_outputs(result)
|
|
|
|
async def execute_payload_async(
|
|
self,
|
|
payload: Mapping[str, Any],
|
|
kwargs: Mapping[str, Any],
|
|
) -> Sequence[Any]:
|
|
"""Execute one payload 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:
|
|
result = await asyncio.to_thread(
|
|
self._consume_stream,
|
|
prepared_payload,
|
|
inputs_payload,
|
|
)
|
|
except asyncio.CancelledError:
|
|
await asyncio.to_thread(self.cancel, prepared_payload, invocation_id)
|
|
raise
|
|
return deserialize_node_outputs(result)
|
|
|
|
def cancel(
|
|
self,
|
|
payload: Mapping[str, Any],
|
|
invocation_id: str,
|
|
) -> bool:
|
|
"""Request cancellation of one active invocation on its assigned worker."""
|
|
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)
|
|
logger.info(
|
|
"Stopped SSH worker to cancel active model staging invocation=%s "
|
|
"environment=%s worker_index=%d stopped=%s.",
|
|
invocation_id,
|
|
active_manager.controller.host.environment_id,
|
|
active_spec.worker_index,
|
|
stopped,
|
|
)
|
|
return stopped
|
|
manager, spec = self._runtime(payload)
|
|
request = encode_json_frame(
|
|
RemoteFrameKind.CANCEL,
|
|
{"invocation_id": invocation_id},
|
|
)
|
|
process = manager.controller.docker_popen(self._relay_arguments(spec))
|
|
try:
|
|
stdout, stderr = process.communicate(request, timeout=30.0)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
process.communicate()
|
|
logger.warning(
|
|
"SSH cancellation timed out environment=%s invocation=%s.",
|
|
manager.controller.host.environment_id,
|
|
invocation_id,
|
|
)
|
|
return False
|
|
if process.returncode != 0:
|
|
logger.warning(
|
|
"SSH cancellation failed environment=%s invocation=%s error=%s.",
|
|
manager.controller.host.environment_id,
|
|
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 to one payload."""
|
|
prepared = dict(payload)
|
|
prepared["execution_provider"] = "ssh_docker"
|
|
prepared.setdefault(
|
|
"invocation_id",
|
|
stable_remote_invocation_id(prepared, inputs_payload),
|
|
)
|
|
prepared.setdefault("capture_remote_outputs", True)
|
|
return prepared
|
|
|
|
def _consume_stream(self, payload: dict[str, Any], inputs_payload: bytes) -> bytes:
|
|
"""Consume SSH frames through the existing 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):
|
|
manager, spec = self._runtime(payload)
|
|
try:
|
|
self._ensure_llm_profiles_staged(manager, spec, payload)
|
|
stream = self._invoke_stream(manager, spec, payload, inputs_payload)
|
|
response = _consume_remote_payload_stream(payload, stream)
|
|
break
|
|
except (
|
|
BrokenPipeError,
|
|
ConnectionResetError,
|
|
OSError,
|
|
RemoteProtocolError,
|
|
SshDockerError,
|
|
SshRemoteTransportError,
|
|
subprocess.SubprocessError,
|
|
) as exc:
|
|
if attempt >= 2:
|
|
raise SshRemoteTransportError(
|
|
f"SSH worker transport failed after one recovery attempt: {exc}"
|
|
) from exc
|
|
logger.warning(
|
|
"Recovering SSH worker transport environment=%s worker_index=%d error=%s.",
|
|
manager.controller.host.environment_id,
|
|
spec.worker_index,
|
|
exc,
|
|
)
|
|
try:
|
|
manager.stop_worker(spec.worker_index)
|
|
except SshDockerError as stop_error:
|
|
logger.warning(
|
|
"Unable to stop failed SSH worker before retry: %s",
|
|
stop_error,
|
|
)
|
|
if response is None:
|
|
raise SshRemoteTransportError("SSH worker returned no execution result.")
|
|
return _materialize_remote_execution_result(
|
|
response,
|
|
settings=self.settings or get_settings(),
|
|
)
|
|
|
|
def _ensure_llm_profiles_staged(
|
|
self,
|
|
manager: SshRuntimeManager,
|
|
spec: SshRuntimeSpec,
|
|
payload: dict[str, Any],
|
|
) -> None:
|
|
"""Resolve, stage, and rewrite LLM model references on one SSH host."""
|
|
model_references = llm_model_references_from_payload(payload)
|
|
if not model_references:
|
|
return
|
|
environment_id = manager.controller.host.environment_id
|
|
with _STAGED_SSH_PROFILES_LOCK:
|
|
missing_references = [
|
|
reference
|
|
for reference in model_references
|
|
if (environment_id, reference) not in _STAGED_SSH_PROFILE_RESULTS
|
|
]
|
|
if missing_references:
|
|
stage_results = self._run_profile_stager(
|
|
manager,
|
|
spec,
|
|
payload,
|
|
missing_references,
|
|
)
|
|
self._cache_staged_profiles(
|
|
environment_id,
|
|
missing_references,
|
|
stage_results,
|
|
)
|
|
resolved_results = {
|
|
reference: _STAGED_SSH_PROFILE_RESULTS[(environment_id, reference)]
|
|
for reference in model_references
|
|
}
|
|
rewrite_llm_model_references(
|
|
payload,
|
|
{
|
|
reference: str(result["profile_id"])
|
|
for reference, result in resolved_results.items()
|
|
},
|
|
)
|
|
logger.info(
|
|
"SSH LLM models are resolved and staged environment=%s component=%s profiles=%s.",
|
|
environment_id,
|
|
payload.get("component_id"),
|
|
sorted(str(result["profile_id"]) for result in resolved_results.values()),
|
|
)
|
|
|
|
def _run_profile_stager(
|
|
self,
|
|
manager: SshRuntimeManager,
|
|
spec: SshRuntimeSpec,
|
|
payload: dict[str, Any],
|
|
model_references: list[str],
|
|
) -> list[dict[str, Any]]:
|
|
"""Run the CPU model stager inside one persistent SSH worker 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="Inspecting and staging LLM on the SSH host",
|
|
)
|
|
arguments = [
|
|
"exec",
|
|
spec.container_name,
|
|
"python",
|
|
"-m",
|
|
"remote.ssh_worker",
|
|
"stage-profiles",
|
|
]
|
|
for model_reference in model_references:
|
|
arguments.extend(("--model-reference", model_reference))
|
|
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)
|
|
try:
|
|
results = self._consume_stager_output(
|
|
process,
|
|
payload,
|
|
_emit_local_llm_staging_progress,
|
|
)
|
|
finally:
|
|
with _ACTIVE_SSH_STAGERS_LOCK:
|
|
_ACTIVE_SSH_STAGERS.pop(invocation_id, None)
|
|
downloaded_gib = sum(
|
|
float(result.get("artifact_bytes") or 0) / 1024**3
|
|
for result in results
|
|
if result.get("downloaded")
|
|
)
|
|
_emit_local_remote_startup_status(
|
|
payload,
|
|
phase="llm_staged",
|
|
status_message=(
|
|
f"SSH LLM staging complete ({downloaded_gib:.1f} GiB downloaded)"
|
|
),
|
|
)
|
|
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'}"
|
|
)
|
|
return results
|
|
|
|
def _cache_staged_profiles(
|
|
self,
|
|
environment_id: str,
|
|
requested_references: list[str],
|
|
stage_results: list[dict[str, Any]],
|
|
) -> None:
|
|
"""Validate and cache SSH staging results by request and immutable ID."""
|
|
confirmed_references: set[str] = set()
|
|
for result in stage_results:
|
|
requested_reference = str(result.get("requested_reference") or "")
|
|
profile_id = str(result.get("profile_id") or "")
|
|
revision = str(result.get("revision") or "")
|
|
if not requested_reference or not profile_id or not revision:
|
|
continue
|
|
_STAGED_SSH_PROFILE_RESULTS[(environment_id, requested_reference)] = result
|
|
_STAGED_SSH_PROFILE_RESULTS[(environment_id, profile_id)] = result
|
|
confirmed_references.add(requested_reference)
|
|
missing_results = set(requested_references) - confirmed_references
|
|
if missing_results:
|
|
raise SshRemoteInvocationError(
|
|
f"SSH model stager did not confirm models {sorted(missing_results)}."
|
|
)
|
|
|
|
def _invoke_stream(
|
|
self,
|
|
manager: SshRuntimeManager,
|
|
spec: SshRuntimeSpec,
|
|
payload: dict[str, Any],
|
|
inputs_payload: bytes,
|
|
) -> Iterator[dict[str, Any]]:
|
|
"""Start one worker relay and yield Modal-compatible 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 = manager.controller.docker_popen(self._relay_arguments(spec))
|
|
if process.stdin is None or process.stdout is None or process.stderr is None:
|
|
process.kill()
|
|
raise SshRemoteInvocationError("SSH worker relay did not expose binary streams.")
|
|
stderr_chunks: list[bytes] = []
|
|
|
|
def collect_stderr() -> None:
|
|
"""Drain bounded diagnostic stderr so the SSH process cannot deadlock."""
|
|
while chunk := process.stderr.read(65536):
|
|
stderr_chunks.append(chunk)
|
|
if sum(len(value) for value in stderr_chunks) > 1024 * 1024:
|
|
del stderr_chunks[:-8]
|
|
|
|
stderr_thread = threading.Thread(target=collect_stderr, 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 SshRemoteInvocationError(
|
|
f"SSH worker {error.get('error_type', 'Error')}: "
|
|
f"{error.get('message', 'remote execution failed')}"
|
|
)
|
|
raise RemoteProtocolError(
|
|
f"Unexpected SSH 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(
|
|
"SSH 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 SshRemoteTransportError(
|
|
f"SSH worker stream ended without a result: {diagnostics or 'no diagnostics'}."
|
|
)
|
|
|
|
def _runtime(
|
|
self,
|
|
payload: Mapping[str, Any],
|
|
) -> tuple[SshRuntimeManager, SshRuntimeSpec]:
|
|
"""Resolve and prepare the worker assigned by the planner."""
|
|
environment_id = str(payload.get("execution_environment_id") or "").strip()
|
|
if not environment_id:
|
|
raise SshRemoteInvocationError(
|
|
"SSH execution payload is missing execution_environment_id."
|
|
)
|
|
raw_host = payload.get("ssh_host_config")
|
|
if isinstance(raw_host, Mapping):
|
|
try:
|
|
host = SshHostConfig.from_dict(raw_host)
|
|
except (TypeError, ValueError) as exc:
|
|
raise SshRemoteInvocationError(
|
|
"SSH execution payload contains an invalid workflow host "
|
|
f"configuration: {exc}"
|
|
) from exc
|
|
if host.environment_id != environment_id:
|
|
raise SshRemoteInvocationError(
|
|
"SSH execution payload environment identity does not match its "
|
|
"workflow host configuration."
|
|
)
|
|
else:
|
|
if self.registry is None:
|
|
raise SshRemoteInvocationError(
|
|
"Legacy SSH execution requires a persistent host registry."
|
|
)
|
|
host = self.registry.get_host(environment_id)
|
|
if not host.enabled or host.draining:
|
|
raise SshRemoteInvocationError(
|
|
f"SSH execution environment {environment_id!r} is not accepting work."
|
|
)
|
|
controller = SshDockerController(host)
|
|
if host.capabilities is None:
|
|
logger.info(
|
|
"Reprobing SSH execution capabilities missing from the queued host "
|
|
"snapshot environment=%s.",
|
|
environment_id,
|
|
)
|
|
try:
|
|
host = replace(host, capabilities=controller.probe_capabilities())
|
|
except (OSError, SshDockerError, ValueError) as exc:
|
|
raise SshRemoteInvocationError(
|
|
f"Unable to refresh SSH execution capabilities for "
|
|
f"{environment_id!r}: {exc}"
|
|
) from exc
|
|
controller = SshDockerController(host)
|
|
resolved_settings = self.settings or get_settings()
|
|
manager = SshRuntimeManager(
|
|
controller=controller,
|
|
repo_root=self.repo_root,
|
|
settings=resolved_settings,
|
|
)
|
|
worker_index = int(payload.get("execution_worker_index", 0))
|
|
return manager, manager.ensure_worker(worker_index)
|
|
|
|
def _relay_arguments(self, spec: SshRuntimeSpec) -> tuple[str, ...]:
|
|
"""Return the Docker exec argv for the worker's binary relay client."""
|
|
return (
|
|
"exec",
|
|
"-i",
|
|
spec.container_name,
|
|
"python",
|
|
"-m",
|
|
"remote.ssh_worker",
|
|
"client",
|
|
)
|
|
|
|
|
|
def _collect_bounded_stream(stream: Any, chunks: list[bytes]) -> None:
|
|
"""Drain a diagnostic stream while retaining at most its recent output."""
|
|
while chunk := stream.read(65536):
|
|
chunks.append(chunk)
|
|
if sum(len(value) for value in chunks) > 1024 * 1024:
|
|
del chunks[:-8]
|
|
|
|
|
|
class _BytesReader:
|
|
"""Minimal binary reader wrapper used for one-shot protocol responses."""
|
|
|
|
def __init__(self, payload: bytes) -> None:
|
|
"""Initialize a read cursor over immutable bytes."""
|
|
self._payload = payload
|
|
self._offset = 0
|
|
|
|
def read(self, length: int = -1) -> bytes:
|
|
"""Read up to ``length`` bytes from the cursor."""
|
|
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
|