tpu-runner 0.1.0__tar.gz

tpu_runner-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 David Hidary
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tpu_runner-0.1.0/PKG-INFO

@@ -0,0 +1,140 @@
+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.

tpu_runner-0.1.0/README.md

@@ -0,0 +1,116 @@
+# 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.

tpu_runner-0.1.0/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "tpu-runner"
+version = "0.1.0"
+description = "Minimal Google Cloud TPU job runner"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.11"
+dependencies = [
+    "google-cloud-firestore>=2.16",
+    "PyYAML>=6.0",
+]
+keywords = ["gcp", "google-cloud", "tpu", "orchestration"]
+classifiers = [
+    "Environment :: Console",
+    "Intended Audience :: Developers",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: System :: Distributed Computing",
+]
+
+[project.optional-dependencies]
+release = ["build>=1,<2", "twine>=6,<7"]
+
+[project.scripts]
+tpu-runner = "tpu_runner.cli:main"
+
+[tool.setuptools.packages.find]
+include = ["tpu_runner*"]
+
+[tool.setuptools.package-data]
+tpu_runner = ["*.sh", "*.txt", "*.yaml", "Dockerfile"]

tpu_runner-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tpu_runner-0.1.0/tpu_runner/Dockerfile

@@ -0,0 +1,10 @@
+FROM gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
+
+WORKDIR /app
+COPY tpu_runner/controller-requirements.txt ./
+COPY tpu_runner ./tpu_runner
+
+RUN python3 -m pip install --no-cache-dir --break-system-packages \
+    -r controller-requirements.txt
+
+ENTRYPOINT ["python3", "-m", "tpu_runner.cli"]

tpu_runner-0.1.0/tpu_runner/__init__.py

@@ -0,0 +1,18 @@
+"""GCP TPU fleet and job orchestration helpers."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .specs import FleetSpec, JobSpec, load_fleet_spec, load_job_specs
+
+try:
+    __version__ = version("tpu-runner")
+except PackageNotFoundError:  # Source tree imported without installation.
+    __version__ = "0+unknown"
+
+__all__ = [
+    "FleetSpec",
+    "JobSpec",
+    "__version__",
+    "load_fleet_spec",
+    "load_job_specs",
+]

tpu_runner-0.1.0/tpu_runner/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+
+raise SystemExit(main())

tpu_runner-0.1.0/tpu_runner/capacity_policy.py

@@ -0,0 +1,248 @@
+"""Pure placement and capacity policy."""
+
+from __future__ import annotations
+
+from .gcp import generated_resource_names
+from .region_pools import region_is_in_pool
+from .runtime import JobRecord, ResourceRecord
+from .specs import FleetSpec, TPUEntry, region_from_zone
+
+
+JOB_PRIORITY_RANK = {"low": 0, "normal": 1, "high": 2}
+
+
+def pending_job_accepts_entry(job: JobRecord, entry) -> bool:
+    if job.status != "pending":
+        return False
+    if job.spec.zone and job.spec.zone != entry.zone:
+        return False
+    entry_region = region_from_zone(entry.zone)
+    if not job.spec.accepts_region(entry_region):
+        return False
+    if job.spec.storage_region and (
+        not job.spec.region or not region_is_in_pool(entry_region, job.spec.region)
+    ):
+        return False
+    if job.spec.tpu_name:
+        if entry.adopted:
+            return (
+                job.spec.accepts_tpu_type(entry.type)
+                and job.spec.tpu_name == entry.existing
+            )
+        return job.spec.accepts_tpu_type(entry.type) and any(
+            generated_resource_names(entry, ordinal)[1] == job.spec.tpu_name
+            for ordinal in range(1, entry.count + 1)
+        )
+    return (entry.adopted or entry.count > 0) and job.spec.accepts_tpu_type(
+        entry.type
+    )
+
+
+def pending_job_accepts_resource(job: JobRecord, resource: ResourceRecord) -> bool:
+    """Match compatible idle capacity while preserving exact TPU pins."""
+
+    if job.status != "pending":
+        return False
+    if not job.spec.accepts_tpu_type(resource.tpu_type):
+        return False
+    if job.spec.zone and job.spec.zone != resource.zone:
+        return False
+    resource_region = region_from_zone(resource.zone)
+    if not job.spec.accepts_region(resource_region):
+        return False
+    if job.spec.storage_region and (
+        not job.spec.region or not region_is_in_pool(resource_region, job.spec.region)
+    ):
+        return False
+    return not job.spec.tpu_name or job.spec.tpu_name == resource.tpu_name
+
+
+def resource_is_busy(resource: ResourceRecord) -> bool:
+    return bool(
+        resource.status == "busy"
+        or resource.current_job_id
+        or resource.current_attempt_id
+    )
+
+
+def pending_job_entry_constraint_key(job: JobRecord, entries: list) -> tuple:
+    compatible = sum(pending_job_accepts_entry(job, entry) for entry in entries)
+    return (
+        compatible == 0,
+        -JOB_PRIORITY_RANK[job.spec.priority],
+        compatible,
+        job.submitted_at,
+        job.spec.id,
+    )
+
+
+def pending_job_resource_constraint_key(
+    job: JobRecord,
+    resources: list[ResourceRecord],
+) -> tuple:
+    compatible = sum(
+        pending_job_accepts_resource(job, resource) for resource in resources
+    )
+    return (
+        compatible == 0,
+        -JOB_PRIORITY_RANK[job.spec.priority],
+        compatible,
+        job.submitted_at,
+        job.spec.id,
+    )
+
+
+def plan_idle_assignments(
+    jobs: list[JobRecord],
+    resources: list[ResourceRecord],
+    *,
+    entries: list[TPUEntry] | None = None,
+) -> list[tuple[JobRecord, ResourceRecord]]:
+    """Match idle resources without displacing an earlier pending job.
+
+    Jobs are considered by priority, constraint count, and FIFO order. An
+    augmenting path may move an earlier flexible job to another idle resource,
+    but it never drops that job merely to fit a later one.
+    """
+
+    idle_resources = sorted(
+        (
+            resource
+            for resource in resources
+            if resource.status == "idle" and not resource_is_busy(resource)
+        ),
+        key=lambda resource: (not resource.adopted, resource.id),
+    )
+    resources_by_id = {resource.id: resource for resource in idle_resources}
+    ordered_jobs = sorted(
+        (job for job in jobs if job.status == "pending"),
+        key=lambda job: (
+            pending_job_entry_constraint_key(job, entries)
+            if entries is not None
+            else pending_job_resource_constraint_key(job, idle_resources)
+        ),
+    )
+    matched_by_resource: dict[str, JobRecord] = {}
+
+    def match(job: JobRecord, visited: set[str]) -> bool:
+        candidates = sorted(
+            (
+                resource
+                for resource in idle_resources
+                if resource.id not in visited
+                and pending_job_accepts_resource(job, resource)
+            ),
+            key=lambda resource: (
+                resource.id in matched_by_resource,
+                not resource.adopted,
+                resource.id,
+            ),
+        )
+        for resource in candidates:
+            visited.add(resource.id)
+            previous = matched_by_resource.get(resource.id)
+            if previous is None or match(previous, visited):
+                matched_by_resource[resource.id] = job
+                return True
+        return False
+
+    for job in ordered_jobs:
+        match(job, set())
+
+    resources_by_job = {
+        job.spec.id: resources_by_id[resource_id]
+        for resource_id, job in matched_by_resource.items()
+    }
+    return [
+        (job, resources_by_job[job.spec.id])
+        for job in ordered_jobs
+        if job.spec.id in resources_by_job
+    ]
+
+
+def allocate_managed_pending_demand(
+    jobs: list[JobRecord],
+    *,
+    fleet: FleetSpec,
+    resources: tuple[ResourceRecord, ...],
+) -> dict[str, int]:
+    """Request every pending job in every compatible Spot entry.
+
+    These are capacity races, not duplicate executions. Assignment remains a
+    Firestore transaction, and the losing entry's demand disappears on the
+    next reconciliation. Every physical entry count remains a hard ceiling.
+    """
+
+    entries = [entry for entry in fleet.tpus if not entry.adopted]
+    busy = {
+        entry.id: sum(
+            1
+            for resource in resources
+            if resource.fleet_entry_id == entry.id
+            and not resource.adopted
+            and resource_is_busy(resource)
+        )
+        for entry in entries
+    }
+    demand = {entry.id: 0 for entry in entries}
+    for job in sorted(
+        jobs,
+        key=lambda candidate: pending_job_entry_constraint_key(candidate, entries),
+    ):
+        for entry in entries:
+            available = max(0, entry.count - busy[entry.id])
+            if demand[entry.id] >= available:
+                continue
+            if pending_job_accepts_entry(job, entry):
+                demand[entry.id] += 1
+    return demand
+
+
+def desired_managed_capacity_counts(
+    jobs: list[JobRecord],
+    *,
+    fleet: FleetSpec,
+    resources: tuple[ResourceRecord, ...],
+) -> dict[str, int]:
+    """Return busy plus raced pending demand, capped by physical ceilings."""
+
+    idle_adopted = [
+        resource
+        for resource in resources
+        if resource.adopted and resource.status == "idle"
+    ]
+    adopted_job_ids = {
+        job.spec.id
+        for job, _ in plan_idle_assignments(
+            jobs,
+            idle_adopted,
+            entries=list(fleet.tpus),
+        )
+    }
+    jobs_requiring_managed = [
+        job
+        for job in jobs
+        if job.status == "pending" and job.spec.id not in adopted_job_ids
+    ]
+
+    pending_demand = allocate_managed_pending_demand(
+        jobs_requiring_managed,
+        fleet=fleet,
+        resources=resources,
+    )
+    desired: dict[str, int] = {}
+    for entry in fleet.tpus:
+        if entry.adopted:
+            continue
+        busy_count = sum(
+            1
+            for resource in resources
+            if resource.fleet_entry_id == entry.id
+            and not resource.adopted
+            and resource_is_busy(resource)
+        )
+        desired[entry.id] = min(
+            entry.count,
+            max(int(entry.keep_warm), busy_count + pending_demand.get(entry.id, 0)),
+        )
+    return desired