yarn-au 0.1.6

Python SDK for the Yarn sovereign AI platform

yarn-au

Python SDK and CLI for Yarn, a sovereign AI compute platform with Australian data residency.

Install

pip install yarn-au
yarn login

Training script

Yarn uses three decorators to turn a standard PyTorch script into a submittable job. The same file runs locally with no changes.

import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import TensorDataset
import yarn.train as yt

@yt.model
def create_model():
    return nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))

@yt.dataset(batch_size=128, shuffle=True)
def load_data():
    return {"train": TensorDataset(...), "val": TensorDataset(...)}

@yt.job
def train(model, data):
    model = yt.prepare(model)
    opt = Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()

    for epoch in range(yt.start_epoch, 10):
        for x, y in data["train"]:
            loss_fn(model(x), y).backward()
            opt.step(); opt.zero_grad()
        yt.report(epoch=epoch)
        yt.checkpoint(model, opt, epoch=epoch)

Pre-flight check

--dry-run analyses your script without submitting:

$ yarn job submit train.py --gpu rtx-4090 --dry-run

Model          3 layers, 203,530 params
Memory         weights 0.78 MB | grads 0.78 MB | optimizer 1.55 MB | activations ~2.1 MB
Total          ~5.2 MB (RTX 4090: 24,576 MB available)
Fit            YES — 0.02% VRAM utilised
Cost           ~$0.12/hr × 10 epochs ≈ $0.03 estimated

CLI quickstart

yarn login                              # authenticate
yarn job submit train.py --gpu rtx-4090 # submit a job
yarn job submit train.py --dry-run      # analyse without submitting
yarn job logs <job-id>                  # stream logs
yarn job list                           # list all jobs

yarn session create --gpu rtx-4090      # interactive GPU session
yarn models                             # available models
yarn gpus                               # GPU types and pricing
yarn balance                            # credit balance
yarn storage upload data/ /datasets/    # upload data
yarn storage download /results/ ./      # download results
yarn storage list /datasets/            # list remote files

Both yarn and yarn-au work as CLI entry points.

Decorator API

Decorator / function	Purpose
@yt.model	Wraps a function returning nn.Module
@yt.dataset(batch_size, shuffle)	Wraps a function returning dict[str, Dataset]
@yt.job(epochs=N)	Wraps the training function; receives (model, data)
yt.prepare(model)	Device placement and distributed wrapping (DDP/FSDP)
yt.data(loader)	Prepare DataLoader for distributed training
yt.report(**metrics)	Log metrics to Yarn
yt.checkpoint(model, opt, epoch=N)	Save a resumable checkpoint
yt.load_checkpoint(model, opt)	Restore from latest checkpoint
yt.start_epoch	Epoch to resume from (0 if fresh)

Requirements

• Python >= 3.10
• Dependencies: requests, click

Documentation

Full docs at docs.yarn.prosodylabs.com.au.

License

Proprietary. Copyright 2025-2026 Prosody Labs Pty Ltd.