245 lines
7.9 KiB
Python
245 lines
7.9 KiB
Python
"""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",
|
|
]
|