532 lines
18 KiB
Python
532 lines
18 KiB
Python
"""Tests for SSH OCI worker image and GPU assignment behavior."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
|
|
def test_worker_lifecycle_reports_build_and_readiness_status(
|
|
ssh_runtime_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""A rebuilt SSH worker should expose every material readiness transition."""
|
|
statuses: list[str] = []
|
|
spec = SimpleNamespace(
|
|
image_tag="comfy-remote:deadbeef",
|
|
storage_volume_name="comfy-remote-lambda",
|
|
)
|
|
controller = SimpleNamespace(
|
|
host=SimpleNamespace(environment_id="lambda"),
|
|
ensure_volume=lambda _name: None,
|
|
)
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=controller,
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(),
|
|
)
|
|
monkeypatch.setattr(manager, "runtime_spec", lambda _worker_index: spec)
|
|
monkeypatch.setattr(manager, "_image_is_current", lambda _spec: False)
|
|
monkeypatch.setattr(manager, "_build_image", lambda _spec: None)
|
|
monkeypatch.setattr(manager, "_remove_stale_worker_containers", lambda _spec: None)
|
|
monkeypatch.setattr(
|
|
manager,
|
|
"_container_is_current_and_running",
|
|
lambda _spec: False,
|
|
)
|
|
monkeypatch.setattr(manager, "_replace_worker_container", lambda _spec: None)
|
|
monkeypatch.setattr(manager, "_wait_until_ready", lambda _spec: None)
|
|
|
|
result = manager.ensure_worker(status_callback=statuses.append)
|
|
|
|
assert result is spec
|
|
assert statuses == [
|
|
"Checking SSH runtime environment=lambda",
|
|
"Building SSH runtime environment=lambda image=comfy-remote:deadbeef",
|
|
"Starting SSH worker environment=lambda",
|
|
"Waiting for SSH worker environment=lambda",
|
|
"Ready for remote execution",
|
|
]
|
|
|
|
|
|
def test_worker_build_loads_image_into_the_remote_daemon(
|
|
ssh_runtime_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""BuildKit docker-container drivers must export the image for docker run."""
|
|
calls: list[tuple[tuple[str, ...], dict[str, Any]]] = []
|
|
|
|
def docker(arguments: tuple[str, ...], **kwargs: Any) -> None:
|
|
"""Record one remote Docker invocation."""
|
|
calls.append((arguments, kwargs))
|
|
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=SimpleNamespace(
|
|
host=SimpleNamespace(environment_id="gpu-host"),
|
|
docker=docker,
|
|
),
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(startup_timeout_seconds=900),
|
|
)
|
|
spec = SimpleNamespace(
|
|
image_tag="comfy-remote:deadbeef",
|
|
identity=SimpleNamespace(fingerprint="deadbeef"),
|
|
)
|
|
monkeypatch.setattr(manager, "_build_context", lambda runtime_spec: b"context")
|
|
|
|
manager._build_image(spec)
|
|
|
|
arguments, kwargs = calls[0]
|
|
assert arguments[:3] == ("build", "--pull", "--load")
|
|
assert kwargs["input_payload"] == b"context"
|
|
|
|
|
|
def test_conflicting_worker_launch_adopts_current_managed_winner(
|
|
ssh_runtime_module: Any,
|
|
ssh_docker_module: Any,
|
|
) -> None:
|
|
"""A concurrent launcher winning the Docker name race should be reused."""
|
|
fingerprint = "ac12dff11abd460ecb642f3f0f19f3c5a8fe3d345f88f0d45e7b05d4036f8fa2"
|
|
container_name = "comfy-remote-lambda-ac12dff11abd460e-w0"
|
|
inspect_count = 0
|
|
|
|
def docker(arguments: tuple[str, ...], **kwargs: Any) -> Any:
|
|
"""Return a name conflict followed by the concurrent winner's state."""
|
|
nonlocal inspect_count
|
|
if arguments[:2] == ("container", "inspect"):
|
|
inspect_count += 1
|
|
if inspect_count == 1:
|
|
return ssh_docker_module.SshCommandResult(b"", b"not found", 1)
|
|
payload = [
|
|
{
|
|
"State": {"Running": True},
|
|
"Config": {
|
|
"Labels": {
|
|
"comfy.remote.runtime-fingerprint": fingerprint,
|
|
"comfy.remote.environment-id": "lambda",
|
|
"comfy.remote.worker-index": "0",
|
|
}
|
|
},
|
|
}
|
|
]
|
|
return ssh_docker_module.SshCommandResult(
|
|
json.dumps(payload).encode(),
|
|
b"",
|
|
0,
|
|
)
|
|
if arguments[0] == "run":
|
|
assert kwargs["check"] is False
|
|
return ssh_docker_module.SshCommandResult(
|
|
b"",
|
|
b"Conflict. The container name is already in use.",
|
|
125,
|
|
)
|
|
raise AssertionError(f"Unexpected Docker arguments: {arguments!r}")
|
|
|
|
controller = SimpleNamespace(
|
|
host=SimpleNamespace(
|
|
environment_id="lambda",
|
|
ssh_target="lambda",
|
|
capabilities=SimpleNamespace(
|
|
gpus=(SimpleNamespace(uuid="GPU-one"),),
|
|
nvidia_container_runtime=True,
|
|
),
|
|
docker_env_file=None,
|
|
),
|
|
docker=docker,
|
|
)
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=controller,
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(),
|
|
)
|
|
spec = ssh_runtime_module.SshRuntimeSpec(
|
|
identity=SimpleNamespace(fingerprint=fingerprint),
|
|
image_tag="comfy-remote:ac12dff11abd460e",
|
|
container_name=container_name,
|
|
storage_volume_name="comfy-remote-lambda",
|
|
worker_index=0,
|
|
)
|
|
|
|
manager._replace_worker_container(spec)
|
|
|
|
assert inspect_count == 2
|
|
|
|
|
|
def test_current_worker_requires_environment_and_slot_ownership_labels(
|
|
ssh_runtime_module: Any,
|
|
ssh_docker_module: Any,
|
|
) -> None:
|
|
"""A matching runtime fingerprint alone must not authorize worker adoption."""
|
|
fingerprint = "deadbeef" * 8
|
|
payload = [
|
|
{
|
|
"State": {"Running": True},
|
|
"Config": {
|
|
"Labels": {
|
|
"comfy.remote.runtime-fingerprint": fingerprint,
|
|
"comfy.remote.environment-id": "different-host",
|
|
"comfy.remote.worker-index": "0",
|
|
}
|
|
},
|
|
}
|
|
]
|
|
controller = SimpleNamespace(
|
|
host=SimpleNamespace(environment_id="lambda"),
|
|
docker=lambda *_args, **_kwargs: ssh_docker_module.SshCommandResult(
|
|
json.dumps(payload).encode(),
|
|
b"",
|
|
0,
|
|
),
|
|
)
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=controller,
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(),
|
|
)
|
|
spec = SimpleNamespace(
|
|
identity=SimpleNamespace(fingerprint=fingerprint),
|
|
container_name="comfy-remote-lambda-deadbeefdeadbeef-w0",
|
|
worker_index=0,
|
|
)
|
|
|
|
assert manager._container_is_current_and_running(spec) is False
|
|
|
|
|
|
def test_worker_context_includes_top_level_comfyui_python_modules(
|
|
ssh_runtime_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""The SSH image must contain modules imported by headless ComfyUI startup."""
|
|
repo_root = tmp_path / "repo"
|
|
comfyui_root = tmp_path / "ComfyUI"
|
|
repo_root.mkdir()
|
|
comfyui_root.mkdir()
|
|
(repo_root / "worker.py").write_text("VALUE = 1\n", encoding="utf-8")
|
|
for module_name in ("execution.py", "folder_paths.py", "nodes.py", "server.py"):
|
|
(comfyui_root / module_name).write_text("VALUE = 1\n", encoding="utf-8")
|
|
(comfyui_root / "requirements.txt").write_text(
|
|
"aiohttp==3.13.3\n",
|
|
encoding="utf-8",
|
|
)
|
|
(comfyui_root / "README.md").write_text("not runtime source\n", encoding="utf-8")
|
|
comfy_package = comfyui_root / "comfy"
|
|
comfy_package.mkdir()
|
|
(comfy_package / "__init__.py").write_text("", encoding="utf-8")
|
|
(comfy_package / "data.json").write_text("{}\n", encoding="utf-8")
|
|
tokenizer_directory = comfy_package / "text_encoders" / "qwen25_tokenizer"
|
|
tokenizer_directory.mkdir(parents=True)
|
|
(tokenizer_directory / "merges.txt").write_text("a b\n", encoding="utf-8")
|
|
(tokenizer_directory / "vocab.json").write_text("{}\n", encoding="utf-8")
|
|
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=SimpleNamespace(),
|
|
repo_root=repo_root,
|
|
settings=SimpleNamespace(comfyui_root=comfyui_root),
|
|
)
|
|
|
|
archive_paths = {
|
|
archive_path for _, archive_path in manager._runtime_context_files()
|
|
}
|
|
|
|
assert {
|
|
"comfyui/execution.py",
|
|
"comfyui/folder_paths.py",
|
|
"comfyui/nodes.py",
|
|
"comfyui/server.py",
|
|
"comfyui/requirements.txt",
|
|
"comfyui/comfy/__init__.py",
|
|
"comfyui/comfy/data.json",
|
|
"comfyui/comfy/text_encoders/qwen25_tokenizer/merges.txt",
|
|
"comfyui/comfy/text_encoders/qwen25_tokenizer/vocab.json",
|
|
}.issubset(archive_paths)
|
|
assert "comfyui/README.md" not in archive_paths
|
|
|
|
|
|
def test_worker_dockerfile_disables_inherited_base_image_healthcheck(
|
|
ssh_runtime_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""The socket worker must not inherit llama-server's HTTP health check."""
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"select_remote_torch_build",
|
|
lambda _gpu: SimpleNamespace(
|
|
install_layers=(),
|
|
validation_command=lambda: "true",
|
|
),
|
|
)
|
|
monkeypatch.setattr(ssh_runtime_module, "remote_apt_packages", lambda: ())
|
|
monkeypatch.setattr(ssh_runtime_module, "remote_runtime_packages", lambda: ())
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"remote_accelerator_packages",
|
|
lambda _gpu: (),
|
|
)
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"remote_accelerator_validation_command",
|
|
lambda _gpu: "true",
|
|
)
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"custom_node_runtime_packages",
|
|
lambda _path: (),
|
|
)
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=SimpleNamespace(),
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(
|
|
modal_gpu="RTX-PRO-6000",
|
|
custom_nodes_dir=Path("/custom_nodes"),
|
|
),
|
|
)
|
|
spec = SimpleNamespace(identity=SimpleNamespace(fingerprint="deadbeef"))
|
|
|
|
dockerfile = manager._dockerfile(spec)
|
|
|
|
assert "HEALTHCHECK NONE" in dockerfile
|
|
assert dockerfile.index("HEALTHCHECK NONE") < dockerfile.index("ENTRYPOINT")
|
|
assert "LD_LIBRARY_PATH=/app:${LD_LIBRARY_PATH}" in dockerfile
|
|
assert "COMFY_MODAL_LLM_EXECUTION_TARGET=ssh_docker" in dockerfile
|
|
assert "from av.video.reformatter import ColorPrimaries" in dockerfile
|
|
|
|
|
|
def test_worker_dockerfile_exposes_triton_compiler_toolchain(
|
|
ssh_runtime_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Generated SSH and Vast images must expose the compiler Triton discovers."""
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"select_remote_torch_build",
|
|
lambda _gpu: SimpleNamespace(
|
|
install_layers=(),
|
|
validation_command=lambda: "true",
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"remote_apt_packages",
|
|
lambda: ("build-essential",),
|
|
)
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"remote_compiler_validation_command",
|
|
lambda: "test -x /usr/bin/gcc",
|
|
)
|
|
monkeypatch.setattr(ssh_runtime_module, "remote_runtime_packages", lambda: ())
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"remote_accelerator_packages",
|
|
lambda _gpu: (),
|
|
)
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"remote_accelerator_validation_command",
|
|
lambda _gpu: "true",
|
|
)
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"custom_node_runtime_packages",
|
|
lambda _path: (),
|
|
)
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=SimpleNamespace(),
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(
|
|
modal_gpu="RTX-PRO-6000",
|
|
custom_nodes_dir=Path("/custom_nodes"),
|
|
),
|
|
)
|
|
spec = SimpleNamespace(identity=SimpleNamespace(fingerprint="deadbeef"))
|
|
|
|
dockerfile = manager._dockerfile(spec)
|
|
|
|
assert "apt-get install -y --no-install-recommends build-essential" in dockerfile
|
|
assert "CC=/usr/bin/gcc CXX=/usr/bin/g++" in dockerfile
|
|
assert "RUN test -x /usr/bin/gcc" in dockerfile
|
|
|
|
|
|
def test_worker_dockerfile_retries_large_accelerator_downloads(
|
|
ssh_runtime_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Large direct wheels should survive transient interrupted downloads."""
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"select_remote_torch_build",
|
|
lambda _gpu: SimpleNamespace(
|
|
install_layers=(),
|
|
validation_command=lambda: "true",
|
|
),
|
|
)
|
|
monkeypatch.setattr(ssh_runtime_module, "remote_apt_packages", lambda: ())
|
|
monkeypatch.setattr(ssh_runtime_module, "remote_runtime_packages", lambda: ())
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"remote_accelerator_packages",
|
|
lambda _gpu: ("https://example.test/large-accelerator.whl",),
|
|
)
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"remote_accelerator_validation_command",
|
|
lambda _gpu: "true",
|
|
)
|
|
monkeypatch.setattr(
|
|
ssh_runtime_module,
|
|
"custom_node_runtime_packages",
|
|
lambda _path: (),
|
|
)
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=SimpleNamespace(),
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(
|
|
modal_gpu="RTX-PRO-6000",
|
|
custom_nodes_dir=Path("/custom_nodes"),
|
|
),
|
|
)
|
|
spec = SimpleNamespace(identity=SimpleNamespace(fingerprint="deadbeef"))
|
|
|
|
dockerfile = manager._dockerfile(spec)
|
|
|
|
assert (
|
|
"RUN python -m pip install --no-cache-dir --resume-retries 20 "
|
|
"--timeout 120 https://example.test/large-accelerator.whl"
|
|
) in dockerfile
|
|
|
|
|
|
def test_worker_indices_are_distributed_across_discovered_gpus(
|
|
ssh_runtime_module: Any,
|
|
remote_hosts_module: Any,
|
|
execution_environments_module: Any,
|
|
) -> None:
|
|
"""Parallel worker containers should deterministically select different GPUs."""
|
|
gpu_type = execution_environments_module.GpuCapability
|
|
capabilities = execution_environments_module.EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=32,
|
|
total_ram_bytes=128 * 1024**3,
|
|
available_ram_bytes=120 * 1024**3,
|
|
available_disk_bytes=1024**4,
|
|
docker_version="28.0.0",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
gpus=(
|
|
gpu_type("GPU-one", "GPU one", 48 * 1024**3),
|
|
gpu_type("GPU-two", "GPU two", 48 * 1024**3),
|
|
),
|
|
)
|
|
host = remote_hosts_module.SshHostConfig(
|
|
environment_id="dual-gpu",
|
|
display_name="Dual GPU",
|
|
ssh_target="dual-gpu",
|
|
capabilities=capabilities,
|
|
)
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=SimpleNamespace(host=host),
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(),
|
|
)
|
|
|
|
assert manager._gpu_arguments(0) == ("--gpus", "device=GPU-one")
|
|
assert manager._gpu_arguments(1) == ("--gpus", "device=GPU-two")
|
|
assert manager._gpu_arguments(2) == ("--gpus", "device=GPU-one")
|
|
|
|
|
|
def test_worker_launch_refuses_missing_gpu_capabilities(
|
|
ssh_runtime_module: Any,
|
|
ssh_docker_module: Any,
|
|
) -> None:
|
|
"""A missing probe must fail instead of silently launching a CPU container."""
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=SimpleNamespace(
|
|
host=SimpleNamespace(capabilities=None),
|
|
),
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(),
|
|
)
|
|
|
|
try:
|
|
manager._gpu_arguments(0)
|
|
except ssh_docker_module.SshDockerError as exc:
|
|
assert "requires probed GPU capabilities" in str(exc)
|
|
else:
|
|
raise AssertionError("Missing GPU capabilities should fail worker launch.")
|
|
|
|
|
|
def test_stale_worker_reconciliation_only_replaces_the_same_slot(
|
|
ssh_runtime_module: Any,
|
|
) -> None:
|
|
"""A new fingerprint should retire its superseded logical worker container."""
|
|
current = SimpleNamespace(
|
|
worker_index=1,
|
|
container_name="comfy-remote-host-new-w1",
|
|
)
|
|
workers = (
|
|
SimpleNamespace(worker_index=0, container_name="comfy-remote-host-old-w0"),
|
|
SimpleNamespace(worker_index=1, container_name="comfy-remote-host-old-w1"),
|
|
SimpleNamespace(worker_index=1, container_name=current.container_name),
|
|
)
|
|
removed: list[str] = []
|
|
controller = SimpleNamespace(
|
|
host=SimpleNamespace(environment_id="host"),
|
|
list_managed_workers=lambda: workers,
|
|
remove_managed_worker=removed.append,
|
|
)
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=controller,
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(),
|
|
)
|
|
|
|
manager._remove_stale_worker_containers(current)
|
|
|
|
assert removed == ["comfy-remote-host-old-w1"]
|
|
|
|
|
|
def test_stop_all_workers_removes_only_controller_managed_names(
|
|
ssh_runtime_module: Any,
|
|
) -> None:
|
|
"""Environment shutdown should delegate every ownership check to the controller."""
|
|
workers = (
|
|
SimpleNamespace(container_name="worker-one"),
|
|
SimpleNamespace(container_name="worker-two"),
|
|
)
|
|
removed: list[str] = []
|
|
controller = SimpleNamespace(
|
|
list_managed_workers=lambda: workers,
|
|
remove_managed_worker=removed.append,
|
|
)
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=controller,
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(),
|
|
)
|
|
|
|
result = manager.stop_all_workers()
|
|
|
|
assert result == ("worker-one", "worker-two")
|
|
assert removed == ["worker-one", "worker-two"]
|
|
|
|
|
|
def test_runtime_passes_only_remote_environment_file_path(
|
|
ssh_runtime_module: Any,
|
|
) -> None:
|
|
"""Worker secrets should remain in an administrator-managed file on the host."""
|
|
manager = ssh_runtime_module.SshRuntimeManager(
|
|
controller=SimpleNamespace(
|
|
host=SimpleNamespace(docker_env_file="/etc/comfy/worker.env")
|
|
),
|
|
repo_root=SimpleNamespace(),
|
|
settings=SimpleNamespace(),
|
|
)
|
|
|
|
assert manager._environment_file_arguments() == (
|
|
"--env-file",
|
|
"/etc/comfy/worker.env",
|
|
)
|