perf: import worker executors lazily via a registry

src/worker/executors/__init__.py

@@ -1,4 +1,6 @@
 import importlib
+from collections.abc import Iterator, Mapping
+from typing import overload
 
 from .base_executor import Executor
 
@@ -17,114 +19,66 @@ def _import_executor(name: str, module: str) -> type[Executor] | None:
     return None
 
 
-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 = _import_executor(
-    "HFTransformersExecutor", ".transformers_executor"
-)
-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 = _import_executor(
-    "DataRetrievalExecutor", ".data_retrieval_executor"
-)
-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 = _import_executor(
-    "OmniText2SpeechExecutor", ".omni_text2speech_executor"
-)
-OmniText2AudioExecutor = _import_executor(
-    "OmniText2AudioExecutor", ".omni_text2audio_executor"
-)
-OmniText2GeneralExecutor = _import_executor(
-    "OmniText2GeneralExecutor", ".omni_text2general_executor"
-)
-
-EXECUTOR_REGISTRY: dict[str, type[Executor] | None] = {
-    "vllm": VLLMExecutor,
-    "vllm_lora": VLLMLoRAExecutor,
-    "ppo": PPOExecutor,
-    "dpo": DPOExecutor,
-    "sft": SFTExecutor,
-    "lora_sft": LoRASFTExecutor,
-    "image_classification_training": ImageClassificationTrainingExecutor,
-    "default": HFTransformersExecutor,
-    "rag": RAGExecutor,
-    "agent": AgentExecutor,
-    "echo": EchoExecutor,
-    "data_profiling": DataProfilingExecutor,
-    "data_retrieval": DataRetrievalExecutor,
-    "diffusers": DiffusersExecutor,
-    "api": APIExecutor,
-    "ssh": SSHExecutor,
-    "omni_text2image": OmniText2ImageExecutor,
-    "omni_text2speech": OmniText2SpeechExecutor,
-    "omni_text2audio": OmniText2AudioExecutor,
-    "omni_text2general": OmniText2GeneralExecutor,
+EXECUTOR_MODULES: dict[str, tuple[str, str]] = {
+    "vllm": ("VLLMExecutor", ".vllm_executor"),
+    "vllm_lora": ("VLLMLoRAExecutor", ".vllm_lora_executor"),
+    "ppo": ("PPOExecutor", ".ppo_executor"),
+    "dpo": ("DPOExecutor", ".dpo_executor"),
+    "sft": ("SFTExecutor", ".sft_executor"),
+    "lora_sft": ("LoRASFTExecutor", ".lora_sft_executor"),
+    "image_classification_training": (
+        "ImageClassificationTrainingExecutor",
+        ".image_classification_executor",
+    ),
+    "default": ("HFTransformersExecutor", ".transformers_executor"),
+    "rag": ("RAGExecutor", ".rag_executor"),
+    "agent": ("AgentExecutor", ".agent_executor"),
+    "echo": ("EchoExecutor", ".echo_executor"),
+    "data_profiling": ("DataProfilingExecutor", ".data_profiling_executor"),
+    "data_retrieval": ("DataRetrievalExecutor", ".data_retrieval_executor"),
+    "diffusers": ("DiffusersExecutor", ".diffusers_executor"),
+    "api": ("APIExecutor", ".api_executor"),
+    "ssh": ("SSHExecutor", ".ssh_executor"),
+    "omni_text2image": ("OmniText2ImageExecutor", ".omni_text2image_executor"),
+    "omni_text2speech": ("OmniText2SpeechExecutor", ".omni_text2speech_executor"),
+    "omni_text2audio": ("OmniText2AudioExecutor", ".omni_text2audio_executor"),
+    "omni_text2general": (
+        "OmniText2GeneralExecutor",
+        ".omni_text2general_executor",
+    ),
 }
 
-EXECUTOR_CLASS_NAMES: dict[str, str] = {
-    "vllm": "VLLMExecutor",
-    "vllm_lora": "VLLMLoRAExecutor",
-    "ppo": "PPOExecutor",
-    "dpo": "DPOExecutor",
-    "sft": "SFTExecutor",
-    "lora_sft": "LoRASFTExecutor",
-    "image_classification_training": "ImageClassificationTrainingExecutor",
-    "default": "HFTransformersExecutor",
-    "rag": "RAGExecutor",
-    "agent": "AgentExecutor",
-    "echo": "EchoExecutor",
-    "data_profiling": "DataProfilingExecutor",
-    "data_retrieval": "DataRetrievalExecutor",
-    "diffusers": "DiffusersExecutor",
-    "api": "APIExecutor",
-    "ssh": "SSHExecutor",
-    "omni_text2image": "OmniText2ImageExecutor",
-    "omni_text2speech": "OmniText2SpeechExecutor",
-    "omni_text2audio": "OmniText2AudioExecutor",
-    "omni_text2general": "OmniText2GeneralExecutor",
-}
 
-IMPORT_ERRORS: dict[str, str] = _IMPORT_ERRORS
+class ExecutorRegistry(Mapping[str, type[Executor] | None]):
+    def __init__(self) -> None:
+        self._executors: dict[str, type[Executor] | None] = {}
+
+    def __getitem__(self, key: str) -> type[Executor] | None:
+        if key in self._executors:
+            return self._executors[key]
+        if key not in EXECUTOR_MODULES:
+            raise KeyError(f"Executor {key!r} not found in registry")
+        name, module = EXECUTOR_MODULES[key]
+        executor = _import_executor(name, module)
+        self._executors[key] = executor
+        return executor
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(EXECUTOR_MODULES)
+
+    def __len__(self) -> int:
+        return len(EXECUTOR_MODULES)
+
 
-__all__ = [
-    name
-    for name, cls in {
-        "VLLMExecutor": VLLMExecutor,
-        "VLLMLoRAExecutor": VLLMLoRAExecutor,
-        "PPOExecutor": PPOExecutor,
-        "DPOExecutor": DPOExecutor,
-        "SFTExecutor": SFTExecutor,
-        "LoRASFTExecutor": LoRASFTExecutor,
-        "ImageClassificationTrainingExecutor": ImageClassificationTrainingExecutor,
-        "HFTransformersExecutor": HFTransformersExecutor,
-        "RAGExecutor": RAGExecutor,
-        "AgentExecutor": AgentExecutor,
-        "EchoExecutor": EchoExecutor,
-        "DataProfilingExecutor": DataProfilingExecutor,
-        "DataRetrievalExecutor": DataRetrievalExecutor,
-        "DiffusersExecutor": DiffusersExecutor,
-        "APIExecutor": APIExecutor,
-        "SSHExecutor": SSHExecutor,
-        "OmniText2ImageExecutor": OmniText2ImageExecutor,
-        "OmniText2SpeechExecutor": OmniText2SpeechExecutor,
-        "OmniText2AudioExecutor": OmniText2AudioExecutor,
-        "OmniText2GeneralExecutor": OmniText2GeneralExecutor,
-    }.items()
-    if cls is not None
-] + ["EXECUTOR_REGISTRY", "IMPORT_ERRORS", "EXECUTOR_CLASS_NAMES"]
+EXECUTOR_REGISTRY = ExecutorRegistry()
+
+
+@overload
+def get_executor_class_name[T](key: str, default: T) -> str | T: ...
+@overload
+def get_executor_class_name(key: str, default: None = None) -> str | None: ...
+def get_executor_class_name[T](key: str, default: T | None = None) -> str | T | None:
+    return mod[0] if (mod := EXECUTOR_MODULES.get(key)) else default
+
+
+IMPORT_ERRORS: dict[str, str] = _IMPORT_ERRORS

src/worker/main.py

@@ -8,7 +8,7 @@
 from shared.tasks.worker_message import WorkerHardware
 
 from .config import WorkerConfig
-from .executors import EXECUTOR_CLASS_NAMES, EXECUTOR_REGISTRY, IMPORT_ERRORS
+from .executors import EXECUTOR_REGISTRY, IMPORT_ERRORS, get_executor_class_name
 from .executors.base_executor import Executor
 from .executors.mp_executor import MPExecutor
 from .hw import collect_hw
@@ -89,7 +89,7 @@ def init_executor(key: str, *, gpu_required: bool = False):
         cls = registry.get(key)
         if cls is None:
             reason = import_errors.get(
-                EXECUTOR_CLASS_NAMES.get(key, key), "dependency missing"
+                get_executor_class_name(key, key), "dependency missing"
             )
             logger.info("Skipping executor %s: %s", key, reason)
             return None

tests/worker/test_executor_registry.py

@@ -10,7 +10,7 @@
 from shared.schemas.result import BaseExecutorResult
 from shared.tasks.task_type import TaskType
 from tests.worker.factories import make_worker_config
-from worker.executors import EXECUTOR_CLASS_NAMES, EXECUTOR_REGISTRY, IMPORT_ERRORS
+from worker.executors import EXECUTOR_MODULES, EXECUTOR_REGISTRY, IMPORT_ERRORS
 from worker.executors.base_executor import Executor, ExecutorTask
 from worker.main import build_capabilities
 
@@ -41,18 +41,12 @@ def test_registry_has_expected_keys(self) -> None:
         }
         assert set(EXECUTOR_REGISTRY.keys()) == expected
 
-    def test_class_names_match_registry(self) -> None:
-        """Every key in EXECUTOR_REGISTRY has a corresponding class name."""
-        assert set(EXECUTOR_CLASS_NAMES.keys()) == set(EXECUTOR_REGISTRY.keys())
-
     def test_unavailable_executors_tracked(self) -> None:
         """Executors that failed to import have their errors recorded."""
         # GPU executors (vllm, etc.) are expected to fail without CUDA
         for key, cls in EXECUTOR_REGISTRY.items():
             if cls is None:
-                assert (
-                    key in IMPORT_ERRORS or EXECUTOR_CLASS_NAMES[key] in IMPORT_ERRORS
-                )
+                assert key in IMPORT_ERRORS or EXECUTOR_MODULES[key][0] in IMPORT_ERRORS
 
     def test_training_executors_are_wrapped_for_isolation(self) -> None:
         """Training executors run in a subprocess for GPU cleanup; ensure the