Repository: dask/distributed

Issue:
Fine performance metrics: measure time spent entering/exiting thread pools
- Part of #7665 
- Related to #7655 
- Related to #5882

When investigating #7655, I got the suspicion that we may have a bottleneck caused by the fact that there's only one worker in the offload executor.
Additionally, there's a known issue (#5882) where, whenever a task calls `rejoin()`, you will have one less thread available then what the WorkerStateMachine believes is available, which will cause `Worker.execute` to get stuck for a long time on the call to `run_in_executor`.

In both cases, as of #7586 this time is displayed as `("execute", <prefix>, "other", "seconds")`.
Wrap these two calls to `run_in_executor` in  `with context_meter.meter("offload")` and `with context_meter.meter("executor")` respectively.


Discussion:

Resolving patch:
diff --git a/distributed/utils.py b/distributed/utils.py
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -67,7 +67,7 @@
 from dask.widgets import get_template
 
 from distributed.compatibility import WINDOWS
-from distributed.metrics import monotonic, time
+from distributed.metrics import context_meter, monotonic, time
 
 try:
     from dask.context import thread_state
@@ -1460,7 +1460,8 @@ def offload(
     run_in_executor_with_context
     https://bugs.python.org/issue34014
     """
-    return run_in_executor_with_context(_offload_executor, func, *args, **kwargs)
+    with context_meter.meter("offload"):
+        return run_in_executor_with_context(_offload_executor, func, *args, **kwargs)
 
 
 class EmptyContext:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -2269,20 +2269,29 @@ async def execute(self, key: str, *, stimulus_id: str) -> StateMachineEvent:
                         self.scheduler_delay,
                     )
                 elif "ThreadPoolExecutor" in str(type(e)):
-                    result = await run_in_executor_with_context(
-                        e,
-                        apply_function,
-                        function,
-                        args2,
-                        kwargs2,
-                        self.execution_state,
-                        key,
-                        self.active_threads,
-                        self.active_threads_lock,
-                        self.scheduler_delay,
-                    )
+                    # The 'executor' time metric should be almost zero most of the time,
+                    # e.g. thread synchronization overhead only, since thread-noncpu and
+                    # thread-cpu inside the thread detract from it. However, it may
+                    # become substantial in case of misalignment between the size of the
+                    # thread pool and the number of running tasks in the worker state
+                    # machine (e.g. https://github.com/dask/distributed/issues/5882)
+                    with context_meter.meter("executor"):
+                        result = await run_in_executor_with_context(
+                            e,
+                            apply_function,
+                            function,
+                            args2,
+                            kwargs2,
+                            self.execution_state,
+                            key,
+                            self.active_threads,
+                            self.active_threads_lock,
+                            self.scheduler_delay,
+                        )
                 else:
-                    # Can't capture contextvars across processes
+                    # Can't capture contextvars across processes. If this is a
+                    # ProcessPoolExecutor, the 'executor' time metric will show the
+                    # whole runtime inside the executor.
                     with context_meter.meter("executor"):
                         result = await self.loop.run_in_executor(
                             e,


Test patch:
diff --git a/distributed/tests/test_worker_metrics.py b/distributed/tests/test_worker_metrics.py
--- a/distributed/tests/test_worker_metrics.py
+++ b/distributed/tests/test_worker_metrics.py
@@ -1,9 +1,8 @@
 from __future__ import annotations
 
 import asyncio
-import secrets
 import time
-from collections.abc import Hashable
+from collections.abc import Collection, Hashable
 from concurrent.futures import ProcessPoolExecutor
 from time import sleep
 
@@ -11,8 +10,8 @@
 
 import dask
 
+import distributed
 from distributed import Event, Worker, wait
-from distributed.comm.utils import OFFLOAD_THRESHOLD
 from distributed.compatibility import WINDOWS
 from distributed.metrics import context_meter, meter
 from distributed.utils_test import (
@@ -25,14 +24,16 @@
 )
 
 
-def get_digests(w: Worker, allow: str | None = None) -> dict[Hashable, float]:
+def get_digests(w: Worker, allow: str | Collection[str] = ()) -> dict[Hashable, float]:
     # import pprint; pprint.pprint(dict(w.digests_total))
+    if isinstance(allow, str):
+        allow = (allow,)
     digests = {
         k: v
         for k, v in w.digests_total.items()
         if k
         not in {"latency", "tick-duration", "transfer-bandwidth", "transfer-duration"}
-        and (allow is None or allow in k)
+        and (any(a in k for a in allow) or not allow)
     }
     assert all(v >= 0 for v in digests.values()), digests
     return digests
@@ -80,6 +81,7 @@ async def test_task_lifecycle(c, s, a, b):
         # -> Run in thread
         ("execute", "z", "thread-cpu", "seconds"),
         ("execute", "z", "thread-noncpu", "seconds"),
+        ("execute", "z", "executor", "seconds"),
         # Spill output; added by _transition_to_memory
         ("execute", "z", "serialize", "seconds"),
         ("execute", "z", "compress", "seconds"),
@@ -150,15 +152,23 @@ async def test_run_spec_deserialization(c, s, a):
 
 
 @gen_cluster(client=True)
-async def test_offload(c, s, a, b):
+async def test_offload(c, s, a, b, monkeypatch):
     """Test that functions wrapped by offload() are metered"""
-    nbytes = int(OFFLOAD_THRESHOLD * 1.1)
-    x = c.submit(secrets.token_bytes, nbytes, key="x", workers=[a.address])
+    monkeypatch.setattr(distributed.comm.utils, "OFFLOAD_THRESHOLD", 1)
+    monkeypatch.setattr(distributed.worker, "OFFLOAD_THRESHOLD", 1)
+
+    x = c.submit(inc, 1, key="x", workers=[a.address])
     y = c.submit(lambda x: None, x, key="y", workers=[b.address])
     await y
 
-    assert 0 < a.digests_total["get-data", "serialize", "seconds"] < 1
-    assert 0 < b.digests_total["gather-dep", "deserialize", "seconds"] < 1
+    assert list(get_digests(b, {"offload", "serialize", "deserialize"})) == [
+        ("gather-dep", "offload", "seconds"),
+        ("gather-dep", "deserialize", "seconds"),
+        ("execute", "y", "offload", "seconds"),
+        ("execute", "y", "deserialize", "seconds"),
+        ("get-data", "offload", "seconds"),
+        ("get-data", "serialize", "seconds"),
+    ]
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])
@@ -315,11 +325,7 @@ async def test_memory_monitor(c, s, a):
     x = c.submit(inc, 1, key="x")
     await async_poll_for(lambda: a.data.disk, timeout=5)
 
-    assert list(get_digests(a)) == [
-        ("execute", "x", "deserialize", "seconds"),
-        ("execute", "x", "thread-cpu", "seconds"),
-        ("execute", "x", "thread-noncpu", "seconds"),
-        ("execute", "x", "other", "seconds"),
+    assert list(get_digests(a, "memory-monitor")) == [
         ("memory-monitor", "serialize", "seconds"),
         ("memory-monitor", "compress", "seconds"),
         ("memory-monitor", "disk-write", "seconds"),
@@ -344,11 +350,13 @@ def f():
         ("execute", "x", "I/O", "seconds"),
         ("execute", "x", "thread-cpu", "seconds"),
         ("execute", "x", "thread-noncpu", "seconds"),
+        ("execute", "x", "executor", "seconds"),
         ("execute", "x", "other", "seconds"),
     ]
     assert get_digests(a)["execute", "x", "I/O", "seconds"] == 5
     assert get_digests(a)["execute", "x", "thread-cpu", "seconds"] == 0
     assert get_digests(a)["execute", "x", "thread-noncpu", "seconds"] == 0
+    assert get_digests(a)["execute", "x", "executor", "seconds"] == 0
     assert get_digests(a)["execute", "x", "other", "seconds"] == 0