Metadata-Version: 2.4
Name: tpu-runner
Version: 0.1.0
Summary: Minimal Google Cloud TPU job runner
License-Expression: MIT
Keywords: gcp,google-cloud,tpu,orchestration
Classifier: Environment :: Console
Classifier: Intended Audience :: Developers
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: System :: Distributed Computing
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: google-cloud-firestore>=2.16
Requires-Dist: PyYAML>=6.0
Provides-Extra: release
Requires-Dist: build<2,>=1; extra == "release"
Requires-Dist: twine<7,>=6; extra == "release"
Dynamic: license-file

# TPU Runner

`tpu-runner` is a job orchestrator designed to **maximize utilization of your Google Cloud TPU allocation.** 

It automatically **races capacity** requests across compatible TPU types and regions, **assigns jobs** by priority + FIFO, and **retries jobs** interrupted by Spot preemptions. It also **minimizes inter-region transfer costs**, cleans the workspace for new jobs, and lets related jobs reuse local caches (e.g. for previously compiled XLA artifacts or software environments).

We use Firestore for queue and orchestration state, GCS for source bundles and job artifacts, and a Cloud Run controller to manage TPU capacity and execution.

## Install

Install `tpu-runner` as a standalone CLI:

```bash
uv tool install tpu-runner
```

For local development, run this from the repository root:

```bash
uv tool install --editable .
```

Configure `gcloud`:

```bash
gcloud auth login
gcloud auth application-default login
gcloud config set project YOUR_PROJECT_ID
gcloud alpha compute tpus --help >/dev/null
```

## Set up the runner

Create an example deployment:

```bash
tpu-runner init
```

Edit `deployment.yaml` with your project ID, existing Secret Manager secret names, and the TPU types, zones, maximum counts, runtime versions, and chip limits you are willing to use. Counts are provisioning ceilings; TPU Runner scales capacity up and down with demand and keeps idle capacity only when `keep_warm` is enabled.

The default `ssh_transport: direct` creates public-IP TPUs; set it to `iap` to create private-IP TPUs and reach them through IAP tunnels instead.

Validate and deploy:

```bash
tpu-runner validate-fleet deployment.yaml
tpu-runner deploy deployment.yaml
```

`deploy` enables the required APIs and creates or updates the runner bucket, Firestore database, service accounts and IAM, SSH identity, worker startup script, controller image, and Cloud Run controller job.

## Submit work

In the root of the code you want to run, create `job.yaml`:

```yaml
jobs:
  - tpu: [v4-64, v6e-64]
    buckets:
      - gs://my-training-us-central2
      - gs://my-training-us-east1
    bundle: .
    priority: high
    caches:
      - key: pip
        path: .cache/pip
    env:
      PIP_CACHE_DIR: .cache/pip
      WANDB_PROJECT: my-project
    command: python3 -m pip install -r requirements.txt && touch "$PIP_CACHE_DIR/.ready" && python3 train.py --data "$JOB_BUCKET/data" --checkpoints "$CHECKPOINT_GCS_DIR"
```

- `id` is optional. When omitted, submission generates and prints one; use that
  ID with `watch`, `logs`, and `cancel`.
- `tpu` may contain one or several compatible TPU types to race.
- Omit `zone` and `tpu_name` to race all compatible fleet capacity. Set `zone`
  to use one zone, or `tpu_name` to use one exact declared TPU.
- Create one listed bucket in each candidate region and mirror required data at
  the same object paths. For a single-region job, list one bucket. The winning
  region's bucket becomes `JOB_BUCKET`, and retries remain pinned to that
  region.
- `bundle` is a local directory relative to `job.yaml`. TPU Runner archives and
  uploads it; use `.tpu-runnerignore` to exclude files.
- `priority` may be `low`, `normal`, or `high`.
- Before a new job starts, TPU Runner clears the previous runner workspace and
  shared memory, then extracts the new bundle into a fresh directory. A directory
  declared under `caches` is preserved when it is marked `.ready` and the next job
  on that worker declares the same `key`. Configure the relevant tool, such as
  pip, to write to the cache path. Incomplete caches are discarded, and all
  caches disappear when the TPU is deleted. Use `CHECKPOINT_GCS_DIR` for
  durable checkpoints.
- `command` runs independently on every TPU worker. Other runner-provided
  variables include `JOB_ID`, `ATTEMPT_ID`, `TPU_WORKER_HOST`,
  `TPU_WORKER_COUNT`, `JOB_DIR`, and `ATTEMPT_GCS_DIR`.

Submit and watch the job:

```bash
tpu-runner validate-jobs job.yaml
tpu-runner submit job.yaml
tpu-runner watch JOB_ID
tpu-runner logs JOB_ID
tpu-runner cancel JOB_ID
```


Spot preemption and recognized infrastructure failures return a job to
`pending` with the same region and checkpoint directory. Application and setup
failures are terminal. TPU Runner schedules higher-priority jobs first. Within
each priority, it considers the most constrained jobs first and moves flexible
jobs to alternative idle TPUs when that allows more jobs to run.

Use `tpu-runner --help` for a list of all commands, or use a specific command such as `tpu-runner submit --help` for its options.

PRs and feature requests are welcome. A big thank you to the [Google TPU Research Cloud (TRC) program](https://sites.research.google/trc/) for inspiring this work.
