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 8b94c3cd141..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,6 +248,30 @@ jobs: --rocm-version "${ROCM_VERSION}" \ --jax-source-dir jax-source + # 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 }} + run: | + python build_tools/github_actions/write_jax_bazel_cache_credentials.py + - name: Build JAX Wheels in manylinux container working-directory: jax-source env: @@ -255,20 +279,36 @@ jobs: JAX_COMMIT: ${{ steps.checkout_jax.outputs.commit }} ML_WHEEL_VERSION_SUFFIX: ${{ steps.resolve_jax_version.outputs.ml_wheel_version_suffix }} JAX_WHEELS: ${{ steps.resolve_jax_version.outputs.wheels }} + 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, staged + # in $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 \ --env ML_WHEEL_VERSION_SUFFIX \ --env JAX_WHEELS \ + ${credentials_mount} \ --env ROCM_VERSION="${ROCM_VERSION}" \ --env PYTHON_VERSION="${PYTHON_VERSION}" \ + --env RELEASE_TYPE="${RELEASE_TYPE}" \ + --env GITHUB_EVENT_NAME \ + --env GITHUB_REF \ + --env JAX_BAZEL_REMOTE_CACHE_URL="${JAX_BAZEL_REMOTE_CACHE_URL}" \ --env PACKAGE_DIST_DIR="/workspace/jax-source/dist" \ --volume "${GITHUB_WORKSPACE}:/workspace" \ --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 \ --wheels="${JAX_WHEELS}" \ --python_version="${PYTHON_VERSION}" \ @@ -278,11 +318,19 @@ 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" ' + - 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/configure_jax_bazel_cache.py b/build_tools/github_actions/configure_jax_bazel_cache.py new file mode 100755 index 00000000000..94eb001f701 --- /dev/null +++ b/build_tools/github_actions/configure_jax_bazel_cache.py @@ -0,0 +1,227 @@ +#!/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 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 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://actinolite.cluster.engflow.com" + +# 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. +SHARED_CACHE_RELEASE_TYPES = frozenset({"ci"}) + +DEFAULT_PROBE_TIMEOUT_SECONDS = 10 +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) + + +def resolve_cache_url(cache_url: str, release_type: str) -> str: + """Returns the cache URL to use, or "" when caching is disabled. + + 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 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]: + """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 check_token(token_file: Path, now: float | None = None): + """Raises CacheUnusableError unless token_file holds an unexpired JWT. + + 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: + 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: + 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 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, 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}", + # 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." + ) + parser.add_argument( + "--cache-url", + default=os.environ.get("JAX_BAZEL_REMOTE_CACHE_URL", ""), + help="Remote cache URL, for release types that may use the cache.", + ) + 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( + "--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( + "--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 + + 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}" + ) + check_token(TOKEN_FILE) + probe_cache(url, args.probe_timeout) + + print(" ".join(bazel_cache_options(url, Path(__file__).resolve(), 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..31c39866c87 --- /dev/null +++ b/build_tools/github_actions/tests/configure_jax_bazel_cache_test.py @@ -0,0 +1,335 @@ +#!/usr/bin/env python3 +# Copyright Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Unit tests for configure_jax_bazel_cache.py""" + +import base64 +import io +import json +import os +import sys +import tempfile +import time +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 ( + HTTP2_CLIENT_PREFACE, + REMOTE_CACHE_URL, + CacheUnusableError, + bazel_cache_options, + check_token, + credential_helper_response, + endpoint_address, + main, + probe_cache, + resolve_cache_url, + upload_allowed, +) + + +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.""" + + 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): + # 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"), "") + + def test_explicit_url_repoints_the_cache(self): + self.assertEqual( + 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.""" + + def test_default_port_is_tls(self): + self.assertEqual( + endpoint_address("grpcs://actinolite.cluster.engflow.com"), + ("actinolite.cluster.engflow.com", 443), + ) + + def test_explicit_port_is_honored(self): + self.assertEqual( + endpoint_address("grpcs://cache.example:8980"), ("cache.example", 8980) + ) + + +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 reach.""" + + 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): + with mock.patch.object( + configure_jax_bazel_cache.ssl, + "create_default_context", + return_value=self.context, + ), mock.patch.object( + configure_jax_bazel_cache.socket, + "create_connection", + side_effect=connect_effect, + ): + probe_cache(url) + + 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_raises_when_the_host_does_not_resolve(self): + with self.assertRaises(CacheUnusableError): + self._probe(connect_effect=OSError("name resolution failed")) + + def test_raises_when_the_connection_times_out(self): + with self.assertRaises(CacheUnusableError): + self._probe(connect_effect=TimeoutError()) + + 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 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.""" + + # 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_credential_helper_is_scoped_to_the_cache_host(self): + # An unscoped helper would send the token to every host Bazel fetches. + self.assertIn( + f"--bazel_options=--credential_helper=cache.example={self.HELPER}", + self._options(allow_upload=False, url="grpcs://cache.example:8980"), + ) + + 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): + 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("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. + 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 probe, mock.patch.dict( + os.environ, github_env, clear=True + ), mock.patch.object(sys, "stdout", stdout), mock.patch.object( + sys, "stdin", io.StringIO(stdin) + ): + main(argv) + return stdout.getvalue().strip() + + def test_prints_options_for_a_reachable_cache(self): + 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) + + 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): + 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_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 new file mode 100644 index 00000000000..ff24bb1797c --- /dev/null +++ b/build_tools/github_actions/tests/write_jax_bazel_cache_credentials_test.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 +# Copyright Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Unit tests for write_jax_bazel_cache_credentials.py""" + +import io +import json +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)) + +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.""" + + 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; a real JWT would trip scanners. + (self.source_dir / TOKEN_NAME).write_text("test token") + + 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 / TOKEN_NAME).read_text(), "test token") + + 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_raises_when_the_token_is_empty(self): + # Present but unusable is a broken runner, not an unprovisioned one. + (self.source_dir / TOKEN_NAME).write_text("") + with self.assertRaises(ValueError): + 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 / 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) + token_mode = stat.S_IMODE((self.dest_dir / TOKEN_NAME).stat().st_mode) + self.assertEqual(dir_mode, 0o700) + self.assertEqual(token_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, 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( + 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" + dest_dir = Path(tmp) / "dest" + source_dir.mkdir() + (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_the_token_is_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..cdd6e87c4fa --- /dev/null +++ b/build_tools/github_actions/write_jax_bazel_cache_credentials.py @@ -0,0 +1,192 @@ +#!/usr/bin/env python3 +# Copyright Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Stage the EngFlow cache token for the JAX wheel build in $RUNNER_TEMP. + +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 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") +# Also the name configure_jax_bazel_cache.py reads inside the build container. +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) + + +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 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 + if not token.is_file(): + _log(f"No cache token at {token}; building without a remote cache.") + return None + if token.stat().st_size == 0: + raise ValueError(f"Cache token is present but empty: {token}") + + dest = _fresh_dest_dir(dest_dir) + shutil.copyfile(token, dest) + _restrict_private(dest, is_dir=False) + + _log(f"Cache token copied to {dest_dir}") + 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="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", + type=Path, + default=Path( + os.environ.get("JAX_BAZEL_CACHE_CREDENTIALS_DIR") or DEFAULT_SOURCE_DIR + ), + help=f"Directory that already holds the reader {TOKEN_NAME} 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}.", + ) + 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") + if not runner_temp: + raise RuntimeError("RUNNER_TEMP is unset; pass --dest-dir") + dest_dir = Path(runner_temp) / DEST_DIRNAME + + 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__": + main() diff --git a/docs/development/ccache_troubleshooting.md b/docs/development/ccache_troubleshooting.md index 356c3282451..16f7c96ad27 100644 --- a/docs/development/ccache_troubleshooting.md +++ b/docs/development/ccache_troubleshooting.md @@ -42,6 +42,17 @@ 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. +### 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#bazel-remote-cache) | + +This page covers ccache only. Builds that use another cache are linked above +rather than documented here. + ### Namespace version `CCACHE_NAMESPACE_VERSION` in `setup_ccache.py` controls the cache diff --git a/external-builds/jax/README.md b/external-builds/jax/README.md index bfe89b0f01c..39c135f8adb 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) +- [Bazel remote cache](#bazel-remote-cache) - [Test instructions](#test-instructions) - [Nightly releases](#nightly-releases) @@ -148,6 +149,50 @@ 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. +## 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)). + +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://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. + +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 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` 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 +cache itself on postsubmit. + ## Test instructions ### Prerequisites