[Feature]: adapt bazel build to use 'build' instead of 'run'

[ksebaz]

Dear developers,

with google/jax or ROCm/jax (certain newer branches), it is possible to add --bazel-options=--nobuild to the build command which will only fetch the dependencies since bazel build is called instead of bazel run. The latter is returning an error with the '--nobuild' option and thus dependencies for multiple wheels cannot be prefetched (if not repeating the build command explicitly for plugin and pjrt) before submitting the build (only; using --nofetch) to some more powerful machine but without network.

Could you maybe also switch to bazel build here? This would easy automated building.

[lucbruni-amd]

Hi @ksebaz, I've linked a PR to address this issue. Please give it a try with your workflow to see if it helps automate building further. Any feedback is welcome. Thanks!

[ksebaz]

Hi @lucbruni-amd, the general approach works, but it fails eventually now with not finding patchelf. I'm not a bazel expert but I have the feeling that the jax_wheel action is run differently under run and build. With the latter the environment is extremely restricted (my feeling is action_env flags are not passed) and none of the methods you implemented in _find_patchelf can actually work (except patchelf is in system's locations).

[ksebaz]

After some investigation, I have the feeling that the only sane way to allow for the find_patchelf functionality to work is to patch jax's jaxlib/jax.bzl to run jax_wheel with use_default_shell_env=True. Otherwise one could add patchelf as a proper dependency in bazel to the targets, but I fear this is a bit more troublesome due to how the whole build is mixing with upstream jax-functionality.

I just tried with use_default_env=True and that works fine. I can also provide a patch on top of your PR for that in case you are interested. But please feel free to also drop this feature request in case it doesn't align.