ci(jax): serve the JAX wheel build from the EngFlow Bazel cache by gulsumgudukbay · Pull Request #7912 · ROCm/TheRock

gulsumgudukbay · 2026-09-03T16:19:52Z

ISSUE ID : #7346

Motivation

The Linux JAX wheel job recompiles jaxlib and the XLA ROCm backend every run: each run gets a fresh container with an empty Bazel action cache.
ccache covers the CMake stages, artifact reuse covers the ROCm graph. The JAX Bazel compile is covered by neither.

Technical Details

Points the CI wheel build at EngFlow (grpcs://actinolite.cluster.engflow.com) as a cache only, not remote execution: --config=rocm_release_wheel keeps compilation on the runner and only action results travel.
Only the ci release type uses the cache. Release builds stay uncached until they have a cache of their own, the way nightly has its own ccache server.
The credential is an EngFlow-issued JWT with the cache-reader role, kept on the runner as /data/engflow-cache.jwt (overridable with the JAX_BAZEL_CACHE_CREDENTIALS_DIR repository variable). The role is enforced by EngFlow, so fork pull requests that can read the file can read the cache but not write it. GitHub secrets are not used.
write_jax_bazel_cache_credentials.py copies the token into $RUNNER_TEMP and the job mounts that copy read-only at /data in the build container.
configure_jax_bazel_cache.py runs inside the build container. A token that was never provisioned is the one quiet case: it prints nothing and the build command is unchanged. A token that is present but empty, malformed or expired, or a cache that does not answer, fails the job.
The token reaches Bazel through configure_jax_bazel_cache.py get, a Bazel credential helper scoped to the cache host. It never appears on the build command line, which build.py --verbose logs, nor in a bazelrc, which JAX's --announce_rc prints, and Bazel never sends it to other hosts.
Only pushes to main in ROCm/TheRock upload; pull requests and manual dispatches pass --remote_upload_local_results=false. That gate is client-side hygiene; what a run can write is decided by its token's role.
The cache-writer token is never on a runner. write_jax_bazel_cache_credentials.py --plan names the therock-jax-cache-writer-main role only for pushes to main in ROCm/TheRock; those builds assume it through GitHub OIDC (credentials in a named AWS profile, 15 minutes), read the token from AWS Secrets Manager into $RUNNER_TEMP for the build container, and delete it after the build. A writer token that cannot be fetched fails the job.
Docs live in external-builds/jax/README.md; ccache_troubleshooting.md only links out.

Test Plan

pytest build_tools/github_actions/tests.
pre-commit run on the touched files, including actionlint; zizmor on the changed workflow.
Against actinolite with a cache-reader token: direct REAPI calls for both reads and both write paths, and a Bazel build with uploads on through the credential helper.
Dispatch the CI workflow from the branch on a fresh azure-linux-scale-rocm runner, which now carries the reader token.
After merge: the first push to main populates the cache with the writer token, and the next pull request reads it.

Test Result

configure_jax_bazel_cache and write_jax_bazel_cache_credentials tests: 48 passed. pre-commit and zizmor pass on the touched files.
cache-reader token on actinolite: GetActionResult returns NOT_FOUND and FindMissingBlobs succeeds; BatchUpdateBlobs and UpdateActionResult return PERMISSION_DENIED (does not allow 'contentaddressablestorage:Write' / 'actioncache:Write'). Bazel with --remote_upload_local_results=true reports it is not authorized to write and builds normally. Bazel calls the helper with get and receives the token headers.
Fresh azure-linux-scale-rocm runner, dispatch 42272479277: the job copied the runner's reader token, skipped the writer role, and Bazel ran against actinolite through the credential helper with --remote_upload_local_results=false and no remote cache errors. It had 0 remote cache hits, as expected: nothing writes TheRock's entries until the first push to main. No token or JWT-shaped string appears in the log.
Two-pass cache measurement, taken earlier on the previous EngFlow endpoint with a read-write certificate, pass 1 33823595745, pass 2 33833980679:

pass 1 (cold)                        pass 2 (warm)
13143 processes: 2722 local          13143 processes: 2722 remote cache hit
  Elapsed: 346.777s                    Elapsed: 102.431s
16872 processes: 14907 local         16872 processes: 14907 remote cache hit
  Elapsed: 2680.469s                   Elapsed: 58.116s

Bazel time 3027s -> 160s, 17629 of 17629 cacheable actions served remotely, no cache errors in either run.
That rate is an upper bound: both passes shared one rocm_package_find_links_url. A real PR builds its own ROCm SDK, so actions consuming ROCm headers will miss. The durable win is the ROCm-independent bulk (LLVM, absl, protobuf, Eigen, XLA core).
The Windows unit-test failure came from this PR (a POSIX mode-bit assertion) and is skipped on Windows now.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/CONTRIBUTING.md.

therock-pr-bot · 2026-09-03T16:20:44Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

therock-pr-bot · 2026-09-03T16:20:46Z

🎉 All checks passed! This PR is ready for review.

ScottTodd · 2026-09-03T16:46:36Z

-      stage_reuse_mode: ${{ inputs.stage_reuse_mode || 'dry-run' }}
+      # Pull requests reuse the stages their diff does not affect. Pushes stay
+      # on dry-run so every main commit builds the stages a PR later reuses.
+      stage_reuse_mode: ${{ inputs.stage_reuse_mode || (github.event_name == 'pull_request' && 'reuse-stage') || 'dry-run' }}


Stage reuse on pull requests. multi_arch_ci.yml passes reuse-stage for
pull_request and keeps every other trigger on dry-run. Pushes are
deliberately left rebuilding: a push that reused stages would stop producing the
artifacts a later pull request reuses. Manual workflow_dispatch still wins.

One change per pull request please. This is a big change independent of JAX that needs substantial testing for rollout to TheRock and other repositories.

Let's focus this PR on just the JAX bazel cache.

See this issue for building/testing just JAX/PyTorch (stage reuse mode is not quite the right tool here, we should be able to reuse rocm packages too):

[CI] Add package reuse mode that jumps straight to PyTorch/JAX builds/tests #7346

Sure, I split them. This pull request is now just the JAX Bazel cache, and the stage-reuse change moved to #7919 as a draft so it can get the rollout testing it needs on its own.

Thanks for the pointer to #7346 as well. Reusing the ROCm packages rather than the artifacts is a better shape for the JAX/PyTorch-only case, so I will follow that

ScottTodd · 2026-09-14T20:53:11Z

+      # 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


Please work with @subodh-dubey-amd and @amd-shiraz to align on a design here that is consistent with how we configure remote cache credentials for ROCm and PyTorch ccache/sccache builds so we aren't maintaining a third cache configuration that shares no code and does not work on PRs from forks without at least a documented reason.

Our other workflows either:

Store base credentials securely on the runners themselves so they are accessible to PRs from forks

Acquire more privileged credentials using OIDC: https://github.com/ROCm/TheRock/blob/main/docs/development/s3_buckets.md#authentication (https://docs.github.com/en/actions/concepts/security/openid-connect)

Related PRs (currently in a private repository that has plans for open sourcing):

https://github.com/ROCm/TheRock-Infra/pull/323

https://github.com/ROCm/TheRock-Infra/pull/416

Docs:

https://github.com/ROCm/TheRock/blob/main/docs/development/ccache_troubleshooting.md

https://github.com/ROCm/TheRock/blob/main/docs/development/s3_buckets.md#cache-buckets

I would strongly prefer for readonly caches to work on all PRs, so github secrets are a non-starter there.

There are also security implications for every github actions secret we add (cc @haribabug and @jayhawk-commits too). We've been yanking secrets that contributors with write access have added to this repository and may do the same here.

i dropped the GH secrets. The job now copies the runner-local /data/ci-cert.crt and /data/ci-cert.key (same path XLA/JAX CI already mount), so fork PRs can read the cache.
upload stays off on pull_request. @amd-shiraz @subodh-dubey-amd, could you please obtain RBE_CI_CERT / RBE_CI_KEY and install them at /data/ci-cert.crt and /data/ci-cert.key on the pool azure-linux-scale-rocm?

ScottTodd · 2026-09-14T20:56:01Z

+          # 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


Please avoid inline bash, particularly for security sensitive code: https://github.com/ROCm/TheRock/blob/main/docs/development/style_guides/github_actions_style_guide.md#prefer-python-scripts-over-inline-bash

Code that directly reads and writes credential files with bash is ripe for vulnerabilities and programmer errors during future debugging.

Moved that into write_jax_bazel_cache_credentials.py with unit tests. The workflow step is just a Python call now

ScottTodd · 2026-09-14T20:58:17Z

+    # Carries the Bazel remote cache credentials into the wheel build.
+    secrets: inherit


We're going to start enforcing https://docs.zizmor.sh/audits/#secrets-inherit soon as part of #7804 and related work (cc @cgoea ). If using secrets (and I don't think we should use secrets given my other comments), we should start applying the suggested code pattern from that page:

# AVOID jobs: pass-secrets-to-workflow: uses: ./.github/workflows/called-workflow.yml secrets: inherit

# Do this instead jobs: pass-secrets-to-workflow: uses: ./.github/workflows/called-workflow.yml secrets: forward-me: ${{ secrets.forward-me }} me-too: ${{ secrets.me-too }}

removed, nothing to forward once the secrets are gone.

ScottTodd · 2026-09-14T21:00:19Z

+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.


This doesn't seem like the right place for these docs, or at least they should go under a #### JAX cache servers subsection.

I would suggest:

Mentions in this document about which builds do use ccache

Links to other documents for builds that use other caches

Details on JAX cache in https://github.com/ROCm/TheRock/blob/main/external-builds/jax/README.md (we should to the same for pytorch @subodh-dubey-amd )

That page now only says which builds use ccache and links out. JAX cache details are in external-builds/jax/README.md.

ScottTodd

Please work with @amd-shiraz and @subodh-dubey-amd on the review here before re-requesting review from me.

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.

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.

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.

…ax-bazel-cache

subodh-dubey-amd

Thanks for splitting out #7919 and moving the credential handling into Python. Requesting changes, mainly on the credential model (inline). This is still inert: the Oct 5 JAX job logged No cache credentials at /data/ci-cert.crt, /data/ci-cert.key (job). But the cache turns on for CI and nightly as soon as certs land, so I'd like the design settled before merge.

Not tied to a changed line: could you link the runs behind the 3027s → 160s numbers, and refresh the CI notes in the description? The Windows unit-test failure comes from this PR.

cc @amd-shiraz

subodh-dubey-amd · 2026-10-06T14:36:32Z

+      # 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:
+          JAX_BAZEL_CACHE_CREDENTIALS_DIR: ${{ vars.JAX_BAZEL_CACHE_CREDENTIALS_DIR }}
+        run: |
+          python build_tools/github_actions/write_jax_bazel_cache_credentials.py


This matches the client side of #7026, but as I noted there, a client-side read-only gate is only hygiene; the boundary has to be server-side. Here that's the scope of this cert on EngFlow, and per the README any PR, including forks, can read it. What is RBE_CI_CERT allowed to do there?

On S3, runner base credentials allow uploads to therock-ci-artifacts-external, and other writes go through OIDC roles that fork PRs can't use (docs); PyTorch sccache has a bucket and role per release type (docs). If an S3-backed cache would work here, we can create the buckets and roles. This job already authenticates to AWS for the wheel upload, so it wouldn't need different runners.

Thanks @subodh-dubey-amd, your other comments are addressed in a9ddda0. On credentials: the certificate I tested with can write to the cache, so I agree it can't sit on the runners as is. Unless you or @amd-shiraz prefer S3, I'd like to stay on EngFlow, with a read-only certificate on the runners and a write certificate that only push-to-main jobs can use.
I am discussing with @charleshofer regarding read-only client certificate.

Agreed: only a read-only cert on the runners. For the write cert, anything on the runner is readable by PR jobs, so it needs a source gated on the ref. On the AWS side we could keep it in Secrets Manager behind an OIDC role that only trusts refs/heads/main, and the job already has id-token: write.

…ci-pr-stage-reuse-and-jax-bazel-cache

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.

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.

@ScottTodd

…rix row (#8524) Builds the latest JAX (`jax-ml/jax@main`) against the latest ROCm, the one that night's rockrel run just built, as a row of the JAX release matrix. Implements Phase 1 of #8201. ## A matrix row, not a canary The original design in #8201 was a separately dispatched, non-blocking canary with artifact, failure, publication and reporting isolation. @ScottTodd's review asked us not to diverge from the PyTorch release precedent, which carries its moving `nightly` ref as an ordinary matrix row, and #8201 was reworked to match. So tip builds through the same workflows as every other row, with no `continue-on-error`: a red tip cell fails the JAX release run and is triaged through ROCM-JAX gardening and therock-hud.amd.com like any other nightly failure. Nothing new is dispatched, so there is no new edge to thread `quartz_tracking_id` through. `configure_jax_release_matrix.py` mirrors `configure_pytorch_release_matrix.py`: stable refs for `prerelease`, stable + `jax-main` otherwise, and CI stays on its pinned ref so no pull request is judged against a tree moving outside the author's control. ``` prerelease 0.10.2 0.11.0 0.11.1 0.11.2 nightly 0.10.2 0.11.0 0.11.1 0.11.2 tip (build only) dev 0.10.2 0.11.0 0.11.1 0.11.2 tip (build only) ci 0.11.2 ``` ## Contents The first three commits are @jayfurmanek's plumbing from `users/jfurmane/jax-tip-plumbing`, unchanged. The rest complete Phase 1: - **The tip row.** `jax-main` was defined but in no default ref list, and `multi_arch_release_linux_jax_wheels.yml` never passed `--release-type`, so the generator ran on its `dev` default whatever the run was. Both fixed, and `wheel_type` is forwarded to the build. - **Versions from the checkout.** PyPI has no release matching `main`, so the tip row builds `jax` and `jaxlib` from the same checkout as the plugin and PJRT. `resolve_jax_version.py` validates `--wheel-type` and emits `ML_WHEEL_VERSION_SUFFIX` and the wheel list: the ROCm suffix alone for `release` (plugin + PJRT, as before), `.dev<date of the HEAD commit>` plus that suffix for `nightly` (all four wheels). Nothing is fetched from an outside index, and the same commit against the same ROCm always gets the same version. The workflows carry no inline version logic and expand no inputs inside `run:` blocks. - **Build only, for now.** The tip row is `publish: false`: every nightly builds all four wheels, which answers "does upstream `main` build against this ROCm", but skips the upload and the tests. The release tests install from `whl-next`, which does not route `jax` or `jaxlib` yet, so a published tip cell would fail at install. Pinned refs publish and test as before. - **Install, test checkout and skip lists.** When the build produced `jaxlib`, the test job installs `jax` and `jaxlib` together from the same place as the plugin. The test job checks out the commit that was built rather than the moving ref. Skip-list lookups are keyed on the release number, so a tip build reads the file of the release it leads to. No skip entries are added. ## The nightly wheel name Versioning a nightly through XLA's `custom` wheel type cannot work. Bazel and `build/build.py` derive the wheel name independently from the same `ML_WHEEL_*` inputs and disagree on one character: XLA emits the `+` only when a git hash is set, `build.py` emits it unconditionally. With no hash, Bazel writes `...dev20260920+devrocm...` while the copy glob looks for `...dev20260920++devrocm...`. `copy_individual_files` globs silently, so nothing is copied and the run dies a step later inside `write_jax_versions.py` as a missing-plugin error rather than a naming one. A nightly is instead a `release` build whose suffix carries the date, which both sides spell identically, for all four wheels. The result is the `<base>.devYYYYMMDD+rocm<ver>` that #8201 specifies, and it is still a PEP 440 pre-release, so `pip install` without `--pre` keeps resolving the pinned releases — the property `torch-2.15.0a0` relies on. ## Trial run One tip cell through the CI workflow, which installs from the run's own wheels, at `medium` on gfx94X: https://github.com/ROCm/TheRock/actions/runs/37328729986 - Built `jax`, `jaxlib`, the plugin and PJRT from upstream `main` in TheRock's container, about 50 minutes (`jaxlib` about 20 of them). - `31492 passed, 22 failed`. All 22 were `tests/thunky/thunky_test.py`, new upstream that day with its FFI handler registered for CUDA only; jax-ml/jax#41305 has since enabled it on ROCm. ## Review notes - #7912 serves the CI JAX wheel build from the Bazel cache; release builds, tip included, stay uncached for now. Both PRs touch the JAX build workflows, so whichever lands second needs `main` merged in. - Publishing and testing tip is a follow-up: route `jax`, `jaxlib` and `opt-einsum` in `rocm_whl_next_ownership.yaml`, publish the supported stable `jax`/`jaxlib` under `jax/whl-next` next to the tip builds (a package routes to one product page), add a RELEASES.md note like PyTorch's, then set the tip row to publish. - `dev-bkc` / `nightly-bkc` include tip, following PyTorch exactly. Say so if BKC runs should be stable-only for JAX. - Phase 2 of #8201 covers the rest: tip skip-list entries, a second GPU family (gfx950 needs `MIOPEN_DEBUG_CONV_WINOGRAD=0`, which `THEROCK_ENV` does not set), and confirming the cell reads sensibly on the HUD. ## Related - #8201: Phase 1 RFC - #7912: Bazel remote cache for the JAX wheel builds - `jax-ml/jax` runs its own ROCm nightly against "TheRock Next" in `.github/workflows/wheel_tests_nightly_release.yml` (gfx950, 1/4/8 GPU, full pytest, jaxlib from their nightly bucket), complementary coverage, not the same job. --------- Co-authored-by: Jason Furmanek <jfurmane@amd.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

…ci-pr-stage-reuse-and-jax-bazel-cache # Conflicts: # .github/workflows/multi_arch_build_linux_jax_wheels_ci.yml

alekstheod · 2026-10-21T13:15:58Z

+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


you can have several --config in one command. Just do --config-rocm_release_wheel --config=rocm_rbe.

alekstheod · 2026-10-21T13:29:06Z

+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


We now do provide self served rbe cluster so in case if there is a fear polluting our upstream rbe cluster(engflow) we can use this.

…ci-pr-stage-reuse-and-jax-bazel-cache

subodh-dubey-amd

Re-checked against daa2ea5: all the code changes from my last review are in. I reran the new probe against a local TLS 1.3 server: a rejected client cert now fails with TLSV1_ALERT_UNKNOWN_CA, an accepted one passes, and a silent server times out.

Approving the code, since it stays inert until certs are provisioned. As you note in the description, only a read-only cert should go on the runners; the write-cert options are in the credentials thread.

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.

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.

+        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}")


+
+def stage_writer_token(token: str, dest_dir: Path) -> Path:
+    dest = _fresh_dest_dir(dest_dir)
+    dest.write_text(token)


gulsumgudukbay requested a review from ScottTodd as a code owner September 3, 2026 16:19

github-project-automation Bot added this to TheRock Triage Sep 3, 2026

github-project-automation Bot moved this to TODO in TheRock Triage Sep 3, 2026

therock-pr-bot Bot added the Not ready to Review PR has unresolved policy failures — reviews blocked label Sep 3, 2026

ScottTodd reviewed Sep 3, 2026

View reviewed changes

gulsumgudukbay force-pushed the users/gulsumgudukbay/ci-pr-stage-reuse-and-jax-bazel-cache branch from 9d5adb0 to b6ea67d Compare September 3, 2026 18:55

gulsumgudukbay changed the title ~~Reuse unaffected build stages on PRs and cache the JAX Bazel build~~ ci(jax): serve the JAX wheel build from the EngFlow Bazel cache Sep 3, 2026

gulsumgudukbay marked this pull request as draft September 3, 2026 18:57

github-advanced-security Bot found potential problems Sep 3, 2026

View reviewed changes

Comment thread build_tools/github_actions/configure_jax_bazel_cache.py Fixed

gulsumgudukbay mentioned this pull request Sep 3, 2026

ci: reuse unaffected build stages on pull requests #7919

Closed

1 task done

therock-pr-bot Bot removed the Not ready to Review PR has unresolved policy failures — reviews blocked label Sep 3, 2026

gulsumgudukbay force-pushed the users/gulsumgudukbay/ci-pr-stage-reuse-and-jax-bazel-cache branch 3 times, most recently from c9e8fb8 to e822257 Compare September 3, 2026 21:10

gulsumgudukbay marked this pull request as ready for review September 4, 2026 00:48

gulsumgudukbay requested a review from ScottTodd September 8, 2026 18:44

gulsumgudukbay force-pushed the users/gulsumgudukbay/ci-pr-stage-reuse-and-jax-bazel-cache branch from e822257 to 119cee4 Compare September 8, 2026 18:54

ScottTodd requested review from amd-shiraz and subodh-dubey-amd September 14, 2026 20:46

ScottTodd reviewed Sep 14, 2026

View reviewed changes

gulsumgudukbay force-pushed the users/gulsumgudukbay/ci-pr-stage-reuse-and-jax-bazel-cache branch from 119cee4 to 6bcba4d Compare September 17, 2026 00:13

gulsumgudukbay requested a review from ScottTodd September 17, 2026 15:58

ScottTodd reviewed Sep 18, 2026

View reviewed changes

gulsumgudukbay mentioned this pull request Sep 25, 2026

ci(jax): build upstream JAX tip against nightly ROCm as a release-matrix row #8524

Merged

gulsumgudukbay force-pushed the users/gulsumgudukbay/ci-pr-stage-reuse-and-jax-bazel-cache branch from 6bcba4d to 237ab6b Compare September 29, 2026 02:35

gulsumgudukbay added 3 commits October 2, 2026 17:19

gulsumgudukbay force-pushed the users/gulsumgudukbay/ci-pr-stage-reuse-and-jax-bazel-cache branch from 237ab6b to 525a491 Compare October 2, 2026 17:19

Merge branch 'main' into users/gulsumgudukbay/ci-pr-stage-reuse-and-j…

4d47079

…ax-bazel-cache

subodh-dubey-amd requested changes Oct 6, 2026

View reviewed changes

gulsumgudukbay added 3 commits October 6, 2026 20:04

Merge remote-tracking branch 'origin/main' into users/gulsumgudukbay/…

96f3336

…ci-pr-stage-reuse-and-jax-bazel-cache

Merge remote-tracking branch 'origin/main' into users/gulsumgudukbay/…

ceb65d9

…ci-pr-stage-reuse-and-jax-bazel-cache # Conflicts: # .github/workflows/multi_arch_build_linux_jax_wheels_ci.yml

alekstheod reviewed Oct 21, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into users/gulsumgudukbay/…

daa2ea5

…ci-pr-stage-reuse-and-jax-bazel-cache

subodh-dubey-amd approved these changes Oct 28, 2026

View reviewed changes

This was referenced Oct 28, 2026

docs(pytorch): document the sccache compiler cache in the PyTorch README #9559

Open

docs(pytorch): document the sccache compiler cache #9560

Open

gulsumgudukbay added 2 commits November 5, 2026 17:25

github-advanced-security Bot found potential problems Nov 6, 2026

View reviewed changes

		# Carries the Bazel remote cache credentials into the wheel build.
		secrets: inherit

Conversation

gulsumgudukbay commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

therock-pr-bot Bot commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ScottTodd left a comment

Choose a reason for hiding this comment

Uh oh!

subodh-dubey-amd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

subodh-dubey-amd Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

subodh-dubey-amd left a comment

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

gulsumgudukbay commented Sep 3, 2026 •

edited

Loading

therock-pr-bot Bot commented Sep 3, 2026 •

edited

Loading

therock-pr-bot Bot commented Sep 3, 2026 •

edited

Loading

subodh-dubey-amd Oct 28, 2026 •

edited

Loading