1
0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00

559 Commits

Author SHA1 Message Date
charleshofer
6451164dc9
Port Jenkins release CI to Actions () 2025-04-15 12:03:18 -05:00
charleshofer
1f93b4b9b5
Add Python and ROCm version matrix for rocm-main () 2025-04-02 11:49:03 -05:00
charleshofer
7b20b7c0b7
Use rebase for upstream sync () 2025-03-27 11:58:59 -05:00
charleshofer
13d88b6340
Add back raw totals in JSON reports () 2025-03-24 11:26:31 -05:00
rocm-repo-management-api-2[bot]
b505df9973
Merge pull request from ROCm/ci-upstream-sync-152_1
CI: 03/19/25 upstream sync
2025-03-19 07:20:19 -05:00
charleshofer
1f2fe33091
Enable upstream CI on release branches () 2025-03-18 17:22:37 -05:00
jax authors
ee0073e605 Merge pull request from vfdev-5:fix-tsan-numpy-install-patch
PiperOrigin-RevId: 738080051
2025-03-18 11:24:15 -07:00
vfdev-5
9145d617b8 Added exit 1 if git patch is failed + other checks 2025-03-18 16:47:34 +01:00
Charles Hofer
c7b407c9f0 Merge branch 'rocm-main' into ci-upstream-sync-151_1 2025-03-18 15:27:35 +00:00
jax authors
b74b16f9b9 Merge pull request from MichaelHudgins:a4-testing
PiperOrigin-RevId: 737733904
2025-03-17 13:36:04 -07:00
Michael Hudgins
ecf7fde714 Add B200 testing to continuous workflow 2025-03-17 20:19:20 +00:00
Nitin Srinivasan
031614c22b Pin numpy~=2.1.0 in workflow file instead of test-requirements.txt
PiperOrigin-RevId: 737632771
2025-03-17 08:59:06 -07:00
GitHub Actions
e275d5cf6c Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync-147_1 2025-03-14 22:42:07 +00:00
charleshofer
022da913e6
Count test totals correctly for dashboards ()
* Account test totals correctly for dashboards

* Add blurb to the dev guide on skipping tests

* Remove extra newline

* Default to 0 if "skipped" isn't found

Co-authored-by: Mathew Odden <1471252+mrodden@users.noreply.github.com>

---------

Co-authored-by: Mathew Odden <1471252+mrodden@users.noreply.github.com>
2025-03-14 16:57:08 -05:00
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
GitHub Actions
a0edd3fbb2 Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync-144_1 2025-03-12 16:57:18 +00: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
Charles Hofer
fb89a4b427 Merge branch 'rocm-main' into ci-upstream-sync-142_1 2025-03-11 16:33:59 +00: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
Charles Hofer
07cd809ba8 Merge branch 'rocm-main' into ci-upstream-sync-135_1 2025-03-04 16:22:17 +00: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
Charles Hofer
45e2060b90 Merge branch 'rocm-main' into ci-upstream-sync-127_1 2025-02-25 19:30:08 +00:00
charleshofer
9cac506bbb
Change CODEOWNERS () 2025-02-25 12:51:24 -06: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