Skip to content

fix(jax): stamp JAX wheels with the real source commit hash - #6979

Open
WBobby wants to merge 1 commit into
mainfrom
users/yanywang/28-7-fix-jax-commit
Open

fix(jax): stamp JAX wheels with the real source commit hash#6979
WBobby wants to merge 1 commit into
mainfrom
users/yanywang/28-7-fix-jax-commit

Conversation

@WBobby

@WBobby WBobby commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Motivation

jaxlib_git_hash was resolved by a git call that ran inside the manylinux
container, where git rejects the runner-owned checkout as dubiously owned. The
command substitution failed silently, so every JAX plugin wheel from these two
workflows was published without a git hash while the job stayed green.

Issue: #6978

Technical Details

This change:

  • Resolves the hash on the host, where the checkout is owned by the user running
    the step, and forwards it with --env JAXLIB_GIT_HASH. Without the --env the
    reference inside the single-quoted script would be an unset container variable,
    i.e. the same empty value by a different route.
  • Validates it as a 40-hex sha and fails the step otherwise, so "no hash" becomes
    a red job instead of a silently unstamped wheel. set -e already covers git
    exiting non-zero; the regex covers git succeeding with unexpected output.
  • Sets safe.directory=* in the container via GIT_CONFIG_COUNT /
    GIT_CONFIG_KEY_0 / GIT_CONFIG_VALUE_0, mirroring
    dockerfiles/build_manylinux_x86_64.Dockerfile. This is defence in depth for any
    other git caller in the image; the hash itself no longer depends on it.

The inner script is otherwise untouched, and stays single-quoted because
$(rocm-sdk path --root) and $(pwd) must be evaluated in the container. Both the
release and CI workflows carried the identical line and are changed identically.

Test Plan

Dispatched multi_arch_build_linux_jax_wheels.yml on this branch for a single
matrix cell (build_mode=manylinux, python_version=3.14,
jax_ref=rocm-jaxlib-v0.10.0, rocm_version=7.15.0a20260728, release_type=dev,
tests skipped), then inspected the published wheel.

Test Result

Run 30425080495,
success in 49m37s. All 15 steps succeeded, the native v0.9.1 step was skipped as
expected, and dubious ownership appears 0 times in the log.
The hash is resolved on the host and reaches Bazel intact:
jaxlib_git_hash=98a16bf22b1cf8384b730e835d785cb0c68033d5
'build' options: ... --//jaxlib/tools:jaxlib_git_hash=98a16bf22b1cf8384b730e835d785cb0c68033d5
That sha is refs/heads/rocm-jaxlib-v0.10.0 in ROCm/jax, confirmed with
git ls-remote.
jax_rocm7_plugin/version.py in the published wheel, before and after:

_release_version _git_hash
before (nightly 20260728) '0.10.0+rocm7.15.0a20260728' None
after (this branch) '0.10.0+rocm7.15.0a20260728' '98a16bf22b1cf8384b730e835d785cb0c68033d5'

Submission Checklist

Resolve the hash on the host instead of inside the container, where git rejects
the runner-owned checkout and silently yields an empty hash.

Signed-off-by: Yanyao Wang <wangyanyao@msn.com>
@therock-pr-bot

therock-pr-bot Bot commented Jul 29, 2026

Copy link
Copy Markdown

✅ All Checks Passed — Ready for Review

Check Status Details
📝 PR Description ✅ Pass
Forbidden Files ✅ Pass
🧪 Unit Test ✅ Pass PR does not contain code files — Unit Test auto-passed
🔎 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

Copy link
Copy Markdown

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

@WBobby
WBobby requested review from ScottTodd and erman-gurses July 29, 2026 20:44
Comment on lines +247 to +254
# Resolved on the host: the container runs as root against a
# runner-owned checkout, where git rejects the repo as dubiously owned.
JAXLIB_GIT_HASH="$(git rev-parse HEAD)"
if [[ ! "${JAXLIB_GIT_HASH}" =~ ^[0-9a-f]{40}$ ]]; then
echo "::error::refusing to build: jax-source HEAD resolved to '${JAXLIB_GIT_HASH}'"
exit 1
fi
echo "jaxlib_git_hash=${JAXLIB_GIT_HASH}"

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The if and echo here look unnecessary, remove?

Comment on lines +247 to +248
# Resolved on the host: the container runs as root against a
# runner-owned checkout, where git rejects the repo as dubiously owned.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could the container itself apply ownership settings for the checkout to fix instead? Again, prefer to keep workflows as simple as possible. Every single line of bash in a workflow should be scrutinized. Treat workflows as recipes for users to run. You wouldn't ask a user to run any of this code... right? Users should just run scripts in sequence. Same for workflows.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: TODO

Development

Successfully merging this pull request may close these issues.

2 participants