pytest: use importlib mode by default #28650
Merged
Conversation
This file contains hidden or 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
Otherwise the default (prepend) mode will add `<jax_src_dir>/tests` and `<jax_src_dir>` to the start of `sys.path`. This is (has always been) fragile, because `<jax_src_dir>` has `jax/` and `jaxlib/` subdirectories, but it recently broke as `setuptools` 80 has a change in behaviour for editable installations, with the result that if `jaxlib` is installed editable then `import jaxlib` with `<jax_src_dir>` in `sys.path` will try to import from `<jax_src_dir>/jaxlib` instead of the editable install location, and fail.
jakevdp
approved these changes
May 9, 2025
google-ml-butler
bot
added
kokoro:force-run
pull ready
|
Hi - it looks like this broke our nightly tests (see #28671) so I think I'm going to roll back this change. |
I guess this is because of https://docs.pytest.org/en/stable/explanation/pythonpath.html#import-modes
and the relevant part is not included in the wheel installation? |
|
The error arises due to missing modules, but they're not test-specific modules. I'm seeing the following three files: I know we had some changes to our wheel build process recently – perhaps these files are missing from the package distribution somehow? |
|
Ah, they're explicitly excluded from the package distribution: https://github.com/jax-ml/jax/blame/76e5bc6f5d5c89c9d56d3211dba1953ae8d20cde/setup.py#L60 This was done deliberately in 8fbe3b1 I'm still not sure why |
The |
copybara-service bot
pushed a commit
that referenced
this pull request
May 15, 2025
Re-land #28650, fixing build failures, I hope! PiperOrigin-RevId: 759244911
copybara-service bot
pushed a commit
that referenced
this pull request
May 15, 2025
Re-land #28650, fixing build failures, I hope! PiperOrigin-RevId: 759244911
copybara-service bot
pushed a commit
that referenced
this pull request
May 15, 2025
Re-land #28650, fixing build failures, I hope! PiperOrigin-RevId: 759244911
copybara-service bot
pushed a commit
that referenced
this pull request
May 15, 2025
Re-land #28650, fixing build failures, I hope! PiperOrigin-RevId: 759244911
copybara-service bot
pushed a commit
that referenced
this pull request
May 16, 2025
Re-land #28650, fixing build failures, I hope! PiperOrigin-RevId: 759244911
copybara-service bot
pushed a commit
that referenced
this pull request
May 16, 2025
This is an attempt to re-land #28650, fixing build failures. **Motivation** #28650 was motivated by recent changes in setuptools, which caused test failures with editable installs, but it also exposes a potentially larger issue with our approach to testing with pytest. As far as I understand it, the default pytest behavior is to prepend the working directory to `sys.path`, meaning that `jax` is imported from the source directory, rather than the installed version. Switching to the `importlib` import mode means that we correctly test against the installed version of `jax`, which seems like what we typically want to do. The catch is that then we need to explicitly package any test utilities into the distribution. We don't currently package test-specific utilities like `internal_test_util` with JAX, but these utilities were still available to tests since they live within the `jax` source tree. This breaks when using `importlib` import mode, and a non-editable install of JAX. **Solutions** The approach that I've taken here is to explicitly package everything needed by the tests into the `jax` distribution. This means that we can correctly test against the _installed_ version of JAX when using pytest. This solution isn't ideal because it means that we're distributing `jax` submodules that aren't actually required except when running the test suite, but this seems like a small price to pay to me. **Alternatives** One different approach that we could take would be to only support using pytest with _editable_ installs of JAX. This would work because the required files would still be discoverable in an editable install because they live within the source tree. In fact, before this change, most of our CI jobs actually did install an editable distribution (which is why the failures in #28650 weren't caught in pre-submit!). The problem with this approach is that we're not actually testing JAX as it is used when installed from a distribution, and it wouldn't catch things like missing submodules. I think it's much better to test against the installed distribution! A more extreme approach would be to switch JAX to a `src/jax` and `src/jaxlib` layout (i.e. moving `jax` and `jaxlib` out of the root directory) as recommended by the Python packaging docs. Unfortunately this would be complicated with the way JAX is distributed internally at Google, so I think that's probably a non-starter. PiperOrigin-RevId: 759244911
copybara-service bot
pushed a commit
that referenced
this pull request
May 22, 2025
This is an attempt to re-land #28650, fixing build failures. **Motivation** #28650 was motivated by recent changes in setuptools, which caused test failures with editable installs, but it also exposes a potentially larger issue with our approach to testing with pytest. As far as I understand it, the default pytest behavior is to prepend the working directory to `sys.path`, meaning that `jax` is imported from the source directory, rather than the installed version. Switching to the `importlib` import mode means that we correctly test against the installed version of `jax`, which seems like what we typically want to do. The catch is that then we need to explicitly package any test utilities into the distribution. We don't currently package test-specific utilities like `internal_test_util` with JAX, but these utilities were still available to tests since they live within the `jax` source tree. This breaks when using `importlib` import mode, and a non-editable install of JAX. **Solutions** The approach that I've taken here is to explicitly package everything needed by the tests into the `jax` distribution. This means that we can correctly test against the _installed_ version of JAX when using pytest. This solution isn't ideal because it means that we're distributing `jax` submodules that aren't actually required except when running the test suite, but this seems like a small price to pay to me. **Alternatives** One different approach that we could take would be to only support using pytest with _editable_ installs of JAX. This would work because the required files would still be discoverable in an editable install because they live within the source tree. In fact, before this change, most of our CI jobs actually did install an editable distribution (which is why the failures in #28650 weren't caught in pre-submit!). The problem with this approach is that we're not actually testing JAX as it is used when installed from a distribution, and it wouldn't catch things like missing submodules. I think it's much better to test against the installed distribution! A more extreme approach would be to switch JAX to a `src/jax` and `src/jaxlib` layout (i.e. moving `jax` and `jaxlib` out of the root directory) as recommended by the Python packaging docs. Unfortunately this would be complicated with the way JAX is distributed internally at Google, so I think that's probably a non-starter. PiperOrigin-RevId: 759244911
copybara-service bot
pushed a commit
that referenced
this pull request
May 22, 2025
This is an attempt to re-land #28650, fixing build failures. **Motivation** #28650 was motivated by recent changes in setuptools, which caused test failures with editable installs, but it also exposes a potentially larger issue with our approach to testing with pytest. As far as I understand it, the default pytest behavior is to prepend the working directory to `sys.path`, meaning that `jax` is imported from the source directory, rather than the installed version. Switching to the `importlib` import mode means that we correctly test against the installed version of `jax`, which seems like what we typically want to do. The catch is that then we need to explicitly package any test utilities into the distribution. We don't currently package test-specific utilities like `internal_test_util` with JAX, but these utilities were still available to tests since they live within the `jax` source tree. This breaks when using `importlib` import mode, and a non-editable install of JAX. **Solutions** The approach that I've taken here is to explicitly package everything needed by the tests into the `jax` distribution. This means that we can correctly test against the _installed_ version of JAX when using pytest. This solution isn't ideal because it means that we're distributing `jax` submodules that aren't actually required except when running the test suite, but this seems like a small price to pay to me. **Alternatives** One different approach that we could take would be to only support using pytest with _editable_ installs of JAX. This would work because the required files would still be discoverable in an editable install because they live within the source tree. In fact, before this change, most of our CI jobs actually did install an editable distribution (which is why the failures in #28650 weren't caught in pre-submit!). The problem with this approach is that we're not actually testing JAX as it is used when installed from a distribution, and it wouldn't catch things like missing submodules. I think it's much better to test against the installed distribution! A more extreme approach would be to switch JAX to a `src/jax` and `src/jaxlib` layout (i.e. moving `jax` and `jaxlib` out of the root directory) as recommended by the Python packaging docs. Unfortunately this would be complicated with the way JAX is distributed internally at Google, so I think that's probably a non-starter. PiperOrigin-RevId: 759244911
copybara-service bot
pushed a commit
that referenced
this pull request
May 23, 2025
This is an attempt to re-land #28650, fixing build failures. **Motivation** #28650 was motivated by recent changes in setuptools, which caused test failures with editable installs, but it also exposes a potentially larger issue with our approach to testing with pytest. As far as I understand it, the default pytest behavior is to prepend the working directory to `sys.path`, meaning that `jax` is imported from the source directory, rather than the installed version. Switching to the `importlib` import mode means that we correctly test against the installed version of `jax`, which seems like what we typically want to do. The catch is that then we need to explicitly package any test utilities into the distribution. We don't currently package test-specific utilities like `internal_test_util` with JAX, but these utilities were still available to tests since they live within the `jax` source tree. This breaks when using `importlib` import mode, and a non-editable install of JAX. **Solutions** The approach that I've taken here is to explicitly package everything needed by the tests into the `jax` distribution. This means that we can correctly test against the _installed_ version of JAX when using pytest. This solution isn't ideal because it means that we're distributing `jax` submodules that aren't actually required except when running the test suite, but this seems like a small price to pay to me. **Alternatives** One different approach that we could take would be to only support using pytest with _editable_ installs of JAX. This would work because the required files would still be discoverable in an editable install because they live within the source tree. In fact, before this change, most of our CI jobs actually did install an editable distribution (which is why the failures in #28650 weren't caught in pre-submit!). The problem with this approach is that we're not actually testing JAX as it is used when installed from a distribution, and it wouldn't catch things like missing submodules. I think it's much better to test against the installed distribution! A more extreme approach would be to switch JAX to a `src/jax` and `src/jaxlib` layout (i.e. moving `jax` and `jaxlib` out of the root directory) as recommended by the Python packaging docs. Unfortunately this would be complicated with the way JAX is distributed internally at Google, so I think that's probably a non-starter. PiperOrigin-RevId: 759244911
copybara-service bot
pushed a commit
that referenced
this pull request
May 23, 2025
This is an attempt to re-land #28650, fixing build failures. **Motivation** #28650 was motivated by recent changes in setuptools, which caused test failures with editable installs, but it also exposes a potentially larger issue with our approach to testing with pytest. As far as I understand it, the default pytest behavior is to prepend the working directory to `sys.path`, meaning that `jax` is imported from the source directory, rather than the installed version. Switching to the `importlib` import mode means that we correctly test against the installed version of `jax`, which seems like what we typically want to do. The catch is that then we need to explicitly package any test utilities into the distribution. We don't currently package test-specific utilities like `internal_test_util` with JAX, but these utilities were still available to tests since they live within the `jax` source tree. This breaks when using `importlib` import mode, and a non-editable install of JAX. **Solutions** The approach that I've taken here is to explicitly package everything needed by the tests into the `jax` distribution. This means that we can correctly test against the _installed_ version of JAX when using pytest. This solution isn't ideal because it means that we're distributing `jax` submodules that aren't actually required except when running the test suite, but this seems like a small price to pay to me. **Alternatives** One different approach that we could take would be to only support using pytest with _editable_ installs of JAX. This would work because the required files would still be discoverable in an editable install because they live within the source tree. In fact, before this change, most of our CI jobs actually did install an editable distribution (which is why the failures in #28650 weren't caught in pre-submit!). The problem with this approach is that we're not actually testing JAX as it is used when installed from a distribution, and it wouldn't catch things like missing submodules. I think it's much better to test against the installed distribution! A more extreme approach would be to switch JAX to a `src/jax` and `src/jaxlib` layout (i.e. moving `jax` and `jaxlib` out of the root directory) as recommended by the Python packaging docs. Unfortunately this would be complicated with the way JAX is distributed internally at Google, so I think that's probably a non-starter. PiperOrigin-RevId: 762419160
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.
Otherwise the default (prepend) mode will add
<jax_src_dir>/testsand<jax_src_dir>to the start ofsys.path.This is (has always been) fragile, because
<jax_src_dir>hasjax/andjaxlib/subdirectories, but it recently broke assetuptools80 has a change in behaviour for editable installations, with the result that ifjaxlibis installed editable thenimport jaxlibwith<jax_src_dir>insys.pathwill try to import from<jax_src_dir>/jaxlibinstead of the editable install location, and fail.See also:
src/directory alongsidetests/at the top level)