[bug] Fix oumi distributed on Slurm to use correct node rank env var (#2020)

Author: wizeng23

src/oumi/cli/distributed_run.py

@@ -422,17 +422,15 @@ def _detect_slurm_process_run_info(env: dict[str, str]) -> Optional[_ProcessRunI
     if len(node_ips) == 0:
         raise RuntimeError("Empty list of nodes in 'SLURM_NODELIST'!")
     gpus_per_node = torch.cuda.device_count()
-    node_rank = _get_optional_int_env_var("SLURM_PROCID", env)
-    if node_rank is None:
-        node_rank = _get_optional_int_env_var("PMI_RANK", env)
 
+    node_rank = _get_optional_int_env_var("SLURM_NODEID", env)
     # If running on a single node, default to 0.
     if node_rank is None and len(node_ips) == 1:
         node_rank = 0
     if node_rank is None:
         raise ValueError(
             "Unable to determine node rank on a multi-node setup. "
-            "Neither 'SLURM_PROCID' nor 'PMI_RANK' is set "
+            "'SLURM_NODEID' is not set."
         )
 
     return _ProcessRunInfo(

tests/unit/cli/test_cli_distributed_run.py

@@ -229,7 +229,7 @@ def test_torchrun_polaris_multi_gpu(
         assert logger.level == logging.DEBUG
 
 
-def test_torchrun_frontier_multi_gpu(
+def test_torchrun_slurm_multi_gpu(
     app,
     mock_os,
     mock_popen,
@@ -239,7 +239,7 @@ def test_torchrun_frontier_multi_gpu(
 ):
     test_env_vars = {
         "SLURM_NODELIST": "frontier[04316-04317]",
-        "PMI_RANK": 1,
+        "SLURM_NODEID": 1,
         "SLURM_JOBID": "123456.frontier",
         # Define the redundant OUMI_ variables to activate consistency checks.
         "OUMI_TOTAL_NUM_GPUS": 16,