Checkpoints & versioning

gpu-train keeps your work reproducible and local-first. Everything it persists lives under ~/.gpu-train (honors GPU_TRAIN_HOME):

~/.gpu-train/
├── registry.sqlite          # jobs / events / metrics / checkpoints / versions
├── artifacts/<job_id>/checkpoints/   # checkpoints pulled back from each run
├── code/<version>.tar.gz    # content-addressed code snapshots
├── data/<name>/<version>.json  # dataset manifests
└── models/<model_id>/       # Hugging Face / local baseline caches

Checkpointing#

Write checkpoints to the directory gpu-train hands you. It sets GPU_TRAIN_CHECKPOINT_DIR for every job and captures whatever lands there back to your machine when the run ends.

import gpu_train
 
ckpt_dir = gpu_train.checkpoint_dir()        # reads GPU_TRAIN_CHECKPOINT_DIR; created for you
model.save_pretrained(ckpt_dir / f"step-{step}")

• On the local provider, GPU_TRAIN_CHECKPOINT_DIR points straight at ~/.gpu-train/artifacts/<job>/checkpoints, so checkpoints land locally with no copy.
• On SSH providers (RunPod, Vast.ai, GCP, Colab) it points at a remote dir that the control plane rsyncs back to that same local directory before the box is torn down.

Either way, captured files are recorded in the registry and listed on the run detail page with one-click download.

Pull-back is on by default; disable it with defaults.fetch_checkpoints = False in your registry if you'd rather leave large artifacts on the box.

Sinks (push checkpoints & logs to a pool)#

On top of the always-on local capture, you can fan artifacts out to one or more sinks. At teardown gpu-train pushes the captured checkpoints and the materialized logs (events.jsonl

• metrics.jsonl) to each configured destination — best-effort, so a failing sink is logged, never fatal.

gpu_train.configure(registry={
    "sinks": {
        "checkpoints": [
            {"type": "local", "dest": "/mnt/nas/runs"},                       # extra local copy
            {"type": "hf", "repo_id": "you/ckpts", "token_ref": "env://HF_TOKEN", "private": True},
            {"type": "s3", "bucket": "my-bucket", "prefix": "gpu-train"},
        ],
        "logs": [
            {"type": "local", "dest": "/mnt/nas/runs"},
            {"type": "gcs", "bucket": "my-gcs-bucket"},
        ],
    },
}, use=[])

Each sink uploads under <dest-or-prefix>/<job_id>/{checkpoints,logs}/. Sink types:

S3/GCS credentials resolve from the sink's secret_ref, the local credential store (the Amazon S3 / Google Cloud Storage cards on the Providers page), or the backend's own default chain (env vars, instance metadata, ~/.aws, GOOGLE_APPLICATION_CREDENTIALS). You can also push on demand:

gpu_train.sinks.push_checkpoints(job_id)
gpu_train.sinks.push_logs(job_id)

Code versioning#

Before each run, gpu-train content-addresses the project tree it pushes (same excludes as the rsync — .git, .venv, __pycache__, …) into a short version id and writes a deduped tarball to ~/.gpu-train/code/<version>.tar.gz. Re-running unchanged code reuses the same id.

If the project is a git repo, the commit SHA and a dirty flag are captured too. The version, commit, and dirty state are stored on the JobRecord and shown on the run detail page.

gpu-train code                                   # list snapshots
gpu-train code --restore 6c77e0025f8c4a14 --dest ./old-run

Restoring extracts the exact tree that ran, so you can always reproduce a past run. Disable snapshotting with defaults.snapshot_code = False.

Data versioning#

Pin exactly which bytes a run trained on — without copying the dataset. register() walks a file or directory, builds a manifest (per-file path, size, and sha256), and derives a content version. Only the manifest is stored (under ~/.gpu-train/data); your data stays put.

import gpu_train
 
dv = gpu_train.data.register("./data/train", name="imdb")
print(dv.version, dv.num_files, dv.size_bytes)
 
# Resolve the same bytes later (latest, or pinned to a version):
path = gpu_train.data.resolve("imdb")
path = gpu_train.data.resolve(f"imdb@{dv.version}")

Registering identical bytes is idempotent (same version id). Browse versions from the CLI or the HTTP API:

gpu-train data                       # list all dataset versions
gpu-train data --register ./data/train --name imdb