Add SLURM-aware exp manager run timer

docs/source/core/exp_manager.rst

@@ -69,6 +69,24 @@ shut down before the procedure has completed. To auto-resume training, set the f
         exp_manager.version: my_experiment_version
 
 
+Wall-clock Time Limits
+----------------------
+
+Set ``max_time_per_run`` to stop training and save the last checkpoint after a wall-clock duration in
+``DD:HH:MM:SS`` format. By default, the timer starts when the training loop starts. To include preprocessing and
+other setup performed earlier in the same SLURM allocation, anchor the timer to the job start:
+
+.. code-block:: yaml
+
+    exp_manager:
+        max_time_per_run: 00:03:45:00
+        max_time_per_run_from_slurm: True
+
+When enabled, the timer reads the SLURM-provided ``SLURM_JOB_START_TIME`` UNIX timestamp and checks the elapsed
+allocation time before training starts and after each configured timer interval. A missing or invalid timestamp
+raises an error instead of silently starting a fresh timer. Leave enough time between ``max_time_per_run`` and the
+SLURM limit for the final checkpoint to finish writing.
+
 Experiment Loggers
 ------------------
 

nemo/utils/exp_manager.py

@@ -281,6 +281,8 @@ class ExpManagerConfig:
     ema: Optional[EMAParams] = field(default_factory=lambda: EMAParams())
     # Wall clock time limit
     max_time_per_run: Optional[str] = None
+    # Count from the SLURM allocation start instead of the training loop start.
+    max_time_per_run_from_slurm: Optional[bool] = False
     # time to sleep non 0 ranks during initialization
     seconds_to_sleep: float = 5
     # Straggler detection
@@ -739,13 +741,17 @@ def exp_manager(trainer: 'lightning.pytorch.Trainer', cfg: Optional[Union[DictCo
                     'Found a PTL Timer callback, replacing with a StatelessTimer callback. '
                     'This will happen if you set trainer.max_time as well as exp_manager.max_time_per_run.'
                 )
-                trainer.callbacks[idx] = StatelessTimer(cfg.max_time_per_run)
+                trainer.callbacks[idx] = StatelessTimer(
+                    cfg.max_time_per_run, max_time_from_slurm=cfg.max_time_per_run_from_slurm
+                )
                 found_ptl_timer = True
                 break
 
         if not found_ptl_timer:
             trainer.max_time = cfg.max_time_per_run
-            trainer.callbacks.append(StatelessTimer(cfg.max_time_per_run))
+            trainer.callbacks.append(
+                StatelessTimer(cfg.max_time_per_run, max_time_from_slurm=cfg.max_time_per_run_from_slurm)
+            )
 
     if cfg.create_straggler_detection_callback:
         if HAVE_STRAGGLER_DET:
@@ -1430,15 +1436,18 @@ def __init__(
         duration: timedelta = None,
         interval: str = Interval.step,
         verbose: bool = True,
+        max_time_from_slurm: bool = False,
     ) -> None:
-        """stateless timer
+        """Create a timer whose elapsed state is reset for every training run.
 
         Args:
-            duration (timedelta, optional): _description_. Defaults to None.
-            interval (str, optional): _description_. Defaults to Interval.step.
-            verbose (bool, optional): _description_. Defaults to True.
+            duration: Maximum elapsed time for this run.
+            interval: Check the time limit after each step or epoch.
+            verbose: Log when the time limit is reached.
+            max_time_from_slurm: Include time elapsed since ``SLURM_JOB_START_TIME``.
         """
         super().__init__(duration, interval, verbose)
+        self._slurm_job_start_time = self._read_slurm_job_start_time() if max_time_from_slurm else None
 
     # Override PTL Timer's state dict to not store elapsed time information so that we can
     # restore and continue training.
@@ -1450,6 +1459,16 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         """load_state_dict"""
         return
 
+    def on_fit_start(self, trainer: lightning.pytorch.Trainer, *args: Any, **kwargs: Any) -> None:
+        """Refresh the SLURM offset before the initial deadline check."""
+        self._update_slurm_time_offset()
+        super().on_fit_start(trainer, *args, **kwargs)
+
+    def on_train_start(self, trainer: lightning.pytorch.Trainer, pl_module: lightning.pytorch.LightningModule) -> None:
+        """Refresh the SLURM offset when the monotonic training clock starts."""
+        self._update_slurm_time_offset()
+        super().on_train_start(trainer, pl_module)
+
     def _check_time_remaining(self, trainer: lightning.pytorch.Trainer) -> None:
         """_check_time_remaining"""
         super()._check_time_remaining(trainer)
@@ -1491,6 +1510,27 @@ def _check_time_remaining(self, trainer: lightning.pytorch.Trainer) -> None:
 
             raise _TunerExitException()
 
+    def _update_slurm_time_offset(self) -> None:
+        """Set the elapsed-time offset to the time used by the current SLURM job."""
+        if self._slurm_job_start_time is not None:
+            self._offset = max(0.0, time.time() - self._slurm_job_start_time)
+
+    @staticmethod
+    def _read_slurm_job_start_time() -> float:
+        """Read and validate the UNIX timestamp exported by SLURM."""
+        value = os.getenv("SLURM_JOB_START_TIME")
+        try:
+            start_time = int(value)
+        except (TypeError, ValueError):
+            raise ValueError(
+                "SLURM-based max_time_per_run requires SLURM_JOB_START_TIME to be a positive UNIX timestamp"
+            ) from None
+        if start_time <= 0:
+            raise ValueError(
+                "SLURM-based max_time_per_run requires SLURM_JOB_START_TIME to be a positive UNIX timestamp"
+            )
+        return float(start_time)
+
 
 def _describe_batch_progress(trainer: lightning.pytorch.Trainer) -> Dict[str, Any]:
     """Return a compact, log-friendly snapshot of Lightning's train batch progress."""

tests/core_ptl/test_ptl_stateless_timer.py

@@ -145,6 +145,100 @@ def test_stateless_timer(self):
         self.cleanup()
 
 
+class TestStatelessTimerSlurmStartTime:
+    @pytest.mark.unit
+    def test_counts_time_since_slurm_job_start(self, monkeypatch):
+        monkeypatch.setenv("SLURM_JOB_ID", "1234")
+        monkeypatch.setenv("SLURM_JOB_START_TIME", "100")
+        clock = {"wall": 1000.0, "monotonic": 50.0}
+        monkeypatch.setattr("nemo.utils.exp_manager.time.time", lambda: clock["wall"])
+        monkeypatch.setattr("nemo.utils.exp_manager.time.monotonic", lambda: clock["monotonic"])
+
+        timer = StatelessTimer(duration="00:00:20:00", max_time_from_slurm=True)
+        timer.on_train_start(None, None)
+        clock["monotonic"] = 55.0
+
+        assert timer.time_elapsed() == pytest.approx(905.0)
+
+    @pytest.mark.unit
+    def test_refreshes_slurm_elapsed_time_before_fit(self, monkeypatch):
+        monkeypatch.setenv("SLURM_JOB_START_TIME", "100")
+        clock = {"wall": 1000.0, "monotonic": 50.0}
+        monkeypatch.setattr("nemo.utils.exp_manager.time.time", lambda: clock["wall"])
+        monkeypatch.setattr("nemo.utils.exp_manager.time.monotonic", lambda: clock["monotonic"])
+        timer = StatelessTimer(duration="00:00:20:00", max_time_from_slurm=True)
+        trainer = MagicMock()
+        trainer.should_stop = False
+        trainer.strategy.broadcast.side_effect = lambda value: value
+
+        clock["wall"] = 1100.0
+        timer.on_fit_start(trainer)
+
+        assert timer.time_elapsed() == pytest.approx(1000.0)
+
+    @pytest.mark.unit
+    def test_expired_slurm_budget_saves_before_training(self, monkeypatch):
+        monkeypatch.setenv("SLURM_JOB_START_TIME", "100")
+        monkeypatch.setattr("nemo.utils.exp_manager.time.time", lambda: 1000.0)
+        checkpoint_callback = MagicMock()
+        checkpoint_callback._monitor_candidates.return_value = {}
+        trainer = SimpleNamespace(
+            strategy=SimpleNamespace(broadcast=lambda value: value),
+            should_stop=False,
+            checkpoint_callback=checkpoint_callback,
+            global_step=0,
+            current_epoch=0,
+        )
+        timer = StatelessTimer(duration="00:00:10:00", max_time_from_slurm=True)
+
+        with pytest.raises(_TunerExitException):
+            timer.on_fit_start(trainer)
+
+        checkpoint_callback._save_last_checkpoint.assert_called_once_with(trainer, {})
+
+    @pytest.mark.unit
+    def test_slurm_timing_is_opt_in(self, monkeypatch):
+        monkeypatch.setenv("SLURM_JOB_START_TIME", "not-a-timestamp")
+        clock = {"monotonic": 50.0}
+        monkeypatch.setattr("nemo.utils.exp_manager.time.monotonic", lambda: clock["monotonic"])
+
+        timer = StatelessTimer(duration="00:00:20:00")
+        timer.on_train_start(None, None)
+        clock["monotonic"] = 55.0
+
+        assert timer.time_elapsed() == pytest.approx(5.0)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("value", [None, "not-a-timestamp", "0", "-1"])
+    def test_rejects_missing_or_invalid_slurm_start_time(self, monkeypatch, value):
+        if value is None:
+            monkeypatch.delenv("SLURM_JOB_START_TIME", raising=False)
+        else:
+            monkeypatch.setenv("SLURM_JOB_START_TIME", value)
+
+        with pytest.raises(ValueError, match="SLURM_JOB_START_TIME"):
+            StatelessTimer(duration="00:00:20:00", max_time_from_slurm=True)
+
+    @pytest.mark.unit
+    def test_exp_manager_enables_slurm_timing(self, monkeypatch, tmp_path):
+        monkeypatch.setenv("SLURM_JOB_START_TIME", "100")
+        trainer = Trainer(accelerator="cpu", logger=False, enable_checkpointing=False)
+        cfg = ExpManagerConfig(
+            explicit_log_dir=str(tmp_path),
+            create_tensorboard_logger=False,
+            create_checkpoint_callback=False,
+            log_step_timing=False,
+            disable_validation_on_resume=False,
+            max_time_per_run="00:00:20:00",
+            max_time_per_run_from_slurm=True,
+        )
+
+        exp_manager(trainer, cfg=OmegaConf.structured(cfg))
+
+        timer = next(callback for callback in trainer.callbacks if isinstance(callback, StatelessTimer))
+        assert timer._slurm_job_start_time == 100.0
+
+
 def _make_trainer_with_batch_progress(batch_progress: _BatchProgress) -> MagicMock:
     trainer = MagicMock()
     trainer.fit_loop.epoch_loop.batch_progress = batch_progress