perf: import worker executors lazily via a registry

[kaiitunnz]

Purpose

Importing worker.executors eagerly imported every executor module at package load, so any process touching the package paid the full executor dependency tree (torch, transformers, vLLM, trl, diffusers) even when it needed none of them. This makes executor imports lazy, which speeds up two paths that don't run a full executor set: the supervisor's per-node hardware probe (worker.main --collect-hw, which never initializes executors) and MPExecutor's spawned subprocess (which re-runs the package import to run a single executor).

Changes

• Worker — replaces the eager module-level imports and the parallel EXECUTOR_REGISTRY / EXECUTOR_CLASS_NAMES dicts with an ExecutorRegistry Mapping that imports each executor on first access and caches it. Class-name lookups move to a module-level get_executor_class_name.
• Tests — registry tests updated for the lazy Mapping contract (lazy resolution, KeyError on unknown keys, class-name lookup).

Design

ExecutorRegistry is a Mapping[str, type[Executor] | None]: __getitem__ imports on demand and caches, raising KeyError for unknown keys (returning None only for a known executor whose import failed). __iter__ / __len__ cover the registered keys without importing them, so capability enumeration stays cheap. Worker startup is unchanged in total cost — initialize_executors still resolves every key — but the import work is deferred from package load to initialization, and processes that never initialize executors skip it entirely.

Performance

Median startup cost, eager baseline vs lazy, on freshly built CPU and GPU worker images:

Path	Image	Eager	Lazy	Improvement
import worker.executors	CPU	14.4 s	0.5 s	−96%
import worker.executors	GPU	21.1 s	0.4 s	−98%
--collect-hw hardware probe	CPU	15.8 s	0.8 s	−95%
--collect-hw hardware probe	GPU	23.1 s	0.9 s	−96%
MPExecutor spawn (light executor)	GPU	21.2 s	3.2 s	−85%
MPExecutor spawn (heavy executor)	GPU	20.9 s	14.1 s	−33%

MPExecutor uses spawn, so each subprocess re-runs the package import; it now loads only the one executor it runs instead of all of them. The win scales inversely with how heavy that executor's own dependencies are — large for light executors, modest for vLLM/training which need most of the stack regardless. Full worker time-to-ready is unchanged, since a fully-initialized worker still imports every executor.

The test suite benefits from the same effect, since collection imports the package without initializing executors: uv run pytest tests/ --ignore=tests/worker/test_mp_executor_cleanup_gpu.py drops from 94.50 s to 28.81 s (~3.3×).

Test Plan

uv run pytest tests/ --ignore=tests/worker/test_mp_executor_cleanup_gpu.py
uv run pre-commit run --all-files

Validated end-to-end on freshly built CPU and GPU images. On a CPU worker: importing the package inside the image loads no executor module or heavy dependency; the worker advertises the expected capability set (supported CPU types present, GPU-only types absent) and runs an echo task to completion; each GPU-gated executor is resolved and logged as skipped with a reason. On a GPU worker: the GPU-only executors are lazily imported and advertised, vLLM initializes without being skipped, and a TinyLlama vLLM inference task runs to completion. The startup-cost figures above were measured on the same images, comparing the eager baseline image against this build.

Test Result

$ uv run pytest tests/ --ignore=tests/worker/test_mp_executor_cleanup_gpu.py   # All passed
$ uv run pre-commit run --all-files          # isort / black / ruff / mypy / codespell passed

End-to-end on CPU + GPU images: 31/31 assertions passed. The CPU worker advertised the expected supported task types and ran the echo task to DONE while logging each GPU-gated executor as skipped; the GPU worker advertised and lazily imported the GPU executor set, initialized vLLM without skipping it, and ran the vLLM inference task to DONE.

---

Pre-submission Checklist

• I have read the contribution guidelines.
• I have run pre-commit on the changed files and fixed any issues.
• I have added or updated tests covering my changes (if applicable).
• I have verified that the affected server/SDK/CLI/shared tests pass locally.
• If I changed shared schemas or proto definitions, I have checked downstream compatibility across Server and Worker.
• If I changed the SDK or CLI, I have verified the affected packages work.
• If this is a breaking change, I have prefixed the PR title with [BREAKING] and described migration steps above.
• I have updated documentation or config examples if user-facing behavior changed.

[timzsu]

One small comment to discuss.

[timzsu]

LGTM.