xrun: smoke-test on a free Kaggle GPU, train on vast.ai, let the agent read the metrics

#ai #automation #cli #rust

I train 3D computer vision models, and classical ones at hackathons. For a long time every run came with the same chores: vastai create instance, scp the data up, start the script, check on it, scp the checkpoint back, remember to destroy the instance. Kaggle had its own version of that with kaggle kernels push. I wanted to hand the whole loop to an agent: launch on whichever platform fits, smoke-test on a free Kaggle GPU before paying for the real training, and collect metrics the same way everywhere, so the agent can read them and decide what to do next.

That became xrun, a Rust CLI with a Python Textual TUI on top. One YAML manifest describes a run: where the GPU comes from, what to upload, what to run and what to bring back. xrun rents the instance, uploads the data, starts training, follows stages and metrics, pulls the checkpoints and destroys the instance when it's over. It supports six vendors: vast.ai, Kaggle, Lightning AI, Google Colab, your own server over SSH, and the local machine.

A manifest

name: resnet50_baseline
vendor: vast                    # vast | kaggle | lightning | colab | ssh | local

vast:
  image: pytorch/pytorch:2.4.1-cuda12.1-cudnn9-devel
  gpu: { type: "RTX 4090", count: 1 }
  price: { max_per_hour: 0.55 }

data:
  - src: data/train.h5
    dst: /workspace/data/train.h5

run:
  cmd: python train.py
  args:
    --lr: 5e-4
    --epochs: 30

artifacts:
  patterns: ["checkpoints/best*.pt"]

policy:
  on_idle_minutes: 30           # no output for 30 minutes: destroy
  early_stop:                   # metric plateau: pull the best, then destroy
    metric: val_f1
    patience: 5

And the day-to-day loop:

xrun launch exp/resnet50.yaml --detach --max-cost 5
xrun events <id> --follow
xrun metrics <id> --key val_f1 --ascii
xrun pull <id> --ckpt best --into models/

The smoke test is a second manifest with vendor: kaggle and a small config. The training script and the metrics code stay the same, so if it survives a few epochs on Kaggle, I launch the vast.ai one.

A copy of the manifest is stored with the run, and its hash is part of the run's identity. I can edit the original freely and xrun rerun <id> still reproduces exactly what ran. For hyperparameters, xrun sweep exp/base.yaml --grid run.args.--lr=1e-3,5e-4 --launch writes out one manifest per variant.

Built for an agent to drive

Every read command (ls, show, events, metrics) has --json, and whatever the TUI can do, the CLI can do too. xrun show prints each metric's last value and step, so a score logged once at the end of a cross-validation run is visible without plotting anything. xrun install skill --claude (or --codex) installs a skill that teaches Claude Code or Codex to use xrun instead of hand-rolled vastai and ssh calls. So the agent can rerun the last experiment with a different learning rate, or pull the checkpoint of the run with the best val_f1 from the last day, through the same commands I would type.

Metrics come from a small hook in the training script:

from xrun_hook import stage, metric, metrics, done, notify

with stage("train"):
    for ep in range(epochs):
        loss = train_one_epoch(model, loader)
        metric("train_loss", loss, step=ep)
        metrics({"val_loss": v.loss, "val_f1": v.f1}, step=ep)

notify("training", f"best val_f1 = {best:.3f}")
done()

The hook writes events.jsonl and metrics.jsonl, and a background poller picks them up. On Kaggle xrun injects the hook into the kernel by itself. An uncaught exception gets recorded as an error event.

A billing instance is never left alone

xrun launch --detach leaves that poller running. It counts the spend and destroys the instance when the run finishes, fails, goes silent, hits a metric plateau or runs into --max-cost / --max-hours. If the destroy call fails, you get a push saying the instance is still billing. The poller can't report its own death, so xrun watchdog runs from the OS scheduler every five minutes: no heartbeat from the poller while the instance is alive means a notification and a respawned poller. It also catches instances that no run record knows about.

Notifications go to ntfy, Telegram, a webhook (Slack, Discord) or desktop toasts. Replying /stop <id> to the Telegram bot kills a run.

The run history lives in a local SQLite database, with no tracking service or account involved. MLflow and W&B plug in as mirrors if you want their charts.

Rough edges

xrun_hook isn't on PyPI yet, so outside Kaggle it travels with your code. Kaggle shows metrics mid-run only when an MLflow server is configured; without one you see the start and the end. Lightning AI was smoke-tested live on a CPU Studio only, and for Lightning and Colab xrun doesn't know the price, so --max-cost doesn't apply there.

Try it

macOS and Linux:

curl -sSf https://raw.githubusercontent.com/Gerrux/xrun/master/install.sh | sh

Windows (PowerShell):

irm https://raw.githubusercontent.com/Gerrux/xrun/master/install.ps1 | iex

xrun doctor and then xrun launch exp/templates/quickstart.yaml check the whole chain locally with no credentials, GPU or data. The code is on GitHub under MIT: Gerrux/xrun.

If you run training on one of these vendors, a report of what broke is the most useful thing you could send me. Each of their APIs has its own habits, and half of the fixes so far came from exactly that.

I built xrun together with my AI sidekick, Claude Code (Opus 5.5), and this post was written with its help too.

DEV Community