"""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", worker_index=0, ) 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, status_callback=None: 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_image_only_preparation_builds_before_storage_helpers_run( ssh_runtime_module: Any, monkeypatch: Any, ) -> None: """R2 setup should make the runtime image available without starting a worker.""" statuses: list[str] = [] ensured_volumes: list[str] = [] built_specs: list[Any] = [] spec = SimpleNamespace( image_tag="comfy-remote:deadbeef", storage_volume_name="comfy-remote-lambda", worker_index=0, ) manager = ssh_runtime_module.SshRuntimeManager( controller=SimpleNamespace( host=SimpleNamespace(environment_id="lambda-image-helper"), ensure_volume=ensured_volumes.append, ), repo_root=SimpleNamespace(), settings=SimpleNamespace(), ) monkeypatch.setattr(manager, "_image_is_current", lambda _spec: False) monkeypatch.setattr( manager, "_build_image", lambda runtime_spec, status_callback=None: built_specs.append(runtime_spec), ) result = manager.ensure_image(spec, status_callback=statuses.append) assert result is spec assert ensured_volumes == ["comfy-remote-lambda"] assert built_specs == [spec] assert statuses == [ "Building SSH runtime environment=lambda-image-helper " "image=comfy-remote:deadbeef" ] 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) -> Any: """Record one remote Docker invocation.""" calls.append((arguments, kwargs)) if arguments[:2] == ("image", "inspect"): return SimpleNamespace(returncode=1, stdout_text="") return None 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, "_dependency_build_context", lambda: b"dependencies", ) monkeypatch.setattr( manager, "_source_overlay_build_context", lambda runtime_spec: b"source-overlay", ) monkeypatch.setattr( ssh_runtime_module, "remote_runtime_dependency_fingerprint", lambda _identity: "cafebabe" * 8, ) monkeypatch.setattr( manager, "_remove_stale_runtime_images", lambda _spec, status_callback=None: (), ) manager._build_image(spec) assert calls[0][0][:2] == ("image", "inspect") dependency_arguments, dependency_kwargs = calls[1] assert dependency_arguments[:6] == ( "buildx", "build", "--builder", "default", "--pull", "--load", ) assert dependency_kwargs["input_payload"] == b"dependencies" overlay_arguments, overlay_kwargs = calls[2] assert overlay_arguments[:5] == ( "buildx", "build", "--builder", "default", "--load", ) assert "--pull" not in overlay_arguments assert overlay_kwargs["input_payload"] == b"source-overlay" def test_worker_build_reuses_retained_dependency_image( ssh_runtime_module: Any, monkeypatch: Any, ) -> None: """A code-only runtime change should build only the source overlay.""" calls: list[tuple[tuple[str, ...], dict[str, Any]]] = [] def docker(arguments: tuple[str, ...], **kwargs: Any) -> Any: """Report a current dependency image and record the overlay build.""" calls.append((arguments, kwargs)) if arguments[:2] == ("image", "inspect"): return SimpleNamespace( returncode=0, stdout_text="cafebabe" * 8, ) return None 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:source-change", identity=SimpleNamespace(fingerprint="source-change"), ) monkeypatch.setattr( ssh_runtime_module, "remote_runtime_dependency_fingerprint", lambda _identity: "cafebabe" * 8, ) def fail_dependency_build() -> bytes: """Fail if a retained dependency image is rebuilt.""" raise AssertionError("dependency context should not be rebuilt") monkeypatch.setattr( manager, "_remove_stale_runtime_images", lambda _spec, status_callback=None: (), ) monkeypatch.setattr(manager, "_dependency_build_context", fail_dependency_build) monkeypatch.setattr( manager, "_source_overlay_build_context", lambda _runtime_spec: b"52MB-source-overlay", ) manager._build_image(spec) assert len(calls) == 2 assert calls[0][0][:2] == ("image", "inspect") assert calls[1][0][:5] == ( "buildx", "build", "--builder", "default", "--load", ) assert calls[1][1]["input_payload"] == b"52MB-source-overlay" def test_worker_build_reclaims_only_superseded_managed_images( ssh_runtime_module: Any, ssh_docker_module: Any, monkeypatch: Any, ) -> None: """Disk reclamation must stay inside extension-owned image repositories.""" removed_references: list[str] = [] def docker(arguments: tuple[str, ...], **_kwargs: Any) -> Any: """Return managed image listings and record targeted removals.""" if arguments[:2] == ("image", "ls"): label_filter = arguments[3] if label_filter == "label=comfy.remote.runtime-fingerprint": output = "\n".join( ( "comfy-remote:current", "comfy-remote:old-source", "unrelated:keep", ) ) else: output = "\n".join( ( "comfy-remote-deps:current-deps", "comfy-remote-deps:old-deps", "comfy-remote:old-source", ) ) return ssh_docker_module.SshCommandResult(output.encode(), b"", 0) if arguments[:2] == ("image", "rm"): removed_references.append(arguments[2]) return ssh_docker_module.SshCommandResult(b"", b"", 0) raise AssertionError(f"Unexpected Docker arguments: {arguments}") manager = ssh_runtime_module.SshRuntimeManager( controller=SimpleNamespace( host=SimpleNamespace(environment_id="gpu-host"), docker=docker, ), repo_root=SimpleNamespace(), settings=SimpleNamespace(), ) spec = SimpleNamespace( image_tag="comfy-remote:current", identity=SimpleNamespace(fingerprint="source-fingerprint"), ) monkeypatch.setattr( manager, "_dependency_image_tag", lambda _spec: "comfy-remote-deps:current-deps", ) removed = manager._remove_stale_runtime_images(spec) assert removed == ( "comfy-remote-deps:old-deps", "comfy-remote:old-source", ) assert removed_references == list(removed) 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_dependency_base_is_separate_from_source_overlay( ssh_runtime_module: Any, monkeypatch: Any, ) -> None: """Frequent source edits must not invalidate heavyweight package layers.""" 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: ("runtime-package==1",), ) monkeypatch.setattr( ssh_runtime_module, "remote_accelerator_packages", lambda _gpu: ("accelerator-package==1",), ) monkeypatch.setattr( ssh_runtime_module, "remote_accelerator_validation_command", lambda _gpu: "true", ) monkeypatch.setattr( ssh_runtime_module, "custom_node_runtime_packages", lambda _path: (), ) monkeypatch.setattr( ssh_runtime_module, "remote_runtime_dependency_fingerprint", lambda _identity: "cafebabe" * 8, ) 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="source-fingerprint")) dependency_dockerfile = manager._dependency_dockerfile() overlay_dockerfile = manager._source_overlay_dockerfile(spec) published_overlay = manager._source_overlay_dockerfile( spec, dependency_image="ghcr.io/example/worker@sha256:" + "a" * 64, ) assert "pip install --no-cache-dir runtime-package==1" in dependency_dockerfile assert "accelerator-package==1" in dependency_dockerfile assert "COPY repo" not in dependency_dockerfile assert "COPY comfyui" not in dependency_dockerfile assert overlay_dockerfile.startswith("FROM comfy-remote-deps:cafebabecafebabe\n") assert "COPY repo /opt/comfy-remote/repo" in overlay_dockerfile assert "COPY comfyui /opt/comfy-remote/ComfyUI" in overlay_dockerfile assert "apt-get" not in overlay_dockerfile assert "pip install" not in overlay_dockerfile assert "COMFY_MODAL_RUNTIME_FINGERPRINT=source-fingerprint" in overlay_dockerfile assert published_overlay.startswith( "FROM ghcr.io/example/worker@sha256:" + "a" * 64 + "\n" ) def test_worker_installs_pinned_torch_before_general_runtime_packages( ssh_runtime_module: Any, monkeypatch: Any, ) -> None: """General requirements must resolve against the selected pinned Torch stack.""" monkeypatch.setattr( ssh_runtime_module, "select_remote_torch_build", lambda _gpu: SimpleNamespace( install_layers=( SimpleNamespace( index_url="https://download.pytorch.org/whl/cu130", extra_options="", packages=("torch==2.13.0",), ), ), validation_command=lambda: "validate-torch", ), ) monkeypatch.setattr(ssh_runtime_module, "remote_apt_packages", lambda: ()) monkeypatch.setattr( ssh_runtime_module, "remote_runtime_packages", lambda: ("runtime-package==1",), ) monkeypatch.setattr( ssh_runtime_module, "remote_accelerator_packages", lambda _gpu: (), ) monkeypatch.setattr( ssh_runtime_module, "remote_accelerator_validation_command", lambda _gpu: "validate-accelerator", ) 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="A100", custom_nodes_dir=Path("/custom_nodes"), ), ) dockerfile = manager._dependency_dockerfile() assert dockerfile.index("torch==2.13.0") < dockerfile.index("runtime-package==1") assert dockerfile.index("validate-torch") < dockerfile.index("runtime-package==1") 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", )