Adds logging [INFO] to TimelineVisualizationCallback and HistoryCallback (#548)

* added info level logging to HistoryCallback to share directory

* add logging to timeline & update mkdir to pathlib style

* moved figure creation into timeline on_compute_end(), removed create_timeline(), removed default logging info level from __init__

* removed history  & timeline logging

* adds more detail to diagnostics.md

* Update docs/user-guide/diagnostics.md

Co-authored-by: Tom White <tom.e.white@gmail.com>

---------

Co-authored-by: Tom White <tom.e.white@gmail.com>

Author: 2 people

cubed/diagnostics/timeline.py

@@ -1,6 +1,6 @@
-import os
 import time
 from dataclasses import asdict
+from pathlib import Path
 from typing import Optional
 
 import matplotlib.patches as mpatches
@@ -16,7 +16,7 @@
 
 
 class TimelineVisualizationCallback(Callback):
-    def __init__(self, format: Optional[str] = None) -> None:
+    def __init__(self, format: Optional[str] = "svg") -> None:
         self.format = format
 
     def on_compute_start(self, event):
@@ -27,75 +27,62 @@ def on_task_end(self, event):
         self.stats.append(asdict(event))
 
     def on_compute_end(self, event):
-        end_tstamp = time.time()
-        dst = f"history/{event.compute_id}"
-        format = self.format
-        create_timeline(self.stats, self.start_tstamp, end_tstamp, dst, format)
+        self.end_tstamp = time.time()
 
-
-# copy of lithops function of the same name, and modified for different field names
-def create_timeline(stats, start_tstamp, end_tstamp, dst=None, format=None):
-    stats_df = pd.DataFrame(stats)
-
-    stats_df = stats_df.sort_values(by=["task_create_tstamp", "name"], ascending=True)
-
-    total_calls = len(stats_df)
-
-    palette = sns.color_palette("deep", 6)
-
-    fig = pylab.figure(figsize=(10, 6))
-    ax = fig.add_subplot(1, 1, 1)
-
-    y = np.arange(total_calls)
-    point_size = 10
-
-    fields = [
-        ("task create", stats_df.task_create_tstamp - start_tstamp),
-        ("function start", stats_df.function_start_tstamp - start_tstamp),
-        ("function end", stats_df.function_end_tstamp - start_tstamp),
-        ("task result", stats_df.task_result_tstamp - start_tstamp),
-    ]
-
-    patches = []
-    for f_i, (field_name, val) in enumerate(fields):
-        ax.scatter(val, y, c=[palette[f_i]], edgecolor="none", s=point_size, alpha=0.8)
-        patches.append(mpatches.Patch(color=palette[f_i], label=field_name))
-
-    ax.set_xlabel("Execution Time (sec)")
-    ax.set_ylabel("Function Call")
-
-    legend = pylab.legend(handles=patches, loc="upper right", frameon=True)
-    legend.get_frame().set_facecolor("#FFFFFF")
-
-    yplot_step = int(np.max([1, total_calls / 20]))
-    y_ticks = np.arange(total_calls // yplot_step + 2) * yplot_step
-    ax.set_yticks(y_ticks)
-    ax.set_ylim(-0.02 * total_calls, total_calls * 1.02)
-    for y in y_ticks:
-        ax.axhline(y, c="k", alpha=0.1, linewidth=1)
-
-    max_seconds = np.max(end_tstamp - start_tstamp) * 1.25
-    xplot_step = max(int(max_seconds / 8), 1)
-    x_ticks = np.arange(max_seconds // xplot_step + 2) * xplot_step
-    ax.set_xlim(0, max_seconds)
-
-    ax.set_xticks(x_ticks)
-    for x in x_ticks:
-        ax.axvline(x, c="k", alpha=0.2, linewidth=0.8)
-
-    ax.grid(False)
-    fig.tight_layout()
-
-    if format is None:
-        format = "svg"
-
-    if dst is None:
-        os.makedirs("plots", exist_ok=True)
-        dst = os.path.join(
-            os.getcwd(), "plots", "{}_{}".format(int(time.time()), f"timeline.{format}")
+        stats_df = pd.DataFrame(self.stats)
+        stats_df = stats_df.sort_values(
+            by=["task_create_tstamp", "name"], ascending=True
         )
-    else:
-        dst = os.path.expanduser(dst) if "~" in dst else dst
-        dst = "{}/{}".format(os.path.realpath(dst), f"timeline.{format}")
-
-    fig.savefig(dst)
+        total_calls = len(stats_df)
+        palette = sns.color_palette("deep", 6)
+
+        fig = pylab.figure(figsize=(10, 6))
+        ax = fig.add_subplot(1, 1, 1)
+
+        y = np.arange(total_calls)
+        point_size = 10
+
+        fields = [
+            ("task create", stats_df.task_create_tstamp - self.start_tstamp),
+            ("function start", stats_df.function_start_tstamp - self.start_tstamp),
+            ("function end", stats_df.function_end_tstamp - self.start_tstamp),
+            ("task result", stats_df.task_result_tstamp - self.start_tstamp),
+        ]
+
+        patches = []
+        for f_i, (field_name, val) in enumerate(fields):
+            ax.scatter(
+                val, y, c=[palette[f_i]], edgecolor="none", s=point_size, alpha=0.8
+            )
+            patches.append(mpatches.Patch(color=palette[f_i], label=field_name))
+
+        ax.set_xlabel("Execution Time (sec)")
+        ax.set_ylabel("Function Call")
+
+        legend = pylab.legend(handles=patches, loc="upper right", frameon=True)
+        legend.get_frame().set_facecolor("#FFFFFF")
+
+        yplot_step = int(np.max([1, total_calls / 20]))
+        y_ticks = np.arange(total_calls // yplot_step + 2) * yplot_step
+        ax.set_yticks(y_ticks)
+        ax.set_ylim(-0.02 * total_calls, total_calls * 1.02)
+        for y in y_ticks:
+            ax.axhline(y, c="k", alpha=0.1, linewidth=1)
+
+        max_seconds = np.max(self.end_tstamp - self.start_tstamp) * 1.25
+        xplot_step = max(int(max_seconds / 8), 1)
+        x_ticks = np.arange(max_seconds // xplot_step + 2) * xplot_step
+        ax.set_xlim(0, max_seconds)
+
+        ax.set_xticks(x_ticks)
+        for x in x_ticks:
+            ax.axvline(x, c="k", alpha=0.2, linewidth=0.8)
+
+        ax.grid(False)
+        fig.tight_layout()
+
+        self.dst = Path(f"history/{event.compute_id}")
+        self.dst.mkdir(parents=True, exist_ok=True)
+        self.dst = self.dst / f"timeline.{self.format}"
+
+        fig.savefig(self.dst)

docs/user-guide/diagnostics.md

@@ -41,7 +41,12 @@ Array `c` is coloured orange, which means it is materialized as a Zarr array. Ar
 
 Similarly, the operation that produces `c` is shown in a lilac colour to signify that it runs tasks to produce the output. Operations `op-001` and `op-002` don't run any tasks since `a` and `b` are just small constant arrays.
 
-## Progress bar
+
+## Callbacks
+
+You can pass callbacks to functions that call `compute`, such as {py:func}`store <cubed.store>` or {py:func}`to_zarr <cubed.to_zarr>`.
+
+### Progress bar
 
 You can display a progress bar to track your computation by passing callbacks to {py:meth}`compute() <cubed.Array.compute()>`:
 
@@ -53,14 +58,37 @@ You can display a progress bar to track your computation by passing callbacks to
   op-003 add    4/4 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:00
 ```
 
+The two current progress bar choice are:
+- `from cubed.diagnostics.rich import RichProgressBar`
+- `from cubed.diagnostics.tqdm import TqdmProgressBar`
+
+
 This will work in Jupyter notebooks, and for all executors.
 
-You can also pass callbacks to functions that call `compute`, such as {py:func}`store <cubed.store>` or {py:func}`to_zarr <cubed.to_zarr>`.
 
-## History and timeline visualization
 
-The history and timeline visualization callbacks can be used to find out how long tasks took to run, and how much memory they used.
+### History
+The history callback can be used to understand how long tasks took to run, and how much memory they used. The history callback will write [`events.csv`, `plan.csv` and `stats.csv`] to a new directory under the current directory with the schema `history/compute-{id}`.
+
 
+```ipython
+>>> from cubed.diagnostics.history import HistoryCallback
+>>> hist = HistoryCallback()
+>>> c.compute(callbacks=[hist])
+```
+
+
+### Timeline
 The timeline visualization is useful to determine how much time was spent in worker startup, as well as how much stragglers affected the overall time of the computation. (Ideally, we want vertical lines on this plot, which would represent perfect horizontal scaling.)
 
+The timeline callback will write a graphic `timeline.svg` to a directory with the schema `history/compute-{id}`.
+
+
+```ipython
+>>> from cubed.diagnostics.timeline import TimelineVisualizationCallback
+>>> timeline_viz = TimelineVisualizationCallback()
+>>> c.compute(callbacks=[timeline_viz])
+```
+
+### Examples in use
 See the [examples](https://github.com/cubed-dev/cubed/blob/main/examples/README.md) for more information about how to use them.