Skip to content

Publish and resolve checkpoints from a job script

CKPT=/lustre/team/finetune
RESUME=$(publish resolve -checkpoint "$CKPT") || exit 1
OUT=$CKPT/exec/$DISPATCH_ATTEMPT

publish watch -checkpoint "$CKPT" -framework hf &
WATCH=$!

python train.py --output_dir "$OUT" ${RESUME:+--resume_from_checkpoint "$RESUME"}
STATUS=$?

kill -TERM "$WATCH"; wait "$WATCH"
publish prune -checkpoint "$CKPT" -keep 3
exit $STATUS

This is the shape of a job script around a Hugging Face Trainer run. resolve prints the latest published save, or nothing on a first run. Training writes into a folder that belongs to this execution alone. watch publishes each save once it's complete, and on SIGTERM looks once more before it exits, so a save finished just before training stopped isn't lost. prune keeps the latest three publications.

The flag-by-flag reference is checkpoint publication.

Before you start

  • The job runs under a dispatch adapter that sets DISPATCH_ATTEMPT and DISPATCH_EPOCH in every rank's environment. The AttemptInjection capability says it does, and Submit refuses a job on a cluster where it isn't Yes.
  • publish is on the compute node's PATH. It's the cmd/publish program in this module; build it static so it runs on any node:
CGO_ENABLED=0 go build -o publish gitlab.com/phpboyscout/go/dispatch/cmd/publish

On Kubeflow the binary comes from the publish image instead, at /opt/dispatch/publish; mount the publish image says how.

  • The checkpoint root, $CKPT here, already exists on a POSIX filesystem every node can reach: local disk, NFSv3 or later, Lustre or GPFS. An s3://BUCKET/PREFIX root works the same way with -creds added to every line; the credential sources are in the reference.
  • Run the publish lines on one rank per execution. Several watch processes over the same folder are harmless, since a publication that already exists counts as done, but they do the same work several times.

Point the framework at this execution's folder

Each execution writes into <Checkpoint>/exec/<DISPATCH_ATTEMPT>/. Pass that to the framework as its output directory. Nothing else writes there, so a previous execution that's still running after its job ended can't touch this one's saves.

Resume from the latest publication

publish resolve -checkpoint DIR prints the absolute path of the save to resume from, and exits 0 whether or not it found one. An empty answer means start from scratch, which is what ${RESUME:+...} handles above.

The answer can be a save in an earlier execution's folder. That's the point: the new execution resumes from it and writes its own saves into its own folder.

Publish saves as they complete

publish watch scans the execution's folder every 10 seconds (-interval changes it) and publishes each save once its completion marker is right:

-framework Save folder Complete when
hf checkpoint-<step>/ trainer_state.json decodes and its global_step is <step>
dcp step-<step>/ .metadata exists and is a regular file

Leave Trainer's save_total_limit unset. With it set, Trainer deletes old checkpoint-<step>/ folders itself, including published ones, and a later resolve can name a folder that's gone.

Publish one save by hand

For a framework with no completion marker, publish each save once your code knows it's complete:

publish publish -checkpoint "$CKPT" "$OUT/save-1200" 1200

Run it a second time for the same save and it says the entry already exists, on stderr, and exits 0 without touching it.

Prune old publications

publish prune -checkpoint DIR -keep N removes all but the latest N publications (at least 1), then the save folders those publications named. A folder stays while a remaining publication names it, a folder inside it or a folder around it. The latest publication is always kept, so resolve still has an answer afterwards.

A save folder reached through a symbolic link below the checkpoint root is never deleted. prune names it on stderr and leaves it, so if you link exec/ or an attempt folder to other storage, clear those folders yourself.

Check what was published

Every publication is an empty file in <Checkpoint>/pub/, named for its epoch, step and save path:

ls "$CKPT/pub"

00000000000001000000-00000000000000000500-MV4GKYZPGEXDIMJQGEXDAL3DNBSWG23QN5UW45BNGUYDA

That one is exec/1.4101.0/checkpoint-500, published at epoch 1000000 (generation 0, submission 1, restart 0) and step 500. The latest is the one with the greatest epoch, then step. A file there whose name doesn't parse stops resolve with an error naming it; it's never skipped.

Reach S3 or Vault through a proxy

publish never uses HTTPS_PROXY, HTTP_PROXY or NO_PROXY from the job's environment, so a node that reaches the store or Vault only through a proxy needs it named explicitly. Give the store's proxy as -proxy on every line, and a Vault source's as proxy= in its location, query-escaped:

PROXY=http://proxy.example:3128
CREDS='vault:https://vault.example:8200/team/s3?auth=kubernetes&role=training&proxy=http%3A%2F%2Fproxy.example%3A3128'
RESUME=$(publish resolve -checkpoint s3://ckpts/team/finetune -creds "$CREDS" -proxy "$PROXY") || exit 1

The two are independent: a store reached directly and a Vault reached through a proxy is only proxy= on the source. Each takes one http:// or https:// URL with no password in it. Run publish probe with the same flags from a compute node to check the store answers through it.