Source code for distributed.diagnostics.memory_sampler

from __future__ import annotations

import uuid
from import AsyncIterator, Iterator
from contextlib import asynccontextmanager, contextmanager
from datetime import datetime
from typing import TYPE_CHECKING, Any

from distributed.compatibility import PeriodicCallback

    # Optional runtime dependencies
    import pandas as pd

    # Circular dependencies
    from distributed.client import Client
    from distributed.scheduler import Scheduler

[docs]class MemorySampler: """Sample cluster-wide memory usage every <interval> seconds. **Usage** .. code-block:: python client = Client() ms = MemorySampler() with ms.sample("run 1"): <run first workflow> with ms.sample("run 2"): <run second workflow> ... ms.plot() or with an asynchronous client: .. code-block:: python client = await Client(asynchronous=True) ms = MemorySampler() async with ms.sample("run 1"): <run first workflow> async with ms.sample("run 2"): <run second workflow> ... ms.plot() """ samples: dict[str, list[tuple[float, int]]] def __init__(self): self.samples = {}
[docs] def sample( self, label: str | None = None, *, client: Client | None = None, measure: str = "process", interval: float = 0.5, ) -> Any: """Context manager that records memory usage in the cluster. This is synchronous if the client is synchronous and asynchronous if the client is asynchronous. The samples are recorded in ``self.samples[<label>]``. Parameters ========== label: str, optional Tag to record the samples under in the self.samples dict. Default: automatically generate a random label client: Client, optional client used to connect to the scheduler. Default: use the global client measure: str, optional One of the measures from :class:`distributed.scheduler.MemoryState`. Default: sample process memory interval: float, optional sampling interval, in seconds. Default: 0.5 """ if not client: from distributed.client import get_client client = get_client() if client.asynchronous: return self._sample_async(label, client, measure, interval) else: return self._sample_sync(label, client, measure, interval)
@contextmanager def _sample_sync( self, label: str | None, client: Client, measure: str, interval: float ) -> Iterator[None]: key = client.sync( client.scheduler.memory_sampler_start,, measure=measure, interval=interval, ) try: yield finally: samples = client.sync(client.scheduler.memory_sampler_stop, key=key) self.samples[label or key] = samples @asynccontextmanager async def _sample_async( self, label: str | None, client: Client, measure: str, interval: float ) -> AsyncIterator[None]: key = await client.scheduler.memory_sampler_start(, measure=measure, interval=interval ) try: yield finally: samples = await client.scheduler.memory_sampler_stop(key=key) self.samples[label or key] = samples
[docs] def to_pandas(self, *, align: bool = False) -> pd.DataFrame: """Return the data series as a pandas.Dataframe. Parameters ========== align : bool, optional If True, change the absolute timestamps into time deltas from the first sample of each series, so that different series can be visualized side by side. If False (the default), use absolute timestamps. """ import pandas as pd ss = {} for label, s_list in self.samples.items(): assert s_list # There's always at least one sample s = pd.DataFrame(s_list).set_index(0)[1] s.index = pd.to_datetime(s.index, unit="s") = label if align: # convert datetime to timedelta from the first sample s.index -= s.index[0] ss[label] = s df = pd.DataFrame(ss) if len(ss) > 1: # Forward-fill NaNs in the middle of a series created either by overlapping # sampling time range or by align=True. Do not ffill series beyond their # last sample. df = df.ffill().where(~pd.isna(df.bfill())) return df
[docs] def plot(self, *, align: bool = False, **kwargs: Any) -> Any: """Plot data series collected so far Parameters ========== align : bool (optional) See :meth:`~distributed.diagnostics.MemorySampler.to_pandas` kwargs Passed verbatim to :meth:`pandas.DataFrame.plot` Returns ======= Output of :meth:`pandas.DataFrame.plot` """ df = self.to_pandas(align=align).resample("1s").nearest() / 2**30 return df.plot( xlabel="time", ylabel="Cluster memory (GiB)", **kwargs, )
class MemorySamplerExtension: """Scheduler extension - server side of MemorySampler""" scheduler: Scheduler samples: dict[str, list[tuple[float, int]]] def __init__(self, scheduler: Scheduler): self.scheduler = scheduler self.scheduler.handlers["memory_sampler_start"] = self.start self.scheduler.handlers["memory_sampler_stop"] = self.stop self.samples = {} def start(self, client: str, measure: str, interval: float) -> str: """Start periodically sampling memory""" assert not measure.startswith("_") assert isinstance(getattr(self.scheduler.memory, measure), int) key = str(uuid.uuid4()) self.samples[key] = [] def sample(): if client in self.scheduler.clients: ts = nbytes = getattr(self.scheduler.memory, measure) self.samples[key].append((ts, nbytes)) else: self.stop(key) pc = PeriodicCallback(sample, interval * 1000) self.scheduler.periodic_callbacks["MemorySampler-" + key] = pc pc.start() # Immediately collect the first sample; this also ensures there's always at # least one sample sample() return key def stop(self, key: str) -> list[tuple[float, int]]: """Stop sampling and return the samples""" pc = self.scheduler.periodic_callbacks.pop("MemorySampler-" + key, None) if pc is not None: # Race condition with scheduler shutdown pc.stop() return self.samples.pop(key)