NVIDIA-NeMo / NVIDIA-NeMo/RL

Anchor TimeoutChecker at SLURM_JOB_START_TIME so checkpoint_must_save_by budgets from allocation start

Open
#3,596 0 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary

`checkpoint_must_save_by` is meant to guarantee "a checkpoint exists before the allocation dies" — but the `TimeoutChecker` that enforces it anchors its budget at **construction time**, which every algorithm reaches only *after* setup (Ray init, cluster spinup, model/vLLM loading). The budget therefore silently excludes setup time: a user who requests a 4h allocation and sets `checkpoint_must_save_by: "00:03:30:00"` actually gets a save at `training_start + 3h30` = `allocation + 3h30 + setup`. With a 20-30 min setup, the timeout save can fire after the allocation is killed — losing exactly the work the knob exists to protect.

Raised as an FYI during review of #3429 (which placed the SC path's `TimeoutChecker` in correct parity with legacy — this is a pre-existing, cross-cutting semantic, not a defect of that PR).

## Mechanism

- `TimeoutChecker.__init__` sets `self.start_time = time.time()` at construction ([timer.py#L433](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/utils/timer.py#L433)); `check_save` compares elapsed-since-construction against the budget ([#L448-L461](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/utils/timer.py#L448-L461)).
- Every path constructs it post-setup: `grpo_train` ([grpo.py#L2615](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/algorithms/grpo.py#L2615)), `async_grpo_train` ([#L3943](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/algorithms/grpo.py#L3943)), `grpo_sync` (grpo_sync.py:408), and the SC actor's `__init__` (#3429) — all reached after `setup()` returns in their run scripts.
- The default `"00:03:45:00"` ([timer.py#L421](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/utils/timer.py#L421)) is clearly sized against a 4h allocation, which invites exactly the allocation-anchored misreading.

## Proposed fix

Anchor at the allocation start when SLURM provides it; otherwise keep today's behavior. Single choke point in `nemo_rl/utils/timer.py`, so all eight algorithms are fixed with zero call-site changes:

```python
def get_slurm_job_start_time() -> Optional[float]:
"""Epoch seconds of the SLURM allocation start, or None outside SLURM.

Malformed/empty values warn and return None — never crash.
"""
raw = os.environ.get("SLURM_JOB_START_TIME")
if not raw:
return None
try:
return float(raw)
except ValueError:
warnings.warn(f"Could not parse SLURM_JOB_START_TIME={raw!r}; "
"anchoring the checkpoint timeout at construction time instead.")
return None

class TimeoutChecker:
def __init__(self, timeout=..., fit_last_save_time=False,
start_time: Optional[float] = None):
...
self.start_time = (start_time
or get_slurm_job_start_time()
or time.time())
```

Plus a loud construction-time log stating the anchor source and effective remaining budget, e.g. `TimeoutChecker anchored at SLURM_JOB_START_TIME (12m34s ago); effective remaining budget 3h17m` — so the semantic shift is visible in every log.

Design points settled during review discussion:

- **Default behavior, not opt-in.** The current behavior is the footgun; an undiscovered opt-in flag fixes nobody. The failure mode of the new default (save slightly earlier than a training-start-anchored user expected) is benign; the old default's failure mode (save after the allocation dies) loses work. Needs a release-note entry.
- **SLURM-only, by name.** `SLURM_JOB_START_TIME` unset (k8s, local, older Slurm) → `None` → today's behavior, a clean no-op. No generic scheduler-agnostic override — k8s doesn't have this problem shape.
- **Function, not import-time constant**, for per-test monkeypatching (implementation freedom).

## Verified facts

- `SLURM_JOB_START_TIME` was added in **Slurm 23.02**: ["Add SLURM_JOB_START_TIME and SLURM_JOB_END_TIME environment variables."](https://github.com/SchedMD/slurm/blob/df46edece3b84b355b4773932b688b444bb5db94/CHANGELOG/slurm-23.02.md) — older Slurm simply doesn't set it, hitting the `None` fallback.
- Format is epoch seconds: slurmd exports it as `%lu` of `env->job_start_time` ([src/common/env.c#L684-L689](https://github.com/SchedMD/slurm/blob/df46edece3b84b355b4773932b688b444bb5db94/src/common/env.c#L684-L689)).
- Propagation note for the SC path: a Ray actor's environment comes from the `ray start` process on its node, not from the submitting shell — this works here because NeMo-RL's `ray.sub` starts Ray *inside* the SLURM allocation, so `SLURM_*` variables are present in the actors' inherited environment on every node of the job.

Related: #3429 (surfaced during review; its `TimeoutChecker` placement is correct legacy parity and needs no change).

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.