---
name: reserve-tpu
description: Reserve an Iris-backed TPU worker for fast debugging with dev_tpu.py.
---

# Skill: Dev TPU

Use this skill for the standard fast TPU debugging loop without wiring a full training job each time.

`scripts/iris/dev_tpu.py` reserves a TPU-backed worker through Iris, waits for the worker VM to come up, and lets you SSH into it or run commands directly against it. It uses `gcloud` SSH and SCP against the worker Iris assigned to the holder job. There is no persistent `~/.ssh/config` alias.

## Critical concurrency rule

Run at most one TPU job at a time on a given dev TPU VM. Do not launch concurrent TPU commands on the same worker from multiple shells, tmux panes, or background jobs.

## Commands

- `allocate`: submit a TPU holder job, resolve the assigned worker VM(s), optionally sync the repo, block until release
- `status`: show the active local session metadata
- `connect`: open an interactive SSH session to one reserved worker
- `setup_env`: sync the repo by default, then install/refresh the remote `uv` environment on one or all reserved workers
- `execute`: sync local files to `~/marin` on the reserved worker(s), then run one command
- `watch`: sync all reserved workers and rerun a command on the selected worker when local files change
- `release`: terminate the holder job and remove the local session file

## Prerequisites

1. Authenticate to GCP and set up the repo.

```bash
gcloud auth login
gcloud config set project hai-gcp-models
gcloud auth application-default login
make dev_setup
```

2. Ensure the Iris controller is running for your cluster. On shared Marin clusters this is usually already true; only start it yourself for a fresh or local cluster.

3. Use a cluster config that can actually provision the TPU type you want.

## Command pattern

All invocations share this shape; only the subcommand and its flags change:

```bash
uv run scripts/iris/dev_tpu.py \
  --config lib/iris/config/marin.yaml \
  --tpu-name "$USER-v5p8" \
  <subcommand> [flags]
```

Subcommands and distinctive flags:

- `allocate --tpu-type v5p-8` — reserves a single-host TPU and holds it until `Ctrl-C`. `--tpu-type` is required (the config may expose many variants; the script does not guess). Add `--zone us-east5-b` to pin the holder job to a zone (the config must expose that TPU family there). Add `--no-setup-env` to skip remote env setup.
- `status` — show the active session.
- `connect` — interactive SSH to the reserved worker. After connecting: `cd ~/marin && source ~/.local/bin/env`.
- `setup_env` — install/refresh the remote env. Run before the first `execute`/`watch` if you allocated with `--no-setup-env`.
- `execute -- <cmd>` — sync local files, then run `<cmd>`. Add `--no-sync` for a no-sync inner loop. Example: `execute -- uv run --package levanter --group test pytest lib/levanter/tests/kernels/test_pallas_fused_cross_entropy_loss.py`.
- `watch -- <cmd>` — rerun `<cmd>` on local file changes.
- `release` — terminate the holder job and clear the session file.

## Worker Selection

Multi-host TPU types reserve more than one worker VM. Use `--worker <index>` with `connect`, `execute`, or `watch` to target a specific worker (`execute` and `watch` default to worker `0`):

```bash
uv run scripts/iris/dev_tpu.py --config lib/iris/config/marin.yaml \
  --tpu-name "$USER-v5p16" connect --worker 1
```

## Observability

Use normal Iris tooling to inspect the backing cluster and holder job:

```bash
uv run iris --config=lib/iris/config/marin.yaml cluster dashboard
uv run iris --config=lib/iris/config/marin.yaml cluster vm status
uv run iris --config=lib/iris/config/marin.yaml job list --prefix /$USER/dev-tpu
uv run iris --config=lib/iris/config/marin.yaml job logs /$USER/dev-tpu-<name>
```

If worker bootstrap fails:

```bash
uv run iris --config=lib/iris/config/marin.yaml cluster vm logs <worker-id>
```

## Session behavior

- Local session state lives under `~/.cache/marin/dev_tpu_iris/`.
- If the `allocate` terminal dies unexpectedly, use `release` to terminate the holder job and clear the stale state file.
- `execute` and `watch` already wrap the remote command in `bash -lc`; do not pass your own `bash -c`.

## Agent Usage

Always pass `--tpu-name` to avoid collisions with other agents:

```bash
export TPU_NAME="${USER}-$(git rev-parse --abbrev-ref HEAD | tr '/' '-')"
uv run scripts/iris/dev_tpu.py --config lib/iris/config/marin.yaml --tpu-name "$TPU_NAME" allocate --tpu-type v5p-8
```

## Cleanup

Normal cleanup is `Ctrl-C` in the `allocate` terminal. To clean up from another shell, run the `release` subcommand.
