From 22e9097b91c54b1d6a08c30df36bf96e4a87654e Mon Sep 17 00:00:00 2001 From: yinlin Date: Mon, 27 Jul 2026 10:17:28 -0700 Subject: [PATCH] ci: automate nightly and stable wheel releases PiperOrigin-RevId: 954696710 --- .github/workflows/nightly_build.yml | 222 ++++++++++++++++++ .github/workflows/release_wheels.yml | 236 ++++++++++++++++++++ .github/workflows/run_benchmarks.yml | 42 ---- .github/workflows/update_d2h_h2d_record.yml | 16 -- RELEASING.md | 93 ++++++++ ci/build_wheel.sh | 114 +--------- ci/build_wheel_impl.sh | 173 ++++++++++++++ torch_tpu.version | 1 + 8 files changed, 732 insertions(+), 165 deletions(-) create mode 100644 .github/workflows/nightly_build.yml create mode 100644 .github/workflows/release_wheels.yml delete mode 100644 .github/workflows/run_benchmarks.yml delete mode 100644 .github/workflows/update_d2h_h2d_record.yml create mode 100644 RELEASING.md create mode 100755 ci/build_wheel_impl.sh create mode 100644 torch_tpu.version diff --git a/.github/workflows/nightly_build.yml b/.github/workflows/nightly_build.yml new file mode 100644 index 00000000..716f0c38 --- /dev/null +++ b/.github/workflows/nightly_build.yml @@ -0,0 +1,222 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Nightly build of the tpu_raiden distributables: +# - tpu_raiden_jax wheel (+ _tpu_raiden_jax.so) +# - tpu_raiden_torch wheel (+ _tpu_raiden_host.so / _tpu_raiden_torch.so) +# +# Wheels are versioned .dev (same scheme as +# torch_tpu) and built inside the ml-build container via +# ci/build_wheel_impl.sh. Every run publishes the wheels and bare .so files as +# GitHub Actions artifacts; scheduled runs (and dispatches that opt in) +# additionally twine-upload the wheels to the raiden Artifact Registry, where +# they are installable with: +# +# pip install --pre tpu_raiden_jax \ +# --extra-index-url https://us-python.pkg.dev/cloud-tpu-inference-test/tpu-raiden/simple/ +# +# The torch wheel is ABI-coupled to torch_tpu: it is built against the +# torch_tpu checkout selected by torch_tpu_ref (default: main HEAD) and the +# exact torch that torch_tpu's per-Python requirements lock pins. The torch_tpu +# commit used is recorded in the job summary and in the +# torch_tpu_commit.txt file inside the wheel artifact. +# +# Required repository secrets/variables: +# secrets.TORCH_TPU_DEPLOY_KEY read-only deploy key for the private +# google-pytorch/torch_tpu repo +# vars.RAIDEN_REGISTRY_URL (optional) overrides the Artifact Registry +# upload URL + +name: Nightly Wheels + +on: + schedule: + - cron: "0 8 * * *" # 08:00 UTC daily + workflow_dispatch: + inputs: + torch_tpu_ref: + description: "torch_tpu ref to build the torch wheel against" + required: false + default: "main" + type: string + upload_to_registry: + description: "Upload wheels to the raiden Artifact Registry?" + type: choice + default: "no" + options: + - "yes" + - "no" + use_remote_cache: + description: "Use the RBE remote bazel cache (--config=ci)?" + type: choice + default: "yes" + options: + - "yes" + - "no" + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }} + cancel-in-progress: false + +env: + RAIDEN_REGISTRY_URL: ${{ vars.RAIDEN_REGISTRY_URL || 'https://us-python.pkg.dev/cloud-tpu-inference-test/tpu-raiden/' }} + +jobs: + version: + name: "Stamp nightly version" + runs-on: ubuntu-latest + outputs: + wheel_version_extras: ${{ steps.stamp.outputs.wheel_version_extras }} + steps: + - id: stamp + # One timestamp for the whole run so the jax and torch wheels of a + # single nightly share a version. + run: echo "wheel_version_extras=.dev$(date -u +%Y%m%d%H%M%S)" >> "$GITHUB_OUTPUT" + + build: + name: "Build ${{ matrix.framework }} wheel" + needs: version + strategy: + fail-fast: false + matrix: + include: + - framework: jax + with_torch: "0" + so_paths: | + tpu-raiden/tpu_raiden/frameworks/jax/_tpu_raiden_jax.so + - framework: torch + with_torch: "1" + so_paths: | + tpu-raiden/tpu_raiden/frameworks/torch/_tpu_raiden_host.so + tpu-raiden/tpu_raiden/frameworks/torch/_tpu_raiden_torch.so + runs-on: linux-x86-n2-32 + container: + image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest + timeout-minutes: 300 + steps: + - name: Checkout tpu-raiden + uses: actions/checkout@v5 + with: + path: tpu-raiden + - name: Checkout torch_tpu + if: matrix.with_torch == '1' + uses: actions/checkout@v5 + with: + repository: google-pytorch/torch_tpu + ref: ${{ inputs.torch_tpu_ref || 'main' }} + ssh-key: ${{ secrets.TORCH_TPU_DEPLOY_KEY }} + path: torch_tpu + # Git security patch (CVE-2022-24765) blocks root from operating on + # runner-owned directories. Tell Git explicitly to trust the workspace + # inside this ephemeral container. + - name: Fix Git Workspace Ownership + run: | + git config --global --add safe.directory "$GITHUB_WORKSPACE/tpu-raiden" + git config --global --add safe.directory "$GITHUB_WORKSPACE/torch_tpu" + - name: Record torch_tpu commit + if: matrix.with_torch == '1' + run: | + TORCH_TPU_SHA=$(git -C "$GITHUB_WORKSPACE/torch_tpu" rev-parse HEAD) + echo "TORCH_TPU_SHA=$TORCH_TPU_SHA" >> "$GITHUB_ENV" + echo "torch_tpu: $TORCH_TPU_SHA (ref: ${{ inputs.torch_tpu_ref || 'main' }})" >> "$GITHUB_STEP_SUMMARY" + - name: Build wheel + working-directory: tpu-raiden + env: + WITH_TORCH: ${{ matrix.with_torch }} + TORCH_TPU_SRC: ${{ github.workspace }}/torch_tpu + WHEEL_VERSION_EXTRAS: ${{ needs.version.outputs.wheel_version_extras }} + EXTRA_BAZEL_FLAGS: ${{ (inputs.use_remote_cache || 'yes') == 'yes' && '--config=ci' || '' }} + run: | + export BAZEL_CACHE_DIR="${RUNNER_TEMP}/bazel_cache" + bash ci/build_wheel_impl.sh + if [[ -n "${TORCH_TPU_SHA:-}" ]]; then + echo "${TORCH_TPU_SHA}" > dist/torch_tpu_commit.txt + fi + - name: Check wheel metadata + working-directory: tpu-raiden + run: uv run --isolated --with twine twine check dist/*${{ needs.version.outputs.wheel_version_extras }}-*.whl + # The JAX-stack coupling (jax/jaxlib/libtpu) is enforced through the + # wheel's own Requires-Dist pins; extract them so each build records the + # stack it must run with, mirroring torch_tpu_commit.txt on the torch side. + - name: Record JAX stack pins + if: matrix.framework == 'jax' + working-directory: tpu-raiden + run: | + python3 - <<'PY' + import glob, re, zipfile + whl = sorted(glob.glob('dist/tpu_raiden_jax-*.whl'))[-1] + with zipfile.ZipFile(whl) as z: + meta = next(n for n in z.namelist() if n.endswith('.dist-info/METADATA')) + text = z.read(meta).decode() + pins = [l.removeprefix('Requires-Dist: ') for l in text.splitlines() + if re.match(r'Requires-Dist: (jax|jaxlib|libtpu)\b', l)] + assert pins, f'no jax/jaxlib/libtpu pins found in {whl} METADATA' + open('dist/jax_pins.txt', 'w').write('\n'.join(pins) + '\n') + print('\n'.join(pins)) + PY + echo "JAX stack pins:" >> "$GITHUB_STEP_SUMMARY" + cat dist/jax_pins.txt >> "$GITHUB_STEP_SUMMARY" + - name: Upload wheel artifact + uses: actions/upload-artifact@v4 + with: + name: nightly-wheel-${{ matrix.framework }} + path: | + tpu-raiden/dist/*${{ needs.version.outputs.wheel_version_extras }}-*.whl + tpu-raiden/dist/torch_tpu_commit.txt + tpu-raiden/dist/jax_pins.txt + if-no-files-found: error + - name: Upload shared libraries + uses: actions/upload-artifact@v4 + with: + name: nightly-so-${{ matrix.framework }} + path: ${{ matrix.so_paths }} + if-no-files-found: error + - name: Summarize + working-directory: tpu-raiden + run: | + echo '```' >> "$GITHUB_STEP_SUMMARY" + ls -lh dist/*${{ needs.version.outputs.wheel_version_extras }}-*.whl >> "$GITHUB_STEP_SUMMARY" + echo '```' >> "$GITHUB_STEP_SUMMARY" + + upload: + name: "Upload wheels to Artifact Registry" + needs: [version, build] + # Scheduled nightlies always publish; manual dispatches only on request. + if: github.event_name == 'schedule' || inputs.upload_to_registry == 'yes' + runs-on: linux-x86-n2-32 + container: + image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest + timeout-minutes: 30 + steps: + - name: Download wheel artifacts + uses: actions/download-artifact@v4 + with: + pattern: nightly-wheel-* + merge-multiple: true + path: dist + # Auth uses the runner's ambient GCP service account via the Artifact + # Registry keyring backend. --skip-existing makes re-runs idempotent. + - name: Upload to Artifact Registry + run: | + uv run --isolated \ + --with twine \ + --with keyrings.google-artifactregistry-auth \ + twine upload --skip-existing \ + --repository-url "${RAIDEN_REGISTRY_URL}" \ + dist/tpu_raiden_*${{ needs.version.outputs.wheel_version_extras }}-*.whl + echo "Uploaded to ${RAIDEN_REGISTRY_URL}:" >> "$GITHUB_STEP_SUMMARY" + ls dist/tpu_raiden_*.whl >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/release_wheels.yml b/.github/workflows/release_wheels.yml new file mode 100644 index 00000000..029978d7 --- /dev/null +++ b/.github/workflows/release_wheels.yml @@ -0,0 +1,236 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Stable release build for tpu_raiden, cut by pushing a vX.Y.Z tag on main +# (see RELEASING.md for the full process). Builds the tpu_raiden_jax and +# tpu_raiden_torch wheels with an EMPTY version suffix — the wheel version is +# exactly the pyproject.toml base version, which must match the tag — then +# uploads them to the raiden Artifact Registry and attaches wheels + .so files +# to a GitHub Release. Stable versions carry no .dev suffix, so +# `pip install tpu_raiden_jax` (without --pre) resolves to the latest release +# while nightlies stay behind --pre. +# +# The torch wheel is built against the torch_tpu commit pinned in the +# torch_tpu.version file at the repo root (falling back to torch_tpu main if +# the file does not exist). Record the pin there before tagging so the release +# is reproducible against a validated torch_tpu build. +# +# workflow_dispatch runs are dry runs: they build and archive the wheels as +# workflow artifacts but never publish. +# +# Required repository secrets/variables: +# secrets.TORCH_TPU_DEPLOY_KEY read-only deploy key for the private +# google-pytorch/torch_tpu repo +# vars.RAIDEN_REGISTRY_URL (optional) overrides the Artifact Registry +# upload URL + +name: Release Wheels + +on: + push: + tags: + - "v[0-9]+.[0-9]+.[0-9]+" + workflow_dispatch: + inputs: + torch_tpu_ref: + description: "torch_tpu ref override for the dry-run build (default: torch_tpu.version file, else main)" + required: false + default: "" + type: string + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + +env: + RAIDEN_REGISTRY_URL: ${{ vars.RAIDEN_REGISTRY_URL || 'https://us-python.pkg.dev/cloud-tpu-inference-test/tpu-raiden/' }} + +jobs: + verify: + name: "Verify tag and version" + runs-on: ubuntu-latest + outputs: + version: ${{ steps.check.outputs.version }} + steps: + - uses: actions/checkout@v5 + with: + fetch-depth: 0 + - id: check + run: | + BASE_VERSION=$(sed -n -E 's/^version = "(.*)"/\1/p' pyproject.toml | head -1) + echo "pyproject.toml version: ${BASE_VERSION}" + if [[ "${GITHUB_REF}" == refs/tags/* ]]; then + TAG_VERSION="${GITHUB_REF_NAME#v}" + if [[ "${TAG_VERSION}" != "${BASE_VERSION}" ]]; then + echo "ERROR: tag ${GITHUB_REF_NAME} does not match pyproject.toml version ${BASE_VERSION}." >&2 + echo " Bump pyproject.toml [project].version on main before tagging." >&2 + exit 1 + fi + if ! git merge-base --is-ancestor "${GITHUB_SHA}" origin/main; then + echo "ERROR: tag ${GITHUB_REF_NAME} is not reachable from origin/main; releases are cut from main only." >&2 + exit 1 + fi + fi + echo "version=${BASE_VERSION}" >> "$GITHUB_OUTPUT" + + build: + name: "Build ${{ matrix.framework }} release wheel" + needs: verify + strategy: + fail-fast: false + matrix: + include: + - framework: jax + with_torch: "0" + so_paths: | + tpu-raiden/tpu_raiden/frameworks/jax/_tpu_raiden_jax.so + - framework: torch + with_torch: "1" + so_paths: | + tpu-raiden/tpu_raiden/frameworks/torch/_tpu_raiden_host.so + tpu-raiden/tpu_raiden/frameworks/torch/_tpu_raiden_torch.so + runs-on: linux-x86-n2-32 + container: + image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest + timeout-minutes: 300 + steps: + - name: Checkout tpu-raiden + uses: actions/checkout@v5 + with: + path: tpu-raiden + - name: Resolve torch_tpu pin + if: matrix.with_torch == '1' + run: | + REF="${{ inputs.torch_tpu_ref || '' }}" + if [[ -z "${REF}" && -f tpu-raiden/torch_tpu.version ]]; then + REF="$(tr -d '\r\n ' < tpu-raiden/torch_tpu.version)" + fi + echo "TORCH_TPU_REF=${REF:-main}" >> "$GITHUB_ENV" + - name: Checkout torch_tpu + if: matrix.with_torch == '1' + uses: actions/checkout@v5 + with: + repository: google-pytorch/torch_tpu + ref: ${{ env.TORCH_TPU_REF }} + ssh-key: ${{ secrets.TORCH_TPU_DEPLOY_KEY }} + path: torch_tpu + # Git security patch (CVE-2022-24765) blocks root from operating on + # runner-owned directories. Tell Git explicitly to trust the workspace + # inside this ephemeral container. + - name: Fix Git Workspace Ownership + run: | + git config --global --add safe.directory "$GITHUB_WORKSPACE/tpu-raiden" + git config --global --add safe.directory "$GITHUB_WORKSPACE/torch_tpu" + - name: Record torch_tpu commit + if: matrix.with_torch == '1' + run: | + TORCH_TPU_SHA=$(git -C "$GITHUB_WORKSPACE/torch_tpu" rev-parse HEAD) + echo "TORCH_TPU_SHA=$TORCH_TPU_SHA" >> "$GITHUB_ENV" + echo "torch_tpu: $TORCH_TPU_SHA (ref: $TORCH_TPU_REF)" >> "$GITHUB_STEP_SUMMARY" + - name: Build wheel + working-directory: tpu-raiden + env: + WITH_TORCH: ${{ matrix.with_torch }} + TORCH_TPU_SRC: ${{ github.workspace }}/torch_tpu + WHEEL_VERSION_EXTRAS: "" + EXTRA_BAZEL_FLAGS: "--config=ci" + run: | + export BAZEL_CACHE_DIR="${RUNNER_TEMP}/bazel_cache" + bash ci/build_wheel_impl.sh + if [[ -n "${TORCH_TPU_SHA:-}" ]]; then + echo "${TORCH_TPU_SHA}" > dist/torch_tpu_commit.txt + fi + - name: Check wheel metadata + working-directory: tpu-raiden + run: uv run --isolated --with twine twine check dist/*.whl + # The JAX-stack coupling (jax/jaxlib/libtpu) is enforced through the + # wheel's own Requires-Dist pins; extract them so the release records the + # stack it must run with, mirroring torch_tpu_commit.txt on the torch side. + - name: Record JAX stack pins + if: matrix.framework == 'jax' + working-directory: tpu-raiden + run: | + python3 - <<'PY' + import glob, re, zipfile + whl = sorted(glob.glob('dist/tpu_raiden_jax-*.whl'))[-1] + with zipfile.ZipFile(whl) as z: + meta = next(n for n in z.namelist() if n.endswith('.dist-info/METADATA')) + text = z.read(meta).decode() + pins = [l.removeprefix('Requires-Dist: ') for l in text.splitlines() + if re.match(r'Requires-Dist: (jax|jaxlib|libtpu)\b', l)] + assert pins, f'no jax/jaxlib/libtpu pins found in {whl} METADATA' + open('dist/jax_pins.txt', 'w').write('\n'.join(pins) + '\n') + print('\n'.join(pins)) + PY + echo "JAX stack pins:" >> "$GITHUB_STEP_SUMMARY" + cat dist/jax_pins.txt >> "$GITHUB_STEP_SUMMARY" + - name: Upload wheel artifact + uses: actions/upload-artifact@v4 + with: + name: release-wheel-${{ matrix.framework }} + path: | + tpu-raiden/dist/*.whl + tpu-raiden/dist/torch_tpu_commit.txt + tpu-raiden/dist/jax_pins.txt + if-no-files-found: error + - name: Upload shared libraries + uses: actions/upload-artifact@v4 + with: + name: release-so-${{ matrix.framework }} + path: ${{ matrix.so_paths }} + if-no-files-found: error + + publish: + name: "Publish release" + needs: [verify, build] + if: startsWith(github.ref, 'refs/tags/') + runs-on: linux-x86-n2-32 + container: + image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest + timeout-minutes: 30 + permissions: + contents: write + steps: + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + pattern: release-* + merge-multiple: true + path: release + # Auth uses the runner's ambient GCP service account via the Artifact + # Registry keyring backend. --skip-existing makes re-runs idempotent. + - name: Upload to Artifact Registry + run: | + uv run --isolated \ + --with twine \ + --with keyrings.google-artifactregistry-auth \ + twine upload --skip-existing \ + --repository-url "${RAIDEN_REGISTRY_URL}" \ + release/tpu_raiden_*.whl + - name: Create GitHub Release + env: + GH_TOKEN: ${{ github.token }} + run: | + NOTES="$(printf 'torch_tpu build pin: %s\n\nJAX stack pins:\n%s\n' \ + "$(cat release/torch_tpu_commit.txt 2>/dev/null || echo 'n/a')" \ + "$(cat release/jax_pins.txt 2>/dev/null || echo 'n/a')")" + gh release create "${GITHUB_REF_NAME}" \ + --repo "${GITHUB_REPOSITORY}" \ + --verify-tag \ + --generate-notes \ + --notes "${NOTES}" \ + release/tpu_raiden_*.whl release/*.so diff --git a/.github/workflows/run_benchmarks.yml b/.github/workflows/run_benchmarks.yml deleted file mode 100644 index f0a58041..00000000 --- a/.github/workflows/run_benchmarks.yml +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -name: Run Benchmarks - -on: - workflow_dispatch: - - push: - branches: - - main - - pull_request: - schedule: - - cron: "0 0 * * *" # Runs at 00:00 UTC every day - -permissions: - contents: read - pull-requests: write - -jobs: - run_benchmarks: - name: Run TPU Raiden H2D D2H Benchmarks - uses: google-ml-infra/bap/.github/workflows/run-benchmarks.yaml@affd28b492f67c242afa8b119e590b25816b7b88 - with: - registry_file: "tpu_raiden/benchmarks/benchmark_registry.pbtxt" - ml_actions_ref: "affd28b492f67c242afa8b119e590b25816b7b88" - runner: "linux-x86-n2-32" - publish_metrics: false - pub_sub_gcp_project_id: "ml-oss-benchmarking-production" - pub_sub_gcp_topic_id: "public-results-prod" diff --git a/.github/workflows/update_d2h_h2d_record.yml b/.github/workflows/update_d2h_h2d_record.yml deleted file mode 100644 index 89211a38..00000000 --- a/.github/workflows/update_d2h_h2d_record.yml +++ /dev/null @@ -1,16 +0,0 @@ -name: Record H2D/D2H Baselines -on: - workflow_dispatch: - -permissions: { contents: read, pull-requests: write } - -jobs: - record: - uses: google-ml-infra/bap/.github/workflows/run-benchmarks.yaml@affd28b492f67c242afa8b119e590b25816b7b88 - with: - registry_file: "tpu_raiden/benchmarks/benchmark_registry_record.pbtxt" - ml_actions_ref: "affd28b492f67c242afa8b119e590b25816b7b88" - runner: "linux-x86-n2-32" - publish_metrics: false - pub_sub_gcp_project_id: "ml-oss-benchmarking-production" - pub_sub_gcp_topic_id: "public-results-prod" diff --git a/RELEASING.md b/RELEASING.md new file mode 100644 index 00000000..e7d22553 --- /dev/null +++ b/RELEASING.md @@ -0,0 +1,93 @@ +# Releasing tpu_raiden + +## Versioning scheme + +The base version is single-sourced from `[project].version` in +[pyproject.toml](pyproject.toml) (read by the Bazel repo rule in +[bazel/wheel_version.bzl](bazel/wheel_version.bzl)). A suffix from the +`WHEEL_VERSION_EXTRAS` environment variable is appended at build time: + +| Channel | Version format | Built by | Install | +|---------|---------------------------|--------------------------------------------|---------| +| nightly | `X.Y.Z.devYYYYMMDDHHMMSS` | `Nightly Wheels` workflow (daily 08:00 UTC) | `pip install --pre tpu_raiden_jax --extra-index-url /simple/` | +| stable | `X.Y.Z` | `Release Wheels` workflow (on `vX.Y.Z` tag) | `pip install tpu_raiden_jax --extra-index-url /simple/` | + +Both channels publish to the same Artifact Registry (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 separate: `.dev` versions are only ever selected with +`--pre`, so a plain `pip install` always resolves to the latest stable +release. + +Two wheels are published per version, one per framework: +`tpu_raiden_jax` (bundles `_tpu_raiden_jax.so`, pulls the jax/jaxlib stack) +and `tpu_raiden_torch` (bundles `_tpu_raiden_host.so` / +`_tpu_raiden_torch.so`, no jax deps). Stable releases additionally attach the +wheels and bare `.so` files to a GitHub Release on the tag. + +## The torch_tpu ABI pin + +`tpu_raiden_torch` is ABI-coupled to `torch_tpu`: both must resolve the same +libtorch symbols at runtime, so raiden is compiled against a specific +torch_tpu checkout and the exact `torch==X.Y.Z+cpu` that torch_tpu's +per-Python requirements lock pins (see [ci/build_wheel_impl.sh](ci/build_wheel_impl.sh)). +torch_tpu publishes nightlies only (`torch_tpu==0.1.1.devYYYYMMDDHHMMSS`), so +every raiden release must record which torch_tpu build it targets: + +- [torch_tpu.version](torch_tpu.version) at the repo root pins the torch_tpu + commit releases are built against (nightlies build against torch_tpu `main` + instead). +- Every built wheel artifact ships a `torch_tpu_commit.txt` beside it, and the + GitHub Release notes carry the same pin. +- The matching `torch_tpu` nightly wheel version must be listed in the + release's CHANGELOG entry so users can install a compatible pair. + +## The JAX stack pin + +`tpu_raiden_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](ci/wheel/BUILD.bazel), kept in lockstep with +[pyproject.toml](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. + +## Cutting a stable release + +1. Pick the release candidate commit on `main` — normally the last-known-good + commit already published in [lkg.version](lkg.version). +2. Open a release PR that: + - updates [torch_tpu.version](torch_tpu.version) to the validated + torch_tpu commit, + - bumps `[project].version` in `pyproject.toml` to `X.Y.Z`, + - adds the `X.Y.Z` section to [CHANGELOG.md](CHANGELOG.md), including the + compatible `torch_tpu` nightly version. +3. (Optional) Dry-run: trigger the `Release Wheels` workflow via + `workflow_dispatch` on the PR branch — it builds stable-versioned wheels as + workflow artifacts without publishing anything. +4. Merge the PR, then tag and push: + + ```bash + git tag vX.Y.Z + git push origin vX.Y.Z + ``` + + The `Release Wheels` workflow verifies the tag (must equal the pyproject + version and be reachable from `main`), rebuilds both wheels with an empty + version suffix, uploads them to the Artifact Registry, and creates the + GitHub Release. +5. Immediately after the release, 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) + +- `TORCH_TPU_DEPLOY_KEY` secret: read-only deploy key for the private + `google-pytorch/torch_tpu` repository (the torch wheel build checks it out). +- The GitHub Actions runner service account needs + `roles/artifactregistry.writer` on the target registry, or set the + `RAIDEN_REGISTRY_URL` repository variable to a registry it can write to. diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh index 5645a12a..48dfc7c2 100755 --- a/ci/build_wheel.sh +++ b/ci/build_wheel.sh @@ -75,120 +75,20 @@ if [[ "${WITH_TORCH}" == "1" ]]; then BUILD_MODE="both" fi -# The in-container build: install clang-18 (XLA .ll targets) + CPU torch (shim -# headers), then drive the existing build.sh for the wheel target. -read -r -d '' INNER <<'INNER_EOF' || true -set -exu -o pipefail -export DEBIAN_FRONTEND=noninteractive -apt-get update -qq -apt-get install -y -qq wget gnupg ca-certificates patchelf >/dev/null -# Add the LLVM jammy-18 apt repo manually (the container's add-apt-repository is -# broken: python apt_pkg is missing for python3.12). -wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | gpg --dearmor -o /usr/share/keyrings/llvm.gpg -echo "deb [signed-by=/usr/share/keyrings/llvm.gpg] http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main" \ - > /etc/apt/sources.list.d/llvm18.list -apt-get update -qq -apt-get install -y -qq clang-18 >/dev/null -ln -sf /usr/bin/clang-18 /usr/bin/clang -ln -sf /usr/bin/clang++-18 /usr/bin/clang++ -clang --version | head -1 - -if [[ "${WITH_TORCH}" == "1" ]]; then - # raiden is built on top of torch_tpu: raiden's _tpu_raiden_torch.so and - # torch_tpu's libpywrap_torch_tpu_common.so must resolve the SAME libtorch - # symbols at runtime. So raiden MUST compile against the EXACT torch that - # torch_tpu was built against — never a floating `torch>=X` specifier, which - # drifts to the latest release and breaks ABI (e.g. torch 2.13.x drops - # `torch::autograd::deleteNode`, which torch_tpu's libpywrap needs → dlopen - # `undefined symbol` at import). torch_tpu's source of truth is its per-Python - # requirements lock, which pins an exact `torch==VERSION+cpu`. - PYTAG="$(python3 -c 'import sys;print(f"{sys.version_info.major}_{sys.version_info.minor}")')" - TORCH_REQ_FILE="/torch_tpu/requirements/requirements_${PYTAG}.txt" - TORCH_PIN="" - if [[ -f "${TORCH_REQ_FILE}" ]]; then - # e.g. line `torch==2.11.0+cpu \` -> `torch==2.11.0+cpu` - TORCH_PIN=$(sed -n -E 's/^(torch==[0-9][0-9A-Za-z.+_-]*).*/\1/p' "${TORCH_REQ_FILE}" | head -1 || true) - fi - if [[ -n "${TORCH_PIN}" ]]; then - echo "Installing torch pinned by torch_tpu (${TORCH_REQ_FILE}): ${TORCH_PIN}" - pip install -q "${TORCH_PIN}" --index-url https://download.pytorch.org/whl/cpu - else - # Fallback: the (looser) specifier from torch_tpu's pyproject.toml. This can - # float to the latest release and may NOT match torch_tpu's ABI, so warn. - TORCH_VERSION="" - if [[ -f /torch_tpu/pyproject.toml ]]; then - TORCH_VERSION=$(sed -n -E 's/.*["'\''`]torch[[:space:]]*([>=<~=]+[0-9.a-zA-Z+-]+)["'\''`].*/\1/p' /torch_tpu/pyproject.toml 2>/dev/null | head -1 || true) - fi - if [[ -z "${TORCH_VERSION}" ]]; then - echo "WARNING: could not determine torch pin from ${TORCH_REQ_FILE} or /torch_tpu/pyproject.toml. Installing latest torch — this may NOT match torch_tpu's ABI." >&2 - pip install -q torch --index-url https://download.pytorch.org/whl/cpu - else - echo "WARNING: no exact pin in ${TORCH_REQ_FILE}; falling back to torch_tpu pyproject specifier 'torch${TORCH_VERSION}', which may float to a torch that does not match torch_tpu's ABI." >&2 - pip install -q "torch${TORCH_VERSION}" --index-url https://download.pytorch.org/whl/cpu - fi - fi - TORCH_SOURCE="$(python3 -c 'import torch,pathlib;print(pathlib.Path(torch.__file__).resolve().parent.parent)')" - export TORCH_SOURCE - export TORCH_TPU_MODULE_PATH=/torch_tpu -fi - -# Persistent, resumable bazel cache + output base on the mounted volume. -export BAZEL_CACHE_DIR=/cache -export BAZEL_OUTPUT_BASE=/cache/output_base - -# Separate per-framework wheels: tpu_raiden_torch (no jax deps) vs -# tpu_raiden_jax. Pick by WITH_TORCH. -if [[ "${WITH_TORCH}" == "1" ]]; then - WHEEL_TARGET="//ci/wheel:raiden_torch_wheel" - WHEEL_DIST="tpu_raiden_torch" -else - WHEEL_TARGET="//ci/wheel:raiden_jax_wheel" - WHEEL_DIST="tpu_raiden_jax" -fi -# Match ONLY the wheel this build just produced. cache/output_base is shared -# across builds, so its bin/ci/wheel/ dir accumulates wheels from earlier runs, -# each with a distinct .dev. A broad "${WHEEL_DIST}-*.whl" glob would -# also match those stale wheels and hand multiple paths to the single-wheel -# patchelf step below (which then fails). WHEEL_VERSION_EXTRAS (.dev) -# is unique per build and appears verbatim in the filename, so scope to it. -WHEEL_GLOB="${WHEEL_DIST}-*${WHEEL_VERSION_EXTRAS}-*.whl" - -cd /workspace -./build.sh "${BUILD_MODE}" "${WHEEL_TARGET}" \ - --repo_env=WHEEL_VERSION_EXTRAS="${WHEEL_VERSION_EXTRAS}" - -mkdir -p /workspace/dist -cp /cache/output_base/execroot/_main/bazel-out/k8-opt/bin/ci/wheel/${WHEEL_GLOB} /workspace/dist/ - -# The bazel-built _tpu_raiden_torch.so does not link libpywrap; the torch -# extension loader (tpu_raiden/api/torch/kv_cache_manager.py) requires it as a -# NEEDED so the torch_tpu symbols (MaterializeAndReturn, AwaitBuffer, ...) -# resolve in RTLD_LOCAL scope at import. build.sh injects this for its -# source-tree copy, but the wheel packages the raw bazel .so -- so inject it -# into the wheel here and repack (which regenerates RECORD with valid hashes). -if [[ "${WITH_TORCH}" == "1" ]]; then - pip install -q wheel - WHL="$(ls /workspace/dist/${WHEEL_GLOB} | head -1)" - UNPACK_DIR="$(mktemp -d)" - wheel unpack "${WHL}" -d "${UNPACK_DIR}" - PKG_DIR="$(ls -d "${UNPACK_DIR}"/*/)" - patchelf --add-needed libpywrap_torch_tpu_common.so \ - "${PKG_DIR}tpu_raiden/frameworks/torch/_tpu_raiden_torch.so" - rm -f "${WHL}" - wheel pack "${PKG_DIR}" -d /workspace/dist - echo "patchelf: injected NEEDED libpywrap_torch_tpu_common.so into wheel .so" -fi -INNER_EOF - +# The in-container build (clang-18 install + pinned torch + build.sh wheel +# target) lives in ci/build_wheel_impl.sh so GitHub Actions jobs that already +# run inside the ml-build container can invoke it without docker-in-docker. echo "===> Building ${BUILD_MODE} wheel in ${CONTAINER_IMAGE}..." docker run --rm \ "${DOCKER_MOUNTS[@]}" \ -w /workspace \ -e WHEEL_VERSION_EXTRAS="${WHEEL_VERSION_EXTRAS}" \ -e WITH_TORCH="${WITH_TORCH}" \ - -e BUILD_MODE="${BUILD_MODE}" \ + -e TORCH_TPU_SRC=/torch_tpu \ + -e BAZEL_CACHE_DIR=/cache \ + -e EXTRA_BAZEL_FLAGS="${EXTRA_BAZEL_FLAGS:-}" \ "${CONTAINER_IMAGE}" \ - bash -c "${INNER}" + bash ci/build_wheel_impl.sh # Scope to THIS build's wheel(s) (.dev); REPO_ROOT/dist and WHEEL_DIR # are persistent and may hold wheels from earlier runs. diff --git a/ci/build_wheel_impl.sh b/ci/build_wheel_impl.sh new file mode 100755 index 00000000..f6c90f2e --- /dev/null +++ b/ci/build_wheel_impl.sh @@ -0,0 +1,173 @@ +#!/bin/bash + +# Copyright 2026 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright 2026 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# In-container build of a tpu_raiden wheel. Must run inside the ml-build +# container (glibc 2.35, matching the TPU runtime) — either via the docker +# wrapper ci/build_wheel.sh, or directly as a GitHub Actions job step whose +# job declares `container: ml-build:latest` (see .github/workflows/ +# nightly_build.yml). Requires root (installs clang-18 via apt). +# +# Produces exactly one wheel per invocation, selected by WITH_TORCH: +# WITH_TORCH=1 -> tpu_raiden_torch (needs a torch_tpu checkout) +# WITH_TORCH=0 -> tpu_raiden_jax +# The wheel lands in ${REPO_ROOT}/dist/, scoped by WHEEL_VERSION_EXTRAS. +# +# Environment: +# WITH_TORCH 1 (default) or 0 +# TORCH_TPU_SRC torch_tpu checkout path (default /torch_tpu) +# WHEEL_VERSION_EXTRAS suffix appended to the pyproject.toml base version +# (default .dev; set to the empty string +# for a stable-release wheel) +# BAZEL_CACHE_DIR bazel disk/repo cache root (default /cache) +# BAZEL_OUTPUT_BASE bazel output base (default ${BAZEL_CACHE_DIR}/output_base) +# EXTRA_BAZEL_FLAGS extra bazel flags, e.g. "--config=ci" for the RBE +# remote cache (space-separated; optional) + +set -exu -o pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "${REPO_ROOT}" + +WITH_TORCH="${WITH_TORCH:-1}" +TORCH_TPU_SRC="${TORCH_TPU_SRC:-/torch_tpu}" +# Default only when UNSET: an empty-but-set WHEEL_VERSION_EXTRAS is meaningful +# (a stable release build, whose wheel version is the bare pyproject version). +WHEEL_VERSION_EXTRAS="${WHEEL_VERSION_EXTRAS-.dev$(date -u +%Y%m%d%H%M%S)}" +export WHEEL_VERSION_EXTRAS +export BAZEL_CACHE_DIR="${BAZEL_CACHE_DIR:-/cache}" +export BAZEL_OUTPUT_BASE="${BAZEL_OUTPUT_BASE:-${BAZEL_CACHE_DIR}/output_base}" +EXTRA_BAZEL_FLAGS="${EXTRA_BAZEL_FLAGS:-}" + +export DEBIAN_FRONTEND=noninteractive +apt-get update -qq +apt-get install -y -qq wget gnupg ca-certificates patchelf >/dev/null +# Add the LLVM jammy-18 apt repo manually (the container's add-apt-repository is +# broken: python apt_pkg is missing for python3.12). +wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | gpg --dearmor -o /usr/share/keyrings/llvm.gpg +echo "deb [signed-by=/usr/share/keyrings/llvm.gpg] http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main" \ + > /etc/apt/sources.list.d/llvm18.list +apt-get update -qq +apt-get install -y -qq clang-18 >/dev/null +ln -sf /usr/bin/clang-18 /usr/bin/clang +ln -sf /usr/bin/clang++-18 /usr/bin/clang++ +clang --version | head -1 + +if [[ "${WITH_TORCH}" == "1" ]]; then + if [[ ! -f "${TORCH_TPU_SRC}/MODULE.bazel" ]]; then + echo "ERROR: torch build needs a torch_tpu checkout at ${TORCH_TPU_SRC}" >&2 + echo " set TORCH_TPU_SRC= or WITH_TORCH=0 for a JAX-only wheel." >&2 + exit 1 + fi + # raiden is built on top of torch_tpu: raiden's _tpu_raiden_torch.so and + # torch_tpu's libpywrap_torch_tpu_common.so must resolve the SAME libtorch + # symbols at runtime. So raiden MUST compile against the EXACT torch that + # torch_tpu was built against — never a floating `torch>=X` specifier, which + # drifts to the latest release and breaks ABI (e.g. torch 2.13.x drops + # `torch::autograd::deleteNode`, which torch_tpu's libpywrap needs → dlopen + # `undefined symbol` at import). torch_tpu's source of truth is its per-Python + # requirements lock, which pins an exact `torch==VERSION+cpu`. + PYTAG="$(python3 -c 'import sys;print(f"{sys.version_info.major}_{sys.version_info.minor}")')" + TORCH_REQ_FILE="${TORCH_TPU_SRC}/requirements/requirements_${PYTAG}.txt" + TORCH_PIN="" + if [[ -f "${TORCH_REQ_FILE}" ]]; then + # e.g. line `torch==2.11.0+cpu \` -> `torch==2.11.0+cpu` + TORCH_PIN=$(sed -n -E 's/^(torch==[0-9][0-9A-Za-z.+_-]*).*/\1/p' "${TORCH_REQ_FILE}" | head -1 || true) + fi + if [[ -n "${TORCH_PIN}" ]]; then + echo "Installing torch pinned by torch_tpu (${TORCH_REQ_FILE}): ${TORCH_PIN}" + pip install -q "${TORCH_PIN}" --index-url https://download.pytorch.org/whl/cpu + else + # Fallback: the (looser) specifier from torch_tpu's pyproject.toml. This can + # float to the latest release and may NOT match torch_tpu's ABI, so warn. + TORCH_VERSION="" + if [[ -f "${TORCH_TPU_SRC}/pyproject.toml" ]]; then + TORCH_VERSION=$(sed -n -E 's/.*["'\''`]torch[[:space:]]*([>=<~=]+[0-9.a-zA-Z+-]+)["'\''`].*/\1/p' "${TORCH_TPU_SRC}/pyproject.toml" 2>/dev/null | head -1 || true) + fi + if [[ -z "${TORCH_VERSION}" ]]; then + echo "WARNING: could not determine torch pin from ${TORCH_REQ_FILE} or ${TORCH_TPU_SRC}/pyproject.toml. Installing latest torch — this may NOT match torch_tpu's ABI." >&2 + pip install -q torch --index-url https://download.pytorch.org/whl/cpu + else + echo "WARNING: no exact pin in ${TORCH_REQ_FILE}; falling back to torch_tpu pyproject specifier 'torch${TORCH_VERSION}', which may float to a torch that does not match torch_tpu's ABI." >&2 + pip install -q "torch${TORCH_VERSION}" --index-url https://download.pytorch.org/whl/cpu + fi + fi + TORCH_SOURCE="$(python3 -c 'import torch,pathlib;print(pathlib.Path(torch.__file__).resolve().parent.parent)')" + export TORCH_SOURCE + export TORCH_TPU_MODULE_PATH="${TORCH_TPU_SRC}" +fi + +# Separate per-framework wheels: tpu_raiden_torch (no jax deps) vs +# tpu_raiden_jax. Pick by WITH_TORCH. +if [[ "${WITH_TORCH}" == "1" ]]; then + BUILD_MODE="both" + WHEEL_TARGET="//ci/wheel:raiden_torch_wheel" + WHEEL_DIST="tpu_raiden_torch" +else + BUILD_MODE="jax" + WHEEL_TARGET="//ci/wheel:raiden_jax_wheel" + WHEEL_DIST="tpu_raiden_jax" +fi +# Match ONLY the wheel this build just produced. The bazel bin dir is shared +# across builds via the persistent cache/output base, so it accumulates wheels +# from earlier runs, each with a distinct .dev. A broad +# "${WHEEL_DIST}-*.whl" glob would also match those stale wheels and hand +# multiple paths to the single-wheel patchelf step below (which then fails). +# WHEEL_VERSION_EXTRAS (.dev) is unique per build and appears +# verbatim in the filename, so scope to it. +WHEEL_GLOB="${WHEEL_DIST}-*${WHEEL_VERSION_EXTRAS}-*.whl" + +./build.sh "${BUILD_MODE}" "${WHEEL_TARGET}" \ + --repo_env=WHEEL_VERSION_EXTRAS="${WHEEL_VERSION_EXTRAS}" \ + ${EXTRA_BAZEL_FLAGS} + +mkdir -p "${REPO_ROOT}/dist" +cp "${REPO_ROOT}"/bazel-bin/ci/wheel/${WHEEL_GLOB} "${REPO_ROOT}/dist/" + +# The bazel-built _tpu_raiden_torch.so does not link libpywrap; the torch +# extension loader (tpu_raiden/api/torch/kv_cache_manager.py) requires it as a +# NEEDED so the torch_tpu symbols (MaterializeAndReturn, AwaitBuffer, ...) +# resolve in RTLD_LOCAL scope at import. build.sh injects this for its +# source-tree copy, but the wheel packages the raw bazel .so -- so inject it +# into the wheel here and repack (which regenerates RECORD with valid hashes). +if [[ "${WITH_TORCH}" == "1" ]]; then + pip install -q wheel + WHL="$(ls "${REPO_ROOT}"/dist/${WHEEL_GLOB} | head -1)" + UNPACK_DIR="$(mktemp -d)" + wheel unpack "${WHL}" -d "${UNPACK_DIR}" + PKG_DIR="$(ls -d "${UNPACK_DIR}"/*/)" + patchelf --add-needed libpywrap_torch_tpu_common.so \ + "${PKG_DIR}tpu_raiden/frameworks/torch/_tpu_raiden_torch.so" + rm -f "${WHL}" + wheel pack "${PKG_DIR}" -d "${REPO_ROOT}/dist" + echo "patchelf: injected NEEDED libpywrap_torch_tpu_common.so into wheel .so" +fi + +echo "===> Wheel(s) built:" +ls -lh "${REPO_ROOT}"/dist/${WHEEL_GLOB} diff --git a/torch_tpu.version b/torch_tpu.version new file mode 100644 index 00000000..c37f47c2 --- /dev/null +++ b/torch_tpu.version @@ -0,0 +1 @@ +5ce45524b3060aeb8597fff325d58d6693892061