fix(slurm): detect srun inside interactive salloc via SLURM_STEP_ID

src/lightning/fabric/CHANGELOG.md

@@ -20,7 +20,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
--
+- Fixed `SLURMEnvironment` raising a spurious "srun is available but not used" warning when running `srun python ...` inside an interactive `salloc` allocation ([#20776](https://github.com/Lightning-AI/pytorch-lightning/issues/20776))
 
 --
 

src/lightning/fabric/plugins/environments/slurm.py

@@ -225,7 +225,14 @@ def _validate_srun_variables() -> None:
 
 
 def _is_srun_used() -> bool:
-    return "SLURM_NTASKS" in os.environ and not _is_slurm_interactive_mode()
+    if "SLURM_NTASKS" not in os.environ:
+        return False
+    # `srun` sets `SLURM_STEP_ID` in the launched process, while `salloc` alone does not. This catches
+    # the case of running `srun python ...` inside an interactive `salloc` allocation, where
+    # `SLURM_JOB_NAME` inherited from the allocation is ``bash`` or ``interactive``.
+    if "SLURM_STEP_ID" in os.environ:
+        return True
+    return not _is_slurm_interactive_mode()
 
 
 def _is_slurm_interactive_mode() -> bool:

tests/tests_fabric/plugins/environments/test_slurm.py

@@ -129,6 +129,15 @@ def test_detect():
     with mock.patch.dict(os.environ, {"SLURM_JOB_NAME": "interactive"}):
         assert not SLURMEnvironment.detect()
 
+    # `srun` launched from inside an `salloc` interactive allocation sets
+    # SLURM_STEP_ID, which indicates srun is actively driving the process.
+    with mock.patch.dict(
+        os.environ,
+        {"SLURM_NTASKS": "2", "SLURM_JOB_NAME": "interactive", "SLURM_STEP_ID": "0"},
+        clear=True,
+    ):
+        assert SLURMEnvironment.detect()
+
 
 @RunIf(skip_windows=True)
 @pytest.mark.skipif(shutil.which("srun") is not None, reason="must run on a machine where srun is not available")
@@ -151,6 +160,36 @@ def test_srun_available_and_not_used(monkeypatch):
         assert not SLURMEnvironment.detect()
 
 
+@RunIf(skip_windows=True)
+@pytest.mark.skipif(shutil.which("srun") is not None, reason="must run on a machine where srun is not available")
+def test_no_srun_warning_inside_salloc_with_srun(monkeypatch):
+    """No `srun` warning when running ``srun python ...`` inside an interactive ``salloc`` allocation.
+
+    Inside an interactive allocation (``SLURM_JOB_NAME`` is ``bash`` or ``interactive``), a nested ``srun`` still
+    sets ``SLURM_STEP_ID``. This distinguishes it from a plain ``python`` invocation in the same shell.
+
+    """
+    monkeypatch.setattr(sys, "argv", ["train.py"])
+    expected = "`srun` .* available .* but is not used"
+
+    with (
+        mock.patch("lightning.fabric.plugins.environments.slurm.shutil.which", return_value="/usr/bin/srun"),
+        mock.patch.dict(
+            os.environ,
+            {
+                "SLURM_NTASKS": "2",
+                "SLURM_NTASKS_PER_NODE": "1",
+                "SLURM_JOB_NAME": "interactive",
+                "SLURM_STEP_ID": "0",
+            },
+            clear=True,
+        ),
+        no_warning_call(PossibleUserWarning, match=expected),
+    ):
+        SLURMEnvironment()
+        assert SLURMEnvironment.detect()
+
+
 def test_srun_variable_validation():
     """Test that we raise useful errors when `srun` variables are misconfigured."""
     with mock.patch.dict(os.environ, {"SLURM_NTASKS": "1"}):