Repository: cubed-dev/cubed
Pull request #548 (2024-08-08 20:01:12)

Adds logging [INFO] to `TimelineVisualizationCallback` and `HistoryCallback` 
Small PR that:

- Adds info level logging to show the directory that `TimelineVisualizationCallback` and `HistoryCallback` write their artifacts to.
- Updates the `TimelineVisualizationCallback` os.mkdir to the Pathlib mkdir used in the HistoryCallback. This seems to fix the `FileNotFoundError` in https://github.com/cubed-dev/cubed/issues/544.

example logging:
```
TimelineVisualizationCallback results saved to directory: history/compute-20240808T135430.657101/
HistoryCallback results saved to directory: history/compute-20240808T135430.657101/
```
Some remaining Q's:
- The `TimelineVisualizationCallback` calls an external function (to the class) `create_timeline` that has args for such as `dst=None`, which are not a part of `TimelineVisualizationCallback`. Is this `create_timeline` function used anywhere else? I ask this b/c the optional `dst` arg creates a few more path creation options. If it is confined to `TimelineVisualizationCallback`, the logic could be more similar to `HistoryCallback` which is a bit cleaner. 

- I added info level logging to `diagnostics/__init__.py` is this OK to set logging here?



TomNicholas: Nice! Would this close either of #543 or #544?
norlandrhagen: I think it should close both!
norlandrhagen: No idea why the Apache Beam tests are failing though 😕 
tomwhite: Thanks for the PR @norlandrhagen!

> * The `TimelineVisualizationCallback` calls an external function (to the class) `create_timeline` that has args for such as `dst=None`, which are not a part of `TimelineVisualizationCallback`. Is this `create_timeline` function used anywhere else? I ask this b/c the optional `dst` arg creates a few more path creation options. If it is confined to `TimelineVisualizationCallback`, the logic could be more similar to `HistoryCallback` which is a bit cleaner.

This code was originally based on code from Lithops for producing a timeline image, so some of it may be redundant. Please go ahead and remove code paths that are not used so it's more similar to `HistoryCallback`.

> * I added info level logging to `diagnostics/__init__.py` is this OK to set logging here?

No, it should go in client code. Perhaps add it to examples/add-random.py and examples/matmul-random.py?

tomwhite: > No idea why the Apache Beam tests are failing though 😕

Not related to this PR. Opened #549

Patch:
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index e6d1ebc9..a4850596 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -756,6 +756,16 @@ def map_direct(
 
 
 def rechunk(x, chunks, target_store=None):
+    if isinstance(chunks, dict):
+        chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()}
+        for i in range(x.ndim):
+            if i not in chunks:
+                chunks[i] = x.chunks[i]
+            elif chunks[i] is None:
+                chunks[i] = x.chunks[i]
+    if isinstance(chunks, (tuple, list)):
+        chunks = tuple(lc if lc is not None else rc for lc, rc in zip(chunks, x.chunks))
+
     normalized_chunks = normalize_chunks(chunks, x.shape, dtype=x.dtype)
     if x.chunks == normalized_chunks:
         return x
diff --git a/cubed/diagnostics/timeline.py b/cubed/diagnostics/timeline.py
index 1c0abc89..64de7551 100644
--- a/cubed/diagnostics/timeline.py
+++ b/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 @@ pylab.switch_backend("Agg")
 
 
 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 @@ class TimelineVisualizationCallback(Callback):
         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)
diff --git a/docs/user-guide/diagnostics.md b/docs/user-guide/diagnostics.md
index 3e013323..278a6f1a 100644
--- a/docs/user-guide/diagnostics.md
+++ b/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.


Test patch:
diff --git a/.github/workflows/array-api-tests.yml b/.github/workflows/array-api-tests.yml
index f41649a4..22ca265f 100644
--- a/.github/workflows/array-api-tests.yml
+++ b/.github/workflows/array-api-tests.yml
@@ -79,15 +79,30 @@ jobs:
           array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__isub__(x1, x2)]
           array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__isub__(x, s)]
 
-          # don't run special cases yet
-          array_api_tests/test_special_cases.py
+          # special cases not implemented
+          array_api_tests/test_special_cases.py::test_empty_arrays[std]
+          array_api_tests/test_special_cases.py::test_empty_arrays[var]
+          array_api_tests/test_special_cases.py::test_nan_propagation[std]
+          array_api_tests/test_special_cases.py::test_nan_propagation[var]
 
-          # don't test signatures or names yet as some are not implemented
-          array_api_tests/test_signatures.py
+          # don't test names yet as some are not implemented
           array_api_tests/test_has_names.py
 
-          # very slow
-          array_api_tests/test_creation_functions.py::test_eye
+          # signatures of items not implemented
+          array_api_tests/test_signatures.py::test_func_signature[std]
+          array_api_tests/test_signatures.py::test_func_signature[var]
+          array_api_tests/test_signatures.py::test_func_signature[unique_all]
+          array_api_tests/test_signatures.py::test_func_signature[unique_counts]
+          array_api_tests/test_signatures.py::test_func_signature[unique_inverse]
+          array_api_tests/test_signatures.py::test_func_signature[unique_values]
+          array_api_tests/test_signatures.py::test_func_signature[nonzero]
+          array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
+          array_api_tests/test_signatures.py::test_func_signature[argsort]
+          array_api_tests/test_signatures.py::test_func_signature[sort]
+          array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
+          array_api_tests/test_signatures.py::test_array_method_signature[__dlpack_device__]
+          array_api_tests/test_signatures.py::test_array_method_signature[__setitem__]
+          array_api_tests/test_signatures.py::test_array_method_signature[to_device]
 
           # edge case failures (https://github.com/cubed-dev/cubed/issues/420)
           array_api_tests/test_linalg.py::test_tensordot
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index f40c7a0f..fdf90ca4 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -246,10 +246,19 @@ def test_multiple_ops(spec, executor):
     )
 
 
-@pytest.mark.parametrize("new_chunks", [(1, 2), {0: 1, 1: 2}])
-def test_rechunk(spec, executor, new_chunks):
+@pytest.mark.parametrize(
+    ("new_chunks", "expected_chunks"),
+    [
+        ((1, 2), ((1, 1, 1), (2, 1))),
+        ({0: 1, 1: 2}, ((1, 1, 1), (2, 1))),
+        ({1: 2}, ((2, 1), (2, 1))),  # dim 0 unchanged
+        ({}, ((2, 1), (1, 1, 1))),  # unchanged
+    ],
+)
+def test_rechunk(spec, executor, new_chunks, expected_chunks):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 1), spec=spec)
     b = a.rechunk(new_chunks)
+    assert b.chunks == expected_chunks
     assert_array_equal(
         b.compute(executor=executor),
         np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),