fix(ray): drop Slurm/PMIx/MPI launcher env in Ray workers so TRT-LLM can import

`import tensorrt_llm` runs MPI_Init. Under srun (ray.sub starts every step with
--mpi=pmix) Open MPI aborts on an inherited SLURM_STEP_ID ("OMPI was not built
with SLURM's PMI support"), and residual OMPI_/MPI_/SLURM_ vars send TRT-LLM to
MPI_Comm_Spawn (MPI_ERR_SPAWN). This breaks every multi-node TRT-LLM run on
such a cluster, including main's own grpo-qwen3-1.7b-2n4g-automodel-trtllm.

A runtime_env cannot delete a variable and TRT-LLM's GPU workers inherit the
raylet's environment, so remove the launcher state inside each worker with a
job-level worker_process_setup_hook. The driver and ray.sub keep their
environment; the OMPI/PRTE root-allow vars set by the TRT-LLM worker are kept.

Verified: grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated passes all
checks on 4 nodes with an unmodified ray.sub.

Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>

Author: hchings

nemo_rl/distributed/virtual_cluster.py

@@ -436,6 +436,36 @@ def init_ray(log_dir: Optional[str] = None) -> None:
         safe_set_span_attributes(span, {RAY_CLUSTER_SOURCE_ATTR: outcome})
 
 
+# Config set by the TRT-LLM worker for Open MPI / PRTE, not launcher state.
+_MPI_ROOT_ALLOW_VARS = frozenset(
+    {
+        "OMPI_ALLOW_RUN_AS_ROOT",
+        "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM",
+        "PRTE_ALLOW_RUN_AS_ROOT",
+        "PRTE_ALLOW_RUN_AS_ROOT_CONFIRM",
+    }
+)
+
+
+def _drop_mpi_launcher_env() -> None:
+    """Ray ``worker_process_setup_hook``: strip Slurm/PMIx/MPI launcher state.
+
+    ``import tensorrt_llm`` runs ``MPI_Init``. Under ``srun`` (``ray.sub`` starts
+    every step with ``--mpi=pmix``) the inherited ``SLURM_STEP_ID``/``PMIX_*``
+    make Open MPI believe it was direct-launched and abort with "OMPI was not
+    built with SLURM's PMI support"; residual ``OMPI_*``/``MPI_*``/``SLURM_*``
+    also push TRT-LLM to ``MPI_Comm_Spawn`` (``MPI_ERR_SPAWN``). Neither an empty
+    value nor an ``env_vars`` override helps, and TRT-LLM's GPU workers are Ray
+    actors whose environment is the raylet's plus a copy, so the variables must
+    be removed inside each worker process before it imports anything.
+    """
+    for key in list(os.environ):
+        if key in _MPI_ROOT_ALLOW_VARS:
+            continue
+        if key.startswith(("PMIX_", "PMI_", "MPI_", "OMPI_", "SLURM_")):
+            os.environ.pop(key, None)
+
+
 def _init_ray(log_dir: Optional[str] = None) -> str:
     """Initialise Ray.
 
@@ -466,6 +496,7 @@ def _init_ray(log_dir: Optional[str] = None) -> str:
 
     runtime_env = {
         "env_vars": env_vars,  # Pass thru all user environment variables
+        "worker_process_setup_hook": _drop_mpi_launcher_env,
     }
 
     cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "ALL")

tests/unit/distributed/test_virtual_cluster.py

@@ -991,3 +991,34 @@ def test_actor_registry_honors_system_flag(use_system_executable):
             "nemo_rl.models.policy.workers.automodel_policy_worker.AutomodelPolicyWorker"
         ].startswith("uv run")
         assert envs["nemo_rl.environments.nemo_gym.NemoGym"].startswith("uv run")
+
+
+def test_init_ray_drops_mpi_launcher_env_in_every_worker():
+    """TRT-LLM's MPI init aborts under srun's launcher env, so workers must not carry it."""
+    from nemo_rl.distributed.virtual_cluster import init_ray
+
+    with (
+        patch("ray.init") as mock_ray_init,
+        patch("ray.cluster_resources") as mock_cluster_resources,
+    ):
+        mock_cluster_resources.return_value = {"nrl_tag_0": 1}
+        with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}, clear=True):
+            init_ray()
+
+        hook = mock_ray_init.call_args_list[0][1]["runtime_env"][
+            "worker_process_setup_hook"
+        ]
+        launcher = {
+            "SLURM_STEP_ID": "14",
+            "SLURM_JOB_ID": "7",
+            "PMIX_NAMESPACE": "x",
+            "OMPI_MCA_plm": "slurm",
+            "MPI_LOCALRANKID": "0",
+        }
+        root_allow = {"OMPI_ALLOW_RUN_AS_ROOT": "1", "PRTE_ALLOW_RUN_AS_ROOT": "1"}
+        unrelated = {"HOME": "/root"}
+        with patch.dict(os.environ, {**launcher, **root_allow, **unrelated}, clear=True):
+            hook()
+            assert not set(launcher) & set(os.environ)
+            assert all(os.environ[k] == v for k, v in {**root_allow, **unrelated}.items())
+            hook()  # a worker that never had any is fine