Switch to a new thread-safe utility for catching warnings. #25626
Conversation
The Python warnings.catch_warnings() functionality is not thread-safe (https://py-free-threading.github.io/porting/#the-warnings-module-is-not-thread-safe), so we cannot use it during tests that use free-threading. This change introduces a private warnings test helper (test_warning_util.py), which hooks the CPython warning infrastructure and uses it to implement thread-safe warnings infrastructure. This requires a handful of small modifications to tests to remove direct uses of the warnings module. We also sadly have to delete one TPU test that checks for a warning raised on another thread; there's no easy way for us to catch that in a thread-safe way, but that test seems like overkill anyway.
|
There are some test failures because there are still some uses of |
| yield | ||
|
|
||
| @contextmanager | ||
| def assertWarnsRegex(self, expected_warning, expected_regex): |
There was a problem hiding this comment.
The existing assertWarnsRegex can be used either as a context manager or as a simple function call, passing the callable and args/kwargs directly. Should we support that in the new version?
| return | ||
| self.fail(f"Expected warning not found {expected_warning}:'{expected_regex}', got " | ||
| f"{ws}") | ||
|
|
There was a problem hiding this comment.
We should also probably define assertWarns(), which is used in a few places.
| addopts = "--doctest-glob='*.rst' --ignore='examples/ffi'" | ||
| # We disable pytest's warning capture system, because neither it not the | ||
| # warnings module is thread-safe. Instead, we use our own utilities to catch | ||
| # and test for warnings in tests, see test_warning_util.py. |
There was a problem hiding this comment.
If I understand correctly, this means that tests which do not import jax._src.test_util will no longer treat warnings as errors. I believe this includes doctests, as well as the array API test. Is there any way we can restore the previous behavior just for those cases?
|
|
||
| def test_warning_raises(self): | ||
| with self.assertRaises(UserWarning, msg="hello"): | ||
| warnings.warn("hello", category=UserWarning) |
There was a problem hiding this comment.
Could we also add a similar test DeprecationWarning? It has some different handing than UserWarning within Python's warning system, and it's important that we still error when we see it.
The Python warnings.catch_warnings() functionality is not thread-safe (https://py-free-threading.github.io/porting/#the-warnings-module-is-not-thread-safe), so we cannot use it during tests that use free-threading. This change introduces a private warnings test helper (test_warning_util.py), which hooks the CPython warning infrastructure and uses it to implement thread-safe warnings infrastructure.
This requires a handful of small modifications to tests to remove direct uses of the warnings module. We also sadly have to delete one TPU test that checks for a warning raised on another thread; there's no easy way for us to catch that in a thread-safe way, but that test seems like overkill anyway.