r2d1 0.1.4

Tiny ML experiment tracking on Cloudflare R2 + D1

r2d1

Tiny ML experiment tracking on Cloudflare R2 + D1.

from r2d1 import Tracker, r2d1

tracker = Tracker.from_env()

job = tracker.start_job("mnist_dit")

for epoch in r2d1(range(10), job=job, checkpoint_every=1, keep_last=2):
    loss = train_step(...)

    epoch.d1(loss=float(loss), lr=float(lr))

    if epoch.should_checkpoint:
        epoch.r2({
            "checkpoint.pt": "ckpt/checkpoint.pt",
            "config.json": {"epoch": epoch.i},
        })

job.complete()

Mental model

epoch.d1({...})  -> small JSON metrics/metadata -> D1
epoch.r2({...})  -> files/artifacts/checkpoints -> R2

Aliases are included:

epoch.log(...)        # same as epoch.d1(...)
epoch.checkpoint(...) # same as epoch.r2(...)

The loop is intentionally lowercase like tqdm:

from r2d1 import r2d1

for epoch in r2d1(range(100), job=job):
    ...

Install locally

pip install -e .

Cloudflare credentials

Tracker.from_env() reads:

export R2D1_ACCOUNT_ID="..."
export R2D1_API_TOKEN="..."
export R2D1_D1_DATABASE_ID="..."
export R2D1_R2_BUCKET="..."
export R2D1_R2_ACCESS_KEY="..."
export R2D1_R2_SECRET_KEY="..."
# optional:
export R2D1_R2_ENDPOINT_URL="https://<account_id>.r2.cloudflarestorage.com"

Or construct directly:

tracker = Tracker(
    account_id="...",
    api_token="...",
    d1_database_id="...",
    r2_bucket="...",
    r2_access_key="...",
    r2_secret_key="...",
)

Checkpoint rotation: keep only last two

Use keep_last=2:

for epoch in r2d1(range(400), job=job, checkpoint_every=10, keep_last=2):
    if epoch.should_checkpoint:
        epoch.r2({"model.safetensors": Path("model.safetensors")})

R2 layout:

jobs/job_3/checkpoints/slot_0/
  manifest.json
  model.safetensors

jobs/job_3/checkpoints/slot_1/
  manifest.json
  model.safetensors

Epoch 0 writes slot_0, epoch 10 writes slot_1, epoch 20 overwrites slot_0, etc. D1 only points at a checkpoint after the files and manifest.json have been uploaded successfully.

D1 schema

r2d1 creates three tables automatically:

• jobs: run lifecycle + latest checkpoint pointer
• epochs: per-epoch metrics, with fixed columns plus metrics_json
• checkpoints: committed checkpoint manifests

epoch.d1(...) stores both convenience columns and a full JSON payload:

epoch.d1({
    "loss": 0.123,
    "lr": 3e-4,
    "fid": 18.2,
    "gpu": "B300",
})

R2 inputs

epoch.r2(...) accepts:

• Path objects — streamed to R2, preferred for large checkpoints
• bytes / bytearray
• str — existing file path if it exists, otherwise UTF-8 text
• dict / list / JSON-ish objects — encoded as JSON
• file-like objects — spooled to a temp file, then uploaded

Example:

epoch.r2({
    "model.safetensors": Path("ckpt/model.safetensors"),
    "config.json": {"hidden_size": 1152, "depth": 28},
    "notes.txt": "loss looked stable",
})

Resume

job = tracker.resume_job(3)
files, manifest = job.load_latest(include_manifest=True)

print(manifest["epoch"])
model_bytes = files["model.safetensors"]

load_latest() validates file SHA256 hashes against the manifest.

Optional serialization helpers

Core r2d1 does not own model serialization. It ships whatever you provide.

There is an optional r2d1.checkpoint module for small/simple PyTorch/JAX cases, but for real training runs prefer framework-native checkpoint files such as safetensors, Orbax, Flax serialization, torch.save, etc., then pass those file paths to epoch.r2(...).