RELEASING.md

Releasing tpu_sync

Versioning scheme

The base version is single-sourced from [project].version in pyproject.toml (read by the Bazel repo rule in bazel/wheel_version.bzl). A suffix from the WHEEL_VERSION_EXTRAS environment variable is appended at build time:

Channel	Version format	Built by	Published to
nightly	X.Y.Z.devYYYYMMDDHHMMSS	Nightly Wheels workflow (daily 08:00 UTC)	Artifact Registry
release	X.Y.Z	Release Wheels workflow (on vX.Y.Z tag)	PyPI + GitHub Release

Two wheels are built per version, one per framework: tpu_sync_jax and tpu_sync_torch. On PyPI they are the projects tpu-sync-jax and tpu-sync-torch.

The jax wheel's build tag is the JAX version it was built against, dots removed (tpu_sync_jax-X.Y.Z-0112-cp312-...whl for jax 0.11.2); see ci/build_wheel_impl.sh.

# JAX: self-contained, pulls the pinned jax/jaxlib/libtpu stack.
pip install tpu-sync-jax

# Torch: requires torch_tpu (and its matching torch) to be installed first,
# from the torch_tpu registry — see "The torch_tpu ABI pin" below for how the
# compatible torch_tpu version is recorded per release.
pip install --pre torch_tpu \
  --index-url "https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/simple/"
pip install tpu-sync-torch

Nightlies live in the tpu-sync Artifact Registry instead (default https://us-python.pkg.dev/cloud-tpu-inference-test/tpu-raiden/, overridable via the RAIDEN_REGISTRY_URL repository variable); pip's pre-release rules keep the channels apart, since .dev versions are only ever selected with --pre:

pip install --pre tpu_sync_jax \
  --extra-index-url "https://us-python.pkg.dev/cloud-tpu-inference-test/tpu-raiden/simple/"

How a release reaches PyPI

The Release Wheels workflow never talks to PyPI directly. It stages the wheels in the project's OSS exit-gate Artifact Registry repository and writes a publishing manifest that names exactly the packages and versions it built to the exit-gate trigger bucket. The gate verifies the staged files, publishes them to PyPI with its own credentials, removes them from the staging repository on success, and emails the outcome (with logs) to the project's notification address. There is no status file to poll: watch for the two emails (one when the release is accepted, one when it finishes), then check https://pypi.org/project/tpu-sync-jax/ and https://pypi.org/project/tpu-sync-torch/.

A failed or abandoned run leaves its files in the staging repository. The workflow's manifest only ever lists the versions it built, so leftovers are never published by a later run; a builder identity can delete them from the repository if they get in the way.

Nightly pre-releases on PyPI

The scheduled Nightly Wheels run also publishes its .dev wheels to PyPI through the exit gate (see .github/workflows/publish_pypi.yml), so pip install --pre tpu-sync-torch tracks main. A nightly skips any wheel above the PyPI file limit instead of failing; the registry still gets every wheel. PyPI's 10 GB project quota holds a few months of daily wheels; request more once the project exists.

The PyPI file-size limit

PyPI accepts files up to 100 MB per file unless the project has been granted a higher limit, and the exit gate uploads a release file by file, so a file over the limit fails the release after the smaller files are already on PyPI (and PyPI releases cannot be replaced). The workflow therefore refuses to stage any wheel above the limit it knows about. Every build job prints the wheel sizes in its summary. Both wheels are compiled with hidden symbol visibility (see ci/build_wheel_impl.sh), which keeps the jax wheel and the torch wheel with one extension per supported torch release under the limit; with default visibility either wheel exceeds it.

Should a wheel outgrow the limit, PyPI raises a project's limit on request once the project has at least one release under the current limit:

1. Run Release Wheels by hand with publish=no and read the wheel sizes off the job summary; a wheel must be under the limit to be published. The frameworks input narrows the run to the wheels that fit, and torch_abis set to a single release (2.11.0) builds the smallest torch wheel the build makes.
2. File a file-size limit increase with PyPI for each project, following https://docs.pypi.org/project-management/storage-limits/. Approval can take weeks.
3. Once granted, record the new limit in the PYPI_FILE_LIMIT_MB repository variable.

The torch_tpu pairing

tpu_sync_torch is ABI-coupled to torch_tpu: the extension compiles against torch_tpu's public tensor-buffer header and binds those symbols at import to the installed torch_tpu wheel, and both must resolve the same libtorch symbols. The build needs no torch_tpu checkout: the wheel-backed module under shims/torch_tpu takes the header and the XLA revision the extension must share with torch_tpu from the installed torch_tpu wheel, which ships both under torch_tpu/include/. The wheel ships one extension variant per torch release in RAIDEN_TORCH_ABIS (see ci/build_wheel_impl.sh), and the loader picks the variant matching the installed torch.

torch_tpu publishes nightlies only (torch_tpu==0.1.1.devYYYYMMDDHHMMSS), so every release records the torch_tpu wheel it was validated with:

• torch_tpu.version at the repo root names that torch_tpu wheel version; the build installs exactly that wheel, so it is the pairing at build time as well. Update it before tagging.
• The GitHub Release notes carry the same version, and the CHANGELOG entry lists it so users can install a compatible pair.

The JAX stack pin

tpu_sync_jax is version-coupled to jax/jaxlib/libtpu, but unlike the torch side this pin is enforced by pip itself: the exact versions are baked into the wheel's Requires-Dist metadata (from JAX_REQUIRES in ci/wheel/BUILD.bazel, kept in lockstep with pyproject.toml), so installing the wheel installs the matching stack. For visibility, both workflows also extract these pins from the built wheel into a jax_pins.txt beside it, and the GitHub Release notes list them next to the torch_tpu pin. Bumping the JAX stack for a release means updating JAX_REQUIRES and pyproject.toml together in the release PR.

Golden commits

A golden commit is a commit on main whose tpu_sync_torch nightly passed the two downstream canaries that consume the wheel. It is what a release is cut from and what the downstream repositories pin:

Record	Meaning
lkg.version	The newest golden commit (the stable checkout the README points users at).
verified_commits.csv	Every golden commit, oldest first: when it was verified, the commit, the tpu-sync-torch nightly version, and the downstream runs that passed.
golden-YYYYMMDD tags	The same commits, for git checkout.

The Golden Commit workflow (.github/workflows/golden_commit.yml, tooling in ci/golden/golden_commit.py) cuts one every Monday and on workflow_dispatch:

1. Candidate: the newest tpu-sync-torch .dev release on PyPI, or the tpu_sync_torch_version input, mapped to its commit through the nightly-manifest-<version> artifact the nightly leaves (older nightlies are matched by their start time).
2. vllm-torchtpu: a build of the Buildkite canary pipeline vllm-torchtpu-torchtpu-nightly with TORCHTPU_NIGHTLY_LEG=tpu-sync and TPU_SYNC_TORCH_VERSION_OVERRIDE=<candidate>: the candidate layered over the image built from vllm-torchtpu's current pins, every lane (tests, integration, perf and eval, verify). It must pass.
3. sglang-torchtpu: sglang-torchtpu runs Nightly Test (TPU - Remote Wheel) with tpu-sync-torch on a Monday 12:00 UTC schedule; sglang_canary finds and waits for that scheduled run read-only when SGLANG_GITHUB_TOKEN is unset, or dispatches a run with wheel=tpu-sync-torch==<candidate> when SGLANG_GITHUB_TOKEN is configured: feature validation and perf against the candidate layered over the latest nightly image. It must succeed.
4. Record: a pull request labeled pull ready that updates lkg.version and appends to verified_commits.csv; Copybara imports it into Piper and exports it back to main. The commit is tagged golden-YYYYMMDD.
5. Downstream pins (opt-in via bump_downstream=yes): pull requests on vllm-torchtpu (pyproject.toml) and sglang-torchtpu (python/pyproject_tpu.toml, requirements/tpu_stack_excludes.txt) that move tpu-sync-torch== to the candidate, for their owners to review.

A failed gate opens an issue in this repository with both runs; the previous golden commit stands. A hosted job ends after six hours, so a canary that is still running when its job times out is waited for by rerunning the workflow with vllm_build_number and sglang_run_id set from the first run (the issue names them). dry_run=yes runs the gates without recording anything, for example to check a candidate before a release.

Cutting a release

1. Pick the release candidate commit on main — normally the newest golden commit, published in lkg.version and recorded with its downstream runs in verified_commits.csv. To release a newer commit, run the Golden Commit workflow on its nightly first.

2. Open a release PR that:

   • updates torch_tpu.version to the torch_tpu wheel the release was validated with,
   • bumps [project].version in pyproject.toml to X.Y.Z,
   • adds the X.Y.Z section to CHANGELOG.md, including the compatible torch_tpu nightly version.

3. (Optional) Dry run: trigger the Release Wheels workflow via workflow_dispatch with publish=no — it builds .dev-versioned wheels as workflow artifacts without publishing anything.

4. Merge the PR, then tag and push:

   git tag vX.Y.Z <merge-commit>
   git push origin vX.Y.Z

   The Release Wheels workflow verifies the tag (must equal the pyproject version and be reachable from main), builds both wheels with an empty version suffix, stages them and triggers the exit gate, and creates the GitHub Release with the wheels attached.

5. Wait for the exit-gate emails and confirm both projects on pypi.org show X.Y.Z.

6. Bump [project].version on main to the next patch version X.Y.(Z+1). This keeps nightlies (X.Y.(Z+1).devN) sorting above the just-released X.Y.Z, so --pre users keep receiving fresh builds.

Prerequisites (one-time repo setup)

• The GitHub Actions runner service account needs roles/artifactregistry.writer on the nightly registry (or set RAIDEN_REGISTRY_URL to a registry it can write to), and must be registered with the OSS exit gate as a builder and triggerer for this project, which grants it write access to the staging repository and the manifest folder.
• PYPI_FILE_LIMIT_MB repository variable: the per-file limit PyPI has granted the project, once it differs from the default 100.
• The Golden Commit workflow uses up to four secrets: BUILDKITE_API_TOKEN (the vllm Buildkite organization, read_builds and write_builds), SGLANG_GITHUB_TOKEN (optional; actions:write on sgl-project/sglang-torchtpu to dispatch runs — when unset, sglang_canary matches and waits for sglang-torchtpu's Monday 12:00 UTC scheduled run read-only), GOLDEN_COMMIT_TOKEN (contents:write and pull_requests:write on this repository, so the record pull request gets CI and Copybara imports it) and DOWNSTREAM_PR_TOKEN (contents:write and pull_requests:write on the two downstream repositories, or on forks named by the DOWNSTREAM_FORK_OWNER repository variable). The header of the workflow file lists them.