Because of C-API changes, this no longer works with upstream numpy without a jaxlib
build, which is too expensive to run as part of a github action. We should replace
with a nightly kokoro job, but in the meantime there's no reason to use these CI
resources for a failing nightly run.
The numpy-dispatch approach has been superseded by the Python Array API (Tracked for JAX in https://github.com/google/jax/issues/18353).
While we're here, we'll reduce the github CI to only two jobs: the oldest and newest supported Python versions. Other versions can be covered by Kokoro.
PiperOrigin-RevId: 587041291
This test is now independent of jax2tf. Move it out and rename it export_harnesses_multi_platform_test.py.
We disable the test in GitHub CI, because it is very large, pending
some changes to ensure it parallelizes well. The test is still
running in internal CI. This is matching the current behavior, since
jax2tf tests are only run internally.
PiperOrigin-RevId: 583603863
- A new matrix option was added to jaxlib-version matrix
- Expected to immediately fail due to strict jaxlib version check (requires libtpu to have the same API verison as jaxlib, fails otherwise)
- Due to expected failure, so chat will be sent when the jaxlib-version matrix is "nightly+oldest_supported_libtpu"
Add a --verbose option that logs all shell() commands run by the script.
Remove some Python 2 backward compatibility logic related to urllib and shutil.
Enable debug logging on Windows wheel builds.
Also include setuptools in the build requirements and test for its presence in build.py.
1. Use `strings` to print the libtpu build date, which is useful if
the libtpu version is stale and causes an API version mismatch
error.
2. Set the default shell to add bash -x command, so all commands are
printed inline. This makes lining up the output with the command
easier.
Notable changes:
* use PEP 585 type names
* use PEP 604 type union syntax where `from __future__ import annotations` is present.
* use f-strings in more places.
* remove redundant arguments to open().
`profiler_test.py:ProfilerTest.test_remote_profiler` fails with the
protobuf upgrade. However, I was seeing mysterious hangs without this,
and in general I think we should be testing with up-to-date deps given
that we don't pin. I'm gonna continue working on getting the Cloud TPU
CI green.
* Add deps to test requirements, including in new
`collect-profile-requirements.txt` (to avoid adding tensorflow to
`test-requirements.txt`).
* Use correct Python executable `ProfilerTest.test_remote_profiler`
(`python` sometimes defaults to python2)
* Run computations for longer in `ProfilerTest.test_remote_profiler`,
othewise `collect_profile` sometimes misses it.
The use of restore-keys allows the pip cache to match stale cache entries. This allows for partial cache reuse, but means the cache also grows without bound.
At the time of writing, the pip cache is 4.7GB, which is much larger than just downloading the wheels we need (20MB)! So it's probably a net positive not to reuse caches even if it requires some redundant downloads.
