Fix annotations and spans leaking between threads #10367

j-bennet · 2023-06-20T21:12:06Z

Switch global annotations to use ContextVar instead of setting values in dask config.

Dependent PR in distributed: dask/distributed#7935

Closes Annotations and spans leak between threads and asyncio tasks #10340
Tests added / passed
Passes pre-commit run --all-files

j-bennet · 2023-06-20T22:12:03Z

Initially, I was going to create a module for this code called dask.annotations. Unfortunately, we already have dask.annotations here:

dask/dask/__init__.py

Line 1 in a6ee367

from __future__ import annotations

The name would be clashing, so for now, I left it in base.py. If you have better suggestions, let me know.

fjetter

Can you please add a test for this as well somewhere? #10340 provides an example that triggers this leakage w/out spans and I'd like to test this behavior.

dask/base.py

crusaderky

Missing unit tests:

the code inside with dask.annotate() raises
annotations in two different asyncio tasks or threads[1] running in parallel do not interfere with each other

[1] only need one. Pick whatever is easier/cleaner to write.

dask/base.py

dask/dataframe/io/parquet/core.py

dask/highlevelgraph.py

dask/base.py

dask/tests/test_highgraph.py

crusaderky · 2023-06-21T16:02:49Z

Please add line to docs/source/api.rst

crusaderky · 2023-06-21T21:31:18Z

dask/tests/test_highgraph.py

+    with dask.annotate(banana=5):
+        with dask.annotate(apple=3):
+            with pytest.raises(ZeroDivisionError):
+                _ = 1 / 0


The exception is being swallowed by pytest.raises and never reaches dask.annotate.

crusaderky · 2023-06-21T21:31:55Z

dask/tests/test_highgraph.py

@@ -186,10 +186,19 @@ def test_annotation_and_config_collision():
    with dask.config.set({"foo": 1}):
        with dask.annotate(foo=2):
            assert dask.config.get("foo") == 1
-            assert dask.config.get("annotations") == {"foo": 2}
+            assert dask.get_annotations() == {"foo": 2}


This test no longer makes sense

crusaderky

Missing unit test:
annotations in two different asyncio tasks or threads[1] running in parallel do not interfere with each other.
[1] only need one. Pick whatever is easier/cleaner to write.

docs/source/api.rst

j-bennet · 2023-06-21T22:08:00Z

Missing unit test: annotations in two different asyncio tasks or threads[1] running in parallel do not interfere with each other. [1] only need one. Pick whatever is easier/cleaner to write.

I added a test called test_annotations_threaded, let me know if that looks right.

j-bennet · 2023-06-21T23:00:20Z

test_datetime_std_with_larger_dataset is failing, but it's unrelated.

crusaderky

This is ready to be merged as soon as you revert the temporary CI changes in dask/distributed#7935

Global annotations use ContextVar.

45d2188

j-bennet marked this pull request as draft June 20, 2023 21:12

github-actions bot added dataframe io labels Jun 20, 2023

j-bennet changed the title ~~Global annotations use ContextVar.~~ Fix annotations and spans leaking between threads Jun 20, 2023

j-bennet force-pushed the j-bennet/10340-annotations-leak branch from e8ac766 to dfe06a4 Compare June 20, 2023 21:51

j-bennet added a commit to j-bennet/distributed that referenced this pull request Jun 20, 2023

Related to change in dask dask/dask#10367.

1a48957

j-bennet mentioned this pull request Jun 20, 2023

Fix annotations and spans leaking between threads dask/distributed#7935

Merged

2 tasks

j-bennet marked this pull request as ready for review June 20, 2023 22:03

j-bennet requested review from crusaderky and fjetter June 20, 2023 22:04

Convenience method to get annotation.

ff86f8f

j-bennet force-pushed the j-bennet/10340-annotations-leak branch from dfe06a4 to ff86f8f Compare June 20, 2023 22:10

fjetter reviewed Jun 21, 2023

View reviewed changes

dask/base.py Outdated Show resolved Hide resolved

crusaderky requested changes Jun 21, 2023

View reviewed changes

Add test for exception in yield.

e7d44a0

j-bennet requested a review from crusaderky June 21, 2023 18:50

Merge branch 'main' into j-bennet/10340-annotations-leak

f39b6d7

crusaderky reviewed Jun 21, 2023

View reviewed changes

github-actions bot added the documentation Improve or add to documentation label Jun 21, 2023

Code review

561f501

crusaderky force-pushed the j-bennet/10340-annotations-leak branch from 0f67385 to 561f501 Compare June 21, 2023 21:42

crusaderky requested changes Jun 21, 2023

View reviewed changes

Test with two threads on single worker. Add doc.

26d6f0a

j-bennet force-pushed the j-bennet/10340-annotations-leak branch from 296ff00 to 26d6f0a Compare June 21, 2023 22:02

crusaderky reviewed Jun 21, 2023

View reviewed changes

docs/source/api.rst Outdated Show resolved Hide resolved

Update docs/source/api.rst

49a8aff

j-bennet requested review from crusaderky and fjetter June 22, 2023 01:09

Review test

48ed0cc

crusaderky approved these changes Jun 22, 2023

View reviewed changes

crusaderky merged commit 0454f34 into dask:main Jun 22, 2023
23 of 24 checks passed

j-bennet deleted the j-bennet/10340-annotations-leak branch June 22, 2023 19:52

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix annotations and spans leaking between threads #10367

Fix annotations and spans leaking between threads #10367

j-bennet commented Jun 20, 2023 •

edited

j-bennet commented Jun 20, 2023

fjetter left a comment

crusaderky left a comment •

edited

crusaderky commented Jun 21, 2023

crusaderky Jun 21, 2023

crusaderky Jun 21, 2023

crusaderky left a comment

j-bennet commented Jun 21, 2023

j-bennet commented Jun 21, 2023

crusaderky left a comment

Fix annotations and spans leaking between threads #10367

Fix annotations and spans leaking between threads #10367

Conversation

j-bennet commented Jun 20, 2023 • edited

j-bennet commented Jun 20, 2023

fjetter left a comment

Choose a reason for hiding this comment

crusaderky left a comment • edited

Choose a reason for hiding this comment

crusaderky commented Jun 21, 2023

crusaderky Jun 21, 2023

Choose a reason for hiding this comment

crusaderky Jun 21, 2023

Choose a reason for hiding this comment

crusaderky left a comment

Choose a reason for hiding this comment

j-bennet commented Jun 21, 2023

j-bennet commented Jun 21, 2023

crusaderky left a comment

Choose a reason for hiding this comment

j-bennet commented Jun 20, 2023 •

edited

crusaderky left a comment •

edited