CI for linux-aarch64

[Rohanjames1997]

Hi Jax community!

I'm interested in understanding more about the current CI capabilities in JAX, especially for linux-aarch64. I am also interested in contributing to any gaps that may currently exist.

1. In ci-build.yaml I see a line in the build jobs section that runs-on: linux-x86-n2-32 alone. Does this mean that CI for every PR is only run on linux-x86 at the moment? What is the effort required to run on a self-hosted arm64 runner as well? Would we have to wait until Github-hosted arm64 runners are GA?

2. At the same time, in bazel_cpu_rbe.yml I see this line that runs some tests on linux-arm64-t2a-16 as well - what tests are they and when are they triggered?

3. Finally, these lines in bazel_cpu_rbe.yml build the test targets on a linux-x86 machine through cross-compilation. When is this triggered? And, is it a gap in Bazel that requires this cross-compilation step instead of native compilation?

Thanks!

[MichaelHudgins]

Hi @Rohanjames1997 thanks for the interest in the arm64 builds for JAX.

Does this mean that CI for every PR is only run on linux-x86 at the moment?

That is the primary one that runs in GitHub actions currently, there are additional tests that run on an internal system that include greater coverage, such as macOS, GPU, and several others. The results for those tests are posted as status contexts in PRs but their configs and several scripts are currently not in the public repo. This will be changing in the upcoming weeks as more moves to be publicly visible and hosted within the JAX repo itself.

What is the effort required to run on a self-hosted arm64 runner as well? Would we have to wait until Github-hosted arm64 runners are GA?

We are in the process of doing that right now! You can expect to see matching runners for arm64 as well as greater testing parity in the next few weeks. It likely wont be 100% equivalent across the two architectures but it will be much greater than it is currently.

At the same time, in bazel_cpu_rbe.yml I see this line that runs some tests on linux-arm64-t2a-16 as well - what tests are they and when are they triggered?

Thats a work in progress job currently. It will be moved to run on PR as a presubmit in the next few weeks. That tests uses a cross complication RBE to run builds on arm64 much faster than we could with a single machine. We cannot yet extend that to all testing on arm64 but you can expect greater coverage as we roll out these additional machines.

Finally, these lines in bazel_cpu_rbe.yml build the test targets on a linux-x86 machine through cross-compilation. When is this triggered? And, is it a gap in Bazel that requires this cross-compilation step instead of native compilation?

Same as above, the trigger is currently manual as we fine tune our rollout of this better coverage. The gap here is not bazel itself but rather a lack of easily utilizing arm64 machines as RBE runners. We cross compile to use RBE which allows for much faster build times in presubmits.

Let us know if you have any other questions!

[Rohanjames1997]

@MichaelHudgins , thanks a lot!
Your reply, along with @hawkinsp's here has clarified most of my doubts 😃

A few more questions out of curiosity:

The results for those tests are posted as status contexts in PRs

This is news to me! Do you happen to have a link to an example 'status context' handy? I couldn't find any from the Github Action logs of some of the PRs that I went through.

You can expect to see matching runners for arm64 as well as greater testing parity in the next few weeks. It likely wont be 100% equivalent across the two architectures but it will be much greater than it is currently.

That is great to know! I'd love to know how JAX can benefit from community contributions at this time. I'm guessing that the task of migrating the internal CI tests to Github Actions can only be handled by the maintainers (?), but increasing testing parity after the Action is released would benefit from OS contribs?

Also, does JAX have a tentative timeline or a tracking issue for when Actions for arm64 CI would be released?

Thank you!

[MichaelHudgins]

This is news to me! Do you happen to have a link to an example 'status context' handy? I couldn't find any from the Github Action logs of some of the PRs that I went through.

Looks like i was wrong on the macOS always running, the ones i was referring to are:
Tests with jaxlib build (Linux CPU, py 3.13, x64=0)
Tests with jaxlib build (Linux GPU, py 3.10, x64=0)
Tests (build-only) with jaxlib build (Linux Arm64 CPU, py 3.13, x64=0)

Some of them do get lost on the PR post submit (they are still on the original commit) so to see a full run you could look for an open, approved PR.

but increasing testing parity after the Action is released would benefit from OS contribs?

That would likely be an area that could use contribution

Also, does JAX have a tentative timeline or a tracking issue for when Actions for arm64 CI would be released?

I don't believe we have a tracker issue for it on GitHub, right now i would say the first or second week of December.

[Rohanjames1997]

Alright, thank you for the clarifications!

[Rohanjames1997]

Hello @MichaelHudgins , does the JAX team have an updated timeline for the release of arm64 Github Actions CI?

[MichaelHudgins]

Hi @Rohanjames1997

Arm64 is now built as part of this workflow https://github.com/jax-ml/jax/blob/main/.github/workflows/bazel_cpu_rbe.yml but those tests have not yet been made the blocking ones for submission. That will likely happen in the next few weeks.

https://github.com/jax-ml/jax/pull/25238/files is currently undergoing review and that will add larger continuous testing for arm64 as well as other configs. .