From 4155359082092ac48a8a625dd1de0bede6bdd7d9 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Thu, 3 Sep 2026 18:48:28 +0000 Subject: [PATCH 1/7] ci(jax): serve the JAX wheel build from the EngFlow Bazel cache The JAX wheel job rebuilds XLA from scratch on every run, because Bazel starts each container with an empty action cache. Point it at the EngFlow cluster and mount the mTLS credentials the ROCm bazelrc already expects under /data, sourced from the RBE_CI_CERT and RBE_CI_KEY secrets. configure_jax_bazel_cache.py emits the --bazel_options flags only when the credentials are present, the endpoint completes a TLS handshake and the release type is one the cache may serve. Every other case prints nothing and leaves the build command untouched, so fork pull requests, unreachable clusters and prerelease builds keep behaving as they do today. Pull requests read from the cache without writing to it, so a proposed change cannot poison what later builds pick up. --- .../multi_arch_build_linux_jax_wheels_ci.yml | 56 +++++ .github/workflows/multi_arch_ci_linux.yml | 2 + .../configure_jax_bazel_cache.py | 190 ++++++++++++++++ .../tests/configure_jax_bazel_cache_test.py | 208 ++++++++++++++++++ docs/development/ccache_troubleshooting.md | 21 ++ 5 files changed, 477 insertions(+) create mode 100644 build_tools/github_actions/configure_jax_bazel_cache.py create mode 100644 build_tools/github_actions/tests/configure_jax_bazel_cache_test.py diff --git a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml index 2a5dc77b306..8cfa615054e 100644 --- a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml +++ b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml @@ -220,17 +220,71 @@ jobs: --rocm-version "${ROCM_VERSION}" + # Fork pull requests receive no secrets. That is not an error: the build + # step mounts nothing and Bazel runs without a remote cache. + - name: Write Bazel remote cache credentials + id: cache_credentials + env: + RBE_CI_CERT: ${{ secrets.RBE_CI_CERT }} + RBE_CI_KEY: ${{ secrets.RBE_CI_KEY }} + run: | + if [[ -z "${RBE_CI_CERT}" || -z "${RBE_CI_KEY}" ]]; then + echo "No cache credentials available; building without a remote cache." + exit 0 + fi + + # Written outside the workspace: that tree is published as artifacts. + credentials_dir="${RUNNER_TEMP}/bazel-cache-credentials" + umask 077 + mkdir -p "${credentials_dir}" + + # The secrets may hold PEM or base64-encoded PEM. Try decoding and + # keep whichever form yields a PEM block, rather than guessing from + # the leading characters. + write_credential() { + if ! printf '%s' "$1" | base64 -d > "$2" 2>/dev/null || + ! grep -q -- "-----BEGIN" "$2"; then + printf '%s\n' "$1" > "$2" + fi + grep -q -- "-----BEGIN" "$2" + } + + # Unusable credentials leave the mount absent so the build runs + # without a cache. Failing here would fail the wheel build itself. + if write_credential "${RBE_CI_CERT}" "${credentials_dir}/ci-cert.crt" && + write_credential "${RBE_CI_KEY}" "${credentials_dir}/ci-cert.key"; then + echo "credentials_dir=${credentials_dir}" >> "${GITHUB_OUTPUT}" + else + echo "Cache credentials are not usable PEM; building without a remote cache." + rm -rf "${credentials_dir}" + fi + - name: Build JAX Wheels in manylinux container working-directory: jax-source env: PYTHON_VERSION: ${{ inputs.python_version }} JAX_COMMIT: ${{ steps.checkout_jax.outputs.commit }} + # Pull requests read the cache without writing to it, so no proposed + # change can put an entry in front of a later build. Pushes fill it. + JAX_BAZEL_CACHE_ALLOW_UPLOAD: ${{ github.event_name != 'pull_request' }} + JAX_BAZEL_REMOTE_CACHE_URL: ${{ vars.JAX_BAZEL_REMOTE_CACHE_URL }} + CACHE_CREDENTIALS_DIR: ${{ steps.cache_credentials.outputs.credentials_dir }} run: | + # rocm.bazelrc reads the client credentials from /data. + credentials_mount="" + if [[ -n "${CACHE_CREDENTIALS_DIR}" ]]; then + credentials_mount="--volume ${CACHE_CREDENTIALS_DIR}:/data:ro" + fi + docker run --rm \ --user root \ --env JAX_COMMIT \ + ${credentials_mount} \ --env ROCM_VERSION="${ROCM_VERSION}" \ --env PYTHON_VERSION="${PYTHON_VERSION}" \ + --env RELEASE_TYPE="${RELEASE_TYPE}" \ + --env JAX_BAZEL_CACHE_ALLOW_UPLOAD="${JAX_BAZEL_CACHE_ALLOW_UPLOAD}" \ + --env JAX_BAZEL_REMOTE_CACHE_URL="${JAX_BAZEL_REMOTE_CACHE_URL}" \ --env ML_WHEEL_VERSION_SUFFIX="${version_suffix}" \ --env PACKAGE_DIST_DIR="/workspace/jax-source/dist" \ --volume "${GITHUB_WORKSPACE}:/workspace" \ @@ -238,6 +292,7 @@ jobs: "${MANYLINUX_IMAGE_TAG}" \ bash -lc ' ROCM_ROOT="$(rocm-sdk path --root)" + bazel_cache_options="$(python /workspace/build_tools/github_actions/configure_jax_bazel_cache.py)" python build/build.py build \ --wheels=jax-rocm-plugin,jax-rocm-pjrt \ --python_version="${PYTHON_VERSION}" \ @@ -247,6 +302,7 @@ jobs: --bazel_options=--repo_env=ML_WHEEL_TYPE=release \ --bazel_options=--repo_env=ML_WHEEL_VERSION_SUFFIX="${ML_WHEEL_VERSION_SUFFIX}" \ --bazel_options=--//jaxlib/tools:jaxlib_git_hash="${JAX_COMMIT}" \ + ${bazel_cache_options} \ --verbose \ --detailed_timestamped_log \ --output_path="$(pwd)/dist" diff --git a/.github/workflows/multi_arch_ci_linux.yml b/.github/workflows/multi_arch_ci_linux.yml index 852edc53278..98f14fc0885 100644 --- a/.github/workflows/multi_arch_ci_linux.yml +++ b/.github/workflows/multi_arch_ci_linux.yml @@ -368,6 +368,8 @@ jobs: matrix: include: ${{ fromJSON(inputs.build_config).jax_build_matrix }} uses: ./.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml + # Carries the Bazel remote cache credentials into the wheel build. + secrets: inherit with: artifact_group: ${{ fromJSON(inputs.build_config).artifact_group }} python_version: ${{ matrix.python_version }} diff --git a/build_tools/github_actions/configure_jax_bazel_cache.py b/build_tools/github_actions/configure_jax_bazel_cache.py new file mode 100644 index 00000000000..18c5cd5bdbb --- /dev/null +++ b/build_tools/github_actions/configure_jax_bazel_cache.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +# Copyright Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Prints the `build/build.py` options that point the JAX wheel build at the +shared Bazel remote cache. + +The JAX wheel job compiles jaxlib, the ROCm plugin and the XLA ROCm backend with +Bazel against an already-installed ROCm. Those compiles repeat across runs +whenever the JAX ref and the ROCm build are unchanged, which is the common case +for a pull request that only touches CI or test-selection files. + +The cache is the EngFlow cluster ROCm's JAX CI already uses, reached with the +mTLS client credentials `rocm.bazelrc` expects at `/data`. Only the cache is +used: the build runs `--config=rocm_release_wheel`, never `rocm_rbe`, so it +never picks up that config's `--remote_executor` and every action still executes +locally. + +This runs *inside* the manylinux build container rather than on the runner. The +container reaches the network through the Docker daemon and receives the +credentials through a mount, so it is the only place that can tell whether Bazel +will actually reach the cache. When the credentials are absent (fork pull +requests get no secrets), the release type may not read shared entries, or the +endpoint does not answer, this prints nothing and the build runs exactly as it +would without a cache. + +Options are printed to stdout so the caller can expand them into the build +command; all logging goes to stderr. +""" + +import argparse +import os +import socket +import ssl +import sys +from pathlib import Path +from urllib.parse import urlsplit + +REMOTE_CACHE_URL = "grpcs://wardite.cluster.engflow.com" + +# Paths `build/rocm/rocm.bazelrc` already expects, populated by mounting the +# runner's decoded credentials at /data. +CLIENT_CERTIFICATE = Path("/data/ci-cert.crt") +CLIENT_KEY = Path("/data/ci-cert.key") + +# Release types that may read shared cache entries, mirroring the ccache policy +# in setup_ccache.py. Stable releases repackage prerelease artifacts, so +# "prerelease" and "nightly-bkc" are absent on purpose. +SHARED_CACHE_RELEASE_TYPES = frozenset({"ci", "dev", "dev-bkc", "nightly"}) + +DEFAULT_PROBE_TIMEOUT_SECONDS = 10 +DEFAULT_REMOTE_TIMEOUT_SECONDS = 60 +DEFAULT_TLS_PORT = 443 + + +def _log(msg: str): + print(f"[jax-bazel-cache] {msg}", file=sys.stderr) + + +def resolve_cache_url(cache_url: str, release_type: str) -> str: + """Returns the cache URL to use, or "" when caching is disabled. + + An explicit `cache_url` wins so infrastructure can repoint the cache + without a code change; otherwise the release type decides. + """ + if cache_url: + return cache_url + if release_type in SHARED_CACHE_RELEASE_TYPES: + return REMOTE_CACHE_URL + return "" + + +def endpoint_address(url: str) -> tuple[str, int]: + """Returns the (host, port) a `grpcs://host[:port]` cache URL points at.""" + split = urlsplit(url) + return split.hostname or "", split.port or DEFAULT_TLS_PORT + + +def probe_cache( + url: str, + certificate: Path, + key: Path, + timeout: float = DEFAULT_PROBE_TIMEOUT_SECONDS, +) -> bool: + """Returns whether the cache completes a TLS handshake from this process. + + This resolves the host, connects, and presents the client credentials, so a + stale or mismatched certificate is reported here rather than as a wall of + Bazel cache warnings during the build. + """ + host, port = endpoint_address(url) + if not host: + _log(f"Cannot parse a host out of {url!r}; building without a remote cache") + return False + try: + context = ssl.create_default_context() + context.minimum_version = ssl.TLSVersion.TLSv1_2 + context.load_cert_chain(certfile=certificate, keyfile=key) + with socket.create_connection((host, port), timeout=timeout) as sock: + with context.wrap_socket(sock, server_hostname=host): + return True + except (OSError, ssl.SSLError) as e: + _log(f"Cache unreachable ({e}); building without a remote cache") + return False + + +def bazel_cache_options( + url: str, certificate: Path, key: Path, allow_upload: bool +) -> list[str]: + """Returns the `--bazel_options=...` arguments for a usable cache.""" + return [ + f"--bazel_options=--remote_cache={url}", + f"--bazel_options=--tls_client_certificate={certificate}", + f"--bazel_options=--tls_client_key={key}", + f"--bazel_options=--remote_timeout={DEFAULT_REMOTE_TIMEOUT_SECONDS}", + f"--bazel_options=--remote_upload_local_results={str(allow_upload).lower()}", + ] + + +def main(argv: list[str] | None = None): + parser = argparse.ArgumentParser( + description="Print build/build.py options for the JAX Bazel remote cache." + ) + parser.add_argument( + "--cache-url", + default=os.environ.get("JAX_BAZEL_REMOTE_CACHE_URL", ""), + help="Remote cache URL. Overrides the release-type default.", + ) + parser.add_argument( + "--release-type", + default=os.environ.get("RELEASE_TYPE", "ci"), + help="Release types that must not read shared entries resolve to no cache.", + ) + parser.add_argument( + "--allow-upload", + default=os.environ.get("JAX_BAZEL_CACHE_ALLOW_UPLOAD", "false"), + help="Whether this run may write results back to the cache.", + ) + parser.add_argument( + "--client-certificate", + type=Path, + default=CLIENT_CERTIFICATE, + help="mTLS client certificate for the cache.", + ) + parser.add_argument( + "--client-key", + type=Path, + default=CLIENT_KEY, + help="mTLS client key for the cache.", + ) + parser.add_argument( + "--probe-timeout", + type=float, + default=DEFAULT_PROBE_TIMEOUT_SECONDS, + help="Seconds to wait for the cache to complete the handshake.", + ) + args = parser.parse_args(argv) + + url = resolve_cache_url(args.cache_url, args.release_type) + if not url: + _log(f"No remote cache for release_type={args.release_type!r}") + return + + missing = [p for p in (args.client_certificate, args.client_key) if not p.exists()] + if missing: + # Fork pull requests get no secrets, so this is a normal outcome and + # not an error: the build simply runs without a cache. + _log(f"No cache credentials at {', '.join(str(p) for p in missing)}") + return + + allow_upload = args.allow_upload.strip().lower() == "true" + _log( + f"Cache mode: url={url} release_type={args.release_type} upload={allow_upload}" + ) + if not probe_cache( + url, args.client_certificate, args.client_key, args.probe_timeout + ): + return + + print( + " ".join( + bazel_cache_options( + url, args.client_certificate, args.client_key, allow_upload + ) + ) + ) + + +if __name__ == "__main__": + main() diff --git a/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py new file mode 100644 index 00000000000..a569b9a5c8b --- /dev/null +++ b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +# Copyright Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Unit tests for configure_jax_bazel_cache.py""" + +import io +import ssl +import sys +import tempfile +import unittest +from pathlib import Path, PurePosixPath +from unittest import mock + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +import configure_jax_bazel_cache +from configure_jax_bazel_cache import ( + REMOTE_CACHE_URL, + bazel_cache_options, + endpoint_address, + main, + probe_cache, + resolve_cache_url, +) + + +class ResolveCacheUrlTest(unittest.TestCase): + """Tests which release types may read the shared cache.""" + + def test_ci_and_dev_use_the_shared_cache(self): + for release_type in ("ci", "dev", "dev-bkc", "nightly"): + self.assertEqual(resolve_cache_url("", release_type), REMOTE_CACHE_URL) + + def test_release_builds_get_no_shared_cache(self): + # Stable releases repackage prerelease artifacts, so they must not read + # entries another build could have written. + for release_type in ("nightly-bkc", "prerelease"): + self.assertEqual(resolve_cache_url("", release_type), "") + + def test_unknown_release_type_gets_no_cache(self): + self.assertEqual(resolve_cache_url("", "something-new"), "") + + def test_explicit_url_overrides_the_release_type_default(self): + self.assertEqual( + resolve_cache_url("grpcs://cache.example", "prerelease"), + "grpcs://cache.example", + ) + + +class EndpointAddressTest(unittest.TestCase): + """Tests parsing the host and port out of a gRPC cache URL.""" + + def test_default_port_is_tls(self): + self.assertEqual( + endpoint_address("grpcs://wardite.cluster.engflow.com"), + ("wardite.cluster.engflow.com", 443), + ) + + def test_explicit_port_is_honored(self): + self.assertEqual( + endpoint_address("grpcs://cache.example:8980"), ("cache.example", 8980) + ) + + +class ProbeCacheTest(unittest.TestCase): + """Tests that the probe reports unreachable rather than raising.""" + + def _probe(self, url=REMOTE_CACHE_URL, connect_effect=None, load_effect=None): + context = mock.MagicMock() + context.load_cert_chain.side_effect = load_effect + with mock.patch.object( + configure_jax_bazel_cache.ssl, + "create_default_context", + return_value=context, + ), mock.patch.object( + configure_jax_bazel_cache.socket, + "create_connection", + side_effect=connect_effect, + ): + return probe_cache(url, Path("cert.crt"), Path("cert.key")) + + def test_reachable_when_the_handshake_completes(self): + self.assertTrue(self._probe()) + + def test_unreachable_when_the_host_does_not_resolve(self): + self.assertFalse(self._probe(connect_effect=OSError("name resolution failed"))) + + def test_unreachable_when_the_connection_times_out(self): + self.assertFalse(self._probe(connect_effect=TimeoutError())) + + def test_unreachable_when_the_credentials_are_rejected(self): + # A stale or mismatched certificate is reported once here instead of as + # a wall of Bazel cache warnings during the build. + self.assertFalse(self._probe(load_effect=ssl.SSLError("bad key"))) + + def test_unreachable_when_the_url_has_no_host(self): + self.assertFalse(self._probe(url="not-a-url")) + + +class BazelCacheOptionsTest(unittest.TestCase): + """Tests the emitted build/build.py options.""" + + def _options(self, allow_upload): + # Paths inside the build container, so they stay POSIX even when the + # test itself runs on Windows. + return bazel_cache_options( + REMOTE_CACHE_URL, + PurePosixPath("/data/ci-cert.crt"), + PurePosixPath("/data/ci-cert.key"), + allow_upload, + ) + + def test_points_bazel_at_the_cache_with_credentials(self): + options = self._options(allow_upload=False) + self.assertIn(f"--bazel_options=--remote_cache={REMOTE_CACHE_URL}", options) + self.assertIn( + "--bazel_options=--tls_client_certificate=/data/ci-cert.crt", options + ) + self.assertIn("--bazel_options=--tls_client_key=/data/ci-cert.key", options) + + def test_never_requests_remote_execution(self): + # The build uses --config=rocm_release_wheel, so it must stay + # cache-only: actions execute on the runner. + self.assertFalse( + [option for option in self._options(True) if "remote_executor" in option] + ) + + def test_upload_is_off_for_read_only_runs(self): + self.assertIn( + "--bazel_options=--remote_upload_local_results=false", + self._options(allow_upload=False), + ) + + def test_upload_is_on_for_trusted_runs(self): + self.assertIn( + "--bazel_options=--remote_upload_local_results=true", + self._options(allow_upload=True), + ) + + +class MainTest(unittest.TestCase): + """Tests that stdout carries options only when the cache is usable.""" + + def setUp(self): + credentials = tempfile.TemporaryDirectory() + self.addCleanup(credentials.cleanup) + self.certificate = Path(credentials.name) / "ci-cert.crt" + self.key = Path(credentials.name) / "ci-cert.key" + # Only existence matters here; the handshake that would read these is + # mocked out, and real PEM headers trip secret scanning. + self.certificate.write_text("test certificate") + self.key.write_text("test key") + + def _main(self, argv, reachable=True, with_credentials=True): + # Always pass explicit paths: a developer machine may happen to have + # credentials at the default location, which would make the + # missing-credentials case pass for the wrong reason. + certificate, key = self.certificate, self.key + if not with_credentials: + certificate = self.certificate.with_name("absent.crt") + key = self.key.with_name("absent.key") + argv = argv + [ + "--client-certificate", + str(certificate), + "--client-key", + str(key), + ] + stdout = io.StringIO() + with mock.patch.object( + configure_jax_bazel_cache, "probe_cache", return_value=reachable + ), mock.patch.object(sys, "stdout", stdout): + main(argv) + return stdout.getvalue().strip() + + def test_prints_options_for_a_reachable_cache(self): + output = self._main(["--release-type", "ci", "--allow-upload", "true"]) + self.assertIn(f"--bazel_options=--remote_cache={REMOTE_CACHE_URL}", output) + self.assertIn("--bazel_options=--remote_upload_local_results=true", output) + + def test_prints_nothing_without_credentials(self): + # Fork pull requests get no secrets, so the mount is absent and the + # build command must come out unchanged. + self.assertEqual( + self._main(["--release-type", "ci"], with_credentials=False), "" + ) + + def test_prints_nothing_when_the_cache_is_unreachable(self): + self.assertEqual(self._main(["--release-type", "ci"], reachable=False), "") + + def test_prints_nothing_when_the_release_type_has_no_cache(self): + self.assertEqual(self._main(["--release-type", "prerelease"]), "") + + def test_probe_is_skipped_when_no_cache_is_configured(self): + with mock.patch.object( + configure_jax_bazel_cache, + "probe_cache", + side_effect=AssertionError("should not be called"), + ): + main(["--release-type", "prerelease"]) + + def test_allow_upload_defaults_to_read_only(self): + output = self._main(["--release-type", "ci"]) + self.assertIn("--bazel_options=--remote_upload_local_results=false", output) + + +if __name__ == "__main__": + unittest.main() diff --git a/docs/development/ccache_troubleshooting.md b/docs/development/ccache_troubleshooting.md index 356c3282451..978a30f0c85 100644 --- a/docs/development/ccache_troubleshooting.md +++ b/docs/development/ccache_troubleshooting.md @@ -42,6 +42,27 @@ remote cache, accessed via ccache's `remote_storage` option with Both servers are on the Kubernetes cluster, accessible without authentication from any pod in the cluster. +The JAX wheel build does not use ccache. It compiles jaxlib, the ROCm plugin +and the XLA ROCm backend with Bazel, and reads Bazel's own remote cache on the +EngFlow cluster that ROCm's JAX CI already uses. `configure_jax_bazel_cache.py` +runs inside the manylinux build container, completes an mTLS handshake with the +cache, and prints no Bazel options at all when the credentials are absent, the +release type may not read shared entries, or the endpoint does not answer, so a +cache that cannot help leaves the build command unchanged. + +The client credentials come from the `RBE_CI_CERT` and `RBE_CI_KEY` secrets +(PEM or base64-encoded PEM). The workflow decodes them under `RUNNER_TEMP` and +mounts that directory read only at `/data`, the path `build/rocm/rocm.bazelrc` +expects. Fork pull requests receive no secrets and therefore no cache. Set the +repository variable `JAX_BAZEL_REMOTE_CACHE_URL` to point the build at a +different cache. + +Only the cache is shared: the build runs `--config=rocm_release_wheel` rather +than `rocm_rbe`, so it never picks up that config's `--remote_executor` and +every action still runs on the runner. Pull requests read the cache without +writing to it, so no proposed change can place an entry in front of a later +build; pushes fill it. + ### Namespace version `CCACHE_NAMESPACE_VERSION` in `setup_ccache.py` controls the cache From e0de8d051c93c774b4e9070ff5ea16e50a6d387a Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Wed, 16 Sep 2026 22:55:09 +0000 Subject: [PATCH 2/7] ci(jax): read EngFlow certs from the runner, not GitHub secrets Fork PRs never see repository secrets, so the previous path could not give them a read-only cache and added a third credential mechanism. Copy the runner-local /data/ci-cert files in Python instead, drop secrets: inherit, and document the cache next to the JAX build. --- .../multi_arch_build_linux_jax_wheels_ci.yml | 44 ++----- .github/workflows/multi_arch_ci_linux.yml | 2 - .../configure_jax_bazel_cache.py | 12 +- .../tests/configure_jax_bazel_cache_test.py | 2 +- .../write_jax_bazel_cache_credentials_test.py | 114 +++++++++++++++++ .../write_jax_bazel_cache_credentials.py | 116 ++++++++++++++++++ docs/development/ccache_troubleshooting.md | 30 ++--- external-builds/jax/README.md | 34 +++++ 8 files changed, 289 insertions(+), 65 deletions(-) create mode 100644 build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py create mode 100644 build_tools/github_actions/write_jax_bazel_cache_credentials.py diff --git a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml index 8cfa615054e..ad8a7a1a073 100644 --- a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml +++ b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml @@ -220,44 +220,15 @@ jobs: --rocm-version "${ROCM_VERSION}" - # Fork pull requests receive no secrets. That is not an error: the build - # step mounts nothing and Bazel runs without a remote cache. - - name: Write Bazel remote cache credentials + # Runner-local files at /data/ci-cert.{crt,key}, the same path XLA and + # JAX CI mount. GitHub secrets are not used, so fork pull requests can + # read the cache. Missing files leave the mount absent. + - name: Locate Bazel remote cache credentials id: cache_credentials env: - RBE_CI_CERT: ${{ secrets.RBE_CI_CERT }} - RBE_CI_KEY: ${{ secrets.RBE_CI_KEY }} + JAX_BAZEL_CACHE_CREDENTIALS_DIR: ${{ vars.JAX_BAZEL_CACHE_CREDENTIALS_DIR }} run: | - if [[ -z "${RBE_CI_CERT}" || -z "${RBE_CI_KEY}" ]]; then - echo "No cache credentials available; building without a remote cache." - exit 0 - fi - - # Written outside the workspace: that tree is published as artifacts. - credentials_dir="${RUNNER_TEMP}/bazel-cache-credentials" - umask 077 - mkdir -p "${credentials_dir}" - - # The secrets may hold PEM or base64-encoded PEM. Try decoding and - # keep whichever form yields a PEM block, rather than guessing from - # the leading characters. - write_credential() { - if ! printf '%s' "$1" | base64 -d > "$2" 2>/dev/null || - ! grep -q -- "-----BEGIN" "$2"; then - printf '%s\n' "$1" > "$2" - fi - grep -q -- "-----BEGIN" "$2" - } - - # Unusable credentials leave the mount absent so the build runs - # without a cache. Failing here would fail the wheel build itself. - if write_credential "${RBE_CI_CERT}" "${credentials_dir}/ci-cert.crt" && - write_credential "${RBE_CI_KEY}" "${credentials_dir}/ci-cert.key"; then - echo "credentials_dir=${credentials_dir}" >> "${GITHUB_OUTPUT}" - else - echo "Cache credentials are not usable PEM; building without a remote cache." - rm -rf "${credentials_dir}" - fi + python build_tools/github_actions/write_jax_bazel_cache_credentials.py - name: Build JAX Wheels in manylinux container working-directory: jax-source @@ -270,7 +241,8 @@ jobs: JAX_BAZEL_REMOTE_CACHE_URL: ${{ vars.JAX_BAZEL_REMOTE_CACHE_URL }} CACHE_CREDENTIALS_DIR: ${{ steps.cache_credentials.outputs.credentials_dir }} run: | - # rocm.bazelrc reads the client credentials from /data. + # rocm.bazelrc reads the client credentials from /data, copied + # from the runner's /data/ci-cert.* into $RUNNER_TEMP. credentials_mount="" if [[ -n "${CACHE_CREDENTIALS_DIR}" ]]; then credentials_mount="--volume ${CACHE_CREDENTIALS_DIR}:/data:ro" diff --git a/.github/workflows/multi_arch_ci_linux.yml b/.github/workflows/multi_arch_ci_linux.yml index 98f14fc0885..852edc53278 100644 --- a/.github/workflows/multi_arch_ci_linux.yml +++ b/.github/workflows/multi_arch_ci_linux.yml @@ -368,8 +368,6 @@ jobs: matrix: include: ${{ fromJSON(inputs.build_config).jax_build_matrix }} uses: ./.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml - # Carries the Bazel remote cache credentials into the wheel build. - secrets: inherit with: artifact_group: ${{ fromJSON(inputs.build_config).artifact_group }} python_version: ${{ matrix.python_version }} diff --git a/build_tools/github_actions/configure_jax_bazel_cache.py b/build_tools/github_actions/configure_jax_bazel_cache.py index 18c5cd5bdbb..12191064d34 100644 --- a/build_tools/github_actions/configure_jax_bazel_cache.py +++ b/build_tools/github_actions/configure_jax_bazel_cache.py @@ -18,9 +18,9 @@ This runs *inside* the manylinux build container rather than on the runner. The container reaches the network through the Docker daemon and receives the -credentials through a mount, so it is the only place that can tell whether Bazel -will actually reach the cache. When the credentials are absent (fork pull -requests get no secrets), the release type may not read shared entries, or the +credentials through a mount of the runner's `/data/ci-cert.*` files, so it is +the only place that can tell whether Bazel will actually reach the cache. When +those files are absent, the release type may not read shared entries, or the endpoint does not answer, this prints nothing and the build runs exactly as it would without a cache. @@ -39,7 +39,7 @@ REMOTE_CACHE_URL = "grpcs://wardite.cluster.engflow.com" # Paths `build/rocm/rocm.bazelrc` already expects, populated by mounting the -# runner's decoded credentials at /data. +# runner's `/data/ci-cert.*` files into the build container. CLIENT_CERTIFICATE = Path("/data/ci-cert.crt") CLIENT_KEY = Path("/data/ci-cert.key") @@ -163,8 +163,8 @@ def main(argv: list[str] | None = None): missing = [p for p in (args.client_certificate, args.client_key) if not p.exists()] if missing: - # Fork pull requests get no secrets, so this is a normal outcome and - # not an error: the build simply runs without a cache. + # Runners without `/data/ci-cert.*` skip the cache. That is a normal + # outcome, not an error: the build simply compiles locally. _log(f"No cache credentials at {', '.join(str(p) for p in missing)}") return diff --git a/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py index a569b9a5c8b..9944f2b68c8 100644 --- a/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py +++ b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py @@ -179,7 +179,7 @@ def test_prints_options_for_a_reachable_cache(self): self.assertIn("--bazel_options=--remote_upload_local_results=true", output) def test_prints_nothing_without_credentials(self): - # Fork pull requests get no secrets, so the mount is absent and the + # Runners without `/data/ci-cert.*` leave the mount absent, so the # build command must come out unchanged. self.assertEqual( self._main(["--release-type", "ci"], with_credentials=False), "" diff --git a/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py b/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py new file mode 100644 index 00000000000..279e6a4b0e9 --- /dev/null +++ b/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +# Copyright Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Unit tests for write_jax_bazel_cache_credentials.py""" + +import os +import stat +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from write_jax_bazel_cache_credentials import ( + CERT_NAME, + KEY_NAME, + main, + materialize_credentials, +) + + +class MaterializeCredentialsTest(unittest.TestCase): + """Tests that runner-local certs are copied privately, or skipped.""" + + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.source_dir = Path(self.tmp.name) / "source" + self.dest_dir = Path(self.tmp.name) / "dest" + self.source_dir.mkdir() + + def _write_source(self) -> None: + # Existence and size are what matter; real PEM headers trip scanners. + (self.source_dir / CERT_NAME).write_text("test certificate") + (self.source_dir / KEY_NAME).write_text("test key") + + def test_copies_both_files_when_present(self): + self._write_source() + copied = materialize_credentials(self.source_dir, self.dest_dir) + self.assertEqual(copied, self.dest_dir) + self.assertEqual((self.dest_dir / CERT_NAME).read_text(), "test certificate") + self.assertEqual((self.dest_dir / KEY_NAME).read_text(), "test key") + + def test_returns_none_when_the_certificate_is_missing(self): + (self.source_dir / KEY_NAME).write_text("test key") + self.assertIsNone(materialize_credentials(self.source_dir, self.dest_dir)) + self.assertFalse(self.dest_dir.exists()) + + def test_returns_none_when_the_key_is_missing(self): + (self.source_dir / CERT_NAME).write_text("test certificate") + self.assertIsNone(materialize_credentials(self.source_dir, self.dest_dir)) + + def test_returns_none_when_the_source_dir_is_empty(self): + self.assertIsNone(materialize_credentials(self.source_dir, self.dest_dir)) + + def test_returns_none_when_a_source_file_is_empty(self): + (self.source_dir / CERT_NAME).write_text("test certificate") + (self.source_dir / KEY_NAME).write_text("") + self.assertIsNone(materialize_credentials(self.source_dir, self.dest_dir)) + + def test_replaces_a_stale_destination(self): + self._write_source() + self.dest_dir.mkdir() + (self.dest_dir / "stale").write_text("leftover") + materialize_credentials(self.source_dir, self.dest_dir) + self.assertFalse((self.dest_dir / "stale").exists()) + self.assertTrue((self.dest_dir / CERT_NAME).is_file()) + + def test_restricts_destination_permissions(self): + self._write_source() + materialize_credentials(self.source_dir, self.dest_dir) + dir_mode = stat.S_IMODE(self.dest_dir.stat().st_mode) + cert_mode = stat.S_IMODE((self.dest_dir / CERT_NAME).stat().st_mode) + key_mode = stat.S_IMODE((self.dest_dir / KEY_NAME).stat().st_mode) + self.assertEqual(dir_mode, 0o700) + self.assertEqual(cert_mode, 0o600) + self.assertEqual(key_mode, 0o600) + + +class MainTest(unittest.TestCase): + """Tests that GITHUB_OUTPUT receives the destination path, or empty.""" + + def _run(self, source_dir: Path, dest_dir: Path) -> str: + with tempfile.TemporaryDirectory() as tmp: + output = Path(tmp) / "github_output" + env = {"GITHUB_OUTPUT": str(output), "RUNNER_TEMP": tmp} + with mock.patch.dict(os.environ, env, clear=False): + main(["--source-dir", str(source_dir), "--dest-dir", str(dest_dir)]) + return output.read_text() + + def test_main_writes_the_destination_when_credentials_exist(self): + with tempfile.TemporaryDirectory() as tmp: + source_dir = Path(tmp) / "source" + dest_dir = Path(tmp) / "dest" + source_dir.mkdir() + (source_dir / CERT_NAME).write_text("test certificate") + (source_dir / KEY_NAME).write_text("test key") + output = self._run(source_dir, dest_dir) + self.assertIn(f"credentials_dir={dest_dir}", output) + + def test_main_writes_empty_when_credentials_are_missing(self): + with tempfile.TemporaryDirectory() as tmp: + source_dir = Path(tmp) / "source" + dest_dir = Path(tmp) / "dest" + source_dir.mkdir() + output = self._run(source_dir, dest_dir) + self.assertRegex(output, r"credentials_dir=\s*$") + + +if __name__ == "__main__": + unittest.main() diff --git a/build_tools/github_actions/write_jax_bazel_cache_credentials.py b/build_tools/github_actions/write_jax_bazel_cache_credentials.py new file mode 100644 index 00000000000..5c46e4fa55a --- /dev/null +++ b/build_tools/github_actions/write_jax_bazel_cache_credentials.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +# Copyright Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Copy runner-local EngFlow client credentials into $RUNNER_TEMP. + +XLA and JAX CI keep the mTLS files at `/data/ci-cert.crt` and +`/data/ci-cert.key` on the runner, where every job can read them — including +pull requests from forks. GitHub Actions secrets are not used: they are +invisible to fork PRs and would be a third credential path next to the +runner-local AWS files and OIDC roles the other caches already use. + +The copy lands under $RUNNER_TEMP rather than the workspace, which CI +publishes as artifacts. Missing source files are not an error: the wheel +build runs without a remote cache. + +Used by `.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml`. +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import stat +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from github_actions_api import gha_set_output + +DEFAULT_SOURCE_DIR = Path("/data") +CERT_NAME = "ci-cert.crt" +KEY_NAME = "ci-cert.key" +DEST_DIRNAME = "bazel-cache-credentials" + + +def _log(msg: str) -> None: + print(f"[jax-bazel-cache] {msg}", file=sys.stderr) + + +def _restrict_private(path: Path, is_dir: bool) -> None: + """Owner-only access; ignore chmod failures on filesystems that lack mode bits.""" + mode = stat.S_IRWXU if is_dir else stat.S_IRUSR | stat.S_IWUSR + try: + os.chmod(path, mode) + except OSError as e: + _log(f"Could not chmod {path}: {e}") + + +def materialize_credentials(source_dir: Path, dest_dir: Path) -> Path | None: + """Copy the runner's cert/key into dest_dir, or return None if they are absent.""" + certificate = source_dir / CERT_NAME + key = source_dir / KEY_NAME + missing = [ + p for p in (certificate, key) if not p.is_file() or p.stat().st_size == 0 + ] + if missing: + _log( + "No cache credentials at " + + ", ".join(str(p) for p in missing) + + "; building without a remote cache." + ) + return None + + if dest_dir.exists(): + shutil.rmtree(dest_dir) + dest_dir.mkdir(parents=True) + _restrict_private(dest_dir, is_dir=True) + + for src, name in ((certificate, CERT_NAME), (key, KEY_NAME)): + dest = dest_dir / name + shutil.copyfile(src, dest) + _restrict_private(dest, is_dir=False) + + _log(f"Cache credentials copied to {dest_dir}") + return dest_dir + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser( + description="Copy runner-local EngFlow credentials into $RUNNER_TEMP." + ) + parser.add_argument( + "--source-dir", + type=Path, + default=Path( + os.environ.get("JAX_BAZEL_CACHE_CREDENTIALS_DIR") or DEFAULT_SOURCE_DIR + ), + help="Directory that already holds ci-cert.crt and ci-cert.key on the runner.", + ) + parser.add_argument( + "--dest-dir", + type=Path, + default=None, + help="Directory to write the copy into. Defaults to $RUNNER_TEMP/" + f"{DEST_DIRNAME}.", + ) + args = parser.parse_args(argv) + + dest_dir = args.dest_dir + if dest_dir is None: + runner_temp = os.environ.get("RUNNER_TEMP") + if not runner_temp: + _log("RUNNER_TEMP is unset; building without a remote cache.") + gha_set_output({"credentials_dir": ""}) + return + dest_dir = Path(runner_temp) / DEST_DIRNAME + + copied = materialize_credentials(args.source_dir, dest_dir) + gha_set_output({"credentials_dir": str(copied) if copied else ""}) + + +if __name__ == "__main__": + main() diff --git a/docs/development/ccache_troubleshooting.md b/docs/development/ccache_troubleshooting.md index 978a30f0c85..606b8b03a21 100644 --- a/docs/development/ccache_troubleshooting.md +++ b/docs/development/ccache_troubleshooting.md @@ -42,26 +42,16 @@ remote cache, accessed via ccache's `remote_storage` option with Both servers are on the Kubernetes cluster, accessible without authentication from any pod in the cluster. -The JAX wheel build does not use ccache. It compiles jaxlib, the ROCm plugin -and the XLA ROCm backend with Bazel, and reads Bazel's own remote cache on the -EngFlow cluster that ROCm's JAX CI already uses. `configure_jax_bazel_cache.py` -runs inside the manylinux build container, completes an mTLS handshake with the -cache, and prints no Bazel options at all when the credentials are absent, the -release type may not read shared entries, or the endpoint does not answer, so a -cache that cannot help leaves the build command unchanged. - -The client credentials come from the `RBE_CI_CERT` and `RBE_CI_KEY` secrets -(PEM or base64-encoded PEM). The workflow decodes them under `RUNNER_TEMP` and -mounts that directory read only at `/data`, the path `build/rocm/rocm.bazelrc` -expects. Fork pull requests receive no secrets and therefore no cache. Set the -repository variable `JAX_BAZEL_REMOTE_CACHE_URL` to point the build at a -different cache. - -Only the cache is shared: the build runs `--config=rocm_release_wheel` rather -than `rocm_rbe`, so it never picks up that config's `--remote_executor` and -every action still runs on the runner. Pull requests read the cache without -writing to it, so no proposed change can place an entry in front of a later -build; pushes fill it. +### Which builds use which cache + +| Build | Cache | Details | +| ------------------------- | --------------------------------- | ----------------------------------------------------------------------- | +| ROCm CMake / Ninja stages | ccache → bazel-remote (this page) | Default for TheRock CI | +| PyTorch wheels | sccache → S3 | [Cache buckets](s3_buckets.md#cache-buckets) | +| JAX wheels | Bazel remote cache on EngFlow | [JAX README](../../external-builds/jax/README.md#ci-bazel-remote-cache) | + +This page covers ccache only. Builds that use another cache are linked above +rather than documented here. ### Namespace version diff --git a/external-builds/jax/README.md b/external-builds/jax/README.md index dd9ea48c5f9..5ddf02c9c0c 100644 --- a/external-builds/jax/README.md +++ b/external-builds/jax/README.md @@ -10,6 +10,7 @@ Table of contents: - [Support status](#support-status) - [Build instructions](#build-instructions) +- [CI Bazel remote cache](#ci-bazel-remote-cache) - [Test instructions](#test-instructions) - [Nightly releases](#nightly-releases) @@ -147,6 +148,39 @@ from the TheRock multi-arch Python package index (the manylinux flow). For more detailed build options, see the `ROCm/jax` repository and the `.github/workflows/multi_arch_build_linux_jax_wheels.yml` workflow in TheRock. +## CI Bazel remote cache + +The Linux JAX wheel job compiles with Bazel inside a fresh manylinux container, +so it cannot use the ccache setup that covers TheRock's CMake stages (see +[CCache Troubleshooting](/docs/development/ccache_troubleshooting.md)). PyTorch +wheels use sccache and S3 instead ([cache buckets](/docs/development/s3_buckets.md#cache-buckets)). + +JAX points Bazel at EngFlow (`grpcs://wardite.cluster.engflow.com`) as a +**cache only**. The build keeps `--config=rocm_release_wheel` and never enables +`rocm_rbe`, so it does not pick up `--remote_executor` and every action still +runs on the runner. + +Credentials are the runner-local files at `/data/ci-cert.crt` and +`/data/ci-cert.key` — the same path XLA and JAX CI already mount, not GitHub +Actions secrets. Fork pull requests see those files, so they can **read** the +cache. `write_jax_bazel_cache_credentials.py` copies them into `$RUNNER_TEMP` +(the workspace is published as artifacts) and the wheel job mounts that copy +read-only at `/data`. Set the repository variable +`JAX_BAZEL_CACHE_CREDENTIALS_DIR` if the runner keeps the files somewhere else. + +`configure_jax_bazel_cache.py` runs inside the build container and prints no +Bazel options when the files are missing, the release type may not read shared +entries (`prerelease`, `nightly-bkc`), or the endpoint does not complete a TLS +handshake. Those cases leave the build command unchanged. + +Pull requests pass `--remote_upload_local_results=false`. Pushes to `main` (and +`workflow_dispatch`) upload. Set `JAX_BAZEL_REMOTE_CACHE_URL` to point the +build at a different cache. + +EngFlow entries from `ROCm/jax` CI do not help TheRock: that CI compiles with +`--config=rocm_rbe`, so the action keys differ. TheRock has to populate the +cache itself on postsubmit. + ## Test instructions ### Prerequisites From 525a491f147712fb75b16a3a85c62abb50bba087 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Mon, 21 Sep 2026 20:13:09 +0000 Subject: [PATCH 3/7] ci(jax): serve the release JAX wheel build from the cache too The cache reached the CI wheel build only, so every release build - and the upstream-tip row that rides in the release matrix - recompiled jaxlib, the plugin and the XLA ROCm backend from scratch nightly, and never left an entry behind for the next one. Both jobs run the same Bazel invocation on the same runner pool, so this is the same three pieces: locate the runner-local credentials, mount them into the build container, and expand whatever configure_jax_bazel_cache.py prints. What differs is the gate. The CI job hardcodes a release type because it has none; here it is an input, so prerelease and nightly-bkc fall out of SHARED_CACHE_RELEASE_TYPES on their own and reach no cache, which is the policy the script already carries. Uploads are unconditional rather than keyed on the event: a release build is the canonical build of its ref and never runs on a pull request, which is the only case that must not write. --- .../multi_arch_build_linux_jax_wheels.yml | 28 +++++++++++++++++++ external-builds/jax/README.md | 17 ++++++----- 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/.github/workflows/multi_arch_build_linux_jax_wheels.yml b/.github/workflows/multi_arch_build_linux_jax_wheels.yml index 853bf7ef7a8..f18d36976e0 100644 --- a/.github/workflows/multi_arch_build_linux_jax_wheels.yml +++ b/.github/workflows/multi_arch_build_linux_jax_wheels.yml @@ -257,17 +257,43 @@ jobs: --rocm-version "${ROCM_VERSION}" + # Runner-local files at /data/ci-cert.{crt,key}, the same path XLA and + # JAX CI mount. Missing files leave the mount absent. + - name: Locate Bazel remote cache credentials + id: cache_credentials + env: + JAX_BAZEL_CACHE_CREDENTIALS_DIR: ${{ vars.JAX_BAZEL_CACHE_CREDENTIALS_DIR }} + run: | + python build_tools/github_actions/write_jax_bazel_cache_credentials.py + - name: Build JAX Wheels in manylinux container working-directory: jax-source env: PYTHON_VERSION: ${{ inputs.python_version }} JAX_COMMIT: ${{ steps.checkout_jax.outputs.commit }} + # The release type decides whether this build shares entries at all, + # so prerelease stays isolated without a rule of its own here. + RELEASE_TYPE: ${{ inputs.release_type }} + JAX_BAZEL_CACHE_ALLOW_UPLOAD: "true" + JAX_BAZEL_REMOTE_CACHE_URL: ${{ vars.JAX_BAZEL_REMOTE_CACHE_URL }} + CACHE_CREDENTIALS_DIR: ${{ steps.cache_credentials.outputs.credentials_dir }} run: | + # rocm.bazelrc reads the client credentials from /data, copied + # from the runner's /data/ci-cert.* into $RUNNER_TEMP. + credentials_mount="" + if [[ -n "${CACHE_CREDENTIALS_DIR}" ]]; then + credentials_mount="--volume ${CACHE_CREDENTIALS_DIR}:/data:ro" + fi + docker run --rm \ --user root \ --env JAX_COMMIT \ + ${credentials_mount} \ --env ROCM_VERSION="${ROCM_VERSION}" \ --env PYTHON_VERSION="${PYTHON_VERSION}" \ + --env RELEASE_TYPE="${RELEASE_TYPE}" \ + --env JAX_BAZEL_CACHE_ALLOW_UPLOAD="${JAX_BAZEL_CACHE_ALLOW_UPLOAD}" \ + --env JAX_BAZEL_REMOTE_CACHE_URL="${JAX_BAZEL_REMOTE_CACHE_URL}" \ --env ML_WHEEL_VERSION_SUFFIX="${version_suffix}" \ --env PACKAGE_DIST_DIR="/workspace/jax-source/dist" \ --volume "${GITHUB_WORKSPACE}:/workspace" \ @@ -275,6 +301,7 @@ jobs: "${MANYLINUX_IMAGE_TAG}" \ bash -lc ' ROCM_ROOT="$(rocm-sdk path --root)" + bazel_cache_options="$(python /workspace/build_tools/github_actions/configure_jax_bazel_cache.py)" python build/build.py build \ --wheels=jax-rocm-plugin,jax-rocm-pjrt \ --python_version="${PYTHON_VERSION}" \ @@ -284,6 +311,7 @@ jobs: --bazel_options=--repo_env=ML_WHEEL_TYPE=release \ --bazel_options=--repo_env=ML_WHEEL_VERSION_SUFFIX="${ML_WHEEL_VERSION_SUFFIX}" \ --bazel_options=--//jaxlib/tools:jaxlib_git_hash="${JAX_COMMIT}" \ + ${bazel_cache_options} \ --verbose \ --detailed_timestamped_log \ --output_path="$(pwd)/dist" diff --git a/external-builds/jax/README.md b/external-builds/jax/README.md index 5ddf02c9c0c..87fc57c9b2a 100644 --- a/external-builds/jax/README.md +++ b/external-builds/jax/README.md @@ -10,7 +10,7 @@ Table of contents: - [Support status](#support-status) - [Build instructions](#build-instructions) -- [CI Bazel remote cache](#ci-bazel-remote-cache) +- [Bazel remote cache](#bazel-remote-cache) - [Test instructions](#test-instructions) - [Nightly releases](#nightly-releases) @@ -148,10 +148,11 @@ from the TheRock multi-arch Python package index (the manylinux flow). For more detailed build options, see the `ROCm/jax` repository and the `.github/workflows/multi_arch_build_linux_jax_wheels.yml` workflow in TheRock. -## CI Bazel remote cache +## Bazel remote cache -The Linux JAX wheel job compiles with Bazel inside a fresh manylinux container, -so it cannot use the ccache setup that covers TheRock's CMake stages (see +Both Linux JAX wheel jobs, CI and release, compile with Bazel inside a fresh +manylinux container, +so they cannot use the ccache setup that covers TheRock's CMake stages (see [CCache Troubleshooting](/docs/development/ccache_troubleshooting.md)). PyTorch wheels use sccache and S3 instead ([cache buckets](/docs/development/s3_buckets.md#cache-buckets)). @@ -173,9 +174,11 @@ Bazel options when the files are missing, the release type may not read shared entries (`prerelease`, `nightly-bkc`), or the endpoint does not complete a TLS handshake. Those cases leave the build command unchanged. -Pull requests pass `--remote_upload_local_results=false`. Pushes to `main` (and -`workflow_dispatch`) upload. Set `JAX_BAZEL_REMOTE_CACHE_URL` to point the -build at a different cache. +Pull requests pass `--remote_upload_local_results=false`, so no proposed change +can put an entry in front of a later build. Everything else uploads, release +builds included: they are the canonical build of a ref, and a release type that +may not share entries never reaches the cache at all. Set +`JAX_BAZEL_REMOTE_CACHE_URL` to point the build at a different cache. EngFlow entries from `ROCm/jax` CI do not help TheRock: that CI compiles with `--config=rocm_rbe`, so the action keys differ. TheRock has to populate the From 91a5d4f291051b19e2e396668cb4bc9dbc57c747 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Tue, 6 Oct 2026 20:06:16 +0000 Subject: [PATCH 4/7] ci(jax): keep the Bazel cache to the CI wheel build for now Release builds would upload whenever the cache is on, and nightly would share CI's entries, where ccache gives nightly a server of its own. Leave the release wheel build uncached and let only the ci release type reach the cache until release builds get one of their own. --- .../multi_arch_build_linux_jax_wheels.yml | 28 ------------------- .../configure_jax_bazel_cache.py | 7 ++--- .../tests/configure_jax_bazel_cache_test.py | 13 ++++----- external-builds/jax/README.md | 20 ++++++------- 4 files changed, 19 insertions(+), 49 deletions(-) diff --git a/.github/workflows/multi_arch_build_linux_jax_wheels.yml b/.github/workflows/multi_arch_build_linux_jax_wheels.yml index f18d36976e0..853bf7ef7a8 100644 --- a/.github/workflows/multi_arch_build_linux_jax_wheels.yml +++ b/.github/workflows/multi_arch_build_linux_jax_wheels.yml @@ -257,43 +257,17 @@ jobs: --rocm-version "${ROCM_VERSION}" - # Runner-local files at /data/ci-cert.{crt,key}, the same path XLA and - # JAX CI mount. Missing files leave the mount absent. - - name: Locate Bazel remote cache credentials - id: cache_credentials - env: - JAX_BAZEL_CACHE_CREDENTIALS_DIR: ${{ vars.JAX_BAZEL_CACHE_CREDENTIALS_DIR }} - run: | - python build_tools/github_actions/write_jax_bazel_cache_credentials.py - - name: Build JAX Wheels in manylinux container working-directory: jax-source env: PYTHON_VERSION: ${{ inputs.python_version }} JAX_COMMIT: ${{ steps.checkout_jax.outputs.commit }} - # The release type decides whether this build shares entries at all, - # so prerelease stays isolated without a rule of its own here. - RELEASE_TYPE: ${{ inputs.release_type }} - JAX_BAZEL_CACHE_ALLOW_UPLOAD: "true" - JAX_BAZEL_REMOTE_CACHE_URL: ${{ vars.JAX_BAZEL_REMOTE_CACHE_URL }} - CACHE_CREDENTIALS_DIR: ${{ steps.cache_credentials.outputs.credentials_dir }} run: | - # rocm.bazelrc reads the client credentials from /data, copied - # from the runner's /data/ci-cert.* into $RUNNER_TEMP. - credentials_mount="" - if [[ -n "${CACHE_CREDENTIALS_DIR}" ]]; then - credentials_mount="--volume ${CACHE_CREDENTIALS_DIR}:/data:ro" - fi - docker run --rm \ --user root \ --env JAX_COMMIT \ - ${credentials_mount} \ --env ROCM_VERSION="${ROCM_VERSION}" \ --env PYTHON_VERSION="${PYTHON_VERSION}" \ - --env RELEASE_TYPE="${RELEASE_TYPE}" \ - --env JAX_BAZEL_CACHE_ALLOW_UPLOAD="${JAX_BAZEL_CACHE_ALLOW_UPLOAD}" \ - --env JAX_BAZEL_REMOTE_CACHE_URL="${JAX_BAZEL_REMOTE_CACHE_URL}" \ --env ML_WHEEL_VERSION_SUFFIX="${version_suffix}" \ --env PACKAGE_DIST_DIR="/workspace/jax-source/dist" \ --volume "${GITHUB_WORKSPACE}:/workspace" \ @@ -301,7 +275,6 @@ jobs: "${MANYLINUX_IMAGE_TAG}" \ bash -lc ' ROCM_ROOT="$(rocm-sdk path --root)" - bazel_cache_options="$(python /workspace/build_tools/github_actions/configure_jax_bazel_cache.py)" python build/build.py build \ --wheels=jax-rocm-plugin,jax-rocm-pjrt \ --python_version="${PYTHON_VERSION}" \ @@ -311,7 +284,6 @@ jobs: --bazel_options=--repo_env=ML_WHEEL_TYPE=release \ --bazel_options=--repo_env=ML_WHEEL_VERSION_SUFFIX="${ML_WHEEL_VERSION_SUFFIX}" \ --bazel_options=--//jaxlib/tools:jaxlib_git_hash="${JAX_COMMIT}" \ - ${bazel_cache_options} \ --verbose \ --detailed_timestamped_log \ --output_path="$(pwd)/dist" diff --git a/build_tools/github_actions/configure_jax_bazel_cache.py b/build_tools/github_actions/configure_jax_bazel_cache.py index 12191064d34..cff76e498ca 100644 --- a/build_tools/github_actions/configure_jax_bazel_cache.py +++ b/build_tools/github_actions/configure_jax_bazel_cache.py @@ -43,10 +43,9 @@ CLIENT_CERTIFICATE = Path("/data/ci-cert.crt") CLIENT_KEY = Path("/data/ci-cert.key") -# Release types that may read shared cache entries, mirroring the ccache policy -# in setup_ccache.py. Stable releases repackage prerelease artifacts, so -# "prerelease" and "nightly-bkc" are absent on purpose. -SHARED_CACHE_RELEASE_TYPES = frozenset({"ci", "dev", "dev-bkc", "nightly"}) +# Release types that may use the cache. Only CI for now: release builds need a +# cache of their own, as nightly has for ccache in setup_ccache.py. +SHARED_CACHE_RELEASE_TYPES = frozenset({"ci"}) DEFAULT_PROBE_TIMEOUT_SECONDS = 10 DEFAULT_REMOTE_TIMEOUT_SECONDS = 60 diff --git a/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py index 9944f2b68c8..f5814920bcf 100644 --- a/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py +++ b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py @@ -28,15 +28,14 @@ class ResolveCacheUrlTest(unittest.TestCase): """Tests which release types may read the shared cache.""" - def test_ci_and_dev_use_the_shared_cache(self): - for release_type in ("ci", "dev", "dev-bkc", "nightly"): - self.assertEqual(resolve_cache_url("", release_type), REMOTE_CACHE_URL) + def test_ci_uses_the_shared_cache(self): + self.assertEqual(resolve_cache_url("", "ci"), REMOTE_CACHE_URL) def test_release_builds_get_no_shared_cache(self): - # Stable releases repackage prerelease artifacts, so they must not read - # entries another build could have written. - for release_type in ("nightly-bkc", "prerelease"): - self.assertEqual(resolve_cache_url("", release_type), "") + # Release builds need a cache of their own, as nightly has for ccache. + for release_type in ("dev", "dev-bkc", "nightly", "nightly-bkc", "prerelease"): + with self.subTest(release_type=release_type): + self.assertEqual(resolve_cache_url("", release_type), "") def test_unknown_release_type_gets_no_cache(self): self.assertEqual(resolve_cache_url("", "something-new"), "") diff --git a/external-builds/jax/README.md b/external-builds/jax/README.md index 87fc57c9b2a..8e5aa551f0a 100644 --- a/external-builds/jax/README.md +++ b/external-builds/jax/README.md @@ -150,12 +150,14 @@ For more detailed build options, see the `ROCm/jax` repository and the ## Bazel remote cache -Both Linux JAX wheel jobs, CI and release, compile with Bazel inside a fresh -manylinux container, -so they cannot use the ccache setup that covers TheRock's CMake stages (see +The Linux JAX wheel job compiles with Bazel inside a fresh manylinux container, +so it cannot use the ccache setup that covers TheRock's CMake stages (see [CCache Troubleshooting](/docs/development/ccache_troubleshooting.md)). PyTorch wheels use sccache and S3 instead ([cache buckets](/docs/development/s3_buckets.md#cache-buckets)). +Only the CI wheel build uses the cache. Release builds stay uncached until they +have a cache of their own, the way nightly has its own ccache server. + JAX points Bazel at EngFlow (`grpcs://wardite.cluster.engflow.com`) as a **cache only**. The build keeps `--config=rocm_release_wheel` and never enables `rocm_rbe`, so it does not pick up `--remote_executor` and every action still @@ -170,15 +172,13 @@ read-only at `/data`. Set the repository variable `JAX_BAZEL_CACHE_CREDENTIALS_DIR` if the runner keeps the files somewhere else. `configure_jax_bazel_cache.py` runs inside the build container and prints no -Bazel options when the files are missing, the release type may not read shared -entries (`prerelease`, `nightly-bkc`), or the endpoint does not complete a TLS -handshake. Those cases leave the build command unchanged. +Bazel options when the files are missing, the release type is not `ci`, or the +endpoint does not complete a TLS handshake. Those cases leave the build command +unchanged. Pull requests pass `--remote_upload_local_results=false`, so no proposed change -can put an entry in front of a later build. Everything else uploads, release -builds included: they are the canonical build of a ref, and a release type that -may not share entries never reaches the cache at all. Set -`JAX_BAZEL_REMOTE_CACHE_URL` to point the build at a different cache. +can put an entry in front of a later build. Set `JAX_BAZEL_REMOTE_CACHE_URL` to +point the build at a different cache. EngFlow entries from `ROCm/jax` CI do not help TheRock: that CI compiles with `--config=rocm_rbe`, so the action keys differ. TheRock has to populate the From a9ddda02ccd7a3fbe540e1e2f879240f76849371 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Tue, 6 Oct 2026 20:09:33 +0000 Subject: [PATCH 5/7] ci(jax): fail fast on an unusable Bazel cache, upload only from main A rejected or empty credential, or a cache that does not answer, used to fall back to an uncached build without a word. Only credentials that were never provisioned do that now; anything else fails the job. The probe also sends the HTTP/2 preface gRPC opens with and reads the reply, since TLS 1.3 reports a rejected client certificate only after the handshake has completed. Uploads come from pushes to main alone, decided in the script rather than an inline expression, and the cache URL override no longer turns the cache on for release types that may not use it. The build script in the container now stops at the first failing command. --- .../multi_arch_build_linux_jax_wheels_ci.yml | 7 +- .../configure_jax_bazel_cache.py | 86 ++++++++----- .../tests/configure_jax_bazel_cache_test.py | 116 ++++++++++++++---- .../write_jax_bazel_cache_credentials_test.py | 7 +- .../write_jax_bazel_cache_credentials.py | 17 ++- docs/development/ccache_troubleshooting.md | 10 +- external-builds/jax/README.md | 16 +-- 7 files changed, 172 insertions(+), 87 deletions(-) diff --git a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml index ad8a7a1a073..64d220a047b 100644 --- a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml +++ b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml @@ -235,9 +235,6 @@ jobs: env: PYTHON_VERSION: ${{ inputs.python_version }} JAX_COMMIT: ${{ steps.checkout_jax.outputs.commit }} - # Pull requests read the cache without writing to it, so no proposed - # change can put an entry in front of a later build. Pushes fill it. - JAX_BAZEL_CACHE_ALLOW_UPLOAD: ${{ github.event_name != 'pull_request' }} JAX_BAZEL_REMOTE_CACHE_URL: ${{ vars.JAX_BAZEL_REMOTE_CACHE_URL }} CACHE_CREDENTIALS_DIR: ${{ steps.cache_credentials.outputs.credentials_dir }} run: | @@ -255,7 +252,8 @@ jobs: --env ROCM_VERSION="${ROCM_VERSION}" \ --env PYTHON_VERSION="${PYTHON_VERSION}" \ --env RELEASE_TYPE="${RELEASE_TYPE}" \ - --env JAX_BAZEL_CACHE_ALLOW_UPLOAD="${JAX_BAZEL_CACHE_ALLOW_UPLOAD}" \ + --env GITHUB_EVENT_NAME \ + --env GITHUB_REF \ --env JAX_BAZEL_REMOTE_CACHE_URL="${JAX_BAZEL_REMOTE_CACHE_URL}" \ --env ML_WHEEL_VERSION_SUFFIX="${version_suffix}" \ --env PACKAGE_DIST_DIR="/workspace/jax-source/dist" \ @@ -263,6 +261,7 @@ jobs: --workdir /workspace/jax-source \ "${MANYLINUX_IMAGE_TAG}" \ bash -lc ' + set -e ROCM_ROOT="$(rocm-sdk path --root)" bazel_cache_options="$(python /workspace/build_tools/github_actions/configure_jax_bazel_cache.py)" python build/build.py build \ diff --git a/build_tools/github_actions/configure_jax_bazel_cache.py b/build_tools/github_actions/configure_jax_bazel_cache.py index cff76e498ca..f751a05179a 100644 --- a/build_tools/github_actions/configure_jax_bazel_cache.py +++ b/build_tools/github_actions/configure_jax_bazel_cache.py @@ -20,9 +20,9 @@ container reaches the network through the Docker daemon and receives the credentials through a mount of the runner's `/data/ci-cert.*` files, so it is the only place that can tell whether Bazel will actually reach the cache. When -those files are absent, the release type may not read shared entries, or the -endpoint does not answer, this prints nothing and the build runs exactly as it -would without a cache. +those files are absent or the release type may not use the cache, this prints +nothing and the build runs exactly as it would without a cache. Credentials +that are present but rejected, or a cache that does not answer, raise instead. Options are printed to stdout so the caller can expand them into the build command; all logging goes to stderr. @@ -51,6 +51,16 @@ DEFAULT_REMOTE_TIMEOUT_SECONDS = 60 DEFAULT_TLS_PORT = 443 +# The HTTP/2 connection preface and an empty SETTINGS frame, which every gRPC +# client sends first and a gRPC server answers with its own SETTINGS frame. +HTTP2_CLIENT_PREFACE = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n" + bytes.fromhex( + "000000040000000000" +) + + +class CacheUnusableError(RuntimeError): + """The cache is configured and credentialed but cannot be used.""" + def _log(msg: str): print(f"[jax-bazel-cache] {msg}", file=sys.stderr) @@ -59,14 +69,18 @@ def _log(msg: str): def resolve_cache_url(cache_url: str, release_type: str) -> str: """Returns the cache URL to use, or "" when caching is disabled. - An explicit `cache_url` wins so infrastructure can repoint the cache - without a code change; otherwise the release type decides. + The release type decides whether to cache at all; an explicit `cache_url` + only repoints where the cache is, so infrastructure can move it without a + code change. """ - if cache_url: - return cache_url - if release_type in SHARED_CACHE_RELEASE_TYPES: - return REMOTE_CACHE_URL - return "" + if release_type not in SHARED_CACHE_RELEASE_TYPES: + return "" + return cache_url or REMOTE_CACHE_URL + + +def upload_allowed(event_name: str, ref: str) -> bool: + """Only pushes to main write; pull requests and dispatches only read.""" + return event_name == "push" and ref == "refs/heads/main" def endpoint_address(url: str) -> tuple[str, int]: @@ -80,27 +94,31 @@ def probe_cache( certificate: Path, key: Path, timeout: float = DEFAULT_PROBE_TIMEOUT_SECONDS, -) -> bool: - """Returns whether the cache completes a TLS handshake from this process. +): + """Raises CacheUnusableError unless the cache accepts the client credentials. - This resolves the host, connects, and presents the client credentials, so a - stale or mismatched certificate is reported here rather than as a wall of - Bazel cache warnings during the build. + Under TLS 1.3 the server checks the client certificate after the handshake + has completed on the client side, so a rejection only arrives on the first + read. Sending the preface gRPC opens with and reading the reply surfaces it + here rather than as a wall of Bazel cache warnings during the build. """ host, port = endpoint_address(url) if not host: - _log(f"Cannot parse a host out of {url!r}; building without a remote cache") - return False + raise CacheUnusableError(f"Cannot parse a host out of {url!r}") + context = ssl.create_default_context() + context.minimum_version = ssl.TLSVersion.TLSv1_2 + context.set_alpn_protocols(["h2"]) try: - context = ssl.create_default_context() - context.minimum_version = ssl.TLSVersion.TLSv1_2 context.load_cert_chain(certfile=certificate, keyfile=key) with socket.create_connection((host, port), timeout=timeout) as sock: - with context.wrap_socket(sock, server_hostname=host): - return True - except (OSError, ssl.SSLError) as e: - _log(f"Cache unreachable ({e}); building without a remote cache") - return False + with context.wrap_socket(sock, server_hostname=host) as tls: + tls.sendall(HTTP2_CLIENT_PREFACE) + if not tls.recv(1): + raise CacheUnusableError(f"{url} closed the connection") + except OSError as e: + raise CacheUnusableError( + f"Cannot use the cache at {url} with {certificate}: {e}" + ) from e def bazel_cache_options( @@ -123,7 +141,7 @@ def main(argv: list[str] | None = None): parser.add_argument( "--cache-url", default=os.environ.get("JAX_BAZEL_REMOTE_CACHE_URL", ""), - help="Remote cache URL. Overrides the release-type default.", + help="Remote cache URL, for release types that may use the cache.", ) parser.add_argument( "--release-type", @@ -131,9 +149,14 @@ def main(argv: list[str] | None = None): help="Release types that must not read shared entries resolve to no cache.", ) parser.add_argument( - "--allow-upload", - default=os.environ.get("JAX_BAZEL_CACHE_ALLOW_UPLOAD", "false"), - help="Whether this run may write results back to the cache.", + "--event-name", + default=os.environ.get("GITHUB_EVENT_NAME", ""), + help="GitHub event that started the run; only a push to main uploads.", + ) + parser.add_argument( + "--ref", + default=os.environ.get("GITHUB_REF", ""), + help="Git ref the run is for, e.g. refs/heads/main.", ) parser.add_argument( "--client-certificate", @@ -167,14 +190,11 @@ def main(argv: list[str] | None = None): _log(f"No cache credentials at {', '.join(str(p) for p in missing)}") return - allow_upload = args.allow_upload.strip().lower() == "true" + allow_upload = upload_allowed(args.event_name, args.ref) _log( f"Cache mode: url={url} release_type={args.release_type} upload={allow_upload}" ) - if not probe_cache( - url, args.client_certificate, args.client_key, args.probe_timeout - ): - return + probe_cache(url, args.client_certificate, args.client_key, args.probe_timeout) print( " ".join( diff --git a/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py index f5814920bcf..9f53c7c9df6 100644 --- a/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py +++ b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py @@ -5,6 +5,7 @@ """Unit tests for configure_jax_bazel_cache.py""" import io +import os import ssl import sys import tempfile @@ -16,12 +17,15 @@ import configure_jax_bazel_cache from configure_jax_bazel_cache import ( + HTTP2_CLIENT_PREFACE, REMOTE_CACHE_URL, + CacheUnusableError, bazel_cache_options, endpoint_address, main, probe_cache, resolve_cache_url, + upload_allowed, ) @@ -40,12 +44,34 @@ def test_release_builds_get_no_shared_cache(self): def test_unknown_release_type_gets_no_cache(self): self.assertEqual(resolve_cache_url("", "something-new"), "") - def test_explicit_url_overrides_the_release_type_default(self): + def test_explicit_url_repoints_the_cache(self): self.assertEqual( - resolve_cache_url("grpcs://cache.example", "prerelease"), - "grpcs://cache.example", + resolve_cache_url("grpcs://cache.example", "ci"), "grpcs://cache.example" ) + def test_explicit_url_does_not_enable_the_cache(self): + # The override only moves the cache; a release type that may not use + # it stays isolated whatever the repository variable says. + self.assertEqual(resolve_cache_url("grpcs://cache.example", "prerelease"), "") + + +class UploadAllowedTest(unittest.TestCase): + """Tests which runs may write results back to the cache.""" + + def test_push_to_main_uploads(self): + self.assertTrue(upload_allowed("push", "refs/heads/main")) + + def test_everything_else_reads_only(self): + for event_name, ref in ( + ("pull_request", "refs/pull/1/merge"), + ("workflow_dispatch", "refs/heads/main"), + ("workflow_dispatch", "refs/heads/users/someone/branch"), + ("push", "refs/heads/multi_arch/experiment"), + ("", ""), + ): + with self.subTest(event_name=event_name, ref=ref): + self.assertFalse(upload_allowed(event_name, ref)) + class EndpointAddressTest(unittest.TestCase): """Tests parsing the host and port out of a gRPC cache URL.""" @@ -63,38 +89,58 @@ def test_explicit_port_is_honored(self): class ProbeCacheTest(unittest.TestCase): - """Tests that the probe reports unreachable rather than raising.""" + """Tests that the probe raises for any cache it cannot use.""" + + def setUp(self): + self.context = mock.MagicMock() + self.tls = self.context.wrap_socket.return_value.__enter__.return_value + self.tls.recv.return_value = b"\x00" - def _probe(self, url=REMOTE_CACHE_URL, connect_effect=None, load_effect=None): - context = mock.MagicMock() - context.load_cert_chain.side_effect = load_effect + def _probe(self, url=REMOTE_CACHE_URL, connect_effect=None): with mock.patch.object( configure_jax_bazel_cache.ssl, "create_default_context", - return_value=context, + return_value=self.context, ), mock.patch.object( configure_jax_bazel_cache.socket, "create_connection", side_effect=connect_effect, ): - return probe_cache(url, Path("cert.crt"), Path("cert.key")) + probe_cache(url, Path("cert.crt"), Path("cert.key")) - def test_reachable_when_the_handshake_completes(self): - self.assertTrue(self._probe()) + def test_usable_when_the_server_answers_the_preface(self): + self._probe() + self.tls.sendall.assert_called_once_with(HTTP2_CLIENT_PREFACE) + self.context.set_alpn_protocols.assert_called_once_with(["h2"]) - def test_unreachable_when_the_host_does_not_resolve(self): - self.assertFalse(self._probe(connect_effect=OSError("name resolution failed"))) + def test_raises_when_the_host_does_not_resolve(self): + with self.assertRaises(CacheUnusableError): + self._probe(connect_effect=OSError("name resolution failed")) - def test_unreachable_when_the_connection_times_out(self): - self.assertFalse(self._probe(connect_effect=TimeoutError())) + def test_raises_when_the_connection_times_out(self): + with self.assertRaises(CacheUnusableError): + self._probe(connect_effect=TimeoutError()) - def test_unreachable_when_the_credentials_are_rejected(self): - # A stale or mismatched certificate is reported once here instead of as - # a wall of Bazel cache warnings during the build. - self.assertFalse(self._probe(load_effect=ssl.SSLError("bad key"))) + def test_raises_when_the_credentials_cannot_be_loaded(self): + self.context.load_cert_chain.side_effect = ssl.SSLError("bad key") + with self.assertRaises(CacheUnusableError): + self._probe() - def test_unreachable_when_the_url_has_no_host(self): - self.assertFalse(self._probe(url="not-a-url")) + def test_raises_when_tls13_rejects_the_certificate_after_the_handshake(self): + # Under TLS 1.3 the handshake completes before the server checks the + # client certificate, so the rejection only shows on the first read. + self.tls.recv.side_effect = ssl.SSLError("TLSV1_ALERT_UNKNOWN_CA") + with self.assertRaises(CacheUnusableError): + self._probe() + + def test_raises_when_the_server_closes_the_connection(self): + self.tls.recv.return_value = b"" + with self.assertRaises(CacheUnusableError): + self._probe() + + def test_raises_when_the_url_has_no_host(self): + with self.assertRaises(CacheUnusableError): + self._probe(url="not-a-url") class BazelCacheOptionsTest(unittest.TestCase): @@ -165,15 +211,30 @@ def _main(self, argv, reachable=True, with_credentials=True): "--client-key", str(key), ] + probe = mock.patch.object( + configure_jax_bazel_cache, + "probe_cache", + side_effect=None if reachable else CacheUnusableError("rejected"), + ) + # The defaults read the runner's own environment; keep it out so a test + # run on a push to main cannot flip the upload decision. + github_env = { + k: v + for k, v in os.environ.items() + if k + not in ("GITHUB_EVENT_NAME", "GITHUB_REF", "JAX_BAZEL_REMOTE_CACHE_URL") + } stdout = io.StringIO() - with mock.patch.object( - configure_jax_bazel_cache, "probe_cache", return_value=reachable + with probe, mock.patch.dict( + os.environ, github_env, clear=True ), mock.patch.object(sys, "stdout", stdout): main(argv) return stdout.getvalue().strip() def test_prints_options_for_a_reachable_cache(self): - output = self._main(["--release-type", "ci", "--allow-upload", "true"]) + output = self._main( + ["--release-type", "ci", "--event-name", "push", "--ref", "refs/heads/main"] + ) self.assertIn(f"--bazel_options=--remote_cache={REMOTE_CACHE_URL}", output) self.assertIn("--bazel_options=--remote_upload_local_results=true", output) @@ -184,8 +245,9 @@ def test_prints_nothing_without_credentials(self): self._main(["--release-type", "ci"], with_credentials=False), "" ) - def test_prints_nothing_when_the_cache_is_unreachable(self): - self.assertEqual(self._main(["--release-type", "ci"], reachable=False), "") + def test_fails_when_the_cache_is_unusable(self): + with self.assertRaises(CacheUnusableError): + self._main(["--release-type", "ci"], reachable=False) def test_prints_nothing_when_the_release_type_has_no_cache(self): self.assertEqual(self._main(["--release-type", "prerelease"]), "") @@ -198,7 +260,7 @@ def test_probe_is_skipped_when_no_cache_is_configured(self): ): main(["--release-type", "prerelease"]) - def test_allow_upload_defaults_to_read_only(self): + def test_upload_defaults_to_read_only(self): output = self._main(["--release-type", "ci"]) self.assertIn("--bazel_options=--remote_upload_local_results=false", output) diff --git a/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py b/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py index 279e6a4b0e9..c5c057137b5 100644 --- a/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py +++ b/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py @@ -56,10 +56,12 @@ def test_returns_none_when_the_key_is_missing(self): def test_returns_none_when_the_source_dir_is_empty(self): self.assertIsNone(materialize_credentials(self.source_dir, self.dest_dir)) - def test_returns_none_when_a_source_file_is_empty(self): + def test_raises_when_a_source_file_is_empty(self): + # Present but unusable is a broken runner, not an unprovisioned one. (self.source_dir / CERT_NAME).write_text("test certificate") (self.source_dir / KEY_NAME).write_text("") - self.assertIsNone(materialize_credentials(self.source_dir, self.dest_dir)) + with self.assertRaises(ValueError): + materialize_credentials(self.source_dir, self.dest_dir) def test_replaces_a_stale_destination(self): self._write_source() @@ -69,6 +71,7 @@ def test_replaces_a_stale_destination(self): self.assertFalse((self.dest_dir / "stale").exists()) self.assertTrue((self.dest_dir / CERT_NAME).is_file()) + @unittest.skipIf(os.name == "nt", "POSIX mode bits") def test_restricts_destination_permissions(self): self._write_source() materialize_credentials(self.source_dir, self.dest_dir) diff --git a/build_tools/github_actions/write_jax_bazel_cache_credentials.py b/build_tools/github_actions/write_jax_bazel_cache_credentials.py index 5c46e4fa55a..e7fb2be97fc 100644 --- a/build_tools/github_actions/write_jax_bazel_cache_credentials.py +++ b/build_tools/github_actions/write_jax_bazel_cache_credentials.py @@ -12,13 +12,11 @@ The copy lands under $RUNNER_TEMP rather than the workspace, which CI publishes as artifacts. Missing source files are not an error: the wheel -build runs without a remote cache. +build runs without a remote cache. Files that are present but empty are. Used by `.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml`. """ -from __future__ import annotations - import argparse import os import shutil @@ -53,9 +51,7 @@ def materialize_credentials(source_dir: Path, dest_dir: Path) -> Path | None: """Copy the runner's cert/key into dest_dir, or return None if they are absent.""" certificate = source_dir / CERT_NAME key = source_dir / KEY_NAME - missing = [ - p for p in (certificate, key) if not p.is_file() or p.stat().st_size == 0 - ] + missing = [p for p in (certificate, key) if not p.is_file()] if missing: _log( "No cache credentials at " @@ -63,6 +59,11 @@ def materialize_credentials(source_dir: Path, dest_dir: Path) -> Path | None: + "; building without a remote cache." ) return None + empty = [p for p in (certificate, key) if p.stat().st_size == 0] + if empty: + raise ValueError( + f"Cache credentials are present but empty: {', '.join(map(str, empty))}" + ) if dest_dir.exists(): shutil.rmtree(dest_dir) @@ -103,9 +104,7 @@ def main(argv: list[str] | None = None) -> None: if dest_dir is None: runner_temp = os.environ.get("RUNNER_TEMP") if not runner_temp: - _log("RUNNER_TEMP is unset; building without a remote cache.") - gha_set_output({"credentials_dir": ""}) - return + raise RuntimeError("RUNNER_TEMP is unset; pass --dest-dir") dest_dir = Path(runner_temp) / DEST_DIRNAME copied = materialize_credentials(args.source_dir, dest_dir) diff --git a/docs/development/ccache_troubleshooting.md b/docs/development/ccache_troubleshooting.md index 606b8b03a21..16f7c96ad27 100644 --- a/docs/development/ccache_troubleshooting.md +++ b/docs/development/ccache_troubleshooting.md @@ -44,11 +44,11 @@ authentication from any pod in the cluster. ### Which builds use which cache -| Build | Cache | Details | -| ------------------------- | --------------------------------- | ----------------------------------------------------------------------- | -| ROCm CMake / Ninja stages | ccache → bazel-remote (this page) | Default for TheRock CI | -| PyTorch wheels | sccache → S3 | [Cache buckets](s3_buckets.md#cache-buckets) | -| JAX wheels | Bazel remote cache on EngFlow | [JAX README](../../external-builds/jax/README.md#ci-bazel-remote-cache) | +| Build | Cache | Details | +| ------------------------- | --------------------------------- | -------------------------------------------------------------------- | +| ROCm CMake / Ninja stages | ccache → bazel-remote (this page) | Default for TheRock CI | +| PyTorch wheels | sccache → S3 | [Cache buckets](s3_buckets.md#cache-buckets) | +| JAX wheels | Bazel remote cache on EngFlow | [JAX README](../../external-builds/jax/README.md#bazel-remote-cache) | This page covers ccache only. Builds that use another cache are linked above rather than documented here. diff --git a/external-builds/jax/README.md b/external-builds/jax/README.md index 8e5aa551f0a..72eee95958f 100644 --- a/external-builds/jax/README.md +++ b/external-builds/jax/README.md @@ -172,13 +172,15 @@ read-only at `/data`. Set the repository variable `JAX_BAZEL_CACHE_CREDENTIALS_DIR` if the runner keeps the files somewhere else. `configure_jax_bazel_cache.py` runs inside the build container and prints no -Bazel options when the files are missing, the release type is not `ci`, or the -endpoint does not complete a TLS handshake. Those cases leave the build command -unchanged. - -Pull requests pass `--remote_upload_local_results=false`, so no proposed change -can put an entry in front of a later build. Set `JAX_BAZEL_REMOTE_CACHE_URL` to -point the build at a different cache. +Bazel options when the files are missing or the release type is not `ci`; the +build command is then unchanged. Credentials that are present but empty or +rejected, or a cache that does not answer, fail the job rather than build +uncached. + +Only pushes to `main` upload. Pull requests and manual dispatches pass +`--remote_upload_local_results=false`. That gate runs on the client, so what a +pull request could write is decided by the certificate's scope on EngFlow. Set +`JAX_BAZEL_REMOTE_CACHE_URL` to point the build at a different cache. EngFlow entries from `ROCm/jax` CI do not help TheRock: that CI compiles with `--config=rocm_rbe`, so the action keys differ. TheRock has to populate the From 53ee88fec89e64859985fc5fb18dd43f2361832e Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Thu, 5 Nov 2026 17:25:41 +0000 Subject: [PATCH 6/7] ci(jax): read the Bazel cache on actinolite with an EngFlow JWT wardite no longer resolves, and actinolite does not accept the runner's /data/ci-cert.* client certificate. Authenticate with an EngFlow-issued JWT instead: a cache-reader token on the runner, whose writes EngFlow rejects, so fork pull requests stay read-only on the server side. configure_jax_bazel_cache.py doubles as the Bazel credential helper (`get`), scoped to the cache host, so the token never reaches the build command line that build.py --verbose logs, nor a bazelrc that --announce_rc prints. A malformed or expired token fails the job. --- .../multi_arch_build_linux_jax_wheels_ci.yml | 10 +- .../configure_jax_bazel_cache.py | 142 +++++++------ .../tests/configure_jax_bazel_cache_test.py | 188 ++++++++++++------ .../write_jax_bazel_cache_credentials_test.py | 45 ++--- .../write_jax_bazel_cache_credentials.py | 55 +++-- external-builds/jax/README.md | 31 +-- 6 files changed, 267 insertions(+), 204 deletions(-) mode change 100644 => 100755 build_tools/github_actions/configure_jax_bazel_cache.py diff --git a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml index 65de7528cd1..0d232246517 100644 --- a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml +++ b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml @@ -248,9 +248,9 @@ jobs: --rocm-version "${ROCM_VERSION}" \ --jax-source-dir jax-source - # Runner-local files at /data/ci-cert.{crt,key}, the same path XLA and - # JAX CI mount. GitHub secrets are not used, so fork pull requests can - # read the cache. Missing files leave the mount absent. + # A runner-local EngFlow cache-reader token, so fork pull requests can + # read the cache; EngFlow rejects its writes. GitHub secrets are not + # used. A missing token leaves the mount absent. - name: Locate Bazel remote cache credentials id: cache_credentials env: @@ -268,8 +268,8 @@ jobs: JAX_BAZEL_REMOTE_CACHE_URL: ${{ vars.JAX_BAZEL_REMOTE_CACHE_URL }} CACHE_CREDENTIALS_DIR: ${{ steps.cache_credentials.outputs.credentials_dir }} run: | - # rocm.bazelrc reads the client credentials from /data, copied - # from the runner's /data/ci-cert.* into $RUNNER_TEMP. + # configure_jax_bazel_cache.py reads the token from /data, a copy + # of the runner's file in $RUNNER_TEMP. credentials_mount="" if [[ -n "${CACHE_CREDENTIALS_DIR}" ]]; then credentials_mount="--volume ${CACHE_CREDENTIALS_DIR}:/data:ro" diff --git a/build_tools/github_actions/configure_jax_bazel_cache.py b/build_tools/github_actions/configure_jax_bazel_cache.py old mode 100644 new mode 100755 index f751a05179a..94eb001f701 --- a/build_tools/github_actions/configure_jax_bazel_cache.py +++ b/build_tools/github_actions/configure_jax_bazel_cache.py @@ -10,38 +10,46 @@ whenever the JAX ref and the ROCm build are unchanged, which is the common case for a pull request that only touches CI or test-selection files. -The cache is the EngFlow cluster ROCm's JAX CI already uses, reached with the -mTLS client credentials `rocm.bazelrc` expects at `/data`. Only the cache is -used: the build runs `--config=rocm_release_wheel`, never `rocm_rbe`, so it -never picks up that config's `--remote_executor` and every action still executes -locally. +The cache is the EngFlow cluster ROCm's JAX CI already uses, reached with an +EngFlow-issued JWT. Its role is fixed when EngFlow issues it: a `cache-reader` +token cannot write, whatever the client asks for. Only the cache is used: the +build runs `--config=rocm_release_wheel`, never `rocm_rbe`, so it never picks up +that config's `--remote_executor` and every action still executes locally. This runs *inside* the manylinux build container rather than on the runner. The -container reaches the network through the Docker daemon and receives the -credentials through a mount of the runner's `/data/ci-cert.*` files, so it is -the only place that can tell whether Bazel will actually reach the cache. When -those files are absent or the release type may not use the cache, this prints -nothing and the build runs exactly as it would without a cache. Credentials -that are present but rejected, or a cache that does not answer, raise instead. +container reaches the network through the Docker daemon and receives the token +through a read-only mount of a copy of the runner's file, so it is the only +place that can tell whether Bazel will actually reach the cache. When the token +is absent or the release type may not use the cache, this prints nothing and +the build runs exactly as it would without a cache. A token that is present but +malformed or expired, or a cache that does not answer, raises instead. Options are printed to stdout so the caller can expand them into the build command; all logging goes to stderr. + +Run as `configure_jax_bazel_cache.py get`, this is also the Bazel credential +helper that hands the token to Bazel, so it never appears on a command line or +in a bazelrc that `--announce_rc` would print. """ import argparse +import base64 +import binascii +import json import os import socket import ssl import sys +import time +from datetime import datetime, timezone from pathlib import Path from urllib.parse import urlsplit -REMOTE_CACHE_URL = "grpcs://wardite.cluster.engflow.com" +REMOTE_CACHE_URL = "grpcs://actinolite.cluster.engflow.com" -# Paths `build/rocm/rocm.bazelrc` already expects, populated by mounting the -# runner's `/data/ci-cert.*` files into the build container. -CLIENT_CERTIFICATE = Path("/data/ci-cert.crt") -CLIENT_KEY = Path("/data/ci-cert.key") +# Where the wheel job mounts the copy of the runner's token inside the build +# container. +TOKEN_FILE = Path("/data/engflow-cache.jwt") # Release types that may use the cache. Only CI for now: release builds need a # cache of their own, as nightly has for ccache in setup_ccache.py. @@ -89,18 +97,32 @@ def endpoint_address(url: str) -> tuple[str, int]: return split.hostname or "", split.port or DEFAULT_TLS_PORT -def probe_cache( - url: str, - certificate: Path, - key: Path, - timeout: float = DEFAULT_PROBE_TIMEOUT_SECONDS, -): - """Raises CacheUnusableError unless the cache accepts the client credentials. +def check_token(token_file: Path, now: float | None = None): + """Raises CacheUnusableError unless token_file holds an unexpired JWT. - Under TLS 1.3 the server checks the client certificate after the handshake - has completed on the client side, so a rejection only arrives on the first - read. Sending the preface gRPC opens with and reading the reply surfaces it - here rather than as a wall of Bazel cache warnings during the build. + EngFlow rejects a bad token per request, and Bazel then builds locally with + only warnings, so an expired token would otherwise pass as a slow uncached + build. Errors name the file, never its contents. + """ + parts = token_file.read_text().strip().split(".") + try: + if len(parts) != 3: + raise ValueError(f"{len(parts)} segments") + payload = parts[1] + "=" * (-len(parts[1]) % 4) + expires = float(json.loads(base64.urlsafe_b64decode(payload))["exp"]) + except (binascii.Error, KeyError, TypeError, ValueError) as e: + raise CacheUnusableError(f"{token_file} does not hold a JWT: {e}") from e + if expires <= (time.time() if now is None else now): + expired_at = datetime.fromtimestamp(expires, timezone.utc).isoformat() + raise CacheUnusableError(f"The token in {token_file} expired at {expired_at}") + + +def probe_cache(url: str, timeout: float = DEFAULT_PROBE_TIMEOUT_SECONDS): + """Raises CacheUnusableError unless the cache answers a gRPC connection. + + Sending the preface gRPC opens with and reading the reply surfaces an + unreachable cache here rather than as a wall of Bazel cache warnings during + the build. """ host, port = endpoint_address(url) if not host: @@ -109,32 +131,46 @@ def probe_cache( context.minimum_version = ssl.TLSVersion.TLSv1_2 context.set_alpn_protocols(["h2"]) try: - context.load_cert_chain(certfile=certificate, keyfile=key) with socket.create_connection((host, port), timeout=timeout) as sock: with context.wrap_socket(sock, server_hostname=host) as tls: tls.sendall(HTTP2_CLIENT_PREFACE) if not tls.recv(1): raise CacheUnusableError(f"{url} closed the connection") except OSError as e: - raise CacheUnusableError( - f"Cannot use the cache at {url} with {certificate}: {e}" - ) from e + raise CacheUnusableError(f"Cannot reach the cache at {url}: {e}") from e + +def credential_helper_response(token: str) -> dict: + """The Bazel credential helper reply carrying an EngFlow-issued JWT.""" + return { + "headers": { + "x-engflow-auth-method": ["jwt-v0"], + "x-engflow-auth-token": [token], + } + } -def bazel_cache_options( - url: str, certificate: Path, key: Path, allow_upload: bool -) -> list[str]: + +def bazel_cache_options(url: str, helper: Path, allow_upload: bool) -> list[str]: """Returns the `--bazel_options=...` arguments for a usable cache.""" + host, _ = endpoint_address(url) return [ f"--bazel_options=--remote_cache={url}", - f"--bazel_options=--tls_client_certificate={certificate}", - f"--bazel_options=--tls_client_key={key}", + # Scoped to the cache host, so downloads from other hosts never see it. + f"--bazel_options=--credential_helper={host}={helper}", f"--bazel_options=--remote_timeout={DEFAULT_REMOTE_TIMEOUT_SECONDS}", f"--bazel_options=--remote_upload_local_results={str(allow_upload).lower()}", ] def main(argv: list[str] | None = None): + argv = sys.argv[1:] if argv is None else argv + if argv[:1] == ["get"]: + # Bazel sends the request URI on stdin; the token is the same for all. + sys.stdin.read() + token = TOKEN_FILE.read_text().strip() + print(json.dumps(credential_helper_response(token))) + return + parser = argparse.ArgumentParser( description="Print build/build.py options for the JAX Bazel remote cache." ) @@ -158,18 +194,6 @@ def main(argv: list[str] | None = None): default=os.environ.get("GITHUB_REF", ""), help="Git ref the run is for, e.g. refs/heads/main.", ) - parser.add_argument( - "--client-certificate", - type=Path, - default=CLIENT_CERTIFICATE, - help="mTLS client certificate for the cache.", - ) - parser.add_argument( - "--client-key", - type=Path, - default=CLIENT_KEY, - help="mTLS client key for the cache.", - ) parser.add_argument( "--probe-timeout", type=float, @@ -183,26 +207,20 @@ def main(argv: list[str] | None = None): _log(f"No remote cache for release_type={args.release_type!r}") return - missing = [p for p in (args.client_certificate, args.client_key) if not p.exists()] - if missing: - # Runners without `/data/ci-cert.*` skip the cache. That is a normal - # outcome, not an error: the build simply compiles locally. - _log(f"No cache credentials at {', '.join(str(p) for p in missing)}") + if not TOKEN_FILE.exists(): + # Runners without the token skip the cache. That is a normal outcome, + # not an error: the build simply compiles locally. + _log(f"No cache token at {TOKEN_FILE}") return allow_upload = upload_allowed(args.event_name, args.ref) _log( f"Cache mode: url={url} release_type={args.release_type} upload={allow_upload}" ) - probe_cache(url, args.client_certificate, args.client_key, args.probe_timeout) - - print( - " ".join( - bazel_cache_options( - url, args.client_certificate, args.client_key, allow_upload - ) - ) - ) + check_token(TOKEN_FILE) + probe_cache(url, args.probe_timeout) + + print(" ".join(bazel_cache_options(url, Path(__file__).resolve(), allow_upload))) if __name__ == "__main__": diff --git a/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py index 9f53c7c9df6..31c39866c87 100644 --- a/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py +++ b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py @@ -4,11 +4,13 @@ """Unit tests for configure_jax_bazel_cache.py""" +import base64 import io +import json import os -import ssl import sys import tempfile +import time import unittest from pathlib import Path, PurePosixPath from unittest import mock @@ -21,6 +23,8 @@ REMOTE_CACHE_URL, CacheUnusableError, bazel_cache_options, + check_token, + credential_helper_response, endpoint_address, main, probe_cache, @@ -29,6 +33,17 @@ ) +def _segment(data: dict) -> str: + return base64.urlsafe_b64encode(json.dumps(data).encode()).decode().rstrip("=") + + +def fake_jwt(expires: float) -> str: + """A JWT-shaped string built at test time; literal tokens trip scanners.""" + return ".".join( + [_segment({"alg": "none"}), _segment({"exp": expires}), "test-signature"] + ) + + class ResolveCacheUrlTest(unittest.TestCase): """Tests which release types may read the shared cache.""" @@ -78,8 +93,8 @@ class EndpointAddressTest(unittest.TestCase): def test_default_port_is_tls(self): self.assertEqual( - endpoint_address("grpcs://wardite.cluster.engflow.com"), - ("wardite.cluster.engflow.com", 443), + endpoint_address("grpcs://actinolite.cluster.engflow.com"), + ("actinolite.cluster.engflow.com", 443), ) def test_explicit_port_is_honored(self): @@ -88,8 +103,47 @@ def test_explicit_port_is_honored(self): ) +class CheckTokenTest(unittest.TestCase): + """Tests that only an unexpired JWT passes.""" + + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.token_file = Path(tmp.name) / "engflow-cache.jwt" + + def _check(self, contents: str, now: float = 1000.0): + self.token_file.write_text(contents) + check_token(self.token_file, now=now) + + def test_unexpired_token_passes(self): + self._check(fake_jwt(expires=2000.0) + "\n") + + def test_expired_token_raises(self): + with self.assertRaisesRegex(CacheUnusableError, "expired"): + self._check(fake_jwt(expires=500.0)) + + def test_malformed_tokens_raise(self): + for contents in ( + "", + "not-a-jwt", + "a.b", + "a.!!!.c", + f"a.{_segment({'no_exp': 1})}.c", + f"a.{base64.urlsafe_b64encode(b'[1]').decode()}.c", + ): + with self.subTest(contents=contents): + with self.assertRaises(CacheUnusableError): + self._check(contents) + + def test_errors_never_include_the_token(self): + token = fake_jwt(expires=500.0) + with self.assertRaises(CacheUnusableError) as raised: + self._check(token) + self.assertNotIn(token.split(".")[1], str(raised.exception)) + + class ProbeCacheTest(unittest.TestCase): - """Tests that the probe raises for any cache it cannot use.""" + """Tests that the probe raises for any cache it cannot reach.""" def setUp(self): self.context = mock.MagicMock() @@ -106,7 +160,7 @@ def _probe(self, url=REMOTE_CACHE_URL, connect_effect=None): "create_connection", side_effect=connect_effect, ): - probe_cache(url, Path("cert.crt"), Path("cert.key")) + probe_cache(url) def test_usable_when_the_server_answers_the_preface(self): self._probe() @@ -121,18 +175,6 @@ def test_raises_when_the_connection_times_out(self): with self.assertRaises(CacheUnusableError): self._probe(connect_effect=TimeoutError()) - def test_raises_when_the_credentials_cannot_be_loaded(self): - self.context.load_cert_chain.side_effect = ssl.SSLError("bad key") - with self.assertRaises(CacheUnusableError): - self._probe() - - def test_raises_when_tls13_rejects_the_certificate_after_the_handshake(self): - # Under TLS 1.3 the handshake completes before the server checks the - # client certificate, so the rejection only shows on the first read. - self.tls.recv.side_effect = ssl.SSLError("TLSV1_ALERT_UNKNOWN_CA") - with self.assertRaises(CacheUnusableError): - self._probe() - def test_raises_when_the_server_closes_the_connection(self): self.tls.recv.return_value = b"" with self.assertRaises(CacheUnusableError): @@ -143,26 +185,43 @@ def test_raises_when_the_url_has_no_host(self): self._probe(url="not-a-url") +class CredentialHelperResponseTest(unittest.TestCase): + """Tests the headers handed to Bazel.""" + + def test_sends_the_token_as_an_engflow_jwt(self): + self.assertEqual( + credential_helper_response("test-token"), + { + "headers": { + "x-engflow-auth-method": ["jwt-v0"], + "x-engflow-auth-token": ["test-token"], + } + }, + ) + + class BazelCacheOptionsTest(unittest.TestCase): """Tests the emitted build/build.py options.""" - def _options(self, allow_upload): - # Paths inside the build container, so they stay POSIX even when the - # test itself runs on Windows. - return bazel_cache_options( - REMOTE_CACHE_URL, - PurePosixPath("/data/ci-cert.crt"), - PurePosixPath("/data/ci-cert.key"), - allow_upload, + # A path inside the build container, so it stays POSIX even when the test + # itself runs on Windows. + HELPER = PurePosixPath("/workspace/configure_jax_bazel_cache.py") + + def _options(self, allow_upload, url=REMOTE_CACHE_URL): + return bazel_cache_options(url, self.HELPER, allow_upload) + + def test_points_bazel_at_the_cache(self): + self.assertIn( + f"--bazel_options=--remote_cache={REMOTE_CACHE_URL}", + self._options(allow_upload=False), ) - def test_points_bazel_at_the_cache_with_credentials(self): - options = self._options(allow_upload=False) - self.assertIn(f"--bazel_options=--remote_cache={REMOTE_CACHE_URL}", options) + def test_credential_helper_is_scoped_to_the_cache_host(self): + # An unscoped helper would send the token to every host Bazel fetches. self.assertIn( - "--bazel_options=--tls_client_certificate=/data/ci-cert.crt", options + f"--bazel_options=--credential_helper=cache.example={self.HELPER}", + self._options(allow_upload=False, url="grpcs://cache.example:8980"), ) - self.assertIn("--bazel_options=--tls_client_key=/data/ci-cert.key", options) def test_never_requests_remote_execution(self): # The build uses --config=rocm_release_wheel, so it must stay @@ -188,33 +247,24 @@ class MainTest(unittest.TestCase): """Tests that stdout carries options only when the cache is usable.""" def setUp(self): - credentials = tempfile.TemporaryDirectory() - self.addCleanup(credentials.cleanup) - self.certificate = Path(credentials.name) / "ci-cert.crt" - self.key = Path(credentials.name) / "ci-cert.key" - # Only existence matters here; the handshake that would read these is - # mocked out, and real PEM headers trip secret scanning. - self.certificate.write_text("test certificate") - self.key.write_text("test key") - - def _main(self, argv, reachable=True, with_credentials=True): - # Always pass explicit paths: a developer machine may happen to have - # credentials at the default location, which would make the - # missing-credentials case pass for the wrong reason. - certificate, key = self.certificate, self.key - if not with_credentials: - certificate = self.certificate.with_name("absent.crt") - key = self.key.with_name("absent.key") - argv = argv + [ - "--client-certificate", - str(certificate), - "--client-key", - str(key), - ] + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + # Never the real default path: a machine that happens to hold a token + # there would make the missing-token case pass for the wrong reason. + self.token_file = Path(tmp.name) / "engflow-cache.jwt" + self.token = fake_jwt(expires=time.time() + 3600) + self.token_file.write_text(self.token + "\n") + patcher = mock.patch.object( + configure_jax_bazel_cache, "TOKEN_FILE", self.token_file + ) + patcher.start() + self.addCleanup(patcher.stop) + + def _main(self, argv, reachable=True, stdin=""): probe = mock.patch.object( configure_jax_bazel_cache, "probe_cache", - side_effect=None if reachable else CacheUnusableError("rejected"), + side_effect=None if reachable else CacheUnusableError("unreachable"), ) # The defaults read the runner's own environment; keep it out so a test # run on a push to main cannot flip the upload decision. @@ -227,7 +277,9 @@ def _main(self, argv, reachable=True, with_credentials=True): stdout = io.StringIO() with probe, mock.patch.dict( os.environ, github_env, clear=True - ), mock.patch.object(sys, "stdout", stdout): + ), mock.patch.object(sys, "stdout", stdout), mock.patch.object( + sys, "stdin", io.StringIO(stdin) + ): main(argv) return stdout.getvalue().strip() @@ -238,12 +290,20 @@ def test_prints_options_for_a_reachable_cache(self): self.assertIn(f"--bazel_options=--remote_cache={REMOTE_CACHE_URL}", output) self.assertIn("--bazel_options=--remote_upload_local_results=true", output) - def test_prints_nothing_without_credentials(self): - # Runners without `/data/ci-cert.*` leave the mount absent, so the - # build command must come out unchanged. - self.assertEqual( - self._main(["--release-type", "ci"], with_credentials=False), "" - ) + def test_options_never_carry_the_token(self): + # build.py --verbose logs the full command line. + self.assertNotIn(self.token, self._main(["--release-type", "ci"])) + + def test_prints_nothing_without_a_token(self): + # Runners without the token leave the mount absent, so the build + # command must come out unchanged. + self.token_file.unlink() + self.assertEqual(self._main(["--release-type", "ci"]), "") + + def test_fails_when_the_token_has_expired(self): + self.token_file.write_text(fake_jwt(expires=time.time() - 60)) + with self.assertRaisesRegex(CacheUnusableError, "expired"): + self._main(["--release-type", "ci"]) def test_fails_when_the_cache_is_unusable(self): with self.assertRaises(CacheUnusableError): @@ -264,6 +324,12 @@ def test_upload_defaults_to_read_only(self): output = self._main(["--release-type", "ci"]) self.assertIn("--bazel_options=--remote_upload_local_results=false", output) + def test_get_answers_the_credential_helper_protocol(self): + output = self._main( + ["get"], stdin='{"uri": "https://actinolite.cluster.engflow.com/x"}' + ) + self.assertEqual(json.loads(output), credential_helper_response(self.token)) + if __name__ == "__main__": unittest.main() diff --git a/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py b/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py index c5c057137b5..1510bef2149 100644 --- a/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py +++ b/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py @@ -15,15 +15,14 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) from write_jax_bazel_cache_credentials import ( - CERT_NAME, - KEY_NAME, + TOKEN_NAME, main, materialize_credentials, ) class MaterializeCredentialsTest(unittest.TestCase): - """Tests that runner-local certs are copied privately, or skipped.""" + """Tests that the runner-local token is copied privately, or skipped.""" def setUp(self): self.tmp = tempfile.TemporaryDirectory() @@ -33,33 +32,22 @@ def setUp(self): self.source_dir.mkdir() def _write_source(self) -> None: - # Existence and size are what matter; real PEM headers trip scanners. - (self.source_dir / CERT_NAME).write_text("test certificate") - (self.source_dir / KEY_NAME).write_text("test key") + # Existence and size are what matter; a real JWT would trip scanners. + (self.source_dir / TOKEN_NAME).write_text("test token") - def test_copies_both_files_when_present(self): + def test_copies_the_token_when_present(self): self._write_source() copied = materialize_credentials(self.source_dir, self.dest_dir) self.assertEqual(copied, self.dest_dir) - self.assertEqual((self.dest_dir / CERT_NAME).read_text(), "test certificate") - self.assertEqual((self.dest_dir / KEY_NAME).read_text(), "test key") + self.assertEqual((self.dest_dir / TOKEN_NAME).read_text(), "test token") - def test_returns_none_when_the_certificate_is_missing(self): - (self.source_dir / KEY_NAME).write_text("test key") + def test_returns_none_when_the_token_is_missing(self): self.assertIsNone(materialize_credentials(self.source_dir, self.dest_dir)) self.assertFalse(self.dest_dir.exists()) - def test_returns_none_when_the_key_is_missing(self): - (self.source_dir / CERT_NAME).write_text("test certificate") - self.assertIsNone(materialize_credentials(self.source_dir, self.dest_dir)) - - def test_returns_none_when_the_source_dir_is_empty(self): - self.assertIsNone(materialize_credentials(self.source_dir, self.dest_dir)) - - def test_raises_when_a_source_file_is_empty(self): + def test_raises_when_the_token_is_empty(self): # Present but unusable is a broken runner, not an unprovisioned one. - (self.source_dir / CERT_NAME).write_text("test certificate") - (self.source_dir / KEY_NAME).write_text("") + (self.source_dir / TOKEN_NAME).write_text("") with self.assertRaises(ValueError): materialize_credentials(self.source_dir, self.dest_dir) @@ -69,18 +57,16 @@ def test_replaces_a_stale_destination(self): (self.dest_dir / "stale").write_text("leftover") materialize_credentials(self.source_dir, self.dest_dir) self.assertFalse((self.dest_dir / "stale").exists()) - self.assertTrue((self.dest_dir / CERT_NAME).is_file()) + self.assertTrue((self.dest_dir / TOKEN_NAME).is_file()) @unittest.skipIf(os.name == "nt", "POSIX mode bits") def test_restricts_destination_permissions(self): self._write_source() materialize_credentials(self.source_dir, self.dest_dir) dir_mode = stat.S_IMODE(self.dest_dir.stat().st_mode) - cert_mode = stat.S_IMODE((self.dest_dir / CERT_NAME).stat().st_mode) - key_mode = stat.S_IMODE((self.dest_dir / KEY_NAME).stat().st_mode) + token_mode = stat.S_IMODE((self.dest_dir / TOKEN_NAME).stat().st_mode) self.assertEqual(dir_mode, 0o700) - self.assertEqual(cert_mode, 0o600) - self.assertEqual(key_mode, 0o600) + self.assertEqual(token_mode, 0o600) class MainTest(unittest.TestCase): @@ -94,17 +80,16 @@ def _run(self, source_dir: Path, dest_dir: Path) -> str: main(["--source-dir", str(source_dir), "--dest-dir", str(dest_dir)]) return output.read_text() - def test_main_writes_the_destination_when_credentials_exist(self): + def test_main_writes_the_destination_when_the_token_exists(self): with tempfile.TemporaryDirectory() as tmp: source_dir = Path(tmp) / "source" dest_dir = Path(tmp) / "dest" source_dir.mkdir() - (source_dir / CERT_NAME).write_text("test certificate") - (source_dir / KEY_NAME).write_text("test key") + (source_dir / TOKEN_NAME).write_text("test token") output = self._run(source_dir, dest_dir) self.assertIn(f"credentials_dir={dest_dir}", output) - def test_main_writes_empty_when_credentials_are_missing(self): + def test_main_writes_empty_when_the_token_is_missing(self): with tempfile.TemporaryDirectory() as tmp: source_dir = Path(tmp) / "source" dest_dir = Path(tmp) / "dest" diff --git a/build_tools/github_actions/write_jax_bazel_cache_credentials.py b/build_tools/github_actions/write_jax_bazel_cache_credentials.py index e7fb2be97fc..6e4306005dc 100644 --- a/build_tools/github_actions/write_jax_bazel_cache_credentials.py +++ b/build_tools/github_actions/write_jax_bazel_cache_credentials.py @@ -2,17 +2,18 @@ # Copyright Advanced Micro Devices, Inc. # SPDX-License-Identifier: MIT -"""Copy runner-local EngFlow client credentials into $RUNNER_TEMP. +"""Copy the runner-local EngFlow cache token into $RUNNER_TEMP. -XLA and JAX CI keep the mTLS files at `/data/ci-cert.crt` and -`/data/ci-cert.key` on the runner, where every job can read them — including -pull requests from forks. GitHub Actions secrets are not used: they are -invisible to fork PRs and would be a third credential path next to the -runner-local AWS files and OIDC roles the other caches already use. +The token is an EngFlow-issued JWT kept on the runner, where every job can read +it — including pull requests from forks. That is safe because the token's +`cache-reader` role is enforced by EngFlow, not by the client. GitHub Actions +secrets are not used: they are invisible to fork PRs and would be a third +credential path next to the runner-local AWS files and OIDC roles the other +caches already use. The copy lands under $RUNNER_TEMP rather than the workspace, which CI -publishes as artifacts. Missing source files are not an error: the wheel -build runs without a remote cache. Files that are present but empty are. +publishes as artifacts. A missing token is not an error: the wheel build runs +without a remote cache. A token file that is present but empty is. Used by `.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml`. """ @@ -29,8 +30,8 @@ from github_actions_api import gha_set_output DEFAULT_SOURCE_DIR = Path("/data") -CERT_NAME = "ci-cert.crt" -KEY_NAME = "ci-cert.key" +# Also the name configure_jax_bazel_cache.py reads inside the build container. +TOKEN_NAME = "engflow-cache.jwt" DEST_DIRNAME = "bazel-cache-credentials" @@ -48,40 +49,30 @@ def _restrict_private(path: Path, is_dir: bool) -> None: def materialize_credentials(source_dir: Path, dest_dir: Path) -> Path | None: - """Copy the runner's cert/key into dest_dir, or return None if they are absent.""" - certificate = source_dir / CERT_NAME - key = source_dir / KEY_NAME - missing = [p for p in (certificate, key) if not p.is_file()] - if missing: - _log( - "No cache credentials at " - + ", ".join(str(p) for p in missing) - + "; building without a remote cache." - ) + """Copy the runner's token into dest_dir, or return None if it is absent.""" + token = source_dir / TOKEN_NAME + if not token.is_file(): + _log(f"No cache token at {token}; building without a remote cache.") return None - empty = [p for p in (certificate, key) if p.stat().st_size == 0] - if empty: - raise ValueError( - f"Cache credentials are present but empty: {', '.join(map(str, empty))}" - ) + if token.stat().st_size == 0: + raise ValueError(f"Cache token is present but empty: {token}") if dest_dir.exists(): shutil.rmtree(dest_dir) dest_dir.mkdir(parents=True) _restrict_private(dest_dir, is_dir=True) - for src, name in ((certificate, CERT_NAME), (key, KEY_NAME)): - dest = dest_dir / name - shutil.copyfile(src, dest) - _restrict_private(dest, is_dir=False) + dest = dest_dir / TOKEN_NAME + shutil.copyfile(token, dest) + _restrict_private(dest, is_dir=False) - _log(f"Cache credentials copied to {dest_dir}") + _log(f"Cache token copied to {dest_dir}") return dest_dir def main(argv: list[str] | None = None) -> None: parser = argparse.ArgumentParser( - description="Copy runner-local EngFlow credentials into $RUNNER_TEMP." + description="Copy the runner-local EngFlow cache token into $RUNNER_TEMP." ) parser.add_argument( "--source-dir", @@ -89,7 +80,7 @@ def main(argv: list[str] | None = None) -> None: default=Path( os.environ.get("JAX_BAZEL_CACHE_CREDENTIALS_DIR") or DEFAULT_SOURCE_DIR ), - help="Directory that already holds ci-cert.crt and ci-cert.key on the runner.", + help=f"Directory that already holds {TOKEN_NAME} on the runner.", ) parser.add_argument( "--dest-dir", diff --git a/external-builds/jax/README.md b/external-builds/jax/README.md index f1beffc389b..adf58655359 100644 --- a/external-builds/jax/README.md +++ b/external-builds/jax/README.md @@ -159,28 +159,31 @@ wheels use sccache and S3 instead ([cache buckets](/docs/development/s3_buckets. Only the CI wheel build uses the cache. Release builds stay uncached until they have a cache of their own, the way nightly has its own ccache server. -JAX points Bazel at EngFlow (`grpcs://wardite.cluster.engflow.com`) as a +JAX points Bazel at EngFlow (`grpcs://actinolite.cluster.engflow.com`) as a **cache only**. The build keeps `--config=rocm_release_wheel` and never enables `rocm_rbe`, so it does not pick up `--remote_executor` and every action still runs on the runner. -Credentials are the runner-local files at `/data/ci-cert.crt` and -`/data/ci-cert.key` — the same path XLA and JAX CI already mount, not GitHub -Actions secrets. Fork pull requests see those files, so they can **read** the -cache. `write_jax_bazel_cache_credentials.py` copies them into `$RUNNER_TEMP` -(the workspace is published as artifacts) and the wheel job mounts that copy -read-only at `/data`. Set the repository variable -`JAX_BAZEL_CACHE_CREDENTIALS_DIR` if the runner keeps the files somewhere else. +The credential is an EngFlow-issued JWT with the `cache-reader` role, kept on +the runner as `/data/engflow-cache.jwt`, not as a GitHub Actions secret. Fork +pull requests see that file, so they can **read** the cache; EngFlow rejects +writes made with it. `write_jax_bazel_cache_credentials.py` copies it into +`$RUNNER_TEMP` (the workspace is published as artifacts) and the wheel job +mounts that copy read-only at `/data`. Set the repository variable +`JAX_BAZEL_CACHE_CREDENTIALS_DIR` if the runner keeps the token somewhere else. `configure_jax_bazel_cache.py` runs inside the build container and prints no -Bazel options when the files are missing or the release type is not `ci`; the -build command is then unchanged. Credentials that are present but empty or -rejected, or a cache that does not answer, fail the job rather than build -uncached. +Bazel options when the token is missing or the release type is not `ci`; the +build command is then unchanged. A token that is present but empty, malformed +or expired, or a cache that does not answer, fails the job rather than build +uncached. The token reaches Bazel through the same script acting as a +credential helper scoped to the cache host, never on the command line, which +`build.py --verbose` logs. Only pushes to `main` upload. Pull requests and manual dispatches pass -`--remote_upload_local_results=false`. That gate runs on the client, so what a -pull request could write is decided by the certificate's scope on EngFlow. Set +`--remote_upload_local_results=false`. That gate runs on the client; what a +run can write is decided by its token's role on EngFlow, so uploads from `main` +need a `cache-writer` token that only `main` can fetch. Set `JAX_BAZEL_REMOTE_CACHE_URL` to point the build at a different cache. EngFlow entries from `ROCm/jax` CI do not help TheRock: that CI compiles with From d0374dbe29ac2f9ac98c396a84ef827ff8d93ba9 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Fri, 6 Nov 2026 18:37:52 +0000 Subject: [PATCH 7/7] ci(jax): fetch the Bazel cache writer token on pushes to main Infra provisioned both EngFlow tokens in AWS Secrets Manager. The cache-reader token is synced to /data/engflow-cache.jwt on the runners; the cache-writer token is not on any runner. Only pushes to main in ROCm/TheRock now assume the therock-jax-cache-writer-main role through GitHub OIDC, read the writer token from Secrets Manager into $RUNNER_TEMP for the build container, and delete it after the build. The role's credentials go to a named AWS profile, so they do not reach the job's later AWS steps. Pull requests stay reader-only, and release builds still never use the cache. --- .../multi_arch_build_linux_jax_wheels_ci.yml | 33 ++++- .../write_jax_bazel_cache_credentials_test.py | 107 ++++++++++++++- .../write_jax_bazel_cache_credentials.py | 122 +++++++++++++++--- external-builds/jax/README.md | 13 +- 4 files changed, 244 insertions(+), 31 deletions(-) diff --git a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml index 0d232246517..a4af356cec2 100644 --- a/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml +++ b/.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml @@ -248,10 +248,24 @@ jobs: --rocm-version "${ROCM_VERSION}" \ --jax-source-dir jax-source - # A runner-local EngFlow cache-reader token, so fork pull requests can - # read the cache; EngFlow rejects its writes. GitHub secrets are not - # used. A missing token leaves the mount absent. - - name: Locate Bazel remote cache credentials + # Pull requests use the runner's EngFlow cache-reader token, whose writes + # EngFlow rejects. Only pushes to main assume the writer role to fetch + # the cache-writer token, which is never pre-mounted on a runner. + - name: Plan Bazel remote cache credentials + id: cache_plan + run: | + python build_tools/github_actions/write_jax_bazel_cache_credentials.py --plan + + - name: Configure AWS credentials for the Bazel cache writer + if: ${{ steps.cache_plan.outputs.writer_role != '' }} + uses: aws-actions/configure-aws-credentials@e1253824e5c10ff9df46874f81ed3ec929e19cfd # v6.3.0 + with: + aws-region: ${{ steps.cache_plan.outputs.aws_region }} + role-to-assume: ${{ steps.cache_plan.outputs.writer_role }} + aws-profile: ${{ steps.cache_plan.outputs.aws_profile }} + role-duration-seconds: 900 + + - name: Stage Bazel remote cache credentials id: cache_credentials env: JAX_BAZEL_CACHE_CREDENTIALS_DIR: ${{ vars.JAX_BAZEL_CACHE_CREDENTIALS_DIR }} @@ -268,8 +282,8 @@ jobs: JAX_BAZEL_REMOTE_CACHE_URL: ${{ vars.JAX_BAZEL_REMOTE_CACHE_URL }} CACHE_CREDENTIALS_DIR: ${{ steps.cache_credentials.outputs.credentials_dir }} run: | - # configure_jax_bazel_cache.py reads the token from /data, a copy - # of the runner's file in $RUNNER_TEMP. + # configure_jax_bazel_cache.py reads the token from /data, staged + # in $RUNNER_TEMP. credentials_mount="" if [[ -n "${CACHE_CREDENTIALS_DIR}" ]]; then credentials_mount="--volume ${CACHE_CREDENTIALS_DIR}:/data:ro" @@ -310,6 +324,13 @@ jobs: --output_path="$(pwd)/dist" ' + - name: Remove Bazel remote cache credentials + if: ${{ always() && steps.cache_credentials.outputs.credentials_dir != '' }} + env: + CACHE_CREDENTIALS_DIR: ${{ steps.cache_credentials.outputs.credentials_dir }} + run: | + rm -rf "${CACHE_CREDENTIALS_DIR}" + - name: Extract JAX versions from built wheels id: write_jax_versions run: | diff --git a/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py b/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py index 1510bef2149..ff24bb1797c 100644 --- a/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py +++ b/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py @@ -4,6 +4,8 @@ """Unit tests for write_jax_bazel_cache_credentials.py""" +import io +import json import os import stat import sys @@ -14,12 +16,74 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) +import write_jax_bazel_cache_credentials from write_jax_bazel_cache_credentials import ( TOKEN_NAME, + WRITER_AWS_PROFILE, + WRITER_ROLE, + WRITER_SECRET_ID, + fetch_writer_token, main, materialize_credentials, + uses_writer_token, ) +MAIN_PUSH = ["--event-name", "push", "--ref", "refs/heads/main"] + + +def _outputs(text: str) -> dict[str, str]: + return dict(line.split("=", 1) for line in text.splitlines() if "=" in line) + + +class UsesWriterTokenTest(unittest.TestCase): + """Tests which runs fetch the cache-writer token.""" + + def test_push_to_main_in_therock_writes(self): + self.assertTrue(uses_writer_token("push", "refs/heads/main", "ROCm/TheRock")) + + def test_everything_else_reads(self): + for event_name, ref, repository in ( + ("pull_request", "refs/pull/1/merge", "ROCm/TheRock"), + ("workflow_dispatch", "refs/heads/main", "ROCm/TheRock"), + ("push", "refs/heads/users/someone/branch", "ROCm/TheRock"), + # A fork's own main must not try the role: it would fail the job. + ("push", "refs/heads/main", "someone/TheRock"), + ): + with self.subTest(event_name=event_name, ref=ref, repository=repository): + self.assertFalse(uses_writer_token(event_name, ref, repository)) + + +class FetchWriterTokenTest(unittest.TestCase): + """Tests reading the writer JWT from Secrets Manager.""" + + def _fetch(self, secret: dict, github_actions: bool = False): + session = mock.MagicMock() + client = session.return_value.client.return_value + client.get_secret_value.return_value = {"SecretString": json.dumps(secret)} + stdout = io.StringIO() + env = {"GITHUB_ACTIONS": "true" if github_actions else "false"} + with mock.patch.object( + write_jax_bazel_cache_credentials.boto3, "Session", session + ), mock.patch.dict(os.environ, env), mock.patch.object(sys, "stdout", stdout): + token = fetch_writer_token() + return token, session, client, stdout.getvalue() + + def test_reads_the_token_property_with_the_writer_profile(self): + token, session, client, _ = self._fetch({"token": "test-writer\n"}) + self.assertEqual(token, "test-writer") + self.assertEqual(session.call_args.kwargs["profile_name"], WRITER_AWS_PROFILE) + client.get_secret_value.assert_called_once_with(SecretId=WRITER_SECRET_ID) + + def test_masks_the_token_in_github_actions(self): + _, _, _, stdout = self._fetch({"token": "test-writer"}, github_actions=True) + self.assertEqual(stdout.strip(), "::add-mask::test-writer") + + def test_raises_without_a_token(self): + for secret in ({}, {"token": ""}): + with self.subTest(secret=secret): + with self.assertRaises(ValueError): + self._fetch(secret) + class MaterializeCredentialsTest(unittest.TestCase): """Tests that the runner-local token is copied privately, or skipped.""" @@ -72,14 +136,53 @@ def test_restricts_destination_permissions(self): class MainTest(unittest.TestCase): """Tests that GITHUB_OUTPUT receives the destination path, or empty.""" - def _run(self, source_dir: Path, dest_dir: Path) -> str: + def _run(self, source_dir: Path, dest_dir: Path, extra=()) -> str: + # Explicit event arguments, so a test run on a push to main cannot + # take the writer path by accident. + argv = ["--event-name", "pull_request", "--ref", "refs/pull/1/merge"] + argv += ["--repository", "ROCm/TheRock", *extra] with tempfile.TemporaryDirectory() as tmp: output = Path(tmp) / "github_output" env = {"GITHUB_OUTPUT": str(output), "RUNNER_TEMP": tmp} with mock.patch.dict(os.environ, env, clear=False): - main(["--source-dir", str(source_dir), "--dest-dir", str(dest_dir)]) + main( + argv + + ["--source-dir", str(source_dir), "--dest-dir", str(dest_dir)] + ) return output.read_text() + def test_plan_names_the_writer_role_only_for_main_pushes(self): + with tempfile.TemporaryDirectory() as tmp: + plan = _outputs(self._run(Path(tmp), Path(tmp), ["--plan", *MAIN_PUSH])) + self.assertEqual(plan["writer_role"], WRITER_ROLE) + self.assertEqual(plan["aws_profile"], WRITER_AWS_PROFILE) + plan = _outputs(self._run(Path(tmp), Path(tmp), ["--plan"])) + self.assertEqual(plan["writer_role"], "") + + def test_main_push_stages_the_writer_not_the_runner_token(self): + with tempfile.TemporaryDirectory() as tmp: + source_dir = Path(tmp) / "source" + dest_dir = Path(tmp) / "dest" + source_dir.mkdir() + (source_dir / TOKEN_NAME).write_text("test reader") + with mock.patch.object( + write_jax_bazel_cache_credentials, + "fetch_writer_token", + return_value="test writer", + ): + output = self._run(source_dir, dest_dir, MAIN_PUSH) + self.assertEqual((dest_dir / TOKEN_NAME).read_text(), "test writer") + self.assertIn(f"credentials_dir={dest_dir}", output) + + def test_writer_fetch_failure_fails_the_job(self): + with tempfile.TemporaryDirectory() as tmp, mock.patch.object( + write_jax_bazel_cache_credentials, + "fetch_writer_token", + side_effect=RuntimeError("AccessDenied"), + ): + with self.assertRaises(RuntimeError): + self._run(Path(tmp), Path(tmp) / "dest", MAIN_PUSH) + def test_main_writes_the_destination_when_the_token_exists(self): with tempfile.TemporaryDirectory() as tmp: source_dir = Path(tmp) / "source" diff --git a/build_tools/github_actions/write_jax_bazel_cache_credentials.py b/build_tools/github_actions/write_jax_bazel_cache_credentials.py index 6e4306005dc..cdd6e87c4fa 100644 --- a/build_tools/github_actions/write_jax_bazel_cache_credentials.py +++ b/build_tools/github_actions/write_jax_bazel_cache_credentials.py @@ -2,31 +2,42 @@ # Copyright Advanced Micro Devices, Inc. # SPDX-License-Identifier: MIT -"""Copy the runner-local EngFlow cache token into $RUNNER_TEMP. +"""Stage the EngFlow cache token for the JAX wheel build in $RUNNER_TEMP. -The token is an EngFlow-issued JWT kept on the runner, where every job can read -it — including pull requests from forks. That is safe because the token's -`cache-reader` role is enforced by EngFlow, not by the client. GitHub Actions -secrets are not used: they are invisible to fork PRs and would be a third -credential path next to the runner-local AWS files and OIDC roles the other +Two EngFlow-issued JWTs exist, and EngFlow enforces each one's role: + +- The `cache-reader` token sits on the runner, where every job can read it — + including pull requests from forks, which is safe because it cannot write. + This copies it. +- The `cache-writer` token is never on a runner. Only ordinary pushes to `main` + in ROCm/TheRock assume `WRITER_ROLE` through GitHub OIDC (`--plan` names the + role for the workflow) and fetch it from AWS Secrets Manager. + +GitHub Actions secrets are not used: they are invisible to fork PRs and would be +a third credential path next to the runner-local files and OIDC roles the other caches already use. The copy lands under $RUNNER_TEMP rather than the workspace, which CI -publishes as artifacts. A missing token is not an error: the wheel build runs -without a remote cache. A token file that is present but empty is. +publishes as artifacts. A missing reader token is not an error: the wheel build +runs without a remote cache. A token that is present but empty is, and so is a +writer token that cannot be fetched. Used by `.github/workflows/multi_arch_build_linux_jax_wheels_ci.yml`. """ import argparse +import json import os import shutil import stat import sys from pathlib import Path +import boto3 + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from configure_jax_bazel_cache import upload_allowed from github_actions_api import gha_set_output DEFAULT_SOURCE_DIR = Path("/data") @@ -34,6 +45,15 @@ TOKEN_NAME = "engflow-cache.jwt" DEST_DIRNAME = "bazel-cache-credentials" +WRITER_REPOSITORY = "ROCm/TheRock" +WRITER_ROLE = "arn:aws:iam::692859939525:role/therock-jax-cache-writer-main" +WRITER_SECRET_ID = "/ossci/ci/ROCm/TheRock/main/jax-bazel-cache-writer" +WRITER_SECRET_KEY = "token" +WRITER_REGION = "us-east-2" +# The workflow stores the role's credentials under this profile, not in the +# environment, so they cannot leak into later AWS steps of the job. +WRITER_AWS_PROFILE = "jax-bazel-cache-writer" + def _log(msg: str) -> None: print(f"[jax-bazel-cache] {msg}", file=sys.stderr) @@ -48,6 +68,19 @@ def _restrict_private(path: Path, is_dir: bool) -> None: _log(f"Could not chmod {path}: {e}") +def uses_writer_token(event_name: str, ref: str, repository: str) -> bool: + """Only the pushes to main that upload get the writer; forks never do.""" + return repository == WRITER_REPOSITORY and upload_allowed(event_name, ref) + + +def _fresh_dest_dir(dest_dir: Path) -> Path: + if dest_dir.exists(): + shutil.rmtree(dest_dir) + dest_dir.mkdir(parents=True) + _restrict_private(dest_dir, is_dir=True) + return dest_dir / TOKEN_NAME + + def materialize_credentials(source_dir: Path, dest_dir: Path) -> Path | None: """Copy the runner's token into dest_dir, or return None if it is absent.""" token = source_dir / TOKEN_NAME @@ -57,12 +90,7 @@ def materialize_credentials(source_dir: Path, dest_dir: Path) -> Path | None: if token.stat().st_size == 0: raise ValueError(f"Cache token is present but empty: {token}") - if dest_dir.exists(): - shutil.rmtree(dest_dir) - dest_dir.mkdir(parents=True) - _restrict_private(dest_dir, is_dir=True) - - dest = dest_dir / TOKEN_NAME + dest = _fresh_dest_dir(dest_dir) shutil.copyfile(token, dest) _restrict_private(dest, is_dir=False) @@ -70,9 +98,38 @@ def materialize_credentials(source_dir: Path, dest_dir: Path) -> Path | None: return dest_dir +def fetch_writer_token() -> str: + """Reads the writer JWT with the credentials of the assumed writer role.""" + session = boto3.Session(profile_name=WRITER_AWS_PROFILE, region_name=WRITER_REGION) + secret = session.client("secretsmanager").get_secret_value( + SecretId=WRITER_SECRET_ID + ) + token = json.loads(secret["SecretString"]).get(WRITER_SECRET_KEY, "").strip() + if not token: + raise ValueError(f"{WRITER_SECRET_ID} has no {WRITER_SECRET_KEY!r} value") + if os.environ.get("GITHUB_ACTIONS") == "true": + # Nothing prints the token; this keeps a stray echo out of the log. + print(f"::add-mask::{token}") + return token + + +def stage_writer_token(token: str, dest_dir: Path) -> Path: + dest = _fresh_dest_dir(dest_dir) + dest.write_text(token) + _restrict_private(dest, is_dir=False) + _log(f"Cache writer token staged in {dest_dir}") + return dest_dir + + def main(argv: list[str] | None = None) -> None: parser = argparse.ArgumentParser( - description="Copy the runner-local EngFlow cache token into $RUNNER_TEMP." + description="Stage the EngFlow cache token in $RUNNER_TEMP." + ) + parser.add_argument( + "--plan", + action="store_true", + help="Only print which AWS role to assume (writer_role, aws_region, " + "aws_profile); empty when the run reads with the runner's token.", ) parser.add_argument( "--source-dir", @@ -80,7 +137,7 @@ def main(argv: list[str] | None = None) -> None: default=Path( os.environ.get("JAX_BAZEL_CACHE_CREDENTIALS_DIR") or DEFAULT_SOURCE_DIR ), - help=f"Directory that already holds {TOKEN_NAME} on the runner.", + help=f"Directory that already holds the reader {TOKEN_NAME} on the runner.", ) parser.add_argument( "--dest-dir", @@ -89,8 +146,34 @@ def main(argv: list[str] | None = None) -> None: help="Directory to write the copy into. Defaults to $RUNNER_TEMP/" f"{DEST_DIRNAME}.", ) + parser.add_argument( + "--event-name", + default=os.environ.get("GITHUB_EVENT_NAME", ""), + help="GitHub event that started the run; only a push to main writes.", + ) + parser.add_argument( + "--ref", + default=os.environ.get("GITHUB_REF", ""), + help="Git ref the run is for, e.g. refs/heads/main.", + ) + parser.add_argument( + "--repository", + default=os.environ.get("GITHUB_REPOSITORY", ""), + help="Repository the run is in; forks never get the writer token.", + ) args = parser.parse_args(argv) + writer = uses_writer_token(args.event_name, args.ref, args.repository) + if args.plan: + gha_set_output( + { + "writer_role": WRITER_ROLE if writer else "", + "aws_region": WRITER_REGION if writer else "", + "aws_profile": WRITER_AWS_PROFILE if writer else "", + } + ) + return + dest_dir = args.dest_dir if dest_dir is None: runner_temp = os.environ.get("RUNNER_TEMP") @@ -98,8 +181,11 @@ def main(argv: list[str] | None = None) -> None: raise RuntimeError("RUNNER_TEMP is unset; pass --dest-dir") dest_dir = Path(runner_temp) / DEST_DIRNAME - copied = materialize_credentials(args.source_dir, dest_dir) - gha_set_output({"credentials_dir": str(copied) if copied else ""}) + if writer: + staged = stage_writer_token(fetch_writer_token(), dest_dir) + else: + staged = materialize_credentials(args.source_dir, dest_dir) + gha_set_output({"credentials_dir": str(staged) if staged else ""}) if __name__ == "__main__": diff --git a/external-builds/jax/README.md b/external-builds/jax/README.md index adf58655359..39c135f8adb 100644 --- a/external-builds/jax/README.md +++ b/external-builds/jax/README.md @@ -180,11 +180,14 @@ uncached. The token reaches Bazel through the same script acting as a credential helper scoped to the cache host, never on the command line, which `build.py --verbose` logs. -Only pushes to `main` upload. Pull requests and manual dispatches pass -`--remote_upload_local_results=false`. That gate runs on the client; what a -run can write is decided by its token's role on EngFlow, so uploads from `main` -need a `cache-writer` token that only `main` can fetch. Set -`JAX_BAZEL_REMOTE_CACHE_URL` to point the build at a different cache. +Only pushes to `main` in ROCm/TheRock upload. Pull requests and manual +dispatches pass `--remote_upload_local_results=false`, but that gate runs on +the client; what a run can write is decided by its token's role on EngFlow. +The `cache-writer` token is never on a runner: those `main` builds assume the +`therock-jax-cache-writer-main` IAM role through GitHub OIDC, fetch the token +from AWS Secrets Manager, stage it for the build container like the reader +token, and delete it after the build. Set `JAX_BAZEL_REMOTE_CACHE_URL` +to point the build at a different cache. EngFlow entries from `ROCm/jax` CI do not help TheRock: that CI compiles with `--config=rocm_rbe`, so the action keys differ. TheRock has to populate the