Metadata-Version: 2.4
Name: yarn-au
Version: 0.1.6
Summary: Python SDK for the Yarn sovereign AI platform
Requires-Python: >=3.10
Requires-Dist: click>=8.0
Requires-Dist: requests>=2.28
Description-Content-Type: text/markdown

# yarn-au

Python SDK and CLI for [Yarn](https://docs.yarn.prosodylabs.com.au), a sovereign AI compute platform with Australian data residency.

## Install

```bash
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.

```python
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

```bash
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](https://docs.yarn.prosodylabs.com.au).

## License

Proprietary. Copyright 2025-2026 Prosody Labs Pty Ltd.
