nemo_rl.distributed.ray_worker_setup

Per-process setup for the Ray workers init_ray launches.

init_ray installs :func:ray_worker_process_setup as the job’s worker_process_setup_hook by module path, so Ray imports this module in every worker process of the job – every actor venv, including the workers vLLM and TRT-LLM start themselves – before any actor code runs. Be careful with imports here: keep to the stdlib. Anything heavier is paid by every worker at startup, and one that a worker’s venv lacks makes that worker fail to start.

Module Contents

Functions

ray_worker_process_setup	The worker_process_setup_hook init_ray installs.
_drop_srun_launcher_env_for_trtllm	Drop the srun launcher env in TRT-LLM workers only.

API

nemo_rl.distributed.ray_worker_setup.ray_worker_process_setup() → None

The worker_process_setup_hook init_ray installs.

nemo_rl.distributed.ray_worker_setup._drop_srun_launcher_env_for_trtllm() → None

Drop the srun launcher env in TRT-LLM workers only.

import tensorrt_llm runs MPI_Init, which the srun launcher env breaks: PMIX_* makes slurm cancel the whole Ray step when the worker exits, and SLURM_* makes a TP>1 engine fail to spawn its MPI workers. env_vars cannot unset a variable, so drop them in-process. TRT-LLM workers set up no telemetry, so nothing is lost.