Checkpoints

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.

Save your progress to the state directory

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.

Choose what is saved

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.

When Nodus checkpoints

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.

Answer checkpoint requests

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.

Hugging Face Trainer

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.

Gang checkpoints (Beta)

Beta

Multi-node Jobs (distributed) are Beta. Their checkpoints use a different format.

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.

Run the example

The example suspends the Job mid-run and resumes it; the second attempt prints resumed from step N and finishes the remaining steps:

$ cd examples/checkpoints/resume
$ nodus run --name checkpoints-resume --cpu 2 --checkpoint /nodus/state -d -- python train.py
$ nodus suspend job/checkpoints-resume
$ nodus resume job/checkpoints-resume
$ nodus logs -f job/checkpoints-resume
resumed from step 12
step 13/120
…

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