Skip to content

fix: anchor checkpoint TimeoutChecker at SLURM allocation start - #3639

Closed
bobby-nandigam wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
bobby-nandigam:fix/timeout-checker-slurm-anchor
Closed

fix: anchor checkpoint TimeoutChecker at SLURM allocation start#3639
bobby-nandigam wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
bobby-nandigam:fix/timeout-checker-slurm-anchor

Conversation

@bobby-nandigam

Copy link
Copy Markdown

What does this PR do ?

checkpoint_must_save_by is meant to guarantee that a checkpoint is written
before the allocation is reclaimed. But TimeoutChecker anchors its budget
at construction time:

self.start_time = time.time()

Every algorithm constructs the checker after setup() returns — Ray init,
cluster spinup, and model/vLLM loading have already happened. The budget
therefore silently excludes 20–30 minutes of setup: 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 enough
setup, the timeout save fires after the allocation is already dead — losing
exactly the work the knob exists to protect.

The fix

Anchor the budget at the allocation start when SLURM provides it, through the
single TimeoutChecker constructor — so all algorithms are corrected with no
call-site changes:

  • Reads SLURM_JOB_START_TIME (exported by Slurm ≥ 23.02) via a small
    get_slurm_job_start_time() helper.
  • Unset / empty / malformed (Kubernetes, local runs, older Slurm) → falls back
    to today's construction-time behavior, a clean no-op. A malformed value warns
    rather than crashing.
  • An explicit start_time argument still wins, for tests and callers that want
    to pin it.
  • A construction-time log records the anchor source and the effective remaining
    budget, so the semantic is visible in every run, e.g.
    TimeoutChecker anchored at SLURM_JOB_START_TIME; effective remaining budget 3:17:00.

The new default's failure mode (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.

Issues

Closes #3596.

Usage

No config change and no new knob. checkpoint_must_save_by now budgets from the
SLURM allocation start automatically; non-SLURM environments are unaffected.

Before your PR is "Ready for review"

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests? (four cases in tests/unit/utils/test_timer.py)
  • Did you run the unit tests locally? (unit only)
  • Did you add or update any necessary documentation? (n/a — behavior fix, no user-facing surface change)

Additional Information

  • SLURM_JOB_START_TIME was added in Slurm 23.02 and is epoch seconds; older
    Slurm simply doesn't set it and hits the None fallback.
  • Surfaced during review of feat(sc): add checkpoint save/restore to SingleController #3429 (whose TimeoutChecker placement is correct
    legacy parity and needs no change — this is a pre-existing, cross-cutting
    semantic).

checkpoint_must_save_by exists to guarantee a checkpoint is written before the
allocation is reclaimed. TimeoutChecker, however, anchored its budget at
construction time, which every algorithm reaches only after setup (Ray init,
cluster spinup, model/vLLM load). The budget therefore silently excluded the
20-30 minutes of setup, so the timeout save could fire after the allocation was
already killed -- losing exactly the work the knob is meant to protect.

Anchor the budget at SLURM_JOB_START_TIME (exported by Slurm >= 23.02) when it
is available, through the single TimeoutChecker constructor, so every algorithm
is corrected with no call-site changes. When the variable is unset or malformed
(Kubernetes, local runs, older Slurm) the checker falls back to the previous
construction-time behavior, and a malformed value warns instead of crashing. A
construction-time log records the anchor source and the effective remaining
budget so the change is visible in every run.

Closes #3596

Signed-off-by: Bobby Nandigam <bobbynandigam.official@gmail.com>
@bobby-nandigam
bobby-nandigam requested review from a team as code owners August 14, 2026 04:45
@copy-pr-bot

copy-pr-bot Bot commented Aug 14, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Aug 16, 2026
@bobby-nandigam bobby-nandigam closed this by deleting the head repository Sep 5, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci removed the waiting-on-maintainers Waiting on maintainers to respond label Sep 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Anchor TimeoutChecker at SLURM_JOB_START_TIME so checkpoint_must_save_by budgets from allocation start

2 participants