Preemption recovery and resume

This note documents the current preemption-handling and resumption design for synthetic-sequences training, with emphasis on the entangled trainer. The goal is to explain the code architecture, the operational contract between trainer and Slurm wrapper, and the design choices that future maintainers should preserve unless they are making a deliberate migration.

Scope

The primary implementation described here is the entangled recovery path in path:experiments/synthetic_sequences/entangled/resumption.py and path:experiments/synthetic_sequences/entangled/train_entangled_model.py.

The disentangled trainer already had most of the software-side recovery primitives. The related disentangled change in this work was mainly to tighten the Slurm wrapper so it launches via srun and uses the same signal-forwarding contract as the entangled wrappers.

Architecture map

• Shared preemption primitives live in path:core/slurm/preemption.py. This is where the trainer-side signal handler and the resumable exit code contract live.
• Shared auto-resume metadata helpers live in path:core/slurm/auto_resume.py. This module owns config_identity.json, run_status.json, and the config-hash helpers used to decide whether an existing run directory belongs to the current setting.
• Entangled-specific recovery logic lives in path:experiments/synthetic_sequences/entangled/resumption.py. That module owns run discovery, legacy-run bootstrap, resume-position arithmetic, shutdown checkpoint writing, and W&B run-id persistence.
• path:experiments/synthetic_sequences/entangled/train_entangled_model.py integrates the recovery flow into startup, checkpoint restore, loop re-entry, and graceful shutdown polling.
• path:experiments/synthetic_sequences/data/baseline_loader.py accepts an optional torch.Generator so shuffled training order can be rebuilt deterministically after restart.
• path:experiments/synthetic_sequences/run_naming.py excludes recovery knobs such as enable_graceful_shutdown, auto_resume, and max_resubmit_attempts from run identity so resubmits map back to the original deterministic run directory.
• The scheduler-facing contract lives in the Slurm templates: path:experiments/synthetic_sequences/slurm/templates/run_entangled_pretrained.job.j2, path:experiments/synthetic_sequences/slurm/templates/run_entangled_task_size.job.j2, and path:experiments/synthetic_sequences/slurm/templates/run_disentangled_beta.job.j2.

End-to-end recovery flow

1. A fresh launch computes the deterministic run directory from the sanitized run config, collection slug, and seed.
2. If --auto-resume is enabled and the user did not provide --resume-from, the entangled trainer checks whether that deterministic run directory already exists and whether it contains checkpoints/last.pt.
3. If the run directory has current metadata, the trainer compares the saved config hash to the current sanitized config hash. If the hashes differ, auto-resume refuses to attach to that run.
4. If the run directory is a legacy run that has last.pt but no config_identity.json, the trainer recomputes the hash from the checkpoint's saved config_snapshot. If that matches, it bootstraps the missing metadata files and resumes the legacy run in place.
5. Model, optimizer, scheduler, and RNG restoration still go through the existing resume_if_requested(...) path. Recovery did not introduce a second state-restore implementation.
6. Resume position is computed separately from checkpoint metadata as (start_epoch, resume_batch_idx). The meaning of resume_batch_idx is "how many effective batches in start_epoch were already fully committed before shutdown."
7. The training loop rebuilds the correct epoch-local dataloader and skips batch_idx <= resume_batch_idx on the first resumed epoch. This is the core invariant that prevents replaying completed work.
8. When a shutdown signal is observed, the trainer saves an immediate step checkpoint, updates last.pt, appends a graceful_shutdown event, updates run_status.json to preempted, increments resubmit_count, finishes the tracker, and exits with code 42.
9. The Slurm wrapper owns the scheduler-side behavior: srun creates the job step that receives signals, the batch shell forwards SIGUSR1 timeout warnings to the child, the batch shell ignores TERM, and #SBATCH --requeue makes preempted jobs eligible for scheduler-native requeue.

Why the design looks like this

Recovery knobs are not part of run identity

Recovery settings such as enable_graceful_shutdown, auto_resume, and max_resubmit_attempts are operational knobs, not experimental setting knobs. They are therefore excluded from run identity in path:experiments/synthetic_sequences/run_naming.py. If they were part of the identity hash, turning on recovery for an existing experiment would create a different run directory and defeat auto-resume.

Entangled resume is defined at effective-batch granularity

The important unit is the effective optimizer batch, not the micro-batch. The checkpoint already contains the optimizer and scheduler state after all committed effective batches. Resuming by replaying any completed effective batch would advance those states twice. This is why the trainer computes resume_batch_idx and explicitly skips already-finished effective batches on the first resumed epoch.

Explicit epoch-local shuffle seeds are easier to reason about than sampler internals

New entangled runs seed the shuffled training dataloader with seed + epoch. That makes the epoch-local record order reconstructible from small, explicit metadata rather than from opaque sampler state. This choice is simpler to inspect and easier to recover correctly after process death.

Legacy entangled runs are supported, but only as far as their old invariants allow

Older entangled checkpoints were created before explicit epoch-local shuffle seeds and before the current metadata files existed. The current code supports them by:

• validating identity from the checkpoint config hash when config_identity.json is missing;
• bootstrapping config_identity.json and run_status.json on first resumed launch;
• treating missing train-order metadata as legacy_random_sampler;
• reconstructing the historical epoch-1 record ordering when resuming a legacy single-long-epoch run.

This is intentionally a compatibility bridge, not a timeless second design. The exactness guarantee is strongest for the current single-long-epoch experiments that motivated the change.

The batch shell must not own the preemption TERM

On this cluster, preemption signals are delivered to job steps started with srun, while timeout warnings configured with #SBATCH --signal=B:SIGUSR1@... are delivered to the batch shell. The templates therefore use two different mechanisms:

• trap '' TERM in the batch shell so the srun step and Python signal handler own the preemption path;
• an explicit USR1 trap in the batch shell so timeout warnings can be forwarded to the child process.

Without srun, or if the batch shell handled TERM itself, the trainer would not reliably receive the signal that is supposed to trigger checkpointing.

run_status.json is mutable control state; checkpoint config is immutable identity evidence

run_status.json answers operational questions such as whether the run is currently running, preempted, or completed, what the last signal was, and how many resumable exits have occurred. The checkpoint's saved config snapshot, by contrast, is used as immutable evidence about what experiment setting produced the run. Keeping these roles separate makes recovery decisions easier to debug.

W&B continuity is opportunistic, not required for filesystem continuity

The recovery path saves wandb_run_id.txt once a tracker exists and reuses it on later launches. If that file is absent, the run can still resume correctly from the filesystem and checkpoint state. W&B continuity is best-effort metadata, not a prerequisite for training continuity.

Submission-layer note: snapshot mode on resubmit

The resubmit path can start from already-generated collection rows. If a previous generation enabled job-code snapshots, those rows may still carry old snapshot metadata fields even after a later resubmit turns snapshot mode off.

The current code handles this in two places:

• path:core/slurmkit/job_code_snapshots.py now clears stale snapshot metadata when job_code_snapshot_enabled is false;
• the templates parse job_code_snapshot_enabled explicitly instead of relying on the truthiness of strings like "false".

This matters operationally because a resubmitted job should not keep cd-ing into a stale detached worktree after snapshot mode has been disabled.

What future developers should preserve

• Auto-resume must stay tied to deterministic run identity, not to ad hoc filesystem scans.
• resume_if_requested(...) should remain the single source of truth for model/optimizer/scheduler/RNG restoration.
• Resume-position arithmetic and dataloader reconstruction should stay explicit and inspectable.
• The trainer and Slurm template together define the recovery contract. Changing one side without the other is risky.
• Exit code 42 is part of the operational contract, not an incidental magic number.
• Legacy support should be removed only as a deliberate migration after the old collections no longer matter.

Debugging checklist

When a run does not resume as expected, inspect these first:

• run_status.json: current status, resubmit_count, and last_signal
• config_identity.json: whether the run directory belongs to the current setting
• checkpoints/last.pt: whether a resumable checkpoint actually exists
• events.jsonl: whether a graceful_shutdown event was written
• wandb_run_id.txt: whether W&B continuity should be expected
• the trainer log line beginning with resumed from checkpoint:: this prints global_step, start_epoch, resume_batch_idx, and train_order_mode

Tests that protect this contract

The most relevant tests are:

• path:experiments/synthetic_sequences/entangled/tests/test_resumption.py
• path:experiments/synthetic_sequences/entangled/tests/test_train_entangled_pretrained_smoke.py
• path:experiments/synthetic_sequences/slurm/tests/test_template_recovery_contracts.py
• path:experiments/synthetic_sequences/slurm/tests/test_sigterm_trap.py
• path:experiments/synthetic_sequences/slurm/tests/test_job_code_snapshots.py
• path:tests/experiments/synthetic_sequences/test_run_naming.py

See also

• Gradient accumulation