refactor: derive worker SSH capability from executor initialization

Advertise capabilities from the initialized executor map rather than a
standalone Docker probe: executors declare an is_available() hook, SSHExecutor
returns docker_available(), and init_executor skips an unavailable executor so
it is never constructed. The worker initializes executors before registering
and reports ssh = "ssh" in executors, so a disabled flag, missing socket,
missing package, or init failure all converge to ssh=False. The registry now
validates entries are Executor subclasses, which types the is_available gate
without a cast. SSHExecutor resolves its container label name lazily so it can
still pick up the worker id once registration completes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

Author: kaiitunnz

src/worker/executors/__init__.py

@@ -1,55 +1,60 @@
 import importlib
 
+from .base_executor import Executor
+
 _IMPORT_ERRORS: dict[str, str] = {}
 
 
-def _safe_import(name: str, module: str) -> type | None:
+def _import_executor(name: str, module: str) -> type[Executor] | None:
     try:
         pkg = importlib.import_module(module, package=__package__)
-        return getattr(pkg, name)
+        if issubclass(cls := getattr(pkg, name), Executor):
+            return cls
+        error = f"{name} is not a subclass of Executor"
     except Exception as exc:
-        _IMPORT_ERRORS[name] = str(exc)
-        return None
+        error = str(exc)
+    _IMPORT_ERRORS[name] = error
+    return None
 
 
-VLLMExecutor = _safe_import("VLLMExecutor", ".vllm_executor")
-VLLMLoRAExecutor = _safe_import("VLLMLoRAExecutor", ".vllm_lora_executor")
-PPOExecutor = _safe_import("PPOExecutor", ".ppo_executor")
-DPOExecutor = _safe_import("DPOExecutor", ".dpo_executor")
-SFTExecutor = _safe_import("SFTExecutor", ".sft_executor")
-LoRASFTExecutor = _safe_import("LoRASFTExecutor", ".lora_sft_executor")
-ImageClassificationTrainingExecutor = _safe_import(
+VLLMExecutor = _import_executor("VLLMExecutor", ".vllm_executor")
+VLLMLoRAExecutor = _import_executor("VLLMLoRAExecutor", ".vllm_lora_executor")
+PPOExecutor = _import_executor("PPOExecutor", ".ppo_executor")
+DPOExecutor = _import_executor("DPOExecutor", ".dpo_executor")
+SFTExecutor = _import_executor("SFTExecutor", ".sft_executor")
+LoRASFTExecutor = _import_executor("LoRASFTExecutor", ".lora_sft_executor")
+ImageClassificationTrainingExecutor = _import_executor(
     "ImageClassificationTrainingExecutor", ".image_classification_executor"
 )
-HFTransformersExecutor = _safe_import(
+HFTransformersExecutor = _import_executor(
     "HFTransformersExecutor", ".transformers_executor"
 )
-RAGExecutor = _safe_import("RAGExecutor", ".rag_executor")
-AgentExecutor = _safe_import("AgentExecutor", ".agent_executor")
-EchoExecutor = _safe_import("EchoExecutor", ".echo_executor")
-DataProfilingExecutor = _safe_import(
+RAGExecutor = _import_executor("RAGExecutor", ".rag_executor")
+AgentExecutor = _import_executor("AgentExecutor", ".agent_executor")
+EchoExecutor = _import_executor("EchoExecutor", ".echo_executor")
+DataProfilingExecutor = _import_executor(
     "DataProfilingExecutor", ".data_profiling_executor"
 )
-DataRetrievalExecutor = _safe_import(
+DataRetrievalExecutor = _import_executor(
     "DataRetrievalExecutor", ".data_retrieval_executor"
 )
-DiffusersExecutor = _safe_import("DiffusersExecutor", ".diffusers_executor")
-APIExecutor = _safe_import("APIExecutor", ".api_executor")
-SSHExecutor = _safe_import("SSHExecutor", ".ssh_executor")
-OmniText2ImageExecutor = _safe_import(
+DiffusersExecutor = _import_executor("DiffusersExecutor", ".diffusers_executor")
+APIExecutor = _import_executor("APIExecutor", ".api_executor")
+SSHExecutor = _import_executor("SSHExecutor", ".ssh_executor")
+OmniText2ImageExecutor = _import_executor(
     "OmniText2ImageExecutor", ".omni_text2image_executor"
 )
-OmniText2SpeechExecutor = _safe_import(
+OmniText2SpeechExecutor = _import_executor(
     "OmniText2SpeechExecutor", ".omni_text2speech_executor"
 )
-OmniText2AudioExecutor = _safe_import(
+OmniText2AudioExecutor = _import_executor(
     "OmniText2AudioExecutor", ".omni_text2audio_executor"
 )
-OmniText2GeneralExecutor = _safe_import(
+OmniText2GeneralExecutor = _import_executor(
     "OmniText2GeneralExecutor", ".omni_text2general_executor"
 )
 
-EXECUTOR_REGISTRY: dict[str, type | None] = {
+EXECUTOR_REGISTRY: dict[str, type[Executor] | None] = {
     "vllm": VLLMExecutor,
     "vllm_lora": VLLMLoRAExecutor,
     "ppo": PPOExecutor,

src/worker/executors/base_executor.py

@@ -81,6 +81,16 @@ def __init__(
         self._hardware = hardware
         self._lifecycle = lifecycle
 
+    @classmethod
+    def is_available(cls, config: WorkerConfig) -> bool:
+        """Whether this executor can run on the worker.
+
+        Override to declare a runtime dependency the worker must satisfy. An unavailable
+        executor is skipped at startup, so it never registers and the worker advertises
+        no capability for it.
+        """
+        return True
+
     def emit_update(self, task_id: str, payload: dict[str, Any]) -> None:
         """Emit a mid-task TASK_UPDATE event.
 

src/worker/executors/ssh_executor.py

@@ -48,7 +48,11 @@
 from shared.utils.manifest import ARTIFACTS_DIR, prepare_output_dir
 from worker.config import WorkerConfig
 from worker.executors.utils.checkpoints import maybe_upload_artifacts
-from worker.executors.utils.docker import DockerUnavailableError, docker_client
+from worker.executors.utils.docker import (
+    DockerUnavailableError,
+    docker_available,
+    docker_client,
+)
 
 from .base_executor import (
     ExecutionError,
@@ -384,19 +388,30 @@ class SSHExecutor(Executor):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         config = self._config
-        lifecycle = self._lifecycle
-        self._worker_name = (
-            config.container_name
-            or config.alias
-            or (lifecycle.worker_id if lifecycle else uuid.uuid4().hex[:8])
-        )
+        self._worker_name = config.container_name or config.alias or None
         self._docker: DockerClient | None = None
         self._docker_gpu_runtime: str | None = config.docker_gpu_runtime
         self._ssh_network: str | None = None
         self._cancel_event = threading.Event()
         self._finish_event = threading.Event()
         self._current_container: Container | None = None
 
+    @classmethod
+    def is_available(cls, config: WorkerConfig) -> bool:
+        return docker_available()
+
+    @property
+    def worker_name(self) -> str:
+        if name := self._worker_name:
+            return name
+        name = (
+            lifecycle.worker_id
+            if (lifecycle := self._lifecycle)
+            else uuid.uuid4().hex[:8]
+        )
+        self._worker_name = name
+        return name
+
     # ------------------------------------------------------------------ #
     # Lifecycle
     # ------------------------------------------------------------------ #
@@ -413,7 +428,7 @@ def teardown(self) -> None:
         client = self._get_docker_client()
         try:
             containers = client.containers.list(
-                filters={"label": f"{_LABEL_WORKER}={self._worker_name}"}
+                filters={"label": f"{_LABEL_WORKER}={self.worker_name}"}
             )
         except Exception as exc:
             logger.warning(
@@ -454,7 +469,8 @@ def run(self, task: ExecutorTask, out_dir: Path) -> SSHResult:
             container_cmd = self._resolve_noninteractive_command(client, cfg)
 
         session_id = new_ssh_session_id()
-        container_name = f"{self._worker_name}_ssh-{task.task_id[:8]}-{session_id[:8]}"
+        worker_name = self.worker_name
+        container_name = f"{worker_name}_ssh-{task.task_id[:8]}-{session_id[:8]}"
 
         prepare_output_dir(out_dir)  # Ensure output dir exists before mounting
         resolved_inputs = self._resolve_inputs(task, cfg)
@@ -463,7 +479,7 @@ def run(self, task: ExecutorTask, out_dir: Path) -> SSHResult:
         )
 
         labels = {
-            _LABEL_WORKER: self._worker_name,
+            _LABEL_WORKER: worker_name,
             _LABEL_TASK: task.task_id,
             _LABEL_SESSION: session_id,
             _LABEL_MANAGED: "true",
@@ -1193,7 +1209,7 @@ def _stage_inputs_in_volume(
         volume = client.volumes.create(
             name=volume_name,
             labels={
-                _LABEL_WORKER: self._worker_name,
+                _LABEL_WORKER: self.worker_name,
                 _LABEL_SESSION: session_id,
                 _LABEL_MANAGED: "true",
             },

src/worker/main.py

@@ -9,7 +9,6 @@
 from .executors import EXECUTOR_CLASS_NAMES, EXECUTOR_REGISTRY, IMPORT_ERRORS
 from .executors.base_executor import Executor
 from .executors.mp_executor import MPExecutor
-from .executors.utils.docker import docker_available
 from .hw import collect_hw
 from .lifecycle import Lifecycle
 from .power import PowerMonitor
@@ -59,7 +58,7 @@ def initialize_executors(
     hardware: WorkerHardware,
     logger: logging.Logger,
     lifecycle: Lifecycle,
-    registry: dict[str, type | None] | None = None,
+    registry: dict[str, type[Executor] | None] | None = None,
     import_errors: dict[str, str] | None = None,
     cuda_available: bool | None = None,
     enable_mp_executors: bool = True,
@@ -97,6 +96,10 @@ def init_executor(key: str, *, gpu_required: bool = False):
             logger.info("Executor %s requires a GPU; unavailable, skipping", key)
             return None
 
+        if not cls.is_available(config):
+            logger.info("Executor %s is unavailable; skipping", key)
+            return None
+
         try:
             if key in configured_wrapped:
                 return MPExecutor(cls, config, hardware)
@@ -197,7 +200,16 @@ def main() -> None:
     )
     hardware = collect_hw(bandwidth_bytes_per_sec=cfg.network_bandwidth_bytes_per_sec)
     logger.info("Collected hardware info: %s", hardware)
-    capabilities = WorkerCapabilities(ssh=docker_available())
+
+    executors, default_executor = initialize_executors(
+        cfg,
+        hardware,
+        logger,
+        lifecycle,
+        enable_mp_executors=cfg.enable_mp_executors,
+    )
+
+    capabilities = WorkerCapabilities(ssh="ssh" in executors)
     ssh_limits = cfg.ssh_limits
     if capabilities.ssh:
         if ssh_limits is None:
@@ -207,8 +219,6 @@ def main() -> None:
             )
         else:
             logger.info("SSH resource cap: %s", ssh_limits.model_dump())
-    else:
-        logger.info("SSH support disabled; Docker is not reachable from this worker.")
     lifecycle.start(
         env={},
         hardware=hardware,
@@ -217,14 +227,6 @@ def main() -> None:
         tags=cfg.tags,
     )
 
-    executors, default_executor = initialize_executors(
-        cfg,
-        hardware,
-        logger,
-        lifecycle,
-        enable_mp_executors=cfg.enable_mp_executors,
-    )
-
     task_stream = supervisor_client.iter_tasks()
     runner = Runner(
         lifecycle,

tests/worker/test_executor_bootstrap.py

@@ -10,9 +10,13 @@
 from pathlib import Path
 from typing import Any
 
+import pytest
+
 from shared.schemas.result import BaseExecutorResult
 from tests.worker.factories import make_live_worker_config, make_worker_hardware
+from worker.executors import ssh_executor as ssh_mod
 from worker.executors.base_executor import Executor, ExecutorTask
+from worker.executors.ssh_executor import SSHExecutor
 from worker.main import initialize_executors
 
 
@@ -28,6 +32,16 @@ def run(self, task: ExecutorTask, out_dir: Path) -> BaseExecutorResult:
         return BaseExecutorResult.model_validate({"ok": True})
 
 
+class _UnavailableExecutor(_PassthroughExecutor):
+    """Executor that declares itself unavailable on this worker."""
+
+    name = "unavailable"
+
+    @classmethod
+    def is_available(cls, config: Any) -> bool:
+        return False
+
+
 class TestInitializeExecutorsHardware:
     def test_executor_receives_hardware_via_passthrough(self, tmp_path: Path) -> None:
         cfg = make_live_worker_config(tmp_path)
@@ -48,3 +62,32 @@ def test_executor_receives_hardware_via_passthrough(self, tmp_path: Path) -> Non
         assert isinstance(default, _PassthroughExecutor)
         assert executors["echo"]._hardware is hw
         assert default._hardware is hw
+
+
+class TestInitializeExecutorsAvailability:
+    def test_unavailable_executor_is_skipped(self, tmp_path: Path) -> None:
+        cfg = make_live_worker_config(tmp_path)
+        executors, _ = initialize_executors(
+            config=cfg,
+            hardware=make_worker_hardware(),
+            logger=logging.getLogger("test"),
+            lifecycle=None,  # type: ignore[arg-type]
+            registry={
+                "default": _PassthroughExecutor,
+                "echo": _UnavailableExecutor,
+            },
+            import_errors={},
+            cuda_available=False,
+            enable_mp_executors=False,
+        )
+        assert "echo" not in executors
+        assert "default" in executors
+
+    def test_ssh_availability_tracks_docker(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        cfg = make_live_worker_config(tmp_path)
+        monkeypatch.setattr(ssh_mod, "docker_available", lambda: False)
+        assert SSHExecutor.is_available(cfg) is False
+        monkeypatch.setattr(ssh_mod, "docker_available", lambda: True)
+        assert SSHExecutor.is_available(cfg) is True

tests/worker/test_executor_registry.py

@@ -1,5 +1,10 @@
 """Tests for the executor registry and safe import mechanism."""
 
+from types import SimpleNamespace
+
+import pytest
+
+import worker.executors as executors_pkg
 from worker.executors import EXECUTOR_CLASS_NAMES, EXECUTOR_REGISTRY, IMPORT_ERRORS
 
 
@@ -51,8 +56,26 @@ def test_training_executors_are_wrapped_for_isolation(self) -> None:
         assert "image_classification_training" in _EXECUTORS_TO_WRAP
         assert "image_classification_training" in EXECUTOR_REGISTRY
 
-    def test_safe_import_does_not_crash(self) -> None:
+    def test_import_executor_does_not_crash(self) -> None:
         """The registry should load without raising, even when deps are missing."""
         # If we got here, the import at module level already succeeded
         assert isinstance(EXECUTOR_REGISTRY, dict)
         assert isinstance(IMPORT_ERRORS, dict)
+
+    def test_import_executor_rejects_non_executor(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """A registered name that resolves to a non-Executor is dropped and recorded."""
+        fake_module = SimpleNamespace(NotAnExecutor=object)
+        monkeypatch.setattr(
+            executors_pkg.importlib, "import_module", lambda *a, **k: fake_module
+        )
+        executors_pkg._IMPORT_ERRORS.pop("NotAnExecutor", None)
+
+        result = executors_pkg._import_executor("NotAnExecutor", ".does_not_matter")
+
+        assert result is None
+        assert (
+            "not a subclass of Executor"
+            in executors_pkg._IMPORT_ERRORS["NotAnExecutor"]
+        )