Visitar URL original
Switch to a new thread-safe utility for catching warnings. by hawkinsp · Pull Request #25626 · jax-ml/jax · GitHub
Skip to content

Switch to a new thread-safe utility for catching warnings. - #25626

Merged
copybara-service[bot] merged 1 commit into
jax-ml:mainfrom
hawkinsp:warnings2
Jan 9, 2025
Merged

copybara-service[bot] merged 1 commit into
jax-ml:mainfrom
hawkinsp:warnings2

Conversation

@hawkinsp

Copy link
Copy Markdown
Collaborator

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.

@hawkinsp
hawkinsp requested a review from jakevdp December 20, 2024 02:32
@hawkinsp

Copy link
Copy Markdown
Collaborator Author

There are some test failures because there are still some uses of warnings.catch_warnings in non-test parts of JAX. I'll need to clean those up (they will certainly break in free-threading mode).

Comment thread jax/_src/test_util.py Outdated
Comment thread jax/_src/test_util.py
Comment thread pyproject.toml Outdated
Comment thread tests/warnings_util_test.py
@nascheme

Copy link
Copy Markdown

FYI: python/cpython#128300 : Add thread-safe context manager to "warnings" module

This is going to take some time to get fixed in CPython but it might be a good idea to make a similar API so that you can switch back to the CPython version in the future, assuming it gets fixed properly. If you have any feedback on my PR, I would be happy to hear about it. My main concerns were to not break existing non-threaded non-async code that is happily using catch_watchings() and add an API that can be used in a backwards compatible way (for code that wants to support both old and newer versions of Python).

@hawkinsp
hawkinsp force-pushed the warnings2 branch 2 times, most recently from 8d0bad1 to 2351fbb Compare January 9, 2025 16:37
@hawkinsp

hawkinsp commented Jan 9, 2025

Copy link
Copy Markdown
Collaborator Author

FYI: python/cpython#128300 : Add thread-safe context manager to "warnings" module

This is going to take some time to get fixed in CPython but it might be a good idea to make a similar API so that you can switch back to the CPython version in the future, assuming it gets fixed properly. If you have any feedback on my PR, I would be happy to hear about it. My main concerns were to not break existing non-threaded non-async code that is happily using catch_watchings() and add an API that can be used in a backwards compatible way (for code that wants to support both old and newer versions of Python).

Thanks! I'm excited to see that this is being fixed properly upstream.

Given I just added a fallback path for non-threaded cases that calls the upstream utilities, it shouldn't be a big lift to switch back to the upstream utilities for the threaded cases also when a thread-safe alternative exists.

I'd assume that your improved warnings will not land until Python 3.14, correct?

Comment thread jax/_src/test_warning_util.py Outdated
@google-ml-butler google-ml-butler Bot added kokoro:force-run pull ready Ready for copybara import and testing labels Jan 9, 2025
@hawkinsp
hawkinsp force-pushed the warnings2 branch 2 times, most recently from 253a908 to be66011 Compare January 9, 2025 16:43
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.
@copybara-service
copybara-service Bot merged commit 7718ac5 into jax-ml:main Jan 9, 2025
dinesh-ad added a commit to Phazeo/truestill that referenced this pull request Aug 21, 2026
…ing a library's words

One `organize` over the format corpus put 866 lines on the user's stderr. It now puts 2 - the
progress line - and the report states what was removed and why.

⚠ THE ENTRY HAD THE SUBJECT BACKWARDS, AND THAT IS THE FINDING. Measured on the same corpus:
478 image files got no near-duplicate check and only 71 of them warned, while 14 warned and
decoded perfectly well. The 131 warnings are a lossy 15% proxy for a gap the product never
mentioned at all. Reporting them would have named 71 of 478 and implied the rest were fine;
suppressing them alone would have made Truestill quieter about a real gap - §4's forty-second
member, a check measuring the cheaper proxy. So what a person is told is derived from the decode
OUTCOME, never from whether a library spoke.

The fact was already in the data and nobody asked the question, which is (aer)'s shape again.
`FileHashes.perceptual_computed` exists because "perceptual=None answers two different questions -
not an image and nobody looked". That settled two of three; the third - an image a pass tried to
decode and could not - was derivable since the field shipped and read by nothing.

⚠ `warnings.catch_warnings` IS DELETED, NOT WIDENED. It assigns process-global module attributes
and scan.py hashes on a ThreadPoolExecutor by default, which CPython documents as undefined
behaviour with two or more threads - so the old one-class suppression was unsound on the common
path and widening it would have widened the race. The upstream fix is out of reach: gh-128384
filed it, the ContextVar implementation gh-130010 is "Changed in version 3.14", and it "defaults
to 1 on free-threaded builds and to 0 otherwise". This project runs 3.13. JAX hit the same wall
(jax-ml/jax#25626) and hooks the warning infrastructure; decode_noise does the same - one write
per process, at a phase boundary, never while workers are in flight.

⚠ SCOPED TO THE EXTERNAL LIBRARY, BOTH DEFENCES PINNED SEPARATELY. We do not silence our own
signals. The filter matches module=PIL\..* and the hook swallows only when the raising file sits
under a PIL directory. Proved by a warning raised from inside truestill_core.hashing during a real
hashing pass, beside a file that makes Pillow warn on the same pass: ours reaches the terminal
with our own path, Pillow's does not.

fd 2 IS PROCESS-WIDE AND THE FIX SAYS SO RATHER THAN PRETENDING. libtiff and libjpeg write there
directly; under a worker pool several decodes are in flight and the lines name no file, so
per-file attribution is impossible and was not attempted. They are counted at phase level and
discarded - safe for one measured reason: grep for the source path over a whole 866-line run
returns 0. Progress survives because _progress_printer writes through the sys.stderr object while
C writes to the descriptor; the object is repointed only when it is actually backed by fd 2.

Installed in two places because there are two: the parent call carries the default thread pool,
`initializer=` carries ProcessPoolExecutor children. The child test is forced onto spawn, because
Linux's fork makes the child inherit everything and the initializer changes nothing - the first
version passed with it deleted. macOS and Windows are spawn, and both are CI lanes.

The action is "always" and that is load-bearing: warn_explicit consults __warningregistry__ before
dispatching, so the default drops repeats - 197 raised, 133 printed. The run now reports 189.

(aew): the unreadable block printed "fix the permission or check the disk" under all five reasons,
and 8 of 8 files named under it on this corpus were UNDECODABLE, where neither is at fault - under
a heading that contradicted its own rows. models.unreadable_remedy beside unreadable_label, the
list grouped by reason, and the heading now names the consequence.

Thirteen mutations, three survived first: a >= against a cumulative counter that leftovers
satisfied; nothing guarding the fd-2 restore, whose test must be a subprocess because pytest has
already replaced fd 2; and the filter's scope, defended only by the hook's filename check.

Verified on both pools: thread and process report an identical 787 = 189 + 598.

Closes (aev).
Closes (aew).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pull ready Ready for copybara import and testing

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants