Add Bazel py_wheel packaging for tpu_raiden (JAX + Torch)

BUILD

@@ -1,6 +1,9 @@
 package(default_visibility = ["//visibility:public"])
 
-exports_files(["LICENSE"])
+exports_files([
+    "LICENSE",
+    "README.md",
+])
 
 config_setting(
     name = "jax_disabled",

MODULE.bazel

@@ -128,3 +128,12 @@ git_override(
     remote = "https://github.com/google/nanobind_abseil.git",
 )
 
+# Single-sources the wheel version from pyproject.toml [project].version and
+# appends WHEEL_VERSION_EXTRAS (e.g. ".dev<timestamp>"). Consumed by
+# //ci/wheel:raiden_wheel via @raiden_version//:version.bzl.
+raiden_version_repo = use_repo_rule("//bazel:wheel_version.bzl", "raiden_version_repo")
+
+raiden_version_repo(
+    name = "raiden_version",
+    pyproject_toml = "//:pyproject.toml",
+)

bazel/BUILD

@@ -0,0 +1,32 @@
+# 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.
+
+# Marks //bazel as a Bazel package so //bazel:wheel_version.bzl is loadable.
+package(default_visibility = ["//visibility:public"])
+
+exports_files(["wheel_version.bzl"])

bazel/wheel_version.bzl

@@ -0,0 +1,76 @@
+# 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.
+
+"""Single-source the wheel version from pyproject.toml.
+
+Mirrors torch_tpu/bazel/wheel_version.bzl: reads the base `version` from the
+[project] section of pyproject.toml and appends the WHEEL_VERSION_EXTRAS env var
+(e.g. ".dev20260613120000") so CI can stamp dated dev builds.
+"""
+
+def _raiden_version_repo_impl(ctx):
+    pyproject_path = ctx.path(ctx.attr.pyproject_toml)
+    content = ctx.read(pyproject_path)
+    version = None
+
+    in_project_section = False
+    for line in content.splitlines():
+        line = line.strip()
+
+        # Skip empty lines and comments.
+        if not line or line.startswith("#"):
+            continue
+
+        # Track the current TOML table/section.
+        if line.startswith("[") and line.endswith("]"):
+            in_project_section = line == "[project]"
+            continue
+
+        # Only parse 'version' while inside the '[project]' section.
+        if in_project_section and line.startswith("version"):
+            parts = line.split("=", 1)
+            if len(parts) == 2 and parts[0].strip() == "version":
+                version = parts[1].strip().strip('"').strip("'")
+                break
+
+    if not version:
+        fail("Version not found in pyproject.toml [project] section")
+
+    suffix = ctx.os.environ.get("WHEEL_VERSION_EXTRAS", "")
+    full_version = version + suffix
+
+    ctx.file("BUILD.bazel", "")
+    ctx.file("version.bzl", "WHEEL_VERSION = '{}'\n".format(full_version))
+
+raiden_version_repo = repository_rule(
+    implementation = _raiden_version_repo_impl,
+    attrs = {
+        "pyproject_toml": attr.label(mandatory = True, allow_single_file = True),
+    },
+    environ = ["WHEEL_VERSION_EXTRAS"],
+)

ci/build_wheel.sh

@@ -0,0 +1,135 @@
+#!/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.
+#
+# Builds the tpu_raiden wheel hermetically inside the ml-build container
+# (glibc 2.35, matching the TPU runtime), mirroring torch_tpu/ci/build_wheel.sh.
+#
+# Unlike torch_tpu, raiden needs (a) clang-18 for XLA's .ll codegen targets and
+# (b) a local torch for the torch_tpu shim headers, plus the torch_tpu module.
+# This script installs clang-18 + CPU torch into the container, mounts a sibling
+# torch_tpu checkout, and runs build_raw_transfer.sh for the wheel target.
+#
+# Usage:
+#   ci/build_wheel.sh                  # JAX + Torch (needs ../torch_tpu)
+#   WITH_TORCH=0 ci/build_wheel.sh     # JAX only
+#   TORCH_TPU_SRC=/path ci/build_wheel.sh
+
+set -exu -o pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "${REPO_ROOT}"
+
+WHEEL_VERSION_EXTRAS=".dev$(date +%Y%m%d%H%M%S)"
+export WHEEL_VERSION_EXTRAS
+echo "WHEEL_VERSION_EXTRAS: ${WHEEL_VERSION_EXTRAS}"
+
+WITH_TORCH="${WITH_TORCH:-1}"
+TORCH_TPU_SRC="${TORCH_TPU_SRC:-${REPO_ROOT}/../torch_tpu}"
+WHEEL_DIR="${KOKORO_ARTIFACTS_DIR:-${HOME}/raiden_artifacts}/dist"
+CACHE_DIR="${RAIDEN_CONTAINER_CACHE:-${HOME}/.bazel_cache_container}"
+mkdir -p "${WHEEL_DIR}" "${REPO_ROOT}/dist" "${CACHE_DIR}"
+
+CONTAINER_IMAGE="us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+echo "===> Pulling ${CONTAINER_IMAGE}..."
+docker pull "${CONTAINER_IMAGE}"
+
+DOCKER_MOUNTS=(
+  -v "${REPO_ROOT}:/workspace"
+  -v "${CACHE_DIR}:/cache"
+)
+BUILD_MODE="jax"
+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=<path> or WITH_TORCH=0 for a JAX-only wheel." >&2
+    exit 1
+  fi
+  TORCH_TPU_SRC="$(cd "${TORCH_TPU_SRC}" && pwd)"
+  DOCKER_MOUNTS+=(-v "${TORCH_TPU_SRC}:/torch_tpu")  # sibling ../torch_tpu == /torch_tpu
+  BUILD_MODE="both"
+fi
+
+# The in-container build: install clang-18 (XLA .ll targets) + CPU torch (shim
+# headers), then drive the existing build_raw_transfer.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 >/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
+  pip install -q torch --index-url https://download.pytorch.org/whl/cpu
+  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
+
+cd /workspace
+./build_raw_transfer.sh "${BUILD_MODE}" //ci/wheel:raiden_wheel \
+  --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/tpu_raiden-*.whl /workspace/dist/
+INNER_EOF
+
+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}" \
+  "${CONTAINER_IMAGE}" \
+  bash -c "${INNER}"
+
+if [[ -n "$(ls -A "${REPO_ROOT}"/dist/*.whl 2>/dev/null)" ]]; then
+  cp "${REPO_ROOT}"/dist/*.whl "${WHEEL_DIR}/"
+  echo "===> Wheel(s) built:"; ls -lh "${WHEEL_DIR}"
+else
+  echo "ERROR: wheel build produced no .whl in dist/" >&2; exit 1
+fi
+
+echo "===> twine check..."
+docker run --rm -v "${WHEEL_DIR}:/dist" "${CONTAINER_IMAGE}" \
+  bash -c "uv run --isolated --with twine twine check /dist/*.whl"
+echo "===> raiden wheel build successful!"

ci/upload_wheel.sh

@@ -0,0 +1,56 @@
+#!/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.
+#
+# Uploads the tpu_raiden wheel to Google Artifact Registry via twine, mirroring
+# torch_tpu/ci/upload_wheel.sh.
+
+set -exu -o pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+WHEEL_DIR="${KOKORO_ARTIFACTS_DIR:-${REPO_ROOT}/../../artifacts}/dist"
+
+# Dedicated raiden Python registry in the cloud-tpu-inference-test project.
+REGISTRY_URL="${RAIDEN_REGISTRY_URL:-https://us-python.pkg.dev/cloud-tpu-inference-test/tpu-raiden/}"
+CONTAINER_IMAGE="us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+
+export UPLOAD_WHEEL_TO_AR="${UPLOAD_WHEEL_TO_AR:-true}"
+if [[ "${UPLOAD_WHEEL_TO_AR}" == "true" ]]; then
+  echo "===> Uploading tpu_raiden wheel(s) to ${REGISTRY_URL} ..."
+  docker run --rm \
+    -v "${WHEEL_DIR}:/dist" \
+    -v "${HOME}/.config/gcloud:/root/.config/gcloud:ro" \
+    "${CONTAINER_IMAGE}" \
+    bash -c "
+      uv run --isolated \
+        --with twine \
+        --with keyrings.google-artifactregistry-auth \
+        twine upload --repository-url ${REGISTRY_URL} /dist/tpu_raiden-*.whl
+    "
+  echo "===> Upload complete."
+fi