Add new scripts and configs for running JAX tests #23677
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
copybara-service
bot
force-pushed
the
test_675316647
branch
from
497539e
to
16bee7c
Compare
This adds new CI scripts and envs for running JAX tests. As with the build artifact scripts, these scripts require that `JAXCI_ENV_FILE` be set to one of the envs inside `ci/envs/run_tests` before invoking `run_bazel_test.sh` or `run_pytest.sh`. Both CPU and GPU commands are in the same script and script behaviors are controlled by `JAXCI` environment variables set by `JAXCI_ENV_FILE`. In order to make the Bazel command concise, test configs are moved to the .bazelrc and are grouped under the `multiaccelerator` and `non_multiaccelerator` configs As an example, for running Bazel CPU tests with RBE, we would run: ``` 1. export JAXCI_ENV_FILE=ci/envs/run_tests/bazel_cpu 2. ./ci/run_bazel_test.sh ``` for running Pytest GPU tests, we would run: ``` 1. export JAXCI_ENV_FILE=ci/envs/run_tests/pytest_gpu 2. ./ci/run_pytest.sh ``` As Pytests are run locally, note that these scripts require the JAX wheels to be present inside the `dist/` folder in the JAX git repository root. PiperOrigin-RevId: 675316647
copybara-service
bot
force-pushed
the
test_675316647
branch
from
16bee7c
to
8c77303
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Add new scripts and configs for running JAX tests
This adds new CI scripts and envs for running JAX tests. As with the build artifact scripts, these scripts require that
JAXCI_ENV_FILEbe set to one of the envs insideci/envs/run_testsbefore invokingrun_bazel_test.shorrun_pytest.sh. Both CPU and GPU commands are in the same script and script behaviors are controlled byJAXCIenvironment variables set byJAXCI_ENV_FILE. In order to make the Bazel command concise, test configs are moved to the .bazelrc and are grouped under themultiacceleratorandnon_multiacceleratorconfigsAs an example, for running Bazel CPU tests with RBE, we would run:
for running Pytest GPU tests, we would run:
As Pytests are run locally, note that these scripts require the JAX wheels to be present inside the
dist/folder in the JAX git repository root.