From 4fc824cc5d9a13f59a366e1b6258f239158f2cb1 Mon Sep 17 00:00:00 2001
From: Dmitri Plotnikov <dplotnikov@google.com>
Date: Tue, 16 Jun 2026 14:21:08 -0700
Subject: [PATCH] Add support for linux/arm64 wheels.

Update pyproject.toml to build for aarch64 on Linux, and update
release_linux.sh to download the corresponding bazelisk binary.
Also download setuptools and wheel on host, and virtualenv resolved for
Python 3.9, and install them with --no-build-isolation to avoid PyPI timeout
issues in container.
Install virtualenv in before-build as well to avoid PyPI timeout during testing.
Increase pip default timeout to 120s to handle slow network/proxy in Kokoro.
Limit Bazel jobs to 4 in setup.py to avoid OOM failures in Kokoro/RBE.

PiperOrigin-RevId: 933297845
---
 release/kokoro/release_linux.sh | 10 ++++++++--
 release/pyproject.toml          |  6 +++++-
 release/setup.py                |  2 +-
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/release/kokoro/release_linux.sh b/release/kokoro/release_linux.sh
index 5292ef6..45199c2 100755
--- a/release/kokoro/release_linux.sh
+++ b/release/kokoro/release_linux.sh
@@ -20,7 +20,7 @@ export VIRTUALENV_NO_DOWNLOAD=1
 export PIP_DISABLE_PIP_VERSION_CHECK=1
 
 # Pass these environment variables to the cibuildwheel Docker container
-export CIBW_ENVIRONMENT="VIRTUALENV_NO_DOWNLOAD=1 PIP_DISABLE_PIP_VERSION_CHECK=1"
+export CIBW_ENVIRONMENT="VIRTUALENV_NO_DOWNLOAD=1 PIP_DISABLE_PIP_VERSION_CHECK=1 PIP_DEFAULT_TIMEOUT=120"
 export CIBW_DEPENDENCY_VERSIONS="latest"
 export CIBW_CONTAINER_ENGINE_EXTRA_ARGS="--network=host"
 
@@ -203,7 +203,13 @@ rm -rf cel_expr_python/*_test.py
 
 echo "Downloading bazelisk on host..."
 curl -LO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-amd64
-chmod +x bazelisk-linux-amd64
+curl -LO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-arm64
+chmod +x bazelisk-linux-amd64 bazelisk-linux-arm64
+
+echo "Downloading build dependencies on host..."
+mkdir -p build_deps
+pip download --only-binary=:all: --dest build_deps "setuptools>=40.8.0" "wheel"
+pip download --only-binary=:all: --dest build_deps --python-version 3.9 --platform manylinux2014_x86_64 "virtualenv"
 
 # Check if pyproject.toml exists before running sed
 if [ -f pyproject.toml ]; then
diff --git a/release/pyproject.toml b/release/pyproject.toml
index c6b322a..d7aa211 100644
--- a/release/pyproject.toml
+++ b/release/pyproject.toml
@@ -44,14 +44,18 @@ test-command = "python {project}/cel_basic_test.py"
 build-verbosity = 1
 
 [tool.cibuildwheel.linux]
+build-frontend = { name = "pip", args = ["--no-build-isolation"] }
+before-build = "pip install --no-index --find-links={project}/build_deps setuptools wheel virtualenv"
+archs = ["x86_64", "aarch64"]
 manylinux-x86_64-image = "manylinux_2_28"
+manylinux-aarch64-image = "manylinux_2_28"
 container-engine = "docker; disable_host_mount: True"
 # Google's internal Kokoro/RBE network uses a secure MITM proxy that resigns HTTPS
 # traffic with an internal Google CA. Since the public manylinux container does not
 # trust this CA, git fetches for external dependencies (like @cel-cpp) will fail
 # with SSL certificate errors. We disable http.sslVerify inside the container to
 # bypass this and allow Bazel to fetch SCM dependencies through the proxy.
-before-all = "git config --global http.sslVerify false && echo 'Installing bazelisk' && cp {project}/bazelisk-linux-amd64 /usr/local/bin/bazel"
+before-all = "git config --global http.sslVerify false && echo 'Installing bazelisk' && if [ $(uname -m) = 'aarch64' ]; then cp {project}/bazelisk-linux-arm64 /usr/local/bin/bazel; else cp {project}/bazelisk-linux-amd64 /usr/local/bin/bazel; fi && python3 -m pip install --no-index --find-links={project}/build_deps virtualenv"
 
 [tool.cibuildwheel.macos]
 archs = ["x86_64", "arm64"]
diff --git a/release/setup.py b/release/setup.py
index 09cfe81..67f3442 100644
--- a/release/setup.py
+++ b/release/setup.py
@@ -71,7 +71,7 @@ def build_extension(self, ext):
 
     # Build with bazel
     # Use --compilation_mode=opt for release builds
-    cmd = ['bazel', 'build', ext.target, '--compilation_mode=opt']
+    cmd = ['bazel', 'build', ext.target, '--compilation_mode=opt', '--jobs=4']
     if sys.platform == 'win32':
       self.platform_config_windows(cmd, python_version)
     if sys.platform == 'darwin':