from __future__ import annotations
import logging
from collections import deque
from typing import Literal
import dask
from dask.utils import format_time, key_split, parse_timedelta
from distributed.client import Client, default_client
from distributed.diagnostics.plugin import SchedulerPlugin
from distributed.diagnostics.progress_stream import color_of
from distributed.metrics import time
logger = logging.getLogger(__name__)
class TaskStreamPlugin(SchedulerPlugin):
name = "task-stream"
def __init__(self, scheduler, maxlen=None):
if maxlen is None:
maxlen = max(
dask.config.get(
"distributed.scheduler.dashboard.status.task-stream-length"
),
dask.config.get(
"distributed.scheduler.dashboard.tasks.task-stream-length"
),
)
self.buffer = deque(maxlen=maxlen)
self.scheduler = scheduler
self.index = 0
def transition(self, key, start, finish, *args, **kwargs):
if start == "processing" and finish in ("memory", "erred"):
assert kwargs["startstops"]
kwargs["key"] = key
self.buffer.append(kwargs)
self.index += 1
def collect(self, start=None, stop=None, count=None, start_index=None):
# ``start_index`` selects records by their position in the monotonically
# increasing append counter (``self.index``) rather than by wall-clock
# time. This is immune to clock differences and latency between the
# client and the workers, which can otherwise cause time-based ``start``
# boundaries to drop tasks that have already completed.
if start_index is not None:
buffer_start = start_index - (self.index - len(self.buffer))
buffer_start = max(0, min(buffer_start, len(self.buffer)))
return [self.buffer[i] for i in range(buffer_start, len(self.buffer))]
def bisect(target, left, right):
while left != right:
mid = (left + right) // 2
stop = max(
startstop["stop"] for startstop in self.buffer[mid]["startstops"]
)
if stop < target:
left = mid + 1
else:
right = mid
return left
if isinstance(start, str):
start = time() - parse_timedelta(start)
if start is not None:
start = bisect(start, 0, len(self.buffer))
if isinstance(stop, str):
stop = time() - parse_timedelta(stop)
if stop is not None:
stop = bisect(stop, 0, len(self.buffer))
if count is not None:
if start is None and stop is None:
stop = len(self.buffer)
start = stop - count
elif start is None and stop is not None:
start = stop - count
elif start is not None and stop is None:
stop = start + count
if stop is None:
stop = len(self.buffer)
if start is None:
start = 0
start = max(0, start)
stop = min(stop, len(self.buffer))
return [self.buffer[i] for i in range(start, stop)]
def rectangles(self, istart, istop=None, workers=None, start_boundary=0):
msgs = []
diff = self.index - len(self.buffer)
if istop is None:
istop = self.index
for i in range(max(0, (istart or 0) - diff), istop - diff if istop else istop):
msg = self.buffer[i]
msgs.append(msg)
return rectangles(msgs, workers=workers, start_boundary=start_boundary)
def rectangles(msgs, workers=None, start_boundary=0):
if workers is None:
workers = {}
L_start = []
L_duration = []
L_duration_text = []
L_key = []
L_name = []
L_color = []
L_alpha = []
L_worker = []
L_worker_thread = []
L_y = []
for msg in msgs:
key = msg["key"]
name = key_split(key)
startstops = msg.get("startstops", [])
try:
worker_thread = f"{msg['worker']}-{msg['thread']}"
except Exception:
continue
if worker_thread not in workers:
workers[worker_thread] = len(workers) / 2
for startstop in startstops:
if startstop["start"] < start_boundary:
continue
color = colors[startstop["action"]]
if type(color) is not str:
color = color(msg)
L_start.append((startstop["start"] + startstop["stop"]) / 2 * 1000)
L_duration.append(1000 * (startstop["stop"] - startstop["start"]))
L_duration_text.append(format_time(startstop["stop"] - startstop["start"]))
L_key.append(key)
L_name.append(prefix[startstop["action"]] + name)
L_color.append(color)
L_alpha.append(alphas[startstop["action"]])
L_worker.append(msg["worker"])
L_worker_thread.append(worker_thread)
L_y.append(workers[worker_thread])
return {
"start": L_start,
"duration": L_duration,
"duration_text": L_duration_text,
"key": L_key,
"name": L_name,
"color": L_color,
"alpha": L_alpha,
"worker": L_worker,
"worker_thread": L_worker_thread,
"y": L_y,
}
def color_of_message(msg):
if msg["status"] == "OK":
split = key_split(msg["key"])
return color_of(split)
else:
return "black"
colors = {
"transfer": "red",
"disk-write": "orange",
"disk-read": "orange",
"deserialize": "gray",
"compute": color_of_message,
}
alphas = {
"transfer": 0.4,
"compute": 1,
"deserialize": 0.4,
"disk-write": 0.4,
"disk-read": 0.4,
}
prefix = {
"transfer": "transfer-",
"disk-write": "disk-write-",
"disk-read": "disk-read-",
"deserialize": "deserialize-",
"compute": "",
}
async def _get_task_stream_impl(
client,
start=None,
stop=None,
count=None,
plot=False,
filename="task-stream.html",
bokeh_resources=None,
start_index=None,
):
"""Asynchronous implementation of Client.get_task_stream and of the
get_task_stream context manager
"""
msgs = await client.scheduler.get_task_stream(
start=start, stop=stop, count=count, start_index=start_index
)
if not plot:
return msgs
from distributed.dashboard.components.scheduler import task_stream_figure
rects = rectangles(msgs)
source, figure = task_stream_figure(sizing_mode="stretch_both")
source.data.update(rects)
if plot == "save":
from bokeh.plotting import output_file, save
output_file(filename=filename, title="Dask Task Stream")
save(figure, filename=filename, resources=bokeh_resources)
return (msgs, figure)
[docs]
class get_task_stream:
"""
Collect task stream within a context block
This provides diagnostic information about every task that was run during
the time when this block was active.
This must be used as a context manager.
Parameters
----------
plot: boolean, str
If true then also return a Bokeh figure
If plot == 'save' then save the figure to a file
filename: str (optional)
The filename to save to if you set ``plot='save'``
Examples
--------
>>> with get_task_stream() as ts:
... x.compute()
>>> ts.data
[...]
Get back a Bokeh figure and optionally save to a file
>>> with get_task_stream(plot='save', filename='task-stream.html') as ts:
... x.compute()
>>> ts.figure
<Bokeh Figure>
To share this file with others you may wish to upload and serve it online.
A common way to do this is to upload the file as a gist, and then serve it
on https://raw.githack.com ::
$ python -m pip install gist
$ gist task-stream.html
https://gist.github.com/8a5b3c74b10b413f612bb5e250856ceb
You can then navigate to that site, click the "Raw" button to the right of
the ``task-stream.html`` file, and then provide that URL to
https://raw.githack.com . This process should provide a sharable link that
others can use to see your task stream plot.
See Also
--------
Client.get_task_stream: Function version of this context manager
"""
data: list[dict]
def __init__(
self,
client: Client | None = None,
plot: bool | Literal["save"] = False,
filename: str = "task-stream.html",
):
self.data = []
self._plot = plot
self._filename = filename
self.figure = None
self.client = client or default_client()
self._start_index = None
def __enter__(self):
return self.client.sync(self.__aenter__)
def __exit__(self, exc_type, exc_value, traceback):
return self.client.sync(self.__aexit__, exc_type, exc_value, traceback)
async def __aenter__(self):
"""Record the scheduler's task-stream cursor on entry and collect
everything appended after it on exit. Using the monotonic index
instead of a wall-clock boundary avoids dropping tasks when there is
latency or clock skew between the client and the workers.
"""
self._start_index = await self.client.scheduler.get_task_stream_index()
return self
async def __aexit__(self, exc_type, exc_value, traceback):
res = await _get_task_stream_impl(
self.client,
start_index=self._start_index,
plot=self._plot,
filename=self._filename,
)
if self._plot:
data, self.figure = res
else:
data = res
self.data.extend(data)
