122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
"""Tests for persistent SSH worker state and relay behavior."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
|
|
def test_top_level_worker_entrypoint_reports_runtime_info() -> None:
|
|
"""The OCI entrypoint must import when the repository is a top-level path."""
|
|
repo_root = Path(__file__).resolve().parents[1]
|
|
environment = dict(os.environ)
|
|
environment["PYTHONPATH"] = str(repo_root)
|
|
environment["COMFY_MODAL_RUNTIME_FINGERPRINT"] = "test-fingerprint"
|
|
|
|
completed = subprocess.run(
|
|
[sys.executable, "-m", "remote.ssh_worker", "runtime-info"],
|
|
cwd=repo_root,
|
|
env=environment,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
check=True,
|
|
text=True,
|
|
timeout=15,
|
|
)
|
|
|
|
runtime_info = json.loads(completed.stdout)
|
|
assert runtime_info["runtime_fingerprint"] == "test-fingerprint"
|
|
assert runtime_info["protocol_version"] > 0
|
|
|
|
|
|
def test_worker_execution_state_registers_and_cancels(ssh_worker_module: Any) -> None:
|
|
"""Cancellation should target a stable invocation identity."""
|
|
state = ssh_worker_module.WorkerExecutionState()
|
|
|
|
cancellation = state.register("RIV_test")
|
|
|
|
assert state.cancel("RIV_test") is True
|
|
assert cancellation.is_set()
|
|
state.unregister("RIV_test")
|
|
assert state.cancel("RIV_test") is False
|
|
|
|
|
|
def test_worker_stage_profiles_streams_progress_and_result(
|
|
ssh_worker_module: Any,
|
|
llm_staging_module: Any,
|
|
monkeypatch: Any,
|
|
capsys: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""The SSH staging command should expose machine-readable progress and metadata."""
|
|
|
|
def stage(model_references: list[str], storage_root: Path, **kwargs: Any) -> Any:
|
|
"""Emit one progress update and return one immutable profile result."""
|
|
assert model_references == ["owner/model"]
|
|
assert storage_root == tmp_path
|
|
kwargs["progress_callback"](
|
|
SimpleNamespace(
|
|
stage="download",
|
|
message="Fetching files",
|
|
value=1,
|
|
maximum=2,
|
|
unit="files",
|
|
indeterminate=False,
|
|
)
|
|
)
|
|
return [
|
|
SimpleNamespace(
|
|
to_dict=lambda: {
|
|
"requested_reference": "owner/model",
|
|
"profile_id": "hf-" + "b" * 64,
|
|
"revision": "8" * 40,
|
|
}
|
|
)
|
|
]
|
|
|
|
monkeypatch.setattr(llm_staging_module, "resolve_and_stage_model_references", stage)
|
|
|
|
results = ssh_worker_module.stage_profiles(["owner/model"], tmp_path)
|
|
|
|
events = [json.loads(line) for line in capsys.readouterr().out.splitlines()]
|
|
assert events[0]["kind"] == "progress"
|
|
assert events[0]["max"] == 2
|
|
assert events[1] == {"kind": "result", "results": results}
|
|
|
|
|
|
def test_worker_request_relay_preserves_framed_binary_payload(
|
|
ssh_worker_module: Any,
|
|
remote_protocol_module: Any,
|
|
) -> None:
|
|
"""The docker-exec relay must preserve request and input frames byte-for-byte."""
|
|
request = remote_protocol_module.encode_json_frame(
|
|
remote_protocol_module.RemoteFrameKind.REQUEST,
|
|
{"invocation_id": "RIV_test", "payload": {"payload_kind": "subgraph"}},
|
|
)
|
|
inputs = remote_protocol_module.encode_frame(
|
|
remote_protocol_module.RemoteFrameKind.INPUTS,
|
|
b"\x00\xfftensor-bytes",
|
|
)
|
|
left, right = socket.socketpair()
|
|
try:
|
|
ssh_worker_module._copy_request_frames(io.BytesIO(request + inputs), left)
|
|
received = right.makefile("rb")
|
|
try:
|
|
assert remote_protocol_module.read_frame(received)[0] is remote_protocol_module.RemoteFrameKind.REQUEST
|
|
assert remote_protocol_module.read_frame(received) == (
|
|
remote_protocol_module.RemoteFrameKind.INPUTS,
|
|
b"\x00\xfftensor-bytes",
|
|
)
|
|
finally:
|
|
received.close()
|
|
finally:
|
|
left.close()
|
|
right.close()
|