TPU agent suite

Autonomous TPU development and onboarding agents for Google Cloud TPUs.

This repository builds and distributes TPU agent suite packages: declarative bundles of multi-agent skills, routing protocols, and tools that compile and install natively across major agent harnesses—including Claude Code and Antigravity, with Codex and Cursor coming soon.

Users talk to the tpu-agent skill. It routes each request to a specialist: onboarding, profiling, a hardware check, cited TPU facts, and kernels. See the design docs for how the framework and the suite are built.

• Python 3.12+ and the host-side packages the agents run with, XProf among them, in a virtualenv (pip install -r package/requirements-client.txt; see docs/install.md). PyYAML alone builds and installs the package, but leaves the agents without XProf.
• Standalone CLI tools (backend-dependent):
  • gcloud (Google Cloud CLI) for Cloud TPU VMs, GCS staging, and GCP credentials (gcloud auth login && gcloud auth application-default login)
  • kubectl (Kubernetes CLI) and gke-gcloud-auth-plugin (gcloud components install gke-gcloud-auth-plugin) for GKE Pod, Job and XPK workloads (gcloud container clusters get-credentials <CLUSTER_NAME> --region=<REGION>)
  • xpk (Accelerated Processing Kit) for multi-node GKE cluster workloads

The installer checks only the host-side Python packages. The first session sets up your backend and checks its tools, credentials and target (see docs/install.md).

Compile the canonical agent package into native distributions for every supported harness:

make build

This writes dist/<harness>/tpu-agent-suite/ for claude-code and antigravity.

See docs/install.md (headless / CI).

Ensure your virtual environment is activated whenever you use this agentic suite. Once inside the session, the tpu-agent skill is activated by any relevant natural language request, so no explicit invocation is needed. For example:

• Onboarding & Benchmarking:

  "Onboard Qwen/Qwen3.8-27B to run as a vLLM inference server on Cloud TPU using the torch_tpu backend. Send it a test completion request, then benchmark it and report TTFT, ITL and throughput in a table."

• Profiling:

  "Analyze the XProf profile at /path/to/profile and identify memory bottlenecks and TPU compute utilization."

• Kernel Optimization:

  "Optimize this custom Pallas flash-attention kernel for TPU v6e and run a self-refinement benchmark loop."

What each request routes to, where workloads run, and where outputs land are in docs/using.md.

Evaluations are orchestrated using Harbor and configured under eval/:

# Set up evaluation virtual environment
./eval/setup.sh
source eval/.venv/bin/activate

# Verify evaluation configuration and dataset schemas
./eval/verify.sh smoke-check

# Run evaluation with Harbor
export GEMINI_API_KEY="your-api-key"
./eval/run.py --config claude-code/agent-suite --dataset tpu-hello-world

Automated PR Evaluations: Evaluations are automatically triggered via GitHub Actions whenever a pull request is opened or updated that modifies relevant parts of the agent packages or evaluation datasets.

For the full list of configs, credentials, and the interactive sandbox, see eval/README.md.

All code contributions, tests, and agent packages live in this repository. Every pull request runs the following GitHub Actions workflows (defined in .github/workflows/):

• make verify: manifest, delegation graph, schema, and knowledge base checks.
• make test: framework and package test suites (translation parity per harness).
• make build: distributions for all harnesses compile cleanly.
• make lint: advisory content checks, reported as annotations.
• pre-commit: formatting and static checks (.pre-commit-config.yaml).

Local setup for these checks, and the short path for common contributions, is in docs/developing.md. Note that external contributions are not being accepted at this time (see CONTRIBUTING.md).

This project is eligible for the Google Open Source Software Vulnerability Rewards Program (OSS VRP).

If you suspect a security vulnerability in this project, please report it via g.co/vulnz. Please do not report security vulnerabilities through public GitHub issues.