Skip to content

Thread JAX build matrix through setup#6270

Merged
ScottTodd merged 6 commits into
mainfrom
users/scotttodd/optional-jax-ci
Jun 30, 2026
Merged

Thread JAX build matrix through setup#6270
ScottTodd merged 6 commits into
mainfrom
users/scotttodd/optional-jax-ci

Conversation

@ScottTodd

@ScottTodd ScottTodd commented Jun 30, 2026

Copy link
Copy Markdown
Member

Motivation

Progress on #5634 and #6218.

This builds on #6117 to provide the JAX equivalent to #6082. Now the dynamic matrix of JAX builds to run will be generated at the start of CI and release pipelines during "setup" / "configure CI" and a table will be included in the summary:
image

Technical Details

The new build_jax option now has these settings:

Workflow build_jax value notes
multi_arch_ci.yml false Matching previous behavior
can later be enabled (opt-in, automatic based on files edited, etc.)
multi_arch_ci_asan.yml false New behavior due to #6218
multi_arch_release.yml inputs.build_jax
(default true)		 Matching previous behavior
multi_arch_release_asan.yml false New behavior due to #6218

Test Plan

Submission Checklist

ScottTodd added 5 commits June 24, 2026 14:33
@ScottTodd
Add optional JAX release build toggle
34e2509 
Mirror the PyTorch release workflow toggle for JAX so manual multi-arch dev releases can skip JAX wheels by default while workflow_call release runs continue to build them. Add a small JAX release matrix script so JAX can share the top-level python_version narrowing behavior.

Also add focused tests for the JAX matrix script and current PyTorch release matrix behavior.

Testing:

- D:/projects/TheRock/.venv/Scripts/python.exe -m pytest github_actions/tests/configure_jax_release_matrix_test.py github_actions/tests/configure_pytorch_release_matrix_test.py

- pre-commit run --files .github/workflows/multi_arch_release.yml .github/workflows/multi_arch_release_linux.yml .github/workflows/multi_arch_release_linux_jax_wheels.yml build_tools/github_actions/configure_jax_release_matrix.py build_tools/github_actions/configure_pytorch_release_matrix.py build_tools/github_actions/tests/configure_jax_release_matrix_test.py build_tools/github_actions/tests/configure_pytorch_release_matrix_test.py

Assisted-by: Codex
@ScottTodd
Merge remote-tracking branch 'upstream/main' into users/scotttodd/opt…
f99c177 
…ional-jax
@ScottTodd
Update JAX release matrix tests
8a776ab 
Avoid exhaustively asserting the expanded default matrix. Check that the default matrix spans multiple Python versions and JAX refs, and that explicit Python versions narrow every generated row.

Testing:

- D:/projects/TheRock/.venv/Scripts/python.exe -m pytest github_actions/tests/configure_jax_release_matrix_test.py

- D:/projects/TheRock/.venv/Scripts/python.exe -m pytest github_actions/tests/configure_pytorch_release_matrix_test.py

- pre-commit run --files build_tools/github_actions/tests/configure_jax_release_matrix_test.py

Assisted-by: Codex
@ScottTodd
Fix bad merge
99c2d48
@ScottTodd
Thread JAX build matrix through setup
a94f141 
Add build_jax to the multi-arch setup configuration path so CI can report that JAX is skipped while release workflows continue to report the JAX matrix they plan to build.

CI and ASAN callers keep JAX disabled. The main release workflow passes its optional JAX toggle into setup, and Linux release dispatch also checks the setup-generated build_config flag.

Testing:

- D:/projects/TheRock/.venv/Scripts/python.exe -m pytest github_actions/tests/configure_multi_arch_ci_test.py github_actions/tests/configure_jax_release_matrix_test.py

- D:/projects/TheRock/.venv/Scripts/python.exe -m pytest github_actions/tests/workflow_dispatch_inputs_test.py

- pre-commit run --files .github/workflows/setup_multi_arch.yml .github/workflows/multi_arch_ci.yml .github/workflows/multi_arch_ci_asan.yml .github/workflows/multi_arch_release.yml .github/workflows/multi_arch_release_asan.yml .github/workflows/multi_arch_release_linux.yml build_tools/github_actions/configure_multi_arch_ci.py build_tools/github_actions/configure_multi_arch_ci_summary.py build_tools/github_actions/tests/configure_multi_arch_ci_test.py

Assisted-by: Codex
@ScottTodd ScottTodd requested a review from erman-gurses June 30, 2026 20:36
erman-gurses
erman-gurses approved these changes Jun 30, 2026
View reviewed changes

@erman-gurses erman-gurses left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LGTM - should we also test it for the true case just in case?

@ScottTodd

ScottTodd commented Jun 30, 2026

Copy link
Copy Markdown
Member Author

(Note that the base branch for this PR is #6117, which is also up for review still)

@ScottTodd

ScottTodd commented Jun 30, 2026

Copy link
Copy Markdown
Member Author

should we also test it for the true case just in case?

Where specifically? I'm testing both this and #6117 at around the same time and I don't want to trigger too many full release builds. I think the main risk would be some of the default values not getting set or used as expected, which could cause jax builds to be skipped or enabled where we don't want them to be. I hope that the test workflow runs I've triggered already have exercised the rest of the plumbing sufficiently.

@erman-gurses

erman-gurses commented Jun 30, 2026

Copy link
Copy Markdown
Contributor

should we also test it for the true case just in case?

Where specifically? I'm testing both this and #6117 at around the same time and I don't want to trigger too many full release builds. I think the main risk would be some of the default values not getting set or used as expected, which could cause jax builds to be skipped or enabled where we don't want them to be. I hope that the test workflow runs I've triggered already have exercised the rest of the plumbing sufficiently.

I forgot about #6117 for a second - reviewing it.

@ScottTodd ScottTodd mentioned this pull request Jun 30, 2026
Open
Base automatically changed from users/scotttodd/optional-jax to main June 30, 2026 23:44
@therock-pr-bot

therock-pr-bot Bot commented Jun 30, 2026 •
edited
Loading

Copy link
Copy Markdown

❌ PR Check — Action Required

Check Status Details
🌿 Branch Name ✅ Pass —
📝 PR Title/Description ❌ Fail Error: Title does not follow Conventional Commits style.
Expected: start with a valid type (feat, fix, docs, …).
Desired format: type(optional-scope): short description
⛔ Forbidden Files ✅ Pass —
🧪 Unit Test ✅ Pass —
🔎 pre-commit ✅ Pass —
🚫 Draft PR 🔜 To Be Enabled —
🚩 Feature Flag 🔜 To Be Enabled —
📊 Code Coverage 🔜 To Be Enabled —

⚠️ 1 policy check(s) failed. Please address the issues above before this PR can be Reviewed.

🚫 Please fix the failed policies

  • ❌ PR Title/Description

The Not ready to Review label was added to this PR. Once all policies pass, the label is removed automatically.

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

@therock-pr-bot

therock-pr-bot Bot commented Jun 30, 2026

Copy link
Copy Markdown

🚫 Please fix the failed policies before requesting reviews.

The following policy checks failed:

  • ❌ PR Title/Description

The Not ready to Review label has been added to this PR.
Once all policies pass, the label will be removed automatically.

@ScottTodd ScottTodd merged commit 297b5cf into main Jun 30, 2026
17 of 18 checks passed
@ScottTodd ScottTodd deleted the users/scotttodd/optional-jax-ci branch June 30, 2026 23:54
@github-project-automation github-project-automation Bot moved this from TODO to Done in TheRock Triage Jul 1, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@erman-gurses erman-gurses erman-gurses approved these changes

Assignees

No one assigned

Labels

None yet

Projects

TheRock Triage
Status: Done

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@ScottTodd @erman-gurses