|
2 | 2 |
|
3 | 3 | import collections |
4 | 4 | import time as timemod |
5 | | -from collections.abc import Callable |
| 5 | +from collections.abc import Callable, Iterator |
| 6 | +from contextlib import contextmanager |
| 7 | +from contextvars import ContextVar |
| 8 | +from dataclasses import dataclass |
6 | 9 | from functools import wraps |
| 10 | +from math import nan |
| 11 | +from typing import Literal |
7 | 12 |
|
8 | 13 | import psutil |
9 | 14 |
|
@@ -103,3 +108,171 @@ def resync(self) -> None: |
103 | 108 | thread_time = timemod.thread_time |
104 | 109 | except (AttributeError, OSError): # pragma: no cover |
105 | 110 | thread_time = process_time |
| 111 | + |
| 112 | + |
| 113 | +@dataclass |
| 114 | +class MeterOutput: |
| 115 | + start: float |
| 116 | + stop: float |
| 117 | + delta: float |
| 118 | + __slots__ = tuple(__annotations__) |
| 119 | + |
| 120 | + |
| 121 | +@contextmanager |
| 122 | +def meter( |
| 123 | + func: Callable[[], float] = timemod.perf_counter, |
| 124 | + floor: float | Literal[False] = 0.0, |
| 125 | +) -> Iterator[MeterOutput]: |
| 126 | + """Convenience context manager which calls func() before and after the wrapped |
| 127 | + code and calculates the delta. |
| 128 | +
|
| 129 | + Parameters |
| 130 | + ---------- |
| 131 | + label: str |
| 132 | + label to pass to the callback |
| 133 | + func: callable |
| 134 | + function to call before and after, which must return a number. |
| 135 | + Besides time, it could return e.g. cumulative network traffic or disk usage. |
| 136 | + Default: :func:`timemod.perf_counter` |
| 137 | + floor: float or False, optional |
| 138 | + Floor the delta to the given value (default: 0). This is useful for strictly |
| 139 | + cumulative functions that can occasionally glitch and go backwards. |
| 140 | + Set to False to disable. |
| 141 | + """ |
| 142 | + out = MeterOutput(func(), nan, nan) |
| 143 | + try: |
| 144 | + yield out |
| 145 | + finally: |
| 146 | + out.stop = func() |
| 147 | + out.delta = out.stop - out.start |
| 148 | + if floor is not False: |
| 149 | + out.delta = max(floor, out.delta) |
| 150 | + |
| 151 | + |
| 152 | +class ContextMeter: |
| 153 | + """Context-based general purpose meter. |
| 154 | +
|
| 155 | + Usage |
| 156 | + ----- |
| 157 | + 1. In high level code, call :meth:`add_callback` to install a hook that defines an |
| 158 | + activity |
| 159 | + 2. In low level code, typically many stack levels below, log quantitative events |
| 160 | + (e.g. elapsed time, transferred bytes, etc.) so that they will be attributed to |
| 161 | + the high-level code calling it, either with :meth:`meter` or |
| 162 | + :meth:`digest_metric`. |
| 163 | +
|
| 164 | + Examples |
| 165 | + -------- |
| 166 | + In the code that e.g. sends a Python object from A to B over the network: |
| 167 | + >>> from distributed.metrics import context_meter |
| 168 | + >>> with context_meter.add_callback(partial(print, "A->B comms:")): |
| 169 | + ... await send_over_the_network(obj) |
| 170 | +
|
| 171 | + In the serialization utilities, called many stack levels below: |
| 172 | + >>> with context_meter.meter("dumps"): |
| 173 | + ... pik = pickle.dumps(obj) |
| 174 | + >>> with context_meter.meter("compress"): |
| 175 | + ... pik = lz4.compress(pik) |
| 176 | +
|
| 177 | + And finally, elsewhere, deep into the TCP stack: |
| 178 | + >>> with context_meter.meter("network-write"): |
| 179 | + ... await comm.write(frames) |
| 180 | +
|
| 181 | + When you call the top-level code, you'll get:: |
| 182 | + A->B comms: dumps 0.012 seconds |
| 183 | + A->B comms: compress 0.034 seconds |
| 184 | + A->B comms: network-write 0.567 seconds |
| 185 | + """ |
| 186 | + |
| 187 | + _callbacks: ContextVar[list[Callable[[str, float, str], None]]] |
| 188 | + default_threshold: float |
| 189 | + |
| 190 | + def __init__(self): |
| 191 | + self._callbacks = ContextVar(f"MetricHook<{id(self)}>._callbacks", default=[]) |
| 192 | + self.default_threshold = 0.001 |
| 193 | + |
| 194 | + @contextmanager |
| 195 | + def add_callback( |
| 196 | + self, callback: Callable[[str, float, str], None] |
| 197 | + ) -> Iterator[None]: |
| 198 | + """Add a callback when entering the context and remove it when exiting it. |
| 199 | + The callback must accept the same parameters as :meth:`digest_metric`. |
| 200 | + """ |
| 201 | + cbs = self._callbacks.get() |
| 202 | + tok = self._callbacks.set(cbs + [callback]) |
| 203 | + try: |
| 204 | + yield |
| 205 | + finally: |
| 206 | + tok.var.reset(tok) |
| 207 | + |
| 208 | + def digest_metric(self, label: str, value: float, unit: str) -> None: |
| 209 | + """Invoke the currently set context callbacks for an arbitrary quantitative |
| 210 | + metric. |
| 211 | + """ |
| 212 | + cbs = self._callbacks.get() |
| 213 | + for cb in cbs: |
| 214 | + cb(label, value, unit) |
| 215 | + |
| 216 | + @contextmanager |
| 217 | + def meter( |
| 218 | + self, |
| 219 | + label: str, |
| 220 | + unit: str = "seconds", |
| 221 | + func: Callable[[], float] = timemod.perf_counter, |
| 222 | + floor: float | Literal[False] = 0.0, |
| 223 | + threshold: float | None = None, |
| 224 | + ) -> Iterator[None]: |
| 225 | + """Convenience context manager which calls func() before and after the wrapped |
| 226 | + code, calculates the delta, and finally calls :meth:`digest_metric`. It also |
| 227 | + subtracts any other calls to :meth:`meter` or :meth:`digest_metric` with the |
| 228 | + same unit performed within the context, so that the total is strictly additive. |
| 229 | +
|
| 230 | + :meth:`digest_metric` is not called in case of exception. |
| 231 | +
|
| 232 | + Parameters |
| 233 | + ---------- |
| 234 | + label: str |
| 235 | + label to pass to the callback |
| 236 | + unit: str, optional |
| 237 | + unit to pass to the callback. Default: seconds |
| 238 | + func: callable |
| 239 | + see :func:`meter` |
| 240 | + floor: bool, optional |
| 241 | + see :func:`meter` |
| 242 | + threshold: float, optional |
| 243 | + Do not call :meth:`digest_metric` if the delta is less than this. |
| 244 | + Default: 1ms |
| 245 | + """ |
| 246 | + offsets = [] |
| 247 | + |
| 248 | + def cb(label2: str, value2: float, unit2: str) -> None: |
| 249 | + if unit2 == unit: |
| 250 | + # This must be threadsafe to support when callbacks are invoked from |
| 251 | + # distributed.utils.offload; '+=' on a float would not be threadsafe! |
| 252 | + offsets.append(value2) |
| 253 | + |
| 254 | + with self.add_callback(cb), meter(func, floor=False) as m: |
| 255 | + yield |
| 256 | + |
| 257 | + delta = m.delta - sum(offsets) |
| 258 | + if floor is not False: |
| 259 | + delta = max(floor, delta) |
| 260 | + if threshold is None: |
| 261 | + threshold = self.default_threshold |
| 262 | + if delta >= threshold: |
| 263 | + self.digest_metric(label, delta, unit) |
| 264 | + |
| 265 | + @contextmanager |
| 266 | + def no_threshold(self) -> Iterator[None]: |
| 267 | + """Temporarily disable default threshold in :meth:`meter()`. |
| 268 | + Useful for unit testing trivial timings. |
| 269 | + """ |
| 270 | + bak = self.default_threshold |
| 271 | + self.default_threshold = 0.0 |
| 272 | + try: |
| 273 | + yield |
| 274 | + finally: |
| 275 | + self.default_threshold = bak |
| 276 | + |
| 277 | + |
| 278 | +context_meter = ContextMeter() |
0 commit comments