mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,234 Commits 501 Branches 263 Tags
Commit Graph

476 Commits

Author SHA1 Message Date
Nitin Srinivasan
5944c9ed65 Install test dependencies from test-requirements.txt instead of requirements.in 
PiperOrigin-RevId: 736878834
 2025-03-14 08:57:20 -07:00
Nitin Srinivasan
12760af236 Add custom job names to group different matrix combinations in the Actions dashboard 
PiperOrigin-RevId: 736481804
 2025-03-13 06:23:04 -07:00
Nitin Srinivasan
d89835acba Fix matrix exclude syntax in TPU tests block 
Also, skip Python 3.13 for now due to missing dependency error.

PiperOrigin-RevId: 736120590
 2025-03-12 07:12:52 -07:00
Nitin Srinivasan
a6ab6bbc20 Ignore Pallas TPU tests when testing with the oldest supported libtpu 
I missed adding this in from https://github.com/jax-ml/jax/blob/main/.github/workflows/cloud-tpu-ci-nightly.yml when I added the TPU jobs to the new CI workflows

PiperOrigin-RevId: 736094492
 2025-03-12 05:20:42 -07:00
Nitin Srinivasan
ff751ecc7b Run single python version for v4-8 and min & max for v5e-8 for TPU tests in nightly/release test workflow 
PiperOrigin-RevId: 735975004
 2025-03-11 20:03:05 -07:00
jax authors
e0545a71eb Remove installation of NVIDIA wheels for CPU tests 
PiperOrigin-RevId: 735875073
 2025-03-11 13:43:13 -07:00
Nitin Srinivasan
7ac6355262 Add TPU test jobs to the new CI continuous and nightly/release test workflows 
Also, modify the TPU presubmit workflow to reuse the `build_artifacts.yml` and `pytest_tpu.yml`

PiperOrigin-RevId: 735832964
 2025-03-11 11:42:21 -07:00
vfdev
1bab037ca0 Add file and zip to tsan.yaml 2025-03-10 16:51:05 +00:00
Nitin Srinivasan
623865fe95 Build JAX wheels instead of installing it from the source repository 
This change allows us to get rid of extra env vars which used to control whether to install `jax` at head. Now, `jax` will be be built and consumed in the same way as the other wheels in the continuous jobs.

PiperOrigin-RevId: 734123590
 2025-03-06 06:48:16 -08:00
Peter Hawkins
40e1a2a561 Remove a TSAN suppression. 
https://github.com/python/cpython/issues/130547 has been marked as fixed and backported to 3.13, so this suppression should no longer be necessary.
 2025-03-05 08:39:58 -05:00
Jake VanderPlas
f0bbd26d03 Update array-api-tests to latest commit 2025-03-04 10:17:51 -08:00
Kanglan Tang
da7c90c4c4 Bump oldest supported libtpu to match the compatibility window (12 weeks) 
PiperOrigin-RevId: 732218061
 2025-02-28 12:13:42 -08:00
Kanglan Tang
55263ce485 Add linux python 3.13t nightly tests 
* Python wheels follow a naming convention: standard wheels use the pattern `*-cp<python_version>-cp<python_version>-*`, while free-threaded wheels use `*-cp<python_version>-cp<python_version>t-*`. Update the pytest workflows to look for free-threaded wheels and ensure that standard wheel tests exclude free-threaded wheels.

* Skip zstandard for python3.13-nogil due to compilation failure https://github.com/indygreg/python-zstandard/issues/231.

PiperOrigin-RevId: 732070585
 2025-02-28 03:13:39 -08:00
Kanglan Tang
d839e441b7 Reduce pytest workers for asan to resolve memory usage causing OOM 
This fixes the current OOM error: https://github.com/jax-ml/jax/actions/runs/13565999206.

PiperOrigin-RevId: 731876781
 2025-02-27 14:54:29 -08:00
Peter Hawkins
6e73637888 Fix a test failure under multi-threading. 
Remove a tsan suppression for a CPython race that is fixed.
 2025-02-27 06:07:05 -05:00
jax authors
615219b1f6 Remove tensorstore dependency from //jax/experimental/array_serialization:serialization in OSS (see https://github.com/google/tensorstore/issues/218) 
Disable serialization_test in OSS.

PiperOrigin-RevId: 731463136
 2025-02-26 14:47:16 -08:00
Nitin Srinivasan
a65de52421 Enable resultstore logging 
Tests logged with resulstore are much easier to read and debug

PiperOrigin-RevId: 731448196
 2025-02-26 14:04:58 -08:00
Dan Foreman-Mackey
b8f236e64d Add --system to uv commands in upstream-nightly workflow. 2025-02-26 13:21:41 -05:00
Peter Hawkins
33bbd5f119 Fix failures in TSAN free threading CI. 2025-02-26 06:04:26 -05:00
Nitin Srinivasan
7566daba68 Use uv instead of pip for installing Python packages 
Missed including these in 4b4f2f9cb9

PiperOrigin-RevId: 731095379
 2025-02-25 17:48:22 -08:00
Nitin Srinivasan
f57c18ad1b Install uv to fix module not found error on Windows 
Ideally, this install should be in the Dockerfile but updating the Windows dockerfile is not straightforward so I'm doing the install here for the time being.

PiperOrigin-RevId: 731055684
 2025-02-25 15:39:07 -08:00
Nitin Srinivasan
771306bab3 Use ${{ !cancelled() }} instead of ${{ always() }} 
`${{ always() }}` makes it difficult to cancel a workflow. See https://github.com/orgs/community/discussions/26303

PiperOrigin-RevId: 731044750
 2025-02-25 15:06:38 -08:00
jax authors
dc1c3f9abd Disable //tests:serialization_test_cpu from TSAN job and remove tensorstore dependency from //jax/experimental/array_serialization:serialization. 
`TSAN CPython` is unable to find a compatible version of `tensorstore` wheel, hence the test can not be executed.

PiperOrigin-RevId: 731027518
 2025-02-25 14:19:02 -08:00
Nitin Srinivasan
cf01fdfe6a Use the 64 core Windows runner to build artifacts 
Now that we have disabled RBE on Windows, we need to use the bigger machine to build fast.

PiperOrigin-RevId: 731012952
 2025-02-25 13:42:16 -08:00
Nitin Srinivasan
2f6f722150 Disable RBE on Windows 
We no longer have a RBE pool with ltsc2019 image and are blocked on upgrading GKE to ltsc2022.

PiperOrigin-RevId: 730997201
 2025-02-25 12:58:45 -08:00
Nitin Srinivasan
7a162f2abc Fix incorrect line separator 
On Windows, we are supposed to use "`" instead of "\"

PiperOrigin-RevId: 730916160
 2025-02-25 09:28:06 -08:00
Peter Hawkins
2325cf3556 Update tsan suppressions. 
Add a new suppression and comment out some old ones.
 2025-02-25 09:23:55 -05:00
Nitin Srinivasan
cc830748bf Remove -y and extra --pre arguments 
Also, remove invalid `install` package from install commands

PiperOrigin-RevId: 730689569
 2025-02-24 19:34:08 -08:00
Peter Hawkins
c8c4cfa04e Update version numbers after 0.5.1 release. 2025-02-24 16:18:25 -05:00
Nitin Srinivasan
cadb531182 Remove extra "--system" argument 
This should fix https://github.com/jax-ml/jax/actions/runs/13504997366/job/37732423493?pr=26378

PiperOrigin-RevId: 730525868
 2025-02-24 11:18:24 -08:00
Nitin Srinivasan
4b4f2f9cb9 Use uv to install Python packages 
PiperOrigin-RevId: 730499307
 2025-02-24 10:13:39 -08:00
Nitin Srinivasan
5089fb01b2 Use new ML Build CUDA images 
Note the CUDA jobs fail in https://github.com/jax-ml/jax/actions/runs/13444028636/job/37565647540 but in the same way as on HEAD.

PiperOrigin-RevId: 729571786
 2025-02-21 10:16:16 -08:00
Michael Hudgins
c664a0cd44 [CI] Enable workflow_dispatch for the continuous workflow 
This enables testing changes that have a high chance of breaking a longer running test.

PiperOrigin-RevId: 729515893
 2025-02-21 06:59:20 -08:00
Nitin Srinivasan
6b39eb6b92 Fix wheel downloads for nightly artifact testing 
This adds an extra filter to make sure that we are only downloading the `jax` python package and uses `call` on Windows so that we can run sequential gsutil commands

PiperOrigin-RevId: 729198507
 2025-02-20 11:59:27 -08:00
Nitin Srinivasan
040b718e44 Bump up timeout for building wheel to 60 minutes 
Avoids running into timeout issues on Windows when we don't see many hits from the RBE cache. E.g: https://github.com/jax-ml/jax/actions/runs/13439584940/job/37550341443

PiperOrigin-RevId: 729151367
 2025-02-20 09:58:48 -08:00
jax authors
91c6e449ae Merge pull request #26461 from ROCm:run-less-rocm-tests 
PiperOrigin-RevId: 726614797
 2025-02-13 13:27:47 -08:00
Nitin Srinivasan
93831bdde7 Download and use jax wheels from GCS bucket for nightly/release test workflows 
Unlike continuous workflows, when testing nightly/release artifacts, we want to download and install the `jax` wheels found in the GCS bucket instead of installing it from HEAD.

It looks like `env` setting in the calling workflow isn't passed over to the called workflows so we define a new workflow input, `install-jax-current-commit`, to control the `jax` install behavior.

PiperOrigin-RevId: 726086522
 2025-02-12 09:32:05 -08:00
Nitin Srinivasan
30acd383fb Run test job irrespective of if the build jobs succeeds or fails 
This lets us avoid losing test coverage if a single unrelated build job fails. E.g Windows build job fails but everything else succeeds. In this case, we still want to run the tests for other platforms.

Also, if a build job fails, its corresponding test job will also report a failure as a result of not being able to download the wheel artifact so we should still be able to tell the source of job failure easily.

PiperOrigin-RevId: 725754098
 2025-02-11 13:37:30 -08:00
Charles Hofer
d50206b830 [ROCm] Only run core tests for ROCm post-build checks 2025-02-10 18:30:44 +00:00
Adam Paszke
c24a3a7766 Bump oldest supported libtpu to match the compatibility window (12 weeks) 
PiperOrigin-RevId: 725193930
 2025-02-10 06:41:07 -08:00
jax authors
1ca8807dca Merge pull request #25810 from ROCm:gh-9948-add-gpu-ci-upstream 
PiperOrigin-RevId: 724378710
 2025-02-07 09:51:59 -08:00
charleshofer
ebf4a54f4f Add AMD ROCm GPU CI post-build check (#137) 2025-02-06 21:41:53 +00:00
Michael Hudgins
be0f3d8f1b Update tsan job to upload to resultstore 2025-02-05 20:14:05 +00:00
Nitin Srinivasan
bc055569df Replace Python 3.12 with Python 3.13 and add Python 3.10 to the matrix 
Expands test coverage to cover the oldest and newest Python versions that we support.

PiperOrigin-RevId: 723520699
 2025-02-05 08:27:22 -08:00
Nitin Srinivasan
b5b913acb6 Run Bazel CPU/CUDA RBE jobs on pushes to main/release branches 
This helps identify breaking commits easily.

PiperOrigin-RevId: 723519953
 2025-02-05 08:24:37 -08:00
dependabot[bot]
333c4e7a0e
Bump actions/setup-python from 5.3.0 to 5.4.0 
Bumps [actions/setup-python](https://github.com/actions/setup-python) from 5.3.0 to 5.4.0.
- [Release notes](https://github.com/actions/setup-python/releases)
- [Commits](0b93645e9f...42375524e2)

---
updated-dependencies:
- dependency-name: actions/setup-python
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
 2025-02-04 21:07:18 +00:00
jax authors
694572bcaf Merge pull request #26310 from dfm:use-uv-in-ci 
PiperOrigin-RevId: 723199573
 2025-02-04 13:06:09 -08:00
jax authors
b1b88a3613 Merge pull request #26307 from dfm:kill-ffi-cache 
PiperOrigin-RevId: 723131658
 2025-02-04 10:07:04 -08:00
jax authors
09ee37a41d Merge pull request #26302 from vfdev-5:readd-missed-sed-part-in-tsan 
PiperOrigin-RevId: 723115097
 2025-02-04 09:23:32 -08:00
Dan Foreman-Mackey
5db5e0d5ca Use uv for dependency resolution on CI. 2025-02-04 11:48:05 -05:00