# Yarn Python SDK — Design Specification

## Overview

Thin Python client for Yarn's REST API. Handles auth, session lifecycle, job submission, storage, and billing queries. Dependencies: `requests` only.

Install: `pip install -e apps/yarn/sdk/` (monorepo) or `pip install git+ssh://git@github.com/prosodylabs/infrastructure.git#subdirectory=apps/yarn/sdk`

## Package structure

```
apps/yarn/sdk/
├── pyproject.toml
├── src/
│   └── yarn/
│       ├── __init__.py       # Re-export top-level convenience functions
│       ├── client.py         # YarnClient — auth, base URL, request helpers
│       ├── jobs.py           # Job submission, status, logs, cancel
│       ├── sessions.py       # Interactive GPU sessions (context manager)
│       ├── notebooks.py      # JupyterLab lifecycle
│       ├── storage.py        # Upload, download, list, delete
│       ├── billing.py        # Balance, usage, cost estimation
│       ├── secrets.py        # Secret CRUD
│       └── exceptions.py     # YarnError, InsufficientCredits, etc.
```

## Authentication

API key is the only auth mechanism for the SDK. The gateway/Kong must resolve API keys to user IDs + roles for `/v1/research/*` and `/v1/data/*` routes (backend item #15 — prerequisite).

```python
client = yarn.Client(api_key="yarn_...")
# or
client = yarn.Client()  # reads YARN_API_KEY env var
```

All requests send `Authorization: Bearer yarn_...` header. No JWT management, no OIDC, no token refresh.

## Base URL

Default: `https://api.au.yarn.prosodylabs.com.au`
Override: `yarn.Client(api_key="...", base_url="...")`

All SDK methods hit this base URL. Research endpoints at `/v1/research/*`, data at `/v1/data/*`, billing at `/v1/billing/*`, inference at `/v1/chat/completions`.

## Top-level convenience API

```python
import yarn

# Configure once
yarn.api_key = "yarn_..."
# or
client = yarn.Client(api_key="yarn_...")
```

### Sessions (interactive GPU)

```python
# Context manager — creates session, waits for ready, tears down on exit
with yarn.session(gpu="rtx-4090", name="experiment-1") as s:
    print(s.ray_address)      # ray://experiment-1.research.prosodylabs.com.au:443
    print(s.dashboard_url)    # https://experiment-1-dash.research.prosodylabs.com.au
    print(s.status)           # "running"

    import ray
    ray.init(address=s.ray_address)
    # ... Ray code ...

# Session is deleted when context exits (including on exception)

# Manual lifecycle
s = client.sessions.create(name="exp-1", gpu="rtx-4090", idle_timeout_minutes=120)
s.wait_until_ready(timeout=300)   # polls every 5s, raises TimeoutError
s.refresh()                       # update status from API
s.stop()                          # DELETE
```

**Session.create() parameters:**
- `name: str` — K8s-safe, 1-63 chars
- `gpu: str` — GPU type (default: "rtx-4090")
- `gpu_count: int` — default 1
- `cpu: str` — K8s resource (default: "4")
- `memory: str` — K8s resource (default: "8Gi")
- `image: str` — container image (default: ray-base)
- `idle_timeout_minutes: int` — auto-shutdown (default: 60)
- `env_vars: dict` — environment variables

**Session properties:**
- `id: str`
- `name: str`
- `status: str` — "creating", "running", "failed", "stopped"
- `ray_address: str` — `ray://{name}.research.prosodylabs.com.au:443`
- `dashboard_url: str`
- `gpu_type: str`
- `created_at: str`

### Jobs (fire-and-forget training)

```python
# From a script file
job = yarn.submit_job(
    script="train.py",                    # local file path — SDK reads + base64 encodes
    name="kairos-batch32",
    gpu="rtx-4090",
    pip_packages=["torch", "wandb"],
    env_vars={"WANDB_PROJECT": "kairos"},
    secret_names=["WANDB_API_KEY"],
    max_runtime_hours=4,
)

# From a string
job = yarn.submit_job(
    code="import torch; print(torch.cuda.is_available())",
    name="gpu-test",
)

# Block until complete
job.wait(timeout=3600, poll_interval=10)  # polls every 10s

# Stream logs
for line in job.stream_logs(tail=100):
    print(line)

# Check status
print(job.status)       # "SUCCEEDED", "FAILED", "RUNNING", "PENDING"
print(job.logs())       # full log output
print(job.cost)         # actual cost in AUD (after completion)

# Cancel
job.cancel()

# List all jobs
jobs = client.jobs.list()

# Cost estimation (before submitting)
estimate = client.jobs.estimate(gpu="rtx-4090", max_runtime_hours=4)
print(f"${estimate.low:.2f} - ${estimate.high:.2f} AUD")
```

**submit_job() parameters:**
- `script: str` — local file path (SDK reads, base64 encodes)
- `code: str` — inline Python code (alternative to script)
- `name: str` — job name (K8s-safe)
- `gpu: str` — GPU type (default: "rtx-4090")
- `gpu_count: int` — default 1
- `cpu: str` — default "2"
- `memory: str` — default "4Gi"
- `image: str` — container image (default: ray-base)
- `entrypoint: str` — default "python /home/ray/code/main.py"
- `pip_packages: list[str]` — installed at runtime
- `env_vars: dict` — environment variables
- `secret_names: list[str]` — K8s secrets to inject
- `max_runtime_hours: float` — 0 = no limit
- `description: str` — optional

Exactly one of `script` or `code` must be provided. If `script`, SDK reads the file and base64-encodes it. If `code`, SDK base64-encodes the string directly.

**Job properties:**
- `id: str`
- `name: str`
- `status: str` — "PENDING", "RUNNING", "SUCCEEDED", "FAILED", "STOPPED"
- `logs() -> str` — fetch latest logs
- `cost: float | None` — AUD, available after completion
- `created_at, start_time, end_time: str`
- `error: str` — error message if failed

### Notebooks (JupyterLab)

```python
nb = client.notebooks.create(
    name="analysis",
    gpu="rtx-4090",
    idle_timeout_minutes=120,
    data_mounts=["datasets/kairos"],
)
nb.wait_until_ready()
print(nb.jupyter_url)       # opens in browser
# ...
nb.stop()
```

### Storage

```python
# Upload
yarn.upload("datasets/train.csv", local_path="./data/train.csv")
yarn.upload("models/checkpoint.pt", local_path="./best.pt")

# Download
yarn.download("results/metrics.json", local_path="./metrics.json")

# List
objects = client.storage.list(prefix="datasets/")
for obj in objects:
    print(f"{obj.key}  {obj.size} bytes")

# Delete
client.storage.delete("datasets/old.csv")

# Usage
usage = client.storage.usage()
print(f"{usage.total_bytes / 1e9:.1f} GB used, {usage.percent_used:.0f}% of quota")
```

### Billing

```python
balance = client.billing.balance()
print(f"${balance.available:.2f} available, ${balance.held:.2f} held")

usage = client.billing.usage(period="daily", days=7)
for day in usage.data:
    print(f"{day.date}: ${day.spend:.2f}")
```

### Inference (chat completions)

```python
# Non-streaming
response = client.chat.completions(
    model="mistral-7b-instruct-v0.2",
    messages=[{"role": "user", "content": "Hello"}],
)
print(response.choices[0].message.content)

# Streaming
for chunk in client.chat.completions(
    model="mistral-7b-instruct-v0.2",
    messages=[{"role": "user", "content": "Hello"}],
    stream=True,
):
    print(chunk.delta.content, end="")
```

### Secrets

```python
client.secrets.create(name="WANDB_API_KEY", value="...", description="W&B tracking")
secrets = client.secrets.list()   # names + descriptions only, never values
client.secrets.delete("WANDB_API_KEY")
```

## Error handling

```python
from yarn.exceptions import (
    YarnError,               # base
    AuthenticationError,     # 401 — invalid API key
    PermissionError,         # 403 — missing researcher role
    InsufficientCredits,     # 402 — not enough balance
    NotFoundError,           # 404
    QuotaExceeded,           # 413 — storage quota
    RateLimited,             # 429
    ServiceUnavailable,      # 502/503
)

try:
    job = yarn.submit_job(script="train.py")
except InsufficientCredits as e:
    print(f"Need ${e.required:.2f}, have ${e.available:.2f}")
except YarnError as e:
    print(f"Error: {e.message} (HTTP {e.status_code})")
```

All API errors raise typed exceptions with `status_code`, `message`, and the raw response body.

## API endpoint mapping

| SDK method | HTTP | Endpoint |
|---|---|---|
| `client.sessions.create()` | POST | `/v1/research/sessions` |
| `client.sessions.get(id)` | GET | `/v1/research/sessions/{id}` |
| `client.sessions.list()` | GET | `/v1/research/sessions` |
| `client.sessions.delete(id)` | DELETE | `/v1/research/sessions/{id}` |
| `client.jobs.submit()` | POST | `/v1/research/jobs` |
| `client.jobs.get(id)` | GET | `/v1/research/jobs/{id}` |
| `client.jobs.list()` | GET | `/v1/research/jobs` |
| `client.jobs.cancel(id)` | DELETE | `/v1/research/jobs/{id}` |
| `client.jobs.logs(id)` | GET | `/v1/research/jobs/{id}/logs` |
| `client.jobs.estimate()` | POST | `/v1/research/jobs/estimate` |
| `client.notebooks.create()` | POST | `/v1/research/notebooks` |
| `client.notebooks.get(id)` | GET | `/v1/research/notebooks/{id}` |
| `client.notebooks.list()` | GET | `/v1/research/notebooks` |
| `client.notebooks.delete(id)` | DELETE | `/v1/research/notebooks/{id}` |
| `client.secrets.create()` | POST | `/v1/research/secrets` |
| `client.secrets.list()` | GET | `/v1/research/secrets` |
| `client.secrets.delete(name)` | DELETE | `/v1/research/secrets/{name}` |
| `client.storage.upload()` | POST | `/v1/data/objects` |
| `client.storage.download(key)` | GET | `/v1/data/objects/{key}` |
| `client.storage.list()` | GET | `/v1/data/objects` |
| `client.storage.delete(key)` | DELETE | `/v1/data/objects/{key}` |
| `client.storage.usage()` | GET | `/v1/data/usage` |
| `client.storage.quota()` | GET | `/v1/data/quota` |
| `client.storage.datasets()` | GET | `/v1/data/datasets` |
| `client.billing.balance()` | GET | `/v1/billing/credits/balance` |
| `client.billing.usage()` | GET | `/v1/billing/credits/usage` |
| `client.billing.packs()` | GET | `/v1/billing/credits/packs` |
| `client.chat.completions()` | POST | `/v1/chat/completions` |
| `client.chat.models()` | GET | `/v1/models` |

## Request/response schemas

Full schemas for every endpoint are documented in the exploration agent's output. Key ones for implementation:

### Job submission request
```json
{
  "name": "str (1-63, K8s-safe)",
  "code": "str (base64-encoded)",
  "entrypoint": "str (default: python /home/ray/code/main.py)",
  "image": "str (allowlist: ghcr.io/jordanlochhill/ray-base:*)",
  "gpu_count": "int (0-8, default: 1)",
  "cpu": "str (default: 2)",
  "memory": "str (default: 4Gi)",
  "env_vars": "dict[str, str]",
  "description": "str (max 500)",
  "max_runtime_hours": "float (0-24)",
  "pip_packages": "list[str]",
  "secret_names": "list[str]"
}
```

### Session response
```json
{
  "id": "str",
  "session_name": "str",
  "status": "created | running | failed",
  "gpu_type": "rtx-4090",
  "gpu_count": 1,
  "ray_endpoint": "ray://{name}.research.prosodylabs.com.au:443",
  "dashboard_url": "https://{name}-dash.research.prosodylabs.com.au",
  "created_at": "ISO datetime",
  "idle_timeout_minutes": 60
}
```

### Credit balance response
```json
{
  "balance": 142.50,
  "available": 97.50,
  "held": 45.00,
  "total_credited": 500.00,
  "total_debited": 357.50,
  "transaction_count": 23,
  "currency": "AUD",
  "active_holds": [],
  "recent_transactions": []
}
```

## Implementation notes

1. **`gpu` parameter mapping:** The SDK accepts friendly names like `"rtx-4090"`, `"a100"`, `"h100"`. Map to the exact strings the API expects. Default is `"rtx-4090"`.

2. **Polling in wait():** Use exponential backoff starting at 2s, capping at 30s. Log status transitions to stderr if verbose mode is on.

3. **Session context manager:** `__enter__` calls `create()` then `wait_until_ready()`. `__exit__` calls `stop()` regardless of exception. If creation fails, don't call stop.

4. **File reading in submit_job:** Read the file, base64-encode, and send as the `code` field. The SDK does the encoding — the researcher passes a file path.

5. **Streaming logs:** `stream_logs()` returns a generator that polls `/jobs/{id}/logs?tail=N` every 2s and yields new lines. Stops when job status is terminal.

6. **Storage upload:** Use `multipart/form-data` with the file. Set `key` query param to the remote path.

7. **Thread safety:** `Client` instances should be thread-safe (each request is independent). Don't share mutable state.

8. **Timeouts:** Default request timeout 30s for normal calls, 300s for uploads/downloads, configurable via `client.timeout`.

## pyproject.toml

```toml
[project]
name = "yarn"
version = "0.1.0"
description = "Python SDK for Yarn — sovereign GPU compute"
requires-python = ">=3.10"
dependencies = ["requests>=2.28"]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["src/yarn"]
```
