from __future__ import annotations
import asyncio
import collections
import sys
import threading
import time as timemod
from collections.abc import Callable, Hashable, Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from functools import wraps
from math import nan
from typing import Literal
import psutil
from distributed.compatibility import WINDOWS
_empty_namedtuple = collections.namedtuple("_empty_namedtuple", ())
def _psutil_caller(method_name, default=_empty_namedtuple):
"""
Return a function calling the given psutil *method_name*,
or returning *default* if psutil fails.
"""
meth = getattr(psutil, method_name)
@wraps(meth)
def wrapper(): # pragma: no cover
try:
return meth()
except RuntimeError:
# This can happen on some systems (e.g. no physical disk in worker)
return default()
return wrapper
disk_io_counters = _psutil_caller("disk_io_counters")
net_io_counters = _psutil_caller("net_io_counters")
class _WindowsTime:
"""Combine time.time() or time.monotonic() with time.perf_counter() to get an
absolute clock with fine resolution.
"""
base_timer: Callable[[], float]
delta: float
previous: float | None
next_resync: float
resync_every: float
def __init__(
self, base: Callable[[], float], is_monotonic: bool, resync_every: float = 600.0
):
self.base_timer = base
self.previous = float("-inf") if is_monotonic else None
self.next_resync = float("-inf")
self.resync_every = resync_every
def time(self) -> float:
cur = timemod.perf_counter()
if cur > self.next_resync:
self.resync()
self.next_resync = cur + self.resync_every
cur += self.delta
if self.previous is not None:
# Monotonic timer
if cur <= self.previous:
cur = self.previous + 1e-9
self.previous = cur
return cur
def resync(self) -> None:
_time = self.base_timer
_perf_counter = timemod.perf_counter
min_samples = 5
while True:
times = [(_time(), _perf_counter()) for _ in range(min_samples * 2)]
abs_times = collections.Counter(t[0] for t in times)
first, nfirst = abs_times.most_common()[0]
if nfirst < min_samples:
# System too noisy? Start again
continue
perf_times = [t[1] for t in times if t[0] == first][:-1]
assert len(perf_times) >= min_samples - 1, perf_times
self.delta = first - sum(perf_times) / len(perf_times)
break
# A high-resolution wall clock timer measuring the seconds since Unix epoch
if WINDOWS and sys.version_info < (3, 13):
time = _WindowsTime(timemod.time, is_monotonic=False).time
monotonic = _WindowsTime(timemod.monotonic, is_monotonic=True).time
else:
# Under modern Unixes, time.time() and time.monotonic() should be good enough
time = timemod.time
monotonic = timemod.monotonic
process_time = timemod.process_time
# Get a per-thread CPU timer function if possible, otherwise
# use a per-process CPU timer function.
try:
# thread_time is not supported on all platforms
thread_time = timemod.thread_time
except (AttributeError, OSError): # pragma: no cover
thread_time = process_time
@dataclass
class MeterOutput:
start: float
stop: float
delta: float
__slots__ = tuple(__annotations__)
[docs]
@contextmanager
def meter(
func: Callable[[], float] = timemod.perf_counter,
floor: float | Literal[False] = 0.0,
) -> Iterator[MeterOutput]:
"""Convenience context manager which calls func() before and after the wrapped
code and calculates the delta.
Parameters
----------
label: str
label to pass to the callback
func: callable
function to call before and after, which must return a number.
Besides time, it could return e.g. cumulative network traffic or disk usage.
Default: :func:`timemod.perf_counter`
floor: float or False, optional
Floor the delta to the given value (default: 0). This is useful for strictly
cumulative functions that can occasionally glitch and go backwards.
Set to False to disable.
Yields
------
:class:`MeterOutput` where the ``start`` attribute is populated straight away, while
``stop`` and ``delta`` are nan until context exit.
"""
out = MeterOutput(func(), nan, nan)
try:
yield out
finally:
out.stop = func()
out.delta = out.stop - out.start
if floor is not False:
out.delta = max(floor, out.delta)
class ContextMeter:
"""Context-based general purpose meter.
Usage
-----
1. In high level code, call :meth:`add_callback` to install a hook that defines an
activity
2. In low level code, typically many stack levels below, log quantitative events
(e.g. elapsed time, transferred bytes, etc.) so that they will be attributed to
the high-level code calling it, either with :meth:`meter`,
:meth:`meter_function`, or :meth:`digest_metric`.
Examples
--------
In the code that e.g. sends a Python object from A to B over the network:
>>> from distributed.metrics import context_meter
>>> with context_meter.add_callback(partial(print, "A->B comms:")):
... await send_over_the_network(obj)
In the serialization utilities, called many stack levels below:
>>> with context_meter.meter("dumps"):
... pik = pickle.dumps(obj)
>>> with context_meter.meter("compress"):
... pik = lz4.compress(pik)
And finally, elsewhere, deep into the TCP stack:
>>> with context_meter.meter("network-write"):
... await comm.write(frames)
When you call the top-level code, you'll get::
A->B comms: dumps 0.012 seconds
A->B comms: compress 0.034 seconds
A->B comms: network-write 0.567 seconds
"""
_callbacks: ContextVar[dict[Hashable, Callable[[Hashable, float, str], None]]]
def __init__(self):
self._callbacks = ContextVar(
f"MetricHook<{id(self)}>._callbacks", default={} # noqa: B039
)
def __reduce__(self):
assert self is context_meter, "Found copy of singleton"
return self._unpickle_singleton, ()
@staticmethod
def _unpickle_singleton():
return context_meter
@contextmanager
def add_callback(
self,
callback: Callable[[Hashable, float, str], None],
*,
key: Hashable | None = None,
allow_offload: bool = False,
) -> Iterator[None]:
"""Add a callback when entering the context and remove it when exiting it.
The callback must accept the same parameters as :meth:`digest_metric`.
Parameters
----------
callback: Callable
``f(label, value, unit)`` to be executed
key: Hashable, optional
Unique key for the callback. If two nested calls to ``add_callback`` use the
same key, suppress the outermost callback.
allow_offload: bool, optional
If set to True, this context must be executed inside a running asyncio
event loop. If a call to :meth:`digest_metric` is performed from a different
thread, e.g. from inside :func:`distributed.utils.offload`, ensure that
the callback is executed in the event loop's thread instead.
"""
if allow_offload:
loop = asyncio.get_running_loop()
tid = threading.get_ident()
def safe_cb(label: Hashable, value: float, unit: str, /) -> None:
if threading.get_ident() == tid:
callback(label, value, unit)
else: # We're inside offload()
loop.call_soon_threadsafe(callback, label, value, unit)
else:
safe_cb = callback
if key is None:
key = object()
cbs = self._callbacks.get()
cbs = cbs.copy()
cbs[key] = safe_cb
tok = self._callbacks.set(cbs)
try:
yield
finally:
tok.var.reset(tok)
@contextmanager
def clear_callbacks(self) -> Iterator[None]:
"""Do not trigger any callbacks set outside of this context"""
tok = self._callbacks.set({})
try:
yield
finally:
tok.var.reset(tok)
def digest_metric(self, label: Hashable, value: float, unit: str) -> None:
"""Invoke the currently set context callbacks for an arbitrary quantitative
metric.
"""
cbs = self._callbacks.get()
for cb in cbs.values():
cb(label, value, unit)
@contextmanager
def meter(
self,
label: Hashable,
unit: str = "seconds",
func: Callable[[], float] = timemod.perf_counter,
floor: float | Literal[False] = 0.0,
) -> Iterator[MeterOutput]:
"""Convenience context manager or decorator which calls func() before and after
the wrapped code, calculates the delta, and finally calls :meth:`digest_metric`.
If unit=='seconds', it also subtracts any other calls to :meth:`meter` or
:meth:`digest_metric` with the same unit performed within the context, so that
the total is strictly additive.
Parameters
----------
label: Hashable
label to pass to the callback
unit: str, optional
unit to pass to the callback. Default: seconds
func: callable
see :func:`meter`
floor: bool, optional
see :func:`meter`
Yields
------
:class:`MeterOutput` where the ``start`` attribute is populated straight away,
while ``stop`` and ``delta`` are nan until context exit. In case of multiple
nested calls to :meth:`meter`, then delta (for seconds only) is reduced by the
inner metrics, to a minimum of ``floor``.
"""
if unit != "seconds":
try:
with meter(func, floor=floor) as m:
yield m
finally:
self.digest_metric(label, m.delta, unit)
return
# If unit=="seconds", subtract time metered from the sub-contexts
offsets = []
def callback(label2: Hashable, value2: float, unit2: str) -> None:
if unit2 == unit:
# This must be threadsafe to support callbacks invoked from
# distributed.utils.offload; '+=' on a float would not be threadsafe!
offsets.append(value2)
try:
with self.add_callback(callback), meter(func, floor=False) as m:
yield m
finally:
delta = m.delta - sum(offsets)
if floor is not False:
delta = max(floor, delta)
m.delta = delta
self.digest_metric(label, delta, unit)
context_meter = ContextMeter()
class DelayedMetricsLedger:
"""Add-on to :class:`ContextMeter` that helps in the case where:
- The code to be metered is not easily expressed as a self-contained code block
e.g. you want to measure latency in the asyncio event loop before and after
running a task
- You want to alter the metrics depending on how the code ends; e.g. you want to
post them differently in case of failure.
Examples
--------
>>> ledger = DelayedMetricsLedger() # Metering starts here
>>> async def wrapper():
... with ledger.record():
... return await metered_function()
>>> task = asyncio.create_task(wrapper())
>>> # (later, elsewhere)
>>> try:
... await task
... coarse_time = False
... except Exception:
... coarse_time = "failed"
... raise
... finally:
... # Metering stops here
... for label, value, unit in ledger.finalize(coarse_time):
... # actually log metrics
"""
func: Callable[[], float]
start: float
metrics: list[tuple[Hashable, float, str]] # (label, value, unit)
def __init__(self, func: Callable[[], float] = timemod.perf_counter):
self.func = func
self.start = func()
self.metrics = []
def _callback(self, label: Hashable, value: float, unit: str) -> None:
self.metrics.append((label, value, unit))
@contextmanager
def record(self, *, key: Hashable | None = None) -> Iterator[None]:
"""Ingest metrics logged with :meth:`ContextMeter.digest_metric` or
:meth:`ContextMeter.meter` and temporarily store them in :ivar:`metrics`.
Parameters
----------
key: Hashable, optional
See :meth:`ContextMeter.add_callback`
"""
with context_meter.add_callback(self._callback, key=key):
yield
def finalize(
self,
coarse_time: str | Literal[False] = False,
floor: float | Literal[False] = 0.0,
) -> Iterator[tuple[Hashable, float, str]]:
"""The metered code is terminated, and we now know how to log it.
Parameters
----------
coarse_time: str | False, optional
False
Yield all acquired metrics, plus an extra time metric, labelled "other",
which is the time between creating the DelayedMetricsLedger and
calling this method, minus any time logged in the metrics.
label
Yield all acquired non-time metrics.
Yield a single metric, labelled <coarse_time>, which is the time
between creating the DelayedMetricsLedger and calling this method.
floor: float | False, optional
Floor either the "other" or the <coarse_time> metric to this value
(default: 0). Set to False to disable.
"""
stop = self.func()
delta = stop - self.start
for label, value, unit in self.metrics:
if unit != "seconds" or not coarse_time:
yield label, value, unit
if unit == "seconds" and not coarse_time:
delta -= value
if floor is not False:
delta = max(floor, delta)
yield coarse_time or "other", delta, "seconds"
