21 Commits

Author SHA1 Message Date
Nitin Srinivasan
8ac2969759 Pass JAXLIB_* env variables to docker container
PiperOrigin-RevId: 730922129
2025-02-25 09:46:30 -08:00
Nitin Srinivasan
e4d63dff3d Pass in "JAX_" variables to the docker container
PiperOrigin-RevId: 730689480
2025-02-24 19:32:29 -08:00
Nitin Srinivasan
4b4f2f9cb9 Use uv to install Python packages
PiperOrigin-RevId: 730499307
2025-02-24 10:13:39 -08:00
Nitin Srinivasan
a073639124 Update Windows docker image to pin it to a sha
PiperOrigin-RevId: 729570957
2025-02-21 10:13:55 -08:00
Adam Paszke
f1ab7514db Make sure we take libTPU version into account in the Pallas lowering
Also, strengthen the presubmit to make sure we catch more errors.

PiperOrigin-RevId: 726061633
2025-02-12 08:15:57 -08:00
Nitin Srinivasan
36679d89e3 Add a retry logic to avoid flakes when downloading Bazel
PiperOrigin-RevId: 720349541
2025-01-27 16:36:01 -08:00
Nitin Srinivasan
89a9c6c244 Add new Bazel remote cache configs
An example run where the cache configs are used: https://github.com/jax-ml/jax/actions/runs/12940123731

PiperOrigin-RevId: 719627011
2025-01-25 05:32:42 -08:00
Nitin Srinivasan
9aad6a6827 Add job that runs Bazel single accelerator and multi-accelerator CUDA tests (non-RBE)
PiperOrigin-RevId: 718637923
2025-01-22 17:51:45 -08:00
Nitin Srinivasan
14029c6762 Replace "gpu" with "cuda" to be specific about the type of gpu tests that are running
Also, make `run_docker_container.sh` executable

PiperOrigin-RevId: 718420712
2025-01-22 09:40:11 -08:00
Nitin Srinivasan
9fb29766a2 Add workflow for testing nightly/release artifacts
This commits adds a Github action workflow that will be used to jobs that test the nightly/release artifacts. These artifacts are built by our internal CI jobs and are uploaded to a transient GCS bucket. After all the wheels have finished uploading, an internal job is run that that will trigger the `wheel_tests_nightly_release.yml` workflow.

PiperOrigin-RevId: 716789482
2025-01-17 13:53:47 -08:00
Nitin Srinivasan
c78487d23d Add Github action workflows for running continuous tests with Pytest
Changes:
- Adds `wheel_tests.yml` that will be used to run continuous jobs that builds artifacts and runs CPU/CUDA tests. Jobs will run by workflow calls to `build_artifacts.yml`/`pytest_cpu.yml`/`pytest_gpu.yml`.
- Adds testing of CUDA tests on H100 gpus
- Make script executable
- Change the name of GPU scripts and workflows to CUDA to be more clear as to what is being tested
PiperOrigin-RevId: 715500412
2025-01-14 13:10:51 -08:00
Sergei Lebedev
90201ce2b7 Removed leftover mentions of xmap from the code
PiperOrigin-RevId: 713202387
2025-01-08 01:39:38 -08:00
Nitin Srinivasan
8552852609 Add utility script and env for running the CI scripts under Docker
This commit adds the scripts and envs that set up a Docker container named "jax". It is used in internal JAX CI jobs that handle building and publishing JAX artifacts to PyPI and/or GCS.

While GitHub action workflows use the same Docker images, they do not run `run_docker_container.sh` script as they leverage built-in containerization features to run jobs within a container.

PiperOrigin-RevId: 706793601
2024-12-16 12:19:23 -08:00
Nitin Srinivasan
d05ab5bb0d Add an experimental Cloud TPU presubmit job
This adds an experimental non-blocking presubmit job that will run a subset of TPU tests, focusing on frequently failing tests. The goal is to achieve comprehensive coverage while keeping the runtime around 10 minutes.

PiperOrigin-RevId: 706064568
2024-12-13 18:06:17 -08:00
Nitin Srinivasan
83c64b2379 Add a flag to enable detailed timestamped logging of subprocess commands.
This adds a new command-line flag, `--detailed_timestamped_log`, that enables detailed logging of Bazel build commands. When disabled (the default), logging mirrors the output you'd see when running the command directly in your terminal.

When this flag is enabled:
- Bazel's output is captured line by line.
- Each line is timestamped for improved traceability.
- The complete log is stored for potential use as an artifact.

The flag is disabled by default and only enabled in the CI builds. If you're running locally and enable `detailed_timestamped_log`, you might notice that Bazel's output is not colored. To force a color output, include `--bazel_options=--color=yes` in your command.

PiperOrigin-RevId: 703581368
2024-12-06 12:32:24 -08:00
Nitin Srinivasan
8c78c1e735 Convert MSYS' Linux-like paths to Windows paths in JAX CI.
This is necessary on Windows because some applications such as Bazel/Docker do not understand/handle Linux-like paths that msys uses.

The script first converts all `JAXCI.*DIR` variables containing msys-like paths to Windows paths using cygpath.  Then it sources those variables into the shell environment so that any tools that use those variables can run correctly.

PiperOrigin-RevId: 702533023
2024-12-03 18:07:23 -08:00
Nitin Srinivasan
031c0acf50 Add new CI scripts for running Pytests
This commit adds the new CI scripts for running Pytests. It makes use of the pytest envs inside the "ci/envs/run_tests" folder to control the build behavior. For e.g: for running the GPU tests with Pytest, we will need to run `./ci/run_pytest.sh ./ci/envs/run_tests/pytest_gpu.env`. Note that Pytests need JAX wheels to be installed on the system to work. The `install_wheels_locally.sh` script installs these wheels in CI builds.

PiperOrigin-RevId: 701331411
2024-11-29 12:08:17 -08:00
Nitin Srinivasan
13726690dd Add new CI script to run Bazel GPU (non-RBE) jobs
This commit adds the CI script needed for running Bazel GPU (non-RBE) tests. These run two Bazel commands: Single accelerator tests with one GPU a piece and multi-accelerator tests with all GPUs

PiperOrigin-RevId: 700523594
2024-11-26 18:42:27 -08:00
Nitin Srinivasan
f7e9f62537 Add new CI scripts for building JAX artifacts
This commit introduces new CI scripts and environment files for building JAX artifacts. It makes use of the artifact envs inside the "ci/envs/build_artifacts" folder to control the build behavior. For e.g: for building jaxlib, we will need to run `./ci/build_artifacts.sh ./ci/envs/build_artifacts/jaxlib.env` from the JAX GitHub root.

PiperOrigin-RevId: 700104283
2024-11-25 14:37:02 -08:00
Nitin Srinivasan
14187399d7 Add new CI script for running Bazel GPU presubmits
PiperOrigin-RevId: 697643622
2024-11-18 08:52:51 -08:00
Nitin Srinivasan
195d407081 Add new CI scripts for running Bazel CPU presubmits
This commit introduces new CI scripts and environment files for running Bazel CPU presubmits.

* Adds a ci directory at the root of the repository to store these files.
* Environment files are located in ci/envs and define new JAXCI_ environment variables to control CI build behavior.
* The build script sources these environment files and set up the build environment before running the build commands.

PiperOrigin-RevId: 695957540
2024-11-12 20:00:35 -08:00