initial integration

Author: gjoseph92

distributed/core.py

@@ -14,7 +14,7 @@
 from collections.abc import Callable, Container, Coroutine, Generator
 from enum import Enum
 from functools import partial
-from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, TypeVar, final
+from typing import TYPE_CHECKING, Any, ClassVar, Hashable, TypedDict, TypeVar, final
 
 import tblib
 from tlz import merge
@@ -943,7 +943,7 @@ async def close(self, timeout=None):
         finally:
             self._event_finished.set()
 
-    def digest_metric(self, name: str, value: float) -> None:
+    def digest_metric(self, name: Hashable, value: float) -> None:
         # Granular data (requires crick)
         if self.digests is not None:
             self.digests[name].add(value)

distributed/protocol/compression.py

@@ -17,6 +17,7 @@
 
 import dask
 
+from distributed.span import meter
 from distributed.utils import ensure_memoryview, nbytes, no_default
 
 compressions: dict[
@@ -150,6 +151,7 @@ def byte_sample(b, size, n):
         return memoryview(b"".join(parts))
 
 
+@meter("compress")
 def maybe_compress(
     payload,
     min_size=10_000,
@@ -195,6 +197,7 @@ def maybe_compress(
     return None, payload
 
 
+@meter("decompress")
 def decompress(header, frames):
     """Decompress frames according to information in the header"""
     return [

distributed/protocol/serialize.py

@@ -24,6 +24,7 @@
     pack_frames_prelude,
     unpack_frames,
 )
+from distributed.span import meter
 from distributed.utils import ensure_memoryview, has_keyword
 
 dask_serialize = dask.utils.Dispatch("dask_serialize")
@@ -426,6 +427,7 @@ def deserialize(header, frames, deserializers=None):
     return loads(header, frames)
 
 
+@meter("serialize")
 def serialize_and_split(
     x, serializers=None, on_error="message", context=None, size=None
 ):
@@ -470,6 +472,7 @@ def serialize_and_split(
     return header, out_frames
 
 
+@meter("deserialize")
 def merge_and_deserialize(header, frames, deserializers=None):
     """Merge and deserialize frames
 

distributed/span.py

@@ -47,6 +47,14 @@ def own_time(self) -> float:
         assert self.stop_time is not None
         return self.total_time - self.other_time
 
+    @property
+    def done(self) -> bool:
+        return self.start_time is not None and self.stop_time is not None
+
+    @property
+    def running(self) -> bool:
+        return self.start_time is not None and self.stop_time is None
+
     def start(self) -> None:
         assert self.start_time is None
         assert self.stop_time is None
@@ -83,14 +91,19 @@ def _unset_current(self) -> None:
         _current_span.reset(self._token)
         self._token = None
 
-    def _subspan(self, label: str | tuple[str, ...]) -> Span:
+    def _subspan(
+        self,
+        label: str | tuple[str, ...],
+        metric: Callable[[], float] = time.perf_counter,
+    ) -> Span:
         assert (
             self.start_time is not None
         ), "Cannot create sub-span for a span that has not started"
         assert (
             self.stop_time is None
         ), "Cannot create sub-span for a span that has already stopped"
-        span = Span(label, self.metric)
+        # TODO allow different metrics, or always use `self.metric`?
+        span = Span(label, metric)
         self.subspans.append(span)
         return span
 
@@ -130,7 +143,7 @@ def __repr__(self) -> str:
         return (
             f"{type(self).__name__}<"
             f"{self.label!r}, "
-            f"total_time={self.total_time}, "
+            f"total_time={self.total_time if self.done else '...'}, "
             f"start_time={self.start_time}, "
             f"stop_time={self.stop_time}, "
             ">"
@@ -155,7 +168,7 @@ def get_span(
     except LookupError:
         span = Span(label, metric)
     else:
-        assert metric is parent.metric, (metric, parent.metric)
+        # assert metric is parent.metric, (metric, parent.metric)
         span = parent._subspan(label)
 
     return span

distributed/spill.py

@@ -14,6 +14,7 @@
 
 from distributed.protocol import deserialize_bytes, serialize_bytelist
 from distributed.sizeof import safe_sizeof
+from distributed.span import meter
 from distributed.utils import RateLimiterFilter
 
 logger = logging.getLogger(__name__)
@@ -239,8 +240,9 @@ def __getitem__(self, key: str) -> Any:
         if key in self.fast:
             # Note: don't log from self.fast.__getitem__, because that's called every
             # time a key is evicted, and we don't want to count those events here.
-            nbytes = cast(int, self.fast.weights[key])
-            self.fast_metrics.log_read(nbytes)
+            with meter("memory-read"):
+                nbytes = cast(int, self.fast.weights[key])
+                self.fast_metrics.log_read(nbytes)
 
         return super().__getitem__(key)
 
@@ -362,7 +364,8 @@ def __init__(self, spill_directory: str, max_weight: int | Literal[False] = Fals
 
     def __getitem__(self, key: str) -> Any:
         t0 = perf_counter()
-        pickled = self.d[key]
+        with meter("disk-read"):
+            pickled = self.d[key]
         assert isinstance(pickled, bytearray if has_zict_230 else bytes)
         t1 = perf_counter()
         out = self.load(pickled)
@@ -408,7 +411,8 @@ def __setitem__(self, key: str, value: Any) -> None:
 
         # Store to disk through File.
         # This may raise OSError, which is caught by SpillBuffer above.
-        self.d[key] = pickled
+        with meter("disk-write"):
+            self.d[key] = pickled
 
         t2 = perf_counter()
 

distributed/tests/test_worker_metrics.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from collections.abc import Hashable
+
+from distributed import Event, Worker
+from distributed.utils_test import async_wait_for, gen_cluster, inc, wait_for_state
+
+
+def get_digests(w: Worker, allow: str | None = None) -> dict[Hashable, float]:
+    # import pprint; pprint.pprint(dict(w.digests_total))
+    digests = {
+        k: v
+        for k, v in w.digests_total.items()
+        if k
+        not in {
+            "latency",
+            "tick-duration",
+            "transfer-bandwidth",
+            "transfer-duration",
+            "compute-duration",
+        }
+        and (allow is None or allow in k)
+    }
+    assert all(v >= 0 for v in digests.values()), digests
+    return digests
+
+
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_basic_execute(c, s, a):
+    await c.submit(inc, 1, key="x")
+    assert list(get_digests(a)) == [
+        ("execute", "x", "thread", "thread-cpu"),
+        ("execute", "x", "thread"),
+        ("execute", "x"),
+    ]
+
+
+# @gen_cluster(client=True, nthreads=[("", 1)])
+# async def test_run_spec_deserialization(c, s, a):
+#     """Test that deserialization of run_spec is metered"""
+#     await c.submit(inc, 1, key="x")
+#     assert 0 < a.digests_total["execute", "x", "deserialize", "seconds"] < 1
+
+
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_cancelled_execute(c, s, a):
+    """cancelled(execute) tasks are metered as a separate lump total"""
+    ev = await Event()
+    x = c.submit(lambda ev: ev.wait(), ev, key="x")
+    await wait_for_state("x", "executing", a)
+    del x
+    await wait_for_state("x", "cancelled", a)
+    await ev.set()
+    await async_wait_for(lambda: not a.state.tasks, timeout=5)
+
+    print(list(get_digests(a)))
+    assert list(get_digests(a)) == [("execute", "x", "cancelled")]

distributed/utils.py

@@ -21,7 +21,11 @@
 from asyncio import TimeoutError
 from collections import deque
 from collections.abc import Callable, Collection, Container, KeysView, ValuesView
-from concurrent.futures import CancelledError, ThreadPoolExecutor  # noqa: F401
+from concurrent.futures import (  # noqa: F401
+    CancelledError,
+    Executor,
+    ThreadPoolExecutor,
+)
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
 from datetime import timedelta
@@ -1410,12 +1414,28 @@ def import_term(name: str) -> AnyType:
     return getattr(module, attr_name)
 
 
-async def offload(fn, *args, **kwargs):
+async def offload(  # type: ignore[valid-type]
+    fn: Callable[P, T],  # FIXME improper use of `ParamSpec`?
+    *args: P.args,
+    executor: Executor | None = None,
+    **kwargs: P.kwargs,
+) -> T:
+    """Wrapper around :meth:`~asyncio.AbstractEventLoop.run_in_executor`, which
+    propagates contextvars.
+    By default, it offloads to a thread pool with a single worker.
+    See also
+    --------
+    https://bugs.python.org/issue34014
+    """
+    if executor is None:
+        # Not the same as defaulting to _offload_executor in the parameters, as this
+        # allows monkey-patching the _offload_executor during unit tests
+        executor = _offload_executor
+
     loop = asyncio.get_running_loop()
-    # Retain context vars while deserializing; see https://bugs.python.org/issue34014
     context = contextvars.copy_context()
     return await loop.run_in_executor(
-        _offload_executor, lambda: context.run(fn, *args, **kwargs)
+        executor, lambda: context.run(fn, *args, **kwargs)
     )
 
 

distributed/utils_test.py

@@ -63,6 +63,7 @@
 from distributed.protocol import deserialize
 from distributed.scheduler import TaskState as SchedulerTaskState
 from distributed.security import Security
+from distributed.span import Span
 from distributed.utils import (
     DequeHandler,
     _offload_executor,
@@ -82,7 +83,7 @@
     StateMachineEvent,
 )
 from distributed.worker_state_machine import TaskState as WorkerTaskState
-from distributed.worker_state_machine import WorkerState
+from distributed.worker_state_machine import TracedEvent, WorkerState
 
 try:
     import dask.array  # register config
@@ -2267,11 +2268,11 @@ def __init__(self, *args, **kwargs):
 
         super().__init__(*args, **kwargs)
 
-    async def execute(self, key: str, *, stimulus_id: str) -> StateMachineEvent:
+    async def execute(self, key: str, *, stimulus_id: str, span: Span) -> TracedEvent:
         self.in_execute.set()
         await self.block_execute.wait()
         try:
-            return await super().execute(key, stimulus_id=stimulus_id)
+            return await super().execute(key, stimulus_id=stimulus_id, span=span)
         finally:
             self.in_execute_exit.set()
             await self.block_execute_exit.wait()