114 lines
3.7 KiB
Python
114 lines
3.7 KiB
Python
"""Tests for SSH worker invocation and local artifact restoration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
|
|
def test_ssh_stream_materializes_remote_output_artifacts(
|
|
ssh_executor_module: Any,
|
|
remote_modal_app_module: Any,
|
|
output_artifacts_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""SSH results should restore output files before deserializing node outputs."""
|
|
output_directory = tmp_path / "output"
|
|
output_directory.mkdir()
|
|
packed_result = output_artifacts_module.pack_remote_execution_result(
|
|
output_artifacts_module.RemoteExecutionResult(
|
|
outputs=b"serialized outputs",
|
|
artifacts=(
|
|
output_artifacts_module.RemoteOutputArtifact(
|
|
relative_path="images/result.png",
|
|
payload=b"png bytes",
|
|
),
|
|
),
|
|
completed_epoch=1_700_000_000,
|
|
)
|
|
)
|
|
settings = SimpleNamespace(
|
|
comfyui_root=tmp_path,
|
|
app_name="test-app",
|
|
)
|
|
client = ssh_executor_module.SshDockerExecutorClient(
|
|
registry=SimpleNamespace(),
|
|
repo_root=tmp_path,
|
|
settings=settings,
|
|
)
|
|
manager = SimpleNamespace()
|
|
spec = SimpleNamespace()
|
|
monkeypatch.setattr(client, "_runtime", lambda _payload: (manager, spec))
|
|
monkeypatch.setattr(
|
|
client,
|
|
"_invoke_stream",
|
|
lambda *_args: iter(({"kind": "result", "outputs": packed_result},)),
|
|
)
|
|
monkeypatch.setattr(
|
|
remote_modal_app_module,
|
|
"_consume_remote_payload_stream",
|
|
lambda _payload, stream: next(stream)["outputs"],
|
|
)
|
|
monkeypatch.setattr(
|
|
remote_modal_app_module,
|
|
"_local_comfy_output_directory",
|
|
lambda _settings: output_directory,
|
|
)
|
|
|
|
outputs = client._consume_stream({"component_id": "1"}, b"inputs")
|
|
|
|
assert outputs == b"serialized outputs"
|
|
materialized = list(output_directory.rglob("*.png"))
|
|
assert len(materialized) == 1
|
|
assert materialized[0].read_bytes() == b"png bytes"
|
|
|
|
|
|
def test_ssh_transport_failure_restarts_worker_and_retries_once(
|
|
ssh_executor_module: Any,
|
|
remote_modal_app_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""A broken relay should recover the exact worker slot with the stable request."""
|
|
stopped: list[int] = []
|
|
manager = SimpleNamespace(
|
|
controller=SimpleNamespace(host=SimpleNamespace(environment_id="host")),
|
|
stop_worker=lambda worker_index: stopped.append(worker_index),
|
|
)
|
|
spec = SimpleNamespace(worker_index=2)
|
|
client = ssh_executor_module.SshDockerExecutorClient(
|
|
registry=SimpleNamespace(),
|
|
repo_root=tmp_path,
|
|
settings=SimpleNamespace(comfyui_root=tmp_path, app_name="app"),
|
|
)
|
|
attempts = 0
|
|
|
|
def invoke(*_args: Any) -> Any:
|
|
"""Fail the first relay attempt and complete the second."""
|
|
nonlocal attempts
|
|
attempts += 1
|
|
if attempts == 1:
|
|
raise ssh_executor_module.SshRemoteTransportError("connection lost")
|
|
return iter(({"kind": "result", "outputs": b"serialized"},))
|
|
|
|
monkeypatch.setattr(client, "_runtime", lambda _payload: (manager, spec))
|
|
monkeypatch.setattr(client, "_invoke_stream", invoke)
|
|
monkeypatch.setattr(
|
|
remote_modal_app_module,
|
|
"_consume_remote_payload_stream",
|
|
lambda _payload, stream: next(stream)["outputs"],
|
|
)
|
|
monkeypatch.setattr(
|
|
remote_modal_app_module,
|
|
"_materialize_remote_execution_result",
|
|
lambda response, settings: response,
|
|
)
|
|
|
|
result = client._consume_stream({"component_id": "1"}, b"inputs")
|
|
|
|
assert result == b"serialized"
|
|
assert attempts == 2
|
|
assert stopped == [2]
|