from __future__ import annotations import copy import uuid import weakref from collections import defaultdict from collections.abc import Hashable, Iterable, Iterator, Mapping from contextlib import contextmanager from itertools import islice from typing import TYPE_CHECKING, Any, TypedDict import dask.config from dask.typing import Key from distributed.collections import sum_mappings from distributed.itertools import ffill from distributed.metrics import time if TYPE_CHECKING: # Needed to avoid Sphinx WARNING: more than one target found for cross-reference # 'TaskState' and 'WorkerState'" # https://github.com/agronholm/sphinx-autodoc-typehints#dealing-with-circular-imports from distributed import Scheduler, Worker from distributed import scheduler as scheduler_module from distributed.client import SourceCode from distributed.scheduler import TaskGroup, TaskStateState CONTEXTS_WITH_SPAN_ID = ("execute", "p2p") class SpanMetadata(TypedDict): collections: list[dict] @contextmanager def span(*tags: str) -> Iterator[str]: """Tag group of tasks to be part of a certain group, called a span. This context manager can be nested, thus creating sub-spans. If you close and re-open a span context manager with the same tag, you'll end up with two separate spans. Every cluster defines a global "default" span when no span has been defined by the client; the default span is automatically closed and reopened when all tasks associated to it have been completed; in other words the cluster is idle save for tasks that are explicitly annotated by a span. Note that, in some edge cases, you may end up with overlapping default spans, e.g. if a worker crashes and all unique tasks that were in memory on it need to be recomputed. You may capture the ID of a span on the client to match it with the :class:`~distributed.spans.Span` objects the scheduler: >>> client = Client() >>> with span("my workflow") as span_id: ... client.submit(lambda: "Hello world!").result() >>> client.cluster.scheduler.extensions["spans"].spans[span_id] Span Notes ----- You may retrieve the current span with ``dask.get_annotations().get("span")``. You can do so in the client code as well as from inside a task. """ if not tags: raise ValueError("Must specify at least one span tag") annotation = dask.get_annotations().get("span") prev_tags = annotation["name"] if annotation else () # You must specify the full history of IDs, not just the parent, because # otherwise you would not be able to uniquely identify grandparents when # they have no tasks of their own. prev_ids = annotation["ids"] if annotation else () ids = tuple(str(uuid.uuid4()) for _ in tags) with dask.annotate(span={"name": prev_tags + tags, "ids": prev_ids + ids}): yield ids[-1] class Span: #: (, , ...) #: Matches ``TaskState.annotations["span"]["name"]``, both on the scheduler and the #: worker. name: tuple[str, ...] #: Unique ID, generated by :func:`~distributed.span` and #: taken from ``TaskState.annotations["span"]["id"][-1]``. #: Matches ``distributed.scheduler.TaskState.group.span_id`` #: and ``distributed.worker_state_machine.TaskState.span_id``. id: str _parent: weakref.ref[Span] | None #: Direct children of this span, sorted by creation time children: list[Span] #: Task groups *directly* belonging to this span. #: #: See Also #: -------- # traverse_groups #: #: Notes #: ----- #: TaskGroups are forgotten by the Scheduler when the last task is forgotten, but #: remain referenced here indefinitely. If a user calls compute() twice on the same #: collection, you'll have more than one group with the same tg.name in this set! #: For the same reason, while the same TaskGroup object is guaranteed to be attached #: to exactly one Span, you may have different TaskGroups with the same key attached #: to different Spans. groups: set[TaskGroup] #: Time when the span first appeared on the scheduler. #: The same property on parent spans is always less than or equal to this. #: #: See Also #: -------- #: start #: stop enqueued: float #: Source code snippets, if it was sent by the client. #: We're using a dict without values as an insertion-sorted set. _code: dict[tuple[SourceCode, ...], None] _metadata: SpanMetadata | None _cumulative_worker_metrics: defaultdict[tuple[Hashable, ...], float] #: reference to SchedulerState.total_nthreads_history _total_nthreads_history: list[tuple[float, int]] #: Length of total_nthreads_history when this span was enqueued _total_nthreads_offset: int # Support for weakrefs to a class with __slots__ __weakref__: Any __slots__ = tuple(__annotations__) _metadata_seen: set[int] = set() def __init__( self, name: tuple[str, ...], id_: str, parent: Span | None, total_nthreads_history: list[tuple[float, int]], ): self.name = name self.id = id_ self._parent = weakref.ref(parent) if parent is not None else None self.enqueued = time() self.children = [] self.groups = set() self._code = {} self._metadata = None # Don't cast int metrics to float self._cumulative_worker_metrics = defaultdict(int) assert len(total_nthreads_history) > 0 self._total_nthreads_history = total_nthreads_history self._total_nthreads_offset = len(total_nthreads_history) - 1 def __repr__(self) -> str: return f"Span" @property def parent(self) -> Span | None: if self._parent: out = self._parent() assert out return out return None def add_metadata(self, metadata: SpanMetadata) -> None: """Add metadata to the span, e.g. code snippets""" id_ = id(metadata) if id_ in self._metadata_seen: return self._metadata_seen.add(id_) if self._metadata is None: self._metadata = copy.deepcopy(metadata) else: self._metadata["collections"].extend(metadata["collections"]) @property def annotation(self) -> dict[str, tuple[str, ...]] | None: """Rebuild the dask graph annotation which contains the full id history Note that this may not match the original annotation in case of TaskGroup collision. """ if self.name == ("default",): return None ids = [] node: Span | None = self while node: ids.append(node.id) node = node.parent return {"name": self.name, "ids": tuple(reversed(ids))} def traverse_spans(self) -> Iterator[Span]: """Top-down recursion of all spans belonging to this branch off span tree, including self """ yield self for child in self.children: yield from child.traverse_spans() def traverse_groups(self) -> Iterator[TaskGroup]: """All TaskGroups belonging to this branch of span tree""" for span in self.traverse_spans(): yield from span.groups @property def start(self) -> float: """Earliest time when a task belonging to this span tree started computing; 0 if no task has *finished* computing yet. Notes ----- This is not updated until at least one task has *finished* computing. It could move backwards as tasks complete. See Also -------- enqueued stop distributed.scheduler.TaskGroup.start """ out = min( (tg.start for tg in self.traverse_groups() if tg.start != 0.0), default=0.0, ) if out: # absorb small errors in worker delay calculation out = max(out, self.enqueued) return out @property def stop(self) -> float: """When this span tree finished computing, or current timestamp if it didn't finish yet. Notes ----- This differs from ``TaskGroup.stop`` when there aren't unfinished tasks; is also will never be zero. See Also -------- enqueued start done distributed.scheduler.TaskGroup.stop """ if self.done: out = max(tg.stop for tg in self.traverse_groups()) else: out = time() # absorb small errors in worker delay calculation, as well as in time() not # being perfectly monotonic return max(out, self.enqueued) @property def metadata(self) -> SpanMetadata | None: return self._metadata @property def states(self) -> dict[TaskStateState, int]: """The number of tasks currently in each state in this span tree; e.g. ``{"memory": 10, "processing": 3, "released": 4, ...}``. See Also -------- distributed.scheduler.TaskGroup.states """ return sum_mappings(tg.states for tg in self.traverse_groups()) @property def done(self) -> bool: """Return True if all tasks in this span tree are completed; False otherwise. Notes ----- This property may transition from True to False, e.g. when a new sub-span is added or when a worker that contained the only replica of a task in memory crashes and the task need to be recomputed. See Also -------- distributed.scheduler.TaskGroup.done """ return all(tg.done for tg in self.traverse_groups()) @property def all_durations(self) -> dict[str, float]: """Cumulative duration of all completed actions in this span tree, by action See Also -------- duration distributed.scheduler.TaskGroup.all_durations """ return sum_mappings(tg.all_durations for tg in self.traverse_groups()) @property def duration(self) -> float: """The total amount of time spent on all tasks in this span tree See Also -------- all_durations distributed.scheduler.TaskGroup.duration """ return sum(tg.duration for tg in self.traverse_groups()) @property def nbytes_total(self) -> int: """The total number of bytes that this span tree has produced See Also -------- distributed.scheduler.TaskGroup.nbytes_total """ return sum(tg.nbytes_total for tg in self.traverse_groups()) @property def code(self) -> list[tuple[SourceCode, ...]]: """Code snippets, sent by the client on compute(), persist(), and submit(). Only populated if ``distributed.diagnostics.computations.nframes`` is non-zero. """ # Deduplicate, but preserve order return list( dict.fromkeys(sc for child in self.traverse_spans() for sc in child._code) ) @property def cumulative_worker_metrics(self) -> dict[tuple[Hashable, ...], float]: """Replica of ``Worker.digests_total`` and ``Scheduler.cumulative_worker_metrics``, but only for the metrics that can be attributed to the current span tree. The span id has been removed from the key. At the moment of writing, all keys are ``("execute", , , )`` or ``("p2p", , , )`` but more may be added in the future with a different format; please test e.g. for ``k[0] == "execute"``. """ out = sum_mappings( child._cumulative_worker_metrics for child in self.traverse_spans() ) known_seconds = sum( v for k, v in out.items() if k[0] == "execute" and k[-1] == "seconds" ) # Besides rounding errors, you may get negative unknown seconds if a user # manually invokes `context_meter.digest_metric`. unknown_seconds = max(0.0, self.active_cpu_seconds - known_seconds) out["execute", "N/A", "idle or other spans", "seconds"] = unknown_seconds return out @staticmethod def merge(*items: Span) -> Span: """Merge multiple spans into a synthetic one. The input spans must not be related with each other. """ if not items: raise ValueError("Nothing to merge") out = Span( name=("(merged)",), id_="(merged)", parent=None, total_nthreads_history=items[0]._total_nthreads_history, ) out._total_nthreads_offset = min( child._total_nthreads_offset for child in items ) out.children.extend(items) out.enqueued = min(child.enqueued for child in items) return out def _nthreads_timeseries(self) -> Iterator[tuple[float, int]]: """Yield (timestamp, number of threads across the cluster), forward-fill""" stop = self.stop if self.done else 0 for t, n in islice( self._total_nthreads_history, self._total_nthreads_offset, None ): if stop and t >= stop: break yield max(self.enqueued, t), n def _active_timeseries(self) -> Iterator[tuple[float, bool]]: """If this span is the output of :meth:`merge`, yield (timestamp, True if at least one input span is active), forward-fill. """ if self.id != "(merged)": yield self.enqueued, True yield self.stop, False return events = [] for child in self.children: events += [(child.enqueued, 1), (child.stop, -1)] # enqueued <= stop by construction. # Occasionally, enqueued == stop, e.g. when the clock is adjusted backwards. # Prevent negative n_active when this happens. events.sort(key=lambda el: el[0]) n_active = 0 for t, delta in events: if not n_active: assert delta == 1 yield t, True n_active += delta if not n_active: yield t, False @property def nthreads_intervals(self) -> list[tuple[float, float, int]]: """ Returns ------- List of tuples: - begin timestamp - end timestamp - Scheduler.total_nthreads during this interval When the Span is the output of :meth:`merge`, the intervals may not be contiguous. See Also -------- enqueued stop active_cpu_seconds distributed.scheduler.SchedulerState.total_nthreads """ nthreads_t, nthreads_count = zip(*self._nthreads_timeseries()) is_active_t, is_active_flag = zip(*self._active_timeseries()) t_interp = sorted({*nthreads_t, *is_active_t}) nthreads_count_interp = ffill(t_interp, nthreads_t, nthreads_count, left=0) is_active_flag_interp = ffill(t_interp, is_active_t, is_active_flag, left=False) return [ (t0, t1, n) for t0, t1, n, active in zip( t_interp, t_interp[1:], nthreads_count_interp, is_active_flag_interp ) if active ] @property def active_cpu_seconds(self) -> float: """Return number of CPU seconds that were made available on the cluster while this Span was running; in other words ``(Span.stop - Span.enqueued) * Scheduler.total_nthreads``. This accounts for workers joining and leaving the cluster while this Span was active. If this Span is the output of :meth:`merge`, do not count gaps between input spans. See Also -------- enqueued stop nthreads_intervals distributed.scheduler.SchedulerState.total_nthreads """ return sum((t1 - t0) * nthreads for t0, t1, nthreads in self.nthreads_intervals) class SpansSchedulerExtension: """Scheduler extension for spans support""" scheduler: Scheduler #: All Span objects by id spans: dict[str, Span] #: Only the spans that don't have any parents, sorted by creation time. #: This is a convenience helper structure to speed up searches. root_spans: list[Span] #: All spans, keyed by their full name and sorted by creation time. #: This is a convenience helper structure to speed up searches. spans_search_by_name: defaultdict[tuple[str, ...], list[Span]] #: All spans, keyed by the individual tags that make up their name and sorted by #: creation time. #: This is a convenience helper structure to speed up searches. #: #: See Also #: -------- #: find_by_tags #: merge_by_tags spans_search_by_tag: defaultdict[str, list[Span]] def __init__(self, scheduler: Scheduler): self.scheduler = scheduler self.spans = {} self.root_spans = [] self.spans_search_by_name = defaultdict(list) self.spans_search_by_tag = defaultdict(list) def observe_tasks( self, tss: Iterable[scheduler_module.TaskState], code: tuple[SourceCode, ...], span_metadata: SpanMetadata, ) -> dict[Key, dict]: """Acknowledge the existence of runnable tasks on the scheduler. These may either be new tasks, tasks that were previously unrunnable, or tasks that were already fed into this method already. Attach newly observed tasks to either the desired span or to ("default", ). Update TaskGroup.span_id and wipe TaskState.annotations["span"]. Returns ------- Updated 'span' annotations: {key: {"name": (..., ...), "ids": (..., ...)}} """ out = {} default_span = None for ts in tss: if ts.annotations is None: ts.annotations = dict() # You may have different tasks belonging to the same TaskGroup but to # different spans. If that happens, arbitrarily force everything onto the # span of the earliest encountered TaskGroup. tg = ts.group if tg.span_id: span = self.spans[tg.span_id] else: ann = ts.annotations.get("span") if ann: span = self._ensure_span(ann["name"], ann["ids"]) else: if not default_span: default_span = self._ensure_default_span() span = default_span tg.span_id = span.id span.groups.add(tg) if code: span._code[code] = None if span_metadata: span.add_metadata(span_metadata) # The span may be completely different from the one referenced by the # annotation, due to the TaskGroup collision issue explained above. if ann := span.annotation: ts.annotations["span"] = out[ts.key] = ann else: ts.annotations.pop("span", None) return out def _ensure_default_span(self) -> Span: """Return the currently active default span, or create one if the previous one terminated. In other words, do not reuse the previous default span if all tasks that were not explicitly annotated with :func:`spans` on the client side are finished. """ defaults = self.spans_search_by_name["default",] if defaults and not defaults[-1].done: return defaults[-1] return self._ensure_span(("default",), (str(uuid.uuid4()),)) def _ensure_span(self, name: tuple[str, ...], ids: tuple[str, ...]) -> Span: """Create Span if it doesn't exist and return it""" try: return self.spans[ids[-1]] except KeyError: pass assert len(name) == len(ids) assert len(name) > 0 parent = None for i in range(1, len(name)): parent = self._ensure_span(name[:i], ids[:i]) span = Span( name=name, id_=ids[-1], parent=parent, total_nthreads_history=self.scheduler.total_nthreads_history, ) self.spans[span.id] = span self.spans_search_by_name[name].append(span) for tag in name: self.spans_search_by_tag[tag].append(span) if parent: parent.children.append(span) else: self.root_spans.append(span) return span def find_by_tags(self, *tags: str) -> Iterator[Span]: """Yield all spans that contain any of the given tags. When a tag is shared both by a span and its (grand)children, only return the parent. """ by_level = defaultdict(list) for tag in tags: for sp in self.spans_search_by_tag[tag]: by_level[len(sp.name)].append(sp) seen = set() for _, level in sorted(by_level.items()): seen.update(level) for sp in level: if sp.parent not in seen: yield sp def merge_all(self) -> Span: """Return a synthetic Span which is the sum of all spans""" return Span.merge(*self.root_spans) def merge_by_tags(self, *tags: str) -> Span: """Return a synthetic Span which is the sum of all spans containing the given tags """ return Span.merge(*self.find_by_tags(*tags)) def heartbeat( self, ws: scheduler_module.WorkerState, data: dict[tuple[Hashable, ...], float] ) -> None: """Triggered by :meth:`SpansWorkerExtension.heartbeat`. Populate :meth:`Span.cumulative_worker_metrics` with data from the worker. See Also -------- SpansWorkerExtension.heartbeat Span.cumulative_worker_metrics """ for (context, span_id, *other), v in data.items(): assert isinstance(span_id, str) span = self.spans[span_id] span._cumulative_worker_metrics[(context, *other)] += v class SpansWorkerExtension: """Worker extension for spans support""" worker: Worker digests_total_since_heartbeat: dict[tuple[Hashable, ...], float] def __init__(self, worker: Worker): self.worker = worker self.digests_total_since_heartbeat = {} def collect_digests( self, digests_total_since_heartbeat: Mapping[Hashable, float] ) -> None: # Note: this method may be called spuriously by Worker._register_with_scheduler, # but when it does it's guaranteed not to find any metrics assert not self.digests_total_since_heartbeat self.digests_total_since_heartbeat = { k: v for k, v in digests_total_since_heartbeat.items() if isinstance(k, tuple) and k[0] in CONTEXTS_WITH_SPAN_ID } def heartbeat(self) -> dict[tuple[Hashable, ...], float]: """Apportion the metrics that do have a span to the Spans on the scheduler Returns ------- ``{(context, span_id, prefix, activity, unit): value}}`` See Also -------- SpansSchedulerExtension.heartbeat Span.cumulative_worker_metrics distributed.worker.Worker.get_metrics """ out = self.digests_total_since_heartbeat self.digests_total_since_heartbeat = {} return out