fix(training): fail fast on malformed launcher rank (#5766)

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/utils/common_utils.py

@@ -22,7 +22,7 @@
 import torch
 import torch.distributed
 from megatron.core import DistributedDataParallel as DDP
-from megatron.core._rank_utils import safe_get_rank as _get_rank_safe
+from megatron.core._rank_utils import safe_get_rank as get_rank_safe  # noqa: F401
 from megatron.core._rank_utils import safe_get_world_size as get_world_size_safe  # noqa: F401
 from megatron.core.transformer.module import Float16Module
 from megatron.core.utils import get_batch_on_this_cp_rank
@@ -43,14 +43,6 @@
     ALL_MODULE_WRAPPER_CLASSNAMES = (DDP, Float16Module)
 
 
-def get_rank_safe() -> int:
-    """Get the distributed rank, falling back to zero for malformed launcher state."""
-    try:
-        return _get_rank_safe()
-    except (TypeError, ValueError):
-        return 0
-
-
 def get_last_rank() -> int:
     """Get the last rank in the distributed group"""
     if not torch.distributed.is_initialized():

tests/unit_tests/utils/test_common_utils.py

@@ -73,12 +73,22 @@ def test_uninitialized_torch_distributed_no_env_var(self, mock_is_initialized):
         mock_is_initialized.assert_called_once()
 
     @patch("torch.distributed.is_initialized")
-    @patch.dict(os.environ, {"RANK": "invalid"})
+    @patch.dict(os.environ, {"RANK": "invalid"}, clear=True)
     def test_invalid_rank_env_var(self, mock_is_initialized):
-        """Test get_rank_safe falls back to 0 when RANK is malformed."""
+        """Test get_rank_safe propagates an actionable error when RANK is malformed."""
         mock_is_initialized.return_value = False
 
-        assert get_rank_safe() == 0
+        with pytest.raises(ValueError, match="invalid literal for int.*invalid"):
+            get_rank_safe()
+
+    @patch("torch.distributed.is_initialized")
+    @patch.dict(os.environ, {"SLURM_NTASKS": "8", "SLURM_PROCID": "invalid"}, clear=True)
+    def test_invalid_slurm_rank_env_var(self, mock_is_initialized):
+        """Test get_rank_safe propagates an actionable error when SLURM_PROCID is malformed."""
+        mock_is_initialized.return_value = False
+
+        with pytest.raises(ValueError, match="invalid literal for int.*invalid"):
+            get_rank_safe()
 
 
 class TestGetWorldSizeSafe: