Skip to content

Save and resume a compute job with checkpoints

View Markdown

A checkpoint is a copy of your Job’s state directory that Nodus takes while the Job runs. When the Job is suspended, or the capacity under it is reclaimed, the next attempt starts with that directory restored and your program carries on from what it saved. Nodus decides when to checkpoint and where to store it; your program decides what to write and how to load it.

Checkpoints hold files, not memory. A restored attempt starts your command again from the beginning, with the state directory as it was at the last checkpoint, so the program must read its own progress back.

Use this guide when a training loop or batch job should continue from saved work after a restart. The example is a small CPU loop using Python’s standard library. It demonstrates the file contract; it is not a model-training benchmark or a measurement of recovery speed.

Write everything you need to continue, such as model weights, optimizer state and the current step, under /nodus/state. The path is also in NODUS_STATE_DIR (NODUS_CHECKPOINT_DIR is the older name for the same directory). On start, load what is there:

examples/checkpoints/resume/train.py
"""A training loop that saves its progress where Nodus checkpoints it and resumes from it.
It uses only the standard library, so it runs in any image. With the Nodus SDK installed,
`nodus.state_dir()`, `nodus.checkpoint.on_request()` and `nodus.restored()` do the same work.
"""
import json
import os
import socket
import threading
import time
from pathlib import Path
STATE_DIR = Path(os.environ.get("NODUS_STATE_DIR", "/nodus/state"))
STATE = STATE_DIR / "progress.json"
STEPS = int(os.environ.get("STEPS", "120"))
SAVE_EVERY = 20 # like Hugging Face Trainer's save_steps: a regular save even when Nodus does not ask
pending = threading.Event() # set while Nodus waits for a consistent checkpoint
request_seq = None
events = None
def save(step):
"""Write the state atomically, so a snapshot never sees a half-written file."""
STATE_DIR.mkdir(parents=True, exist_ok=True)
tmp = STATE.with_suffix(".tmp")
tmp.write_text(json.dumps({"step": step}))
os.replace(tmp, STATE)
def listen():
"""Subscribe to checkpoint requests on the events socket and flag each one for the loop."""
global events, request_seq
path = os.environ.get("NODUS_EVENTS_SOCKET", "/run/nodus/events.sock")
try:
events = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
events.connect(path)
except OSError:
return # outside Nodus there is nobody to ask for checkpoints
events.sendall(b'{"type":"checkpoint.subscribe"}\n')
for line in events.makefile("r"):
message = json.loads(line)
if message.get("type") == "checkpoint.request":
request_seq = message["seq"]
pending.set()
def ack():
"""Tell Nodus the files in the state directory are complete."""
events.sendall((json.dumps({"type": "checkpoint.ack", "seq": request_seq}) + "\n").encode())
pending.clear()
def main():
start = 0
if STATE.exists():
start = json.loads(STATE.read_text())["step"]
print(f"resumed from step {start}", flush=True)
threading.Thread(target=listen, daemon=True).start()
for step in range(start + 1, STEPS + 1):
time.sleep(1) # one step of work
print(f"step {step}/{STEPS}", flush=True)
if pending.is_set():
save(step)
ack()
elif step % SAVE_EVERY == 0:
save(step)
print("training complete", flush=True)
if __name__ == "__main__":
main()

Two habits keep checkpoints useful:

  • Write atomically. Write to a temporary file and rename it over the old one, as save() does. A checkpoint can then never capture a half-written file.
  • Keep outputs separate. Results you want to download go to /nodus/outputs. The state directory is recovery state: it is restored into the next attempt, not offered as a download.

NODUS_RESTORED=1 is set in an attempt that started from a checkpoint, if your program wants to log the difference. An empty state directory never replaces an earlier checkpoint that had files in it, so an attempt that fails before it saves anything cannot erase progress.

recovery.checkpoint in the Job spec controls the checkpoint. The defaults suit most programs:

recovery:
continuity: Checkpointed # restore the latest checkpoint into each new attempt
checkpoint:
paths: [/nodus/state] # up to 64 absolute paths, all saved in one checkpoint
interval: auto # or a fixed interval from 1m to 6h
maxSize: 1Ti # larger checkpoints fail with CheckpointTooLarge
retainAfterFinish: 168h # after the Job finishes, keep only the final checkpoint

On the command line, nodus run --checkpoint /nodus/state sets paths. Saving a whole folder such as your working directory is possible by listing it, but it makes every checkpoint larger and slower; list only what you need to continue.

continuity: Restartable is the lighter choice for programs that track their position as a counter: Nodus keeps NODUS_CURSOR_COMPLETED and NODUS_CURSOR_TOTAL from your progress reports and skips the file copy. continuity: Ephemeral starts every attempt from scratch.

With interval: auto, Nodus sets the cadence from how often the capacity your Job runs on is interrupted and how long a checkpoint takes to save. A Job gets at least four checkpoints over its expected run time, and checkpointing takes no more than about a tenth of it. Capacity that is rarely interrupted is checkpointed less often.

Nodus also checkpoints, whatever the interval:

  • when you run nodus suspend job/NAME, before the compute is released, so nodus resume continues from there;
  • when the capacity gives notice that it is about to be reclaimed, so the replacement attempt loses as little work as possible.

nodus describe job/NAME shows the latest checkpoint, and GET …/jobs/NAME/checkpoints lists each one with its sequence number, attempt, time, size and file count.

A checkpoint taken while your program is halfway through writing its files would restore a broken state. To avoid that, a program can ask to be told before each checkpoint and say when its files are complete. This is the request/ack handshake, and the example above implements it with the standard library.

It runs over the events socket at /run/nodus/events.sock (NODUS_EVENTS_SOCKET), one JSON object per line:

  1. Your program connects and sends {"type": "checkpoint.subscribe"} once.
  2. Before each checkpoint, Nodus sends {"type": "checkpoint.request", "seq": 7, "urgent": false}.
  3. Your program finishes the current step, writes its state and replies {"type": "checkpoint.ack", "seq": 7}.
  4. Nodus copies the state directory, then your program carries on. It does not need to pause while the copy runs.

urgent: true means the capacity is about to go away: save at the next safe point and skip optional work. Nodus waits for the ack for as long as the shutdown allows and then takes the checkpoint anyway, so a program that hangs cannot block a suspend.

You choose whether to use the handshake with recovery.checkpoint.integration:

Value Behaviour
Auto (default) Use the handshake when the program subscribes; otherwise checkpoint without asking
None Never ask; checkpoint the paths as they are
HFTrainer Answer requests from inside Hugging Face Trainer, with no change to your image (below)

With the Python SDK installed (pip install nodus-compute), nodus.checkpoint.on_request(save) registers a callback and sends the ack after it returns, nodus.checkpoint.requested() lets a loop poll instead, and nodus.state_dir() returns the directory. All of them do nothing outside Nodus, so the same script runs on your laptop.

Trainer already saves and resumes; point it at the state directory and it works with Nodus checkpoints:

  • set output_dir to the state directory (os.environ["NODUS_STATE_DIR"]), or a folder inside it;
  • set save_steps to how often Trainer saves on its own, and save_total_limit (for example 2) so old checkpoints do not fill the directory;
  • call trainer.train(resume_from_checkpoint=True) when NODUS_RESTORED is 1, and trainer.train() on a first start, because Trainer refuses to resume from an empty directory.

To also save when Nodus asks, add nodus.checkpoint.HFTrainerCallback() to the Trainer’s callbacks, or set integration: HFTrainer and Nodus registers the same callback for you, even in an image without the SDK. Trainer then saves at the end of the current step and Nodus checkpoints once the save is written.

A Job with distributed set defaults to recovery.checkpoint.format: Dcp. Every rank writes its shard with torch.distributed.checkpoint to $NODUS_CHECKPOINT_URI, and rank 0 commits it, which nodus.checkpoint.dcp.save and .load do for you. After a restart, $NODUS_RESTORE_URI names the latest committed checkpoint. A rank that lost its place in the gang cannot commit, so a restored gang always loads a checkpoint every rank finished.

Save the Python example above as train.py, or download the same source. In a terminal, use a fresh directory for state:

Terminal window
export NODUS_STATE_DIR="$(mktemp -d)"
STEPS=120 python3 -u train.py

After step 21/120 appears, press Ctrl+C. The loop saves every 20 steps, so cat "$NODUS_STATE_DIR/progress.json" should show {"step": 20}. Restart the same script with the same directory and a shorter finishing point:

Terminal window
STEPS=24 python3 -u train.py

Observed output from a local stop/restart validation on October 8, 2026:

resumed from step 20
step 21/24
step 22/24
step 23/24
step 24/24
training complete

Step 21 runs again because it was not in the last saved state. The validation terminated the first process after it printed step 21 and restarted the unmodified example with the same directory. It verifies loading saved files; it does not exercise Nodus storage, scheduling, or the checkpoint handshake.

For a real training loop, save the model, optimizer, scheduler, random-number state, and data position needed by your framework. A step counter alone cannot reconstruct a trained model.

With the CLI installed and signed in, run from the directory containing train.py. This uses paid compute and sets a $0.50 spending cap; the cap is not a quoted run cost.

Terminal window
$ nodus run --name checkpoints-resume --cpu 2 --memory 4Gi --checkpoint /nodus/state --max-cost 0.50 -d -- python train.py
$ nodus wait job/checkpoints-resume --for=jsonpath='{.status.phase}'=Running --timeout=8m
$ nodus logs -f job/checkpoints-resume

Once the log shows the loop progressing, stop following logs with Ctrl+C. The Job keeps running. Suspend it and wait for the save to finish before resuming:

Terminal window
$ nodus suspend job/checkpoints-resume
$ nodus wait job/checkpoints-resume --for=jsonpath='{.status.phase}'=Suspended --timeout=5m
$ nodus describe job/checkpoints-resume
$ nodus resume job/checkpoints-resume
$ nodus wait job/checkpoints-resume --for=condition=Succeeded --timeout=8m
$ nodus logs -f job/checkpoints-resume

The resumed attempt should print resumed from step N; N depends on when the checkpoint was saved. Inspect Runs → Details for attempts and the actual execution cost. The local result above is not evidence of a cloud run or a promised recovery time. The downloadable source and the code displayed in this guide come from the same example, so you can inspect the exact save and load behavior.

If the Job does not reach Suspended, inspect describe and its events before continuing. If it reaches the spending cap, review the cost before raising it. Download any outputs you need before deleting the Job with nodus delete job/checkpoints-resume.

Checkpoints are deleted with their Job. After a Job finishes, Nodus keeps only its final checkpoint once retainAfterFinish (seven days by default) has passed.