Skip to content

Checkpoint publication

cmd/publish is the helper a job script runs on the compute node to find the checkpoint to resume from, publish saves as training writes them, and prune old ones. It wraps the checkpoint package. The design is spec 0001 D16 in the project wiki.

A checkpoint root is a POSIX directory (local disk, NFSv3 and later, Lustre, GPFS) or an S3-compatible bucket and prefix, s3://BUCKET/PREFIX, on a store that honours conditional writes. Any other scheme is refused: a -checkpoint that starts with word: and is not s3:// (gs://, s3:/bucket, file:/data) exits 1 rather than being read as a relative path, so put ./ in front of a relative path whose first element holds a colon. A relative POSIX -checkpoint is made absolute against the working directory. A POSIX root must already exist as a directory, for every verb: one that does not is an error, never an empty log.

Layout

<Checkpoint>/
  exec/<DISPATCH_ATTEMPT>/            one folder per execution; pass it to the framework as its output dir
  pub/<epoch>-<step>-<path>           one empty file per publication
  • <epoch> and <step> are decimal, zero-padded to 20 digits.
  • <path> is the save's path relative to <Checkpoint>, encoded as RFC 4648 base32 (the alphabet A–Z, 2–7), with no padding.
  • Everything about a publication is in its name, so an entry is either there or not. Entries are created with O_CREAT|O_EXCL on POSIX, or on S3 as an empty object put with If-None-Match: *, and never renamed or rewritten.
  • On S3 the same layout is a key prefix: PREFIX/exec/... and PREFIX/pub/.... The latest entry is found by a LIST of PREFIX/pub/ and nothing else.
  • The latest checkpoint is the entry with the greatest (epoch, step). A save from an older execution, published late, has a lower epoch and is never chosen.
  • A file in pub/ whose name does not parse is an error naming it. It is never skipped.

Environment

Variable Flag Used by
DISPATCH_EPOCH -epoch watch, publish
DISPATCH_ATTEMPT -attempt watch with no directory

The flag wins when both are set. -checkpoint ROOT is required by every verb.

S3 credentials

S3 credentials never pass through the job script or the job's environment. -creds names where to find them: an ordered, comma-separated list of kind:location sources. The helper tries them in order on the node it runs on and uses the first that answers.

publish resolve -checkpoint s3://ckpts/run-7 \
  -creds secret:/var/run/secrets/ckpt-s3,file:$HOME/.config/ckpt-s3
Kind Location Holds
file an absolute directory, mode 0700, every file in it 0600 (or stricter) one file per field
secret the absolute directory a Kubernetes Secret is mounted at as a volume one file per field
vault a URL: the Vault server and the KV v2 secret's path, the login in its query one field per field

The fields are access_key_id and secret_access_key, and optionally session_token, endpoint (the store's http(s)://host[:port], AWS's regional endpoint when absent) and region (us-east-1 when absent). Any other file or field is refused, as is a value that is empty or holds whitespace. One trailing newline on a file is dropped.

A source does not answer when its location is absent from this node: no directory, no service-account token, no AppRole directory, or a Vault server that cannot be reached. The next source is then tried. A source that is present but wrong (a directory others can read, an unknown field, a denied Vault login, a missing secret) stops the run with an error naming it, rather than being passed over. When no source answers, the error names each one and why.

A vault location:

vault:https://vault.example:8200/team/s3?mount=secret&auth=kubernetes&role=training
vault:https://vault.example:8200/team/s3?auth=approle&approle=/home/u/.config/vault-approle
Query key Meaning
mount the KV v2 mount, secret by default
auth kubernetes or approle, required
auth-mount where the auth method is mounted, by default its name
role the Kubernetes auth role; auth=kubernetes only
token the projected service-account token, /var/run/secrets/kubernetes.io/serviceaccount/token by default
approle a directory holding role_id and secret_id, held to the file kind's modes
ca a PEM file to verify the server's certificate against
proxy the HTTP proxy every request to this Vault server goes through, http(s)://host[:port], query-escaped

A key not in this table, a key given twice or with no value, and a query that does not parse (a bad % escape, a ;) are refused rather than dropped.

No VAULT_* variable is used: the client is configured from the location alone. A token in VAULT_TOKEN is never sent, since the helper always logs in, and VAULT_ADDR, VAULT_AGENT_ADDR, VAULT_SKIP_VERIFY, VAULT_CACERT, VAULT_NAMESPACE, VAULT_HEADERS, VAULT_WRAP_TTL and the rest change neither where the login goes nor how the server is trusted; only ca= does that. The HashiCorp client still parses those variables as it is built, so one holding a malformed value (a VAULT_SKIP_VERIFY that is not a boolean, a VAULT_CACERT naming no file) refuses the source rather than being ignored.

Proxies

HTTPS_PROXY, HTTP_PROXY and NO_PROXY are ignored, for Vault and for S3 alike. A proxy is configured only explicitly, and separately for each thing reached:

  • a vault source names its own in its location's proxy= key, for example vault:https://vault.example:8200/team/s3?auth=kubernetes&role=training&proxy=http%3A%2F%2Fproxy.example%3A3128;
  • the S3 store's is the -proxy flag, -proxy http://proxy.example:3128.

Each is one URL, http:// or https:// and a host with an optional port. Userinfo, a path, a query or a fragment is refused, as is any other scheme: a proxy password would sit in the job script. Without one, the connection is direct. There is no exclusion list: a proxy given is used for every request to that server or store. A malformed proxy= refuses the source and exits 1; a malformed -proxy is a usage error and exits 2.

Flags for S3 roots

The flags -creds, -proxy and -timeout (default 5m, the time an S3 verb may take, or one watch scan) apply to s3:// roots; -creds or -proxy with a POSIX root is a usage error, and so is resolve -probe. -timeout with a POSIX root is accepted, and bounds only the scan watch makes on its way out.

Verbs

Flags follow the verb and come before any positional argument.

resolve

publish resolve -checkpoint ROOT [-creds LIST [-proxy URL]] [-probe]

Prints the latest publication's save location (a path, or an s3:// URL), or nothing if none exists, and exits 0 either way. A POSIX root that does not exist (an unmounted filesystem, a typo) exits 1 instead, so a job never restarts from step 0 because it looked in the wrong place. Make it the job script's first action: on S3 it resolves the credential sources there, so a source only compute nodes can see fails the job before training. -probe runs the probe checks first, for a job whose preflight could not reach its credentials, and fails the same way probe does. With a POSIX root, -probe is a usage error and exits 2.

watch

publish watch -checkpoint ROOT [-creds LIST [-proxy URL]] -framework hf|dcp [-interval 10s] [-epoch N] [-attempt A] [SAVES]

Runs beside training and publishes each save once it is complete, printing published <entry> (epoch E, step S, <save>) for each. SAVES defaults to <Checkpoint>/exec/$DISPATCH_ATTEMPT; on S3 it is that prefix, and a SAVES given as an argument is an s3:// URL under the root. Every -interval it lists the save folders, publishes the complete ones in step order, each once, and passes over an incomplete one until a later look. On SIGTERM or SIGINT it scans once more, then exits.

-framework Save folder Complete when
hf (Hugging Face Trainer) checkpoint-<step>/ trainer_state.json decodes and its global_step equals <step>
dcp (PyTorch DCP) step-<step>/ .metadata exists; on POSIX it must be a regular file, and anything else stops watch with checkpoint.ErrMalformed

On S3, a save folder is a common prefix of a LIST delimited by /, trainer_state.json is read with a GET (an object over 64 MiB is refused), and .metadata is checked with a HEAD. -timeout bounds each scan rather than the whole run, and the scan on the way out.

A run can outlive its credentials. When the store refuses a request with ExpiredToken (an expired session token, such as one from a Vault secret that has since been renewed) or InvalidAccessKeyId (a rotated key), watch resolves -creds again from the start, says so on stderr (publish: the store rejected the credentials (<code>); resolving the sources again, then the source that answered), and sends the request once more. If that is refused too, or no source answers, or the source now names another endpoint or region, the scan fails and watch exits 1. Credentials carry no expiry time, so they are renewed when refused, never ahead of time. The other verbs are short-lived and fail on the first refusal.

trainer_state.json is read as Python's json module writes it, so the bare NaN, Infinity and -Infinity that a diverged loss puts in log_history are accepted. A file cut off before its end is a save still being written and is passed over. Any other fault in it (a syntax error, data after the document, no global_step, or a global_step that disagrees with its folder) stops watch with an error, as does any error from the store.

publish

publish publish -checkpoint ROOT [-creds LIST [-proxy URL]] [-epoch N] SAVE STEP

Publishes one save, for frameworks without a completion marker, and prints the same published ... line as watch. Run it only once the save is complete. SAVE is a path, or an s3:// URL under the root holding at least one object. If the entry already exists, it says so on stderr, leaves the entry untouched and exits 0. On S3, a put whose connection dropped is sent again; if the first put had landed, the retry reports the entry as already published. A put S3 refuses with 409 ConditionalRequestConflict, because another conditional put for the same entry is in flight, is sent up to four times in all, with jittered waits of 50 to 100, 100 to 200 and 200 to 400 ms between them. Once the other put settles, this one either lands or gets 412 and reports the entry as already published; a 409 on every attempt fails the publication.

prune

publish prune -checkpoint ROOT [-creds LIST [-proxy URL]] [-keep 3]

Removes all but the latest -keep entries (at least 1), then the folders those entries named, printing removed <folder> for each. The latest entry is always kept. A folder is left in place while any remaining entry names it, a folder inside it, or a folder around it. On S3 a folder is every object under its prefix. On POSIX, a folder with a symbolic link anywhere on its path below the root (the folder itself or a parent such as exec/<attempt>) is never removed: prune leaves it in place, prints publish: prune: left <folder> in place: <link> is a symbolic link on stderr, and carries on with the rest, so a link can't make it delete outside the root. The root itself may be a link.

probe

publish probe -checkpoint s3://BUCKET/PREFIX -creds LIST [-proxy URL]

The preflight for an S3 root, run before a job is accepted. It resolves the credential sources, then puts a probe object beside pub/, reads it back, puts it again with If-None-Match: * and expects 412, lists pub/, and deletes the probe. On success it prints the root and the source that answered.

On failure it exits 1. An error that starts credentials: is about -creds or what it found: a list that does not parse, a source that is refused or none that answered, or an endpoint field the helper cannot use. One that starts store: probe: names the step that failed: write, read, conditional write, list or delete. A -checkpoint that is not a valid s3:// root, or names another scheme, also exits 1, before any source is tried. A missing -creds, a malformed -proxy, or a POSIX -checkpoint, is a usage error and exits 2.

Exit status

Code Meaning
0 success, including "nothing published" and "already published"
1 the operation failed, the checkpoint store is not supported, or a value is malformed (an epoch, a step, a -creds list)
2 usage error: an unknown verb or flag, a required flag or argument missing, or a flag for the other kind of root

Limits

  • Framework rotation. With Trainer's save_total_limit set, Trainer deletes old checkpoint-<step>/ folders itself, including ones already published. Leave it unset and let prune rotate.
  • The marker is not a data sync. A publication is durable when publish returns, but the save's own files are only as durable as the framework made them. DCP's file writer fsyncs by default; Trainer writes trainer_state.json without one.
  • A build for compute nodes uses CGO_ENABLED=0, so the binary is static.
  • S3 addressing is path-style (ENDPOINT/BUCKET/KEY), which every S3-compatible server and AWS accept. Requests are signed with SigV4 for service s3.
  • A secret source's modes are not checked. The kubelet sets them from the pod's defaultMode. On Soperator a mounted Secret is readable by every job on the worker.