Merge branch 'main' into shuffle-shutdown

hendrikmakait · hendrikmakait · commit 74599ee44be8 · 2023-07-13T08:55:45.000+02:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -58,7 +58,7 @@ repos:
           - types-psutil
           - types-setuptools
           # Typed libraries
-          - click
+          - click!=8.1.4  # https://github.com/pallets/click/issues/2558
           - numpy
           - pytest
           - tornado
diff --git a/distributed/diagnostics/progress.py b/distributed/diagnostics/progress.py
@@ -2,6 +2,7 @@
 
 import asyncio
 import logging
+import warnings
 from collections import defaultdict
 from timeit import default_timer
 from typing import ClassVar
@@ -140,9 +141,19 @@ def stop(self, exception=None, key=None):
 class MultiProgress(Progress):
     """Progress variant that keeps track of different groups of keys
 
-    See Progress for most details.  This only adds a function ``func=``
-    that splits keys.  This defaults to ``key_split`` which aligns with naming
-    conventions chosen in the dask project (tuples, hyphens, etc..)
+    See Progress for most details.
+
+    Parameters
+    ----------
+
+    func : Callable (deprecated)
+        Function that splits keys. This defaults to ``key_split`` which
+        aligns with naming conventions chosen in the dask project (tuples,
+        hyphens, etc..)
+
+    group_by : Callable | Literal["spans"] | Literal["prefix"], default: "prefix"
+        How to group keys to display multiple bars. Defaults to "prefix",
+        which uses ``key_split`` from dask project
 
     State
     -----
@@ -161,10 +172,24 @@ class MultiProgress(Progress):
     """
 
     def __init__(
-        self, keys, scheduler=None, func=key_split, minimum=0, dt=0.1, complete=False
+        self,
+        keys,
+        scheduler=None,
+        *,
+        func=None,
+        group_by="prefix",
+        minimum=0,
+        dt=0.1,
+        complete=False,
     ):
-        self.func = func
-        name = f"multi-progress-{tokenize(keys, func, minimum, dt, complete)}"
+        if func is not None:
+            warnings.warn(
+                "`func` is deprecated, use `group_by`", category=DeprecationWarning
+            )
+            group_by = func
+        self.group_by = key_split if group_by in (None, "prefix") else group_by
+        self.func = None
+        name = f"multi-progress-{tokenize(keys, group_by, minimum, dt, complete)}"
         super().__init__(
             keys, scheduler, minimum=minimum, dt=dt, complete=complete, name=name
         )
@@ -191,6 +216,22 @@ async def setup(self):
         if not self.keys:
             self.stop(exception=None, key=None)
 
+        if self.group_by == "spans":
+            spans_ext = self.scheduler.extensions["spans"]
+            span_defs = spans_ext.spans if spans_ext else None
+
+            def group_key(k):
+                span_id = self.scheduler.tasks[k].group.span_id
+                span_name = ", ".join(span_defs[span_id].name) if span_defs else span_id
+                return span_name, span_id
+
+            group_keys = {k: group_key(k) for k in self.all_keys}
+            self.func = group_keys.get
+        elif self.group_by == "prefix":
+            self.func = key_split
+        else:
+            self.func = self.group_by
+
         # Group keys by func name
         self.keys = valmap(set, groupby(self.func, self.keys))
         self.all_keys = valmap(set, groupby(self.func, self.all_keys))
diff --git a/distributed/diagnostics/progressbar.py b/distributed/diagnostics/progressbar.py
@@ -7,6 +7,7 @@
 import weakref
 from contextlib import suppress
 from timeit import default_timer
+from typing import Callable
 
 from tlz import valmap
 from tornado.ioloop import IOLoop
@@ -243,7 +244,9 @@ def __init__(
         self,
         keys,
         scheduler=None,
-        func=key_split,
+        *,
+        func=None,
+        group_by="prefix",
         interval="100ms",
         complete=False,
         **kwargs,
@@ -256,8 +259,17 @@ def __init__(
                 self.client = weakref.ref(key.client)
                 break
 
+        if func is not None:
+            warnings.warn(
+                "`func` is deprecated, use `group_by` instead",
+                category=DeprecationWarning,
+            )
+            group_by = func
+        elif group_by in (None, "prefix"):
+            group_by = key_split
+
         self.keys = {k.key if hasattr(k, "key") else k for k in keys}
-        self.func = func
+        self.group_by = group_by
         self.interval = interval
         self.complete = complete
         self._start_time = default_timer()
@@ -269,10 +281,15 @@ def elapsed(self):
     async def listen(self):
         complete = self.complete
         keys = self.keys
-        func = self.func
+        group_by = self.group_by
 
         async def setup(scheduler):
-            p = MultiProgress(keys, scheduler, complete=complete, func=func)
+            p = MultiProgress(
+                keys,
+                scheduler,
+                complete=complete,
+                group_by=group_by,
+            )
             await p.setup()
             return p
 
@@ -339,29 +356,31 @@ def __init__(
         keys,
         scheduler=None,
         minimum=0,
-        interval=0.1,
-        func=key_split,
-        complete=False,
         **kwargs,
     ):
-        super().__init__(keys, scheduler, func, interval, complete)
+        super().__init__(keys, scheduler, **kwargs)
         from ipywidgets import VBox
 
         self.widget = VBox([])
 
     def make_widget(self, all):
         from ipywidgets import HTML, FloatProgress, HBox, VBox
 
+        def make_label(key):
+            if isinstance(key, tuple):
+                # tuple of (group_name, group_id)
+                key = key[0]
+            key = key.decode() if isinstance(key, bytes) else key
+            return html.escape(key)
+
         self.elapsed_time = HTML("")
         self.bars = {key: FloatProgress(min=0, max=1, description="") for key in all}
         self.bar_texts = {key: HTML("") for key in all}
         self.bar_labels = {
             key: HTML(
                 '<div style="padding: 0px 10px 0px 10px;'
                 " text-align:left; word-wrap: "
-                'break-word;">'
-                + html.escape(key.decode() if isinstance(key, bytes) else key)
-                + "</div>"
+                'break-word;">' + make_label(key) + "</div>"
             )
             for key in all
         }
@@ -429,7 +448,9 @@ def _draw_bar(self, remaining, all, status, **kwargs):
             )
 
 
-def progress(*futures, notebook=None, multi=True, complete=True, **kwargs):
+def progress(
+    *futures, notebook=None, multi=True, complete=True, group_by="prefix", **kwargs
+):
     """Track progress of futures
 
     This operates differently in the notebook and the console
@@ -448,6 +469,9 @@ def progress(*futures, notebook=None, multi=True, complete=True, **kwargs):
     complete : bool (optional)
         Track all keys (True) or only keys that have not yet run (False)
         (defaults to True)
+    group_by : Callable | Literal["spans"] | Literal["prefix"]
+        Use spans instead of task key names for grouping tasks
+        (defaults to "prefix")
 
     Notes
     -----
@@ -465,9 +489,18 @@ def progress(*futures, notebook=None, multi=True, complete=True, **kwargs):
         futures = [futures]
     if notebook is None:
         notebook = is_kernel()  # often but not always correct assumption
+    if kwargs.get("func", None) is not None:
+        warnings.warn(
+            "`func` is deprecated, use `group_by` instead", category=DeprecationWarning
+        )
+        group_by = kwargs.pop("func")
+    if group_by not in ("spans", "prefix") and not isinstance(group_by, Callable):
+        raise ValueError("`group_by` should be 'spans', 'prefix', or a Callable")
     if notebook:
         if multi:
-            bar = MultiProgressWidget(futures, complete=complete, **kwargs)
+            bar = MultiProgressWidget(
+                futures, complete=complete, group_by=group_by, **kwargs
+            )
         else:
             bar = ProgressWidget(futures, complete=complete, **kwargs)
         return bar
diff --git a/distributed/diagnostics/tests/test_progress.py b/distributed/diagnostics/tests/test_progress.py
@@ -4,6 +4,7 @@
 
 import pytest
 
+import distributed
 from distributed import Nanny
 from distributed.client import wait
 from distributed.compatibility import LINUX
@@ -14,7 +15,16 @@
     Progress,
     SchedulerPlugin,
 )
-from distributed.utils_test import dec, div, gen_cluster, inc, nodebug, slowdec, slowinc
+from distributed.utils_test import (
+    dec,
+    div,
+    gen_cluster,
+    inc,
+    nodebug,
+    slowdec,
+    slowinc,
+    wait_for_state,
+)
 
 
 def f(*args):
@@ -72,6 +82,49 @@ async def test_multiprogress(c, s, a, b):
     assert p.status == "finished"
 
 
+@gen_cluster(client=True)
+async def test_multiprogress_cancel(c, s, a, b):
+    lock = distributed.Lock()
+    await lock.acquire()
+
+    async def wait_and_raise(*args, **kwargs):
+        async with lock:
+            raise RuntimeError()
+
+    f = c.submit(wait_and_raise, key="cancel", workers=[a.address])
+    p = MultiProgress([f], scheduler=s, complete=True)
+    await p.setup()
+    await wait_for_state(f.key, "executing", a)
+    f.release()
+    await wait_for_state(f.key, "cancelled", a)
+    assert p.status == "error"
+    assert p.all_keys.keys() == {"cancel"}
+
+
+@gen_cluster(client=True)
+async def test_multiprogress_with_spans(c, s, a, b):
+    x = c.submit(inc, 1)
+    p = MultiProgress([x], scheduler=s, complete=True, group_by="spans")
+    await p.setup()
+    group_names = {k[0] for k in p.all_keys}
+    assert group_names == {"default"}
+
+
+@gen_cluster(client=True)
+async def test_multiprogress_with_prefix(c, s, a, b):
+    x = c.submit(inc, 1)
+    p = MultiProgress([x], scheduler=s, complete=True, group_by="prefix")
+    await p.setup()
+    group_names = {k for k in p.all_keys}
+    assert group_names == {"inc"}
+
+
+def test_multiprogress_warns():
+    with pytest.warns(DeprecationWarning, match="func` is deprecated, use `group_by"):
+        p = MultiProgress([], complete=True, func="spans")
+        assert p.group_by == "spans"
+
+
 @gen_cluster(client=True)
 async def test_robust_to_bad_plugin(c, s, a, b):
     class Bad(SchedulerPlugin):
diff --git a/distributed/diagnostics/tests/test_progressbar.py b/distributed/diagnostics/tests/test_progressbar.py
@@ -77,6 +77,16 @@ def test_progress_function_w_kwargs(client, capsys):
     check_bar_completed(capsys)
 
 
+def test_progress_function_warns(client):
+    with pytest.warns(DeprecationWarning, match="`func` is deprecated"):
+        progress(None, func="prefix")
+
+
+def test_progress_function_raises():
+    with pytest.raises(ValueError, match="`group_by` should be "):
+        progress(None, group_by="incorrect")
+
+
 @gen_cluster(client=True, nthreads=[])
 async def test_deprecated_loop_properties(c, s):
     class ExampleTextProgressBar(TextProgressBar):
diff --git a/distributed/diagnostics/tests/test_widgets.py b/distributed/diagnostics/tests/test_widgets.py
@@ -8,7 +8,10 @@
 import pytest
 from packaging.version import parse as parse_version
 
+from dask.utils import key_split
+
 from distributed.client import wait
+from distributed.spans import span
 from distributed.utils_test import dec, gen_cluster, gen_tls_cluster, inc, throws
 
 ipywidgets = pytest.importorskip("ipywidgets")
@@ -207,6 +210,7 @@ async def test_multibar_complete(c, s, a, b):
     y2 = c.submit(dec, y1, key="y-2")
     e = c.submit(throws, y2, key="e")
     other = c.submit(inc, 123, key="other")
+    await other.cancel()
 
     p = MultiProgressWidget([e.key], scheduler=s.address, complete=True)
     await p.listen()
@@ -227,6 +231,51 @@ def test_fast(client):
     assert set(p._last_response["all"]) == {"inc", "dec", "add"}
 
 
+@mock_widget()
+def test_multibar_with_spans(client):
+    """Test progress(group_by='spans'"""
+    with span("span 1"):
+        L = client.map(inc, range(100))
+    with span("span 2"):
+        L2 = client.map(dec, L)
+    with span("span 3"):
+        L3 = client.map(add, L, L2)
+    with span("other span"):
+        _ = client.submit(inc, 123)
+    e = client.submit(throws, L3)
+
+    p = progress(e, complete=True, multi=True, notebook=True, group_by="spans")
+    client.sync(p.listen)
+
+    # keys are tuples of (group_name, group_id), just get names
+    bar_items = {k[0]: v.value for k, v in p.bars.items()}
+    bar_texts = {k[0]: v.value for k, v in p.bar_texts.items()}
+    bar_labels = {k[0]: v.value for k, v in p.bar_labels.items()}
+
+    assert bar_items == {"span 1": 1, "span 2": 1, "span 3": 1, "default": 0}
+    assert bar_texts.keys() == {"span 1", "span 2", "span 3", "default"}
+    assert all("100 / 100" in v for k, v in bar_texts.items() if k != "default")
+    assert bar_labels.keys() == {"span 1", "span 2", "span 3", "default"}
+    assert all(f">{k}<" in v for k, v in bar_labels.items())
+
+
+@mock_widget()
+def test_multibar_func_warns(client):
+    """Deprecate `func`, use `group_by`"""
+    L = client.map(inc, range(100))
+    L2 = client.map(dec, L)
+    L3 = client.map(add, L, L2)
+
+    # ensure default value if nothing is set
+    p = MultiProgressWidget(L3)
+    assert p.group_by == key_split
+
+    with pytest.warns(
+        DeprecationWarning, match="`func` is deprecated, use `group_by` instead"
+    ):
+        MultiProgressWidget(L3, func="foo")
+
+
 @mock_widget()
 @gen_cluster(client=True, client_kwargs={"serializers": ["msgpack"]})
 async def test_serializers(c, s, a, b):
diff --git a/distributed/distributed.yaml b/distributed/distributed.yaml
diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py
