331 lines
12 KiB
Python
331 lines
12 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
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
if __package__:
|
|
from .durable_state import stable_remote_invocation_id
|
|
from .remote_hosts import RemoteHostRegistry
|
|
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 remote_hosts import RemoteHostRegistry
|
|
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__)
|
|
|
|
|
|
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
|
|
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."""
|
|
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:
|
|
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 _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."
|
|
)
|
|
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."
|
|
)
|
|
resolved_settings = self.settings or get_settings()
|
|
manager = SshRuntimeManager(
|
|
controller=SshDockerController(host),
|
|
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",
|
|
)
|
|
|
|
|
|
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
|