Draining a Training Node Without Losing a Checkpoint

A drain request is a countdown, not a suggestion to save soon. Write the hard termination time into the event and subtract fixed allowances for signal delivery, process shutdown, checkpoint verification, and a safety margin. The remainder is the save budget. If the largest observed durable-save time does not fit, the correct action is to preserve the last committed checkpoint and exit cleanly. Starting a save that cannot commit before termination can create an attractive but unusable directory.

Measure the components independently. Record notice receipt time, training quiesce time, bytes staged, durable write completion, manifest publication, and process exit. Use a high percentile from real saves rather than the median. Kubernetes counts the preStop hook inside terminationGracePeriodSeconds and eventually sends SIGKILL to remaining processes. Slurm can signal a job before its time limit, but its documentation says delivery may occur up to 60 seconds earlier than requested. The application therefore needs an absolute deadline, not an assumed interval.

First stop new work from landing on the node. Then tell the distributed job that one member is leaving. These are separate control paths. A scheduler drain prevents replacement work from arriving, while a job-level signal coordinates the current ranks. If the training framework treats one disappearing rank as an unexplained fault, the remaining ranks may wait in a collective until their own grace periods expire.

Choose one coordinator, usually rank zero or a fault-tolerance controller, to decide whether a new checkpoint is possible. Every rank should acknowledge the same completed optimizer step, finish or discard partial accumulation according to a documented rule, and enter the checkpoint collective once. Repeated notices must be idempotent. A second signal may shorten the deadline, but it must not launch a competing save into the same destination.

Distributed checkpoints often contain multiple rank-local files plus shared metadata. Write them under a unique staging identifier. Each rank reports its file name, byte count, and checksum to the coordinator. Only after all expected entries verify should the coordinator write the manifest and publish a small commit marker or atomically move a pointer where the storage system supports that operation. Restart logic ignores any generation without that final evidence.

Asynchronous saving can shorten the time training is paused, but it does not erase the drain deadline. In PyTorch's documented async_save pattern, state is staged from device memory to CPU memory and the storage write finishes asynchronously; the example waits for the prior checkpoint future before issuing another save and waits for the final future before cleanup. During a drain, wait for the save future, verify the expected files, and then publish the application's durable commit evidence. Completion of CPU staging alone is insufficient, and exhausting host memory while staging can make the evacuation slower.

Consider an illustrative 150-second drain window. The team reserves 5 seconds for handler latency, 20 seconds to align ranks, 15 seconds to verify and publish, and 15 seconds to exit. The available durable-write budget is 150 - 5 - 20 - 15 - 15 = 95 seconds. Recent checkpoint write times at comparable load have a 95th percentile of 82 seconds and a maximum of 103 seconds. Starting a full save has only 13 seconds of percentile headroom and cannot cover the observed maximum.

The decision rule is explicit after the handler receives the notice: start a drain checkpoint only when time remaining until the hard stop exceeds quiesce p95 + durable-save p99 + commit p95 + exit margin. Otherwise terminate from the most recent verified generation. In this illustrative case, if the adopted p99 is 101 seconds, required time is 20 + 101 + 15 + 15 = 151 seconds, so the job skips the new save. The separate 5-second handler allowance has already elapsed when this comparison is made. The job may lose recent work, but it does not advertise a partial checkpoint as recoverable.

A successful commit means the expected bytes became durable. Restartability adds discovery, strict loading, optimizer and scheduler restoration, data-position restoration, and one valid subsequent step. The drain controller should record the committed checkpoint ID, training step, world size, parallelism layout, container digest, dataset revision, and configuration digest in the requeue request. A replacement job must consume that exact record rather than search for the lexicographically newest directory.

Use a fast gate during evacuation: manifest completeness, checksums, expected rank count, and a format compatibility check. Then run periodic clean-worker restores outside emergencies. A drain is the wrong moment to discover that a checkpoint format cannot load on the available topology. When matching replacement capacity matters, ClusterBid can be evaluated against the measured GPU count, topology, storage bandwidth, and restore-time requirement recorded by those drills.

Notices can arrive late or never reach the training process. AWS describes Spot interruption notices as best effort, which is a useful general warning against making notice handling the only protection. Regular checkpoints set the maximum work at risk. A node heartbeat, scheduler state watcher, and job-level timeout should detect disappearance without waiting forever for a graceful callback.

On an abrupt loss, fence the failed attempt so it cannot later publish a stale active pointer. The new attempt selects only a committed generation, verifies it, and resumes with a fresh attempt ID. If the old process returns, its credentials or lease should no longer authorize publication. Measure missed-notice rate, checkpoint age at failure, requeue delay, restore duration, and the fraction of restarts that reach a healthy step.

Test voluntary maintenance, a scheduler time-limit signal, a short grace period, a notice delivered during an existing save, a rank stuck in a collective, and abrupt node loss. For each case, capture the event timeline and inspect storage for abandoned staging data. Cleanup can remove uncommitted objects later, but recovery must be safe before cleanup runs. Also verify that a second drain request does not overwrite the first checkpoint or extend the hard deadline.

Pass when no work is scheduled onto the drained node, all ranks converge on one save decision, every published checkpoint has a complete manifest, the process exits inside the budget, and the replacement job reaches a valid optimizer step. This runbook cannot guarantee zero lost work after hardware failure or network partition. It limits loss to the age of the last verified checkpoint and makes graceful opportunities useful without depending on them.

Kubernetes Documentation: Pod Lifecycle

Slurm Workload Manager: sbatch

PyTorch Tutorial: Asynchronous Saving with Distributed Checkpoint

Amazon EC2 Documentation: Spot Instance Interruption Notices

GB300 NVL72 Rack-Scale Buyers Guide 2026: Power, Cooling, and Real Cost

Academic GPU Compute Grants: NSF, NIH, and University AI Funding in 2026