Paper-matrix harness, Tunix backend, GCP scaffolding

.github/workflows/ci.yml

@@ -0,0 +1,29 @@
+name: ci
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v3
+        with:
+          enable-cache: true
+      - name: Set up Python
+        run: uv python install 3.11
+      - name: Sync deps
+        run: uv sync --extra dev --extra ot
+      - name: Lint
+        run: uv run ruff check src tests scripts
+      - name: Format check
+        run: uv run ruff format --check src tests scripts
+      - name: Pyright
+        run: uv run pyright
+        continue-on-error: true
+      - name: Test
+        run: uv run pytest -q -m "not slow"

.gitignore

@@ -10,4 +10,40 @@ __pycache__/
 .pyright/
 .DS_Store
 *.orbax-checkpoint-tmp*
-.env*
+.env*
+
+# Build / dist
+/build/
+/dist/
+*.whl
+
+# Coverage
+.coverage
+.coverage.*
+coverage.xml
+htmlcov/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Logs
+*.log
+
+# Python tooling
+.python-version
+.tox/
+.nox/
+
+# GCP credentials — never commit service-account keys
+*-credentials.json
+*-key.json
+service-account*.json
+gcp-key*.json
+.gcloud/
+application_default_credentials.json

.pre-commit-config.yaml

@@ -0,0 +1,15 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.6.9
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-added-large-files
+        args: ["--maxkb=1024"]

AGENTS.md

@@ -15,6 +15,22 @@
 - Tests: `uv run pytest -q`. Mark slow paths with `@pytest.mark.slow`.
 - No emojis or decorative comments. Comments only when *why* is non-obvious.
 
+## Make targets
+
+- `make setup` / `make setup-all` — sync deps (core / + hf, teacher, gcp).
+- `make fmt` / `make lint` / `make type` / `make test` — the usual.
+- `make smoke` — `scripts/smoke_all_methods.sh`.
+- `make matrix-local` — run the paper matrix on CPU fixtures + aggregate.
+- `make teacher-data` — generate the full replay shards (needs API keys).
+- `make docker-cpu` / `make docker-tpu` — build deploy images locally.
+
+## `deploy/` boundary
+
+Anything under `deploy/` is launchers, Dockerfiles, and Cloud Build config — **no
+algorithm code**. The CPU image runs CI smoke; the TPU image is the actual training
+target. `deploy/launch_tpu.sh` and `deploy/launch_vertex.py` are the only entry
+points; both honor a `DRY_RUN=1` env to print the underlying `gcloud` invocations.
+
 ## Common workflows
 
 Generate teacher data into a replay shard:

Makefile

@@ -0,0 +1,51 @@
+.PHONY: setup fmt lint type test smoke matrix-local matrix-tpu teacher-data clean docker-cpu docker-tpu
+
+UV ?= uv
+PY ?= $(UV) run python
+OUT ?= artifacts/results/smoke
+SEEDS ?= 2
+BACKEND ?= local_nano
+TASKS ?= math
+
+setup:
+	$(UV) sync --extra dev --extra ot
+
+setup-all:
+	$(UV) sync --extra dev --extra ot --extra hf --extra teacher --extra gcp
+
+fmt:
+	$(UV) run ruff format src tests scripts
+
+lint:
+	$(UV) run ruff check src tests scripts
+
+type:
+	$(UV) run pyright
+
+test:
+	$(UV) run pytest -q -m "not slow"
+
+test-all:
+	$(UV) run pytest -q
+
+smoke:
+	bash scripts/smoke_all_methods.sh
+
+matrix-local:
+	$(PY) scripts/run_paper_matrix.py --backend $(BACKEND) --tasks $(TASKS) --seeds $(SEEDS) --out $(OUT)
+	$(PY) scripts/aggregate_results.py $(OUT) --base B3 --out $(OUT)/table.md
+
+matrix-tpu:
+	bash deploy/launch_tpu.sh
+
+teacher-data:
+	bash scripts/gen_all_teacher_data.sh
+
+docker-cpu:
+	docker build -f deploy/Dockerfile.cpu -t cts:cpu .
+
+docker-tpu:
+	docker build -f deploy/Dockerfile.tpu -t cts:tpu .
+
+clean:
+	rm -rf artifacts/results/smoke artifacts/runs .pytest_cache .ruff_cache

README.md

@@ -17,6 +17,38 @@ JAX · Flax NNX · Optax · Orbax · OTT-JAX (entropic / unbalanced Sinkhorn) ·
 `LocalModelAPI` — the smoke / unit-test path (nano-LM, optional HF Flax). Lets every loss and method run on a laptop.
 `ScaleBackendAdapter` — the sharding-aware path. Two impls: `TunixAdapter` (M8) and `MaxTextAdapter` (M9). Methods are written against one side or the other; we deliberately do **not** force them through a single protocol.
 
+## Run the paper
+
+The full experiment matrix (B0–B5 × {gsm8k, apps} × seeds) is driven by one script:
+
+```bash
+# Local CPU smoke (fixtures, 2 seeds, math only)
+make matrix-local                                # → artifacts/results/smoke/table.md
+
+# Real run (after generating teacher data)
+make teacher-data                                # honors $TEACHER, $TRAIN_LIMIT, etc.
+uv run python scripts/run_paper_matrix.py \
+  --backend tunix --tasks both --seeds 5 \
+  --max-steps 2000 --out artifacts/results/paper/
+uv run python scripts/aggregate_results.py \
+  artifacts/results/paper/ --base B3 --out artifacts/results/paper/table.md
+```
+
+The matrix orchestrator asserts the fairness fingerprint matches across cells before
+launching a single training job (see `cts.train.fairness.assert_consistent`).
+
+## Deploy to GCP
+
+See [deploy/README.md](deploy/README.md). TL;DR:
+
+```bash
+gcloud builds submit --config deploy/cloudbuild.yaml .          # build + push images
+PROJECT=… ZONE=… TPU_NAME=… IMAGE_TAG=… GCS_BUCKET=… \
+  bash deploy/launch_tpu.sh                                     # reserve TPU + run matrix
+```
+
+`DRY_RUN=1` prints the `gcloud` invocations without executing.
+
 ## Quickstart
 
 ```bash

deploy/Dockerfile.cpu

@@ -0,0 +1,23 @@
+# CPU image — used by CI smoke and Vertex eval jobs.
+FROM python:3.11-slim
+
+ENV PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    UV_SYSTEM_PYTHON=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        git curl ca-certificates build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN pip install --no-cache-dir uv==0.4.18
+
+WORKDIR /workspace
+COPY pyproject.toml uv.lock README.md ./
+COPY src ./src
+COPY tests ./tests
+COPY scripts ./scripts
+
+RUN uv sync --frozen --extra dev --extra ot --extra teacher --extra gcp
+
+ENV PATH="/workspace/.venv/bin:${PATH}"
+CMD ["bash"]

deploy/Dockerfile.tpu

@@ -0,0 +1,20 @@
+# TPU image — used for actual training runs on Cloud TPU VMs.
+# Base matches Tunix's required JAX/libtpu pin; bump together when upgrading JAX.
+FROM us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.34-rev1
+
+ENV PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1
+
+RUN pip install --no-cache-dir uv==0.4.18
+
+WORKDIR /workspace
+COPY pyproject.toml uv.lock README.md ./
+COPY src ./src
+COPY scripts ./scripts
+COPY deploy ./deploy
+
+# --extra scale pulls Tunix; --extra ot pulls OTT-JAX for the OT backend.
+# Use --inexact-resolve so JAX from the base image isn't replaced.
+RUN uv pip install --system -e ".[ot,scale,gcp,teacher]"
+
+CMD ["python", "-m", "cts.train.main"]

deploy/README.md

@@ -0,0 +1,66 @@
+# Deploying CTS to GCP
+
+Five commands. Nothing here runs automatically — invoke when you're ready.
+
+## 0. One-time setup
+
+```bash
+gcloud auth login
+gcloud config set project $PROJECT
+gcloud artifacts repositories create cts \
+  --repository-format=docker --location=us-central1
+gsutil mb gs://my-cts-runs
+```
+
+## 1. Build & push images
+
+```bash
+gcloud builds submit --config deploy/cloudbuild.yaml \
+  --substitutions=_TAG=$(git rev-parse --short HEAD) .
+```
+
+Pushes both `…/cts/cpu:<sha>` and `…/cts/tpu:<sha>`.
+
+## 2. Launch the paper matrix on a TPU
+
+```bash
+PROJECT=my-proj ZONE=us-east5-a TPU_NAME=cts-run-1 \
+IMAGE_TAG=$(git rev-parse --short HEAD) GCS_BUCKET=gs://my-cts-runs \
+TPU_TYPE=v5p-8 bash deploy/launch_tpu.sh
+```
+
+Use `DRY_RUN=1` to print the `gcloud` invocations without executing.
+
+Alternative (Vertex AI, with managed retries + Tensorboard):
+
+```bash
+uv run python deploy/launch_vertex.py \
+  --project $PROJECT --region us-central1 \
+  --image us-central1-docker.pkg.dev/$PROJECT/cts/tpu:latest \
+  --bucket gs://my-cts-runs --tpu-type v5litepod-8 \
+  -- --tasks both --seeds 5 --max-steps 2000
+```
+
+## 3. Pull & aggregate results
+
+```bash
+gsutil -m rsync -r gs://my-cts-runs/runs/cts-run-1/ artifacts/results/cts-run-1/
+uv run python scripts/aggregate_results.py artifacts/results/cts-run-1/ \
+  --base B3 --out artifacts/results/cts-run-1/table.md
+```
+
+## 4. Tear down
+
+```bash
+gcloud alpha compute tpus queued-resources delete cts-run-1 \
+  --project=$PROJECT --zone=$ZONE --force
+```
+
+## Notes
+
+- No Terraform / GKE / secrets manager intentionally. A single SA + manual
+  `gcloud auth` is enough for this research repo.
+- The TPU image base pin (`jax0.4.34-rev1`) must match the Tunix `[scale]`
+  extra. Bump them together.
+- Algorithm code never lives under `deploy/`; this dir is launchers + Dockerfiles
+  only.

deploy/cloudbuild.yaml

@@ -0,0 +1,40 @@
+# Build + push CPU and TPU images to Artifact Registry.
+# Trigger: gcloud builds submit --config deploy/cloudbuild.yaml .
+substitutions:
+  _REGION: us-central1
+  _REPO: cts
+  _TAG: ${SHORT_SHA}
+
+steps:
+  - id: build-cpu
+    name: gcr.io/cloud-builders/docker
+    args:
+      - build
+      - -f
+      - deploy/Dockerfile.cpu
+      - -t
+      - ${_REGION}-docker.pkg.dev/${PROJECT_ID}/${_REPO}/cpu:${_TAG}
+      - -t
+      - ${_REGION}-docker.pkg.dev/${PROJECT_ID}/${_REPO}/cpu:latest
+      - .
+  - id: build-tpu
+    name: gcr.io/cloud-builders/docker
+    args:
+      - build
+      - -f
+      - deploy/Dockerfile.tpu
+      - -t
+      - ${_REGION}-docker.pkg.dev/${PROJECT_ID}/${_REPO}/tpu:${_TAG}
+      - -t
+      - ${_REGION}-docker.pkg.dev/${PROJECT_ID}/${_REPO}/tpu:latest
+      - .
+  - id: push-cpu
+    name: gcr.io/cloud-builders/docker
+    args: [push, --all-tags, "${_REGION}-docker.pkg.dev/${PROJECT_ID}/${_REPO}/cpu"]
+  - id: push-tpu
+    name: gcr.io/cloud-builders/docker
+    args: [push, --all-tags, "${_REGION}-docker.pkg.dev/${PROJECT_ID}/${_REPO}/tpu"]
+
+options:
+  machineType: E2_HIGHCPU_8
+timeout: 1800s