nemo_rl.distributed.actor_environments#

Which uv extras each Ray actor needs. The single source of truth.

Two readers:

  • nemo_rl.distributed.ray_actor_environment_registry imports ACTOR_ENVIRONMENTS and turns it into py_executable strings at runtime.

  • docker/Dockerfile runs this file as a script to list the venvs it must pre-build, so the image ships one venv per actor.

DO NOT IMPORT ANYTHING FROM nemo_rl HERE, AND KEEP IT STDLIB-ONLY. The Dockerfile runs this from the dependency layer, where pyproject.toml, uv.lock and a handful of standalone files exist, but the rest of the nemo_rl package does not (see the COPY directives in docker/Dockerfile). Running it as a script rather than importing it is what keeps nemo_rl/__init__.py – which is copied in, and does real work at import – from executing there. An import added here breaks the image build in its most expensive layer. tests/unit/distributed/test_actor_environments.py enforces this.

Module Contents#

Functions#

_build_stage

Which image layer can finish this venv.

main

Print “\t\t” for uv-managed actors.

Data#

ACTOR_ENVIRONMENTS

API#

nemo_rl.distributed.actor_environments.ACTOR_ENVIRONMENTS: dict[str, list[str] | None]#

None

nemo_rl.distributed.actor_environments._build_stage(extras: list[str]) → str[source]#

Which image layer can finish this venv.

The tensorrt_llm wheel is only built in the TRT-LLM layer, so those venvs get their third-party packages in two steps; everything else finishes in the dependency layer.

nemo_rl.distributed.actor_environments.main(argv: list[str]) → int[source]#

Print “\t\t” for uv-managed actors.

Usage: actor_environments.py [] [ …]

<stage> is “deps” or “trtllm” (omit for all). Any extras listed after it are skipped, which is how the Dockerfile honors SKIP_VLLM_BUILD and friends. Filtering on the declared extras – rather than on a substring of the actor name – is what makes SKIP_VLLM_BUILD also skip actors like AsyncTrajectoryCollector, whose name contains no “vllm”.