Repository: cubed-dev/cubed (python)

Issue:
[Docs] Usage instructions for progressbars 
I've been inspired by @thodson-usgs's Virtualizarr / Cubed rechunking example work and figured it's time to give Cubed a shot. 

Wondering if it's possible to add some documentation on how to use progress bars to monitor a cubed job. I've tried piecing together some from the examples, but haven't had much luck. 

RichProgressBar:

``` python

from cubed import Spec
import xarray as xr 
spec = Spec(work_dir='tmp', allowed_mem='2GB')

from cubed.extensions.rich import RichProgressBar

ds = xr.tutorial.open_dataset('air_temperature', chunked_array_type='cubed',
     from_array_kwargs={'spec': spec},chunks={})
ds = ds.chunk({'time':1})

ds.compute(callbacks=RichProgressBar())

# OR

ds.to_zarr("air_temp.zarr", mode='w',consolidated=True, chunkmanager_store_kwargs={'from_array_kwargs': {'spec': spec}, 'callbacks':RichProgressBar() })



```

Throws a TypeError: `zip() argument after * must be an iterable, not RichProgressBar`

Full traceback:
<details>

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[4], line 16
     14 ds.air.data.visualize()
     15 # ds.compute(callbacks=RichProgressBar())
---> 16 ds.to_zarr("air_temp.zarr", mode='w',consolidated=True, chunkmanager_store_kwargs={'from_array_kwargs': {'spec': spec}, 'callbacks':RichProgressBar() })

File ~/micromamba/envs/cubed/lib/python3.11/site-packages/xarray/core/dataset.py:2553, in Dataset.to_zarr(self, store, chunk_store, mode, synchronizer, group, encoding, compute, consolidated, append_dim, region, safe_chunks, storage_options, zarr_version, write_empty_chunks, chunkmanager_store_kwargs)
   2406 """Write dataset contents to a zarr group.
   2407 
   2408 Zarr chunks are determined in the following way:
   (...)
   2549     The I/O user guide, with more details and examples.
   2550 """
   2551 from xarray.backends.api import to_zarr
-> 2553 return to_zarr(  # type: ignore[call-overload,misc]
   2554     self,
   2555     store=store,
   2556     chunk_store=chunk_store,
   2557     storage_options=storage_options,
   2558     mode=mode,
   2559     synchronizer=synchronizer,
   2560     group=group,
   2561     encoding=encoding,
   2562     compute=compute,
   2563     consolidated=consolidated,
   2564     append_dim=append_dim,
   2565     region=region,
   2566     safe_chunks=safe_chunks,
   2567     zarr_version=zarr_version,
   2568     write_empty_chunks=write_empty_chunks,
   2569     chunkmanager_store_kwargs=chunkmanager_store_kwargs,
   2570 )

File ~/micromamba/envs/cubed/lib/python3.11/site-packages/xarray/backends/api.py:1721, in to_zarr(dataset, store, chunk_store, mode, synchronizer, group, encoding, compute, consolidated, append_dim, region, safe_chunks, storage_options, zarr_version, write_empty_chunks, chunkmanager_store_kwargs)
   1719 # TODO: figure out how to properly handle unlimited_dims
   1720 dump_to_store(dataset, zstore, writer, encoding=encoding)
-> 1721 writes = writer.sync(
   1722     compute=compute, chunkmanager_store_kwargs=chunkmanager_store_kwargs
   1723 )
   1725 if compute:
   1726     _finalize_store(writes, zstore)

File ~/micromamba/envs/cubed/lib/python3.11/site-packages/xarray/backends/common.py:267, in ArrayWriter.sync(self, compute, chunkmanager_store_kwargs)
    264 if chunkmanager_store_kwargs is None:
    265     chunkmanager_store_kwargs = {}
--> 267 delayed_store = chunkmanager.store(
    268     self.sources,
    269     self.targets,
    270     lock=self.lock,
    271     compute=compute,
    272     flush=True,
    273     regions=self.regions,
    274     **chunkmanager_store_kwargs,
    275 )
    276 self.sources = []
    277 self.targets = []

File ~/micromamba/envs/cubed/lib/python3.11/site-packages/xarray/namedarray/daskmanager.py:249, in DaskManager.store(self, sources, targets, **kwargs)
    241 def store(
    242     self,
    243     sources: Any | Sequence[Any],
    244     targets: Any,
    245     **kwargs: Any,
    246 ) -> Any:
    247     from dask.array import store
--> 249     return store(
    250         sources=sources,
    251         targets=targets,
    252         **kwargs,
    253     )

File ~/micromamba/envs/cubed/lib/python3.11/site-packages/dask/array/core.py:1229, in store(***failed resolving arguments***)
   1227 elif compute:
   1228     store_dsk = HighLevelGraph(layers, dependencies)
-> 1229     compute_as_if_collection(Array, store_dsk, map_keys, **kwargs)
   1230     return None
   1232 else:

File ~/micromamba/envs/cubed/lib/python3.11/site-packages/dask/base.py:403, in compute_as_if_collection(cls, dsk, keys, scheduler, get, **kwargs)
    401 schedule = get_scheduler(scheduler=scheduler, cls=cls, get=get)
    402 dsk2 = optimization_function(cls)(dsk, keys, **kwargs)
--> 403 return schedule(dsk2, keys, **kwargs)

File ~/micromamba/envs/cubed/lib/python3.11/site-packages/dask/threaded.py:90, in get(dsk, keys, cache, num_workers, pool, **kwargs)
     87     elif isinstance(pool, multiprocessing.pool.Pool):
     88         pool = MultiprocessingPoolExecutor(pool)
---> 90 results = get_async(
     91     pool.submit,
     92     pool._max_workers,
     93     dsk,
     94     keys,
     95     cache=cache,
     96     get_id=_thread_get_id,
     97     pack_exception=pack_exception,
     98     **kwargs,
     99 )
    101 # Cleanup pools associated to dead threads
    102 with pools_lock:

File ~/micromamba/envs/cubed/lib/python3.11/site-packages/dask/local.py:427, in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
    425 dsk = dict(dsk)
    426 with local_callbacks(callbacks) as callbacks:
--> 427     _, _, pretask_cbs, posttask_cbs, _ = unpack_callbacks(callbacks)
    428     started_cbs = []
    429     succeeded = False

File ~/micromamba/envs/cubed/lib/python3.11/site-packages/dask/callbacks.py:89, in unpack_callbacks(cbs)
     87 """Take an iterable of callbacks, return a list of each callback."""
     88 if cbs:
---> 89     return [[i for i in f if i] for f in zip(*cbs)]
     90 else:
     91     return [(), (), (), (), ()]

TypeError: zip() argument after * must be an iterable, not RichProgressBar

</details>

Fingers crossed it's a user error!

Pull request description:
Diagnostics documentation
Fixes #538

Adds a page of docs covering `visualize` and progress bars.

Patch:
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
index e91070bf..9dc9f22b 100644
--- a/cubed/core/plan.py
+++ b/cubed/core/plan.py
@@ -1,4 +1,5 @@
 import atexit
+import dataclasses
 import inspect
 import shutil
 import tempfile
@@ -30,6 +31,8 @@ CONTEXT_ID = f"cubed-{datetime.now().strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4()}"
 # Delete local context dirs when Python exits
 CONTEXT_DIRS = set()
 
+Decorator = Callable
+
 
 def delete_on_exit(context_dir: str) -> None:
     if context_dir not in CONTEXT_DIRS and is_local_path(context_dir):
@@ -200,13 +203,45 @@ class Plan:
 
         return dag
 
+    def _compile_blockwise(self, dag, compile_function: Decorator) -> nx.MultiDiGraph:
+        """Compiles functions from all blockwise ops by mutating the input dag."""
+        # Recommended: make a copy of the dag before calling this function.
+
+        compile_with_config = 'config' in inspect.getfullargspec(compile_function).kwonlyargs
+
+        for n in dag.nodes:
+            node = dag.nodes[n]
+
+            if "primitive_op" not in node:
+                continue
+
+            if not isinstance(node["pipeline"].config, BlockwiseSpec):
+                continue
+
+            if compile_with_config:
+                compiled = compile_function(node["pipeline"].config.function, config=node["pipeline"].config)
+            else:
+                compiled = compile_function(node["pipeline"].config.function)
+
+            # node is a blockwise primitive_op.
+            # maybe we should investigate some sort of optics library for frozen dataclasses...
+            new_pipeline = dataclasses.replace(
+                node["pipeline"],
+                config=dataclasses.replace(node["pipeline"].config, function=compiled)
+            )
+            node["pipeline"] = new_pipeline
+
+        return dag
+
     @lru_cache
     def _finalize_dag(
-        self, optimize_graph: bool = True, optimize_function=None
+        self, optimize_graph: bool = True, optimize_function=None, compile_function: Optional[Decorator] = None,
     ) -> nx.MultiDiGraph:
         dag = self.optimize(optimize_function).dag if optimize_graph else self.dag
         # create a copy since _create_lazy_zarr_arrays mutates the dag
         dag = dag.copy()
+        if callable(compile_function):
+            dag = self._compile_blockwise(dag, compile_function)
         dag = self._create_lazy_zarr_arrays(dag)
         return nx.freeze(dag)
 
@@ -216,11 +251,12 @@ class Plan:
         callbacks=None,
         optimize_graph=True,
         optimize_function=None,
+        compile_function=None,
         resume=None,
         spec=None,
         **kwargs,
     ):
-        dag = self._finalize_dag(optimize_graph, optimize_function)
+        dag = self._finalize_dag(optimize_graph, optimize_function, compile_function)
 
         compute_id = f"compute-{datetime.now().strftime('%Y%m%dT%H%M%S.%f')}"
 
diff --git a/docs/images/cubed-add.svg b/docs/images/cubed-add.svg
new file mode 100644
index 00000000..58d2f711
--- /dev/null
+++ b/docs/images/cubed-add.svg
@@ -0,0 +1,132 @@
+<?xml version="1.0" encoding="UTF-8" standalone="no"?>
+<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
+ "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
+<!-- Generated by graphviz version 10.0.1 (20240210.2158)
+ -->
+<!-- Pages: 1 -->
+<svg width="230pt" height="319pt"
+ viewBox="0.00 0.00 229.75 318.75" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
+<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 314.75)">
+<polygon fill="white" stroke="none" points="-4,4 -4,-314.75 225.75,-314.75 225.75,4 -4,4"/>
+<text text-anchor="start" x="8" y="-39.5" font-family="Times,serif" font-size="10.00">num tasks: 5</text>
+<text text-anchor="start" x="8" y="-28.25" font-family="Times,serif" font-size="10.00">max projected memory: 100.0 MB</text>
+<text text-anchor="start" x="8" y="-17" font-family="Times,serif" font-size="10.00">total nbytes written: 72 bytes</text>
+<text text-anchor="start" x="8" y="-5.75" font-family="Times,serif" font-size="10.00">optimized: True</text>
+<!-- op&#45;001 -->
+<g id="node1" class="node">
+<title>op&#45;001</title>
+<g id="a_node1"><a xlink:title="name: op&#45;001&#10;op: asarray&#10;calls: &lt;module&gt; &#45;&gt; asarray&#10;line: 1 in &lt;module&gt;">
+<path fill="none" stroke="black" d="M43.25,-310.75C43.25,-310.75 13.25,-310.75 13.25,-310.75 7.25,-310.75 1.25,-304.75 1.25,-298.75 1.25,-298.75 1.25,-286.75 1.25,-286.75 1.25,-280.75 7.25,-274.75 13.25,-274.75 13.25,-274.75 43.25,-274.75 43.25,-274.75 49.25,-274.75 55.25,-280.75 55.25,-286.75 55.25,-286.75 55.25,-298.75 55.25,-298.75 55.25,-304.75 49.25,-310.75 43.25,-310.75"/>
+<text text-anchor="middle" x="28.25" y="-294.5" font-family="Helvetica,sans-Serif" font-size="10.00">op&#45;001</text>
+<text text-anchor="middle" x="28.25" y="-283.25" font-family="Helvetica,sans-Serif" font-size="10.00">asarray</text>
+</a>
+</g>
+</g>
+<!-- array&#45;001 -->
+<g id="node2" class="node">
+<title>array&#45;001</title>
+<g id="a_node2"><a xlink:title="name: array&#45;001&#10;variable: a&#10;shape: (3, 3)&#10;chunks: (2, 2)&#10;dtype: int64&#10;chunk memory: 32 bytes">
+<polygon fill="none" stroke="black" points="56.5,-238.75 0,-238.75 0,-202.75 56.5,-202.75 56.5,-238.75"/>
+<text text-anchor="middle" x="28.25" y="-222.5" font-family="Helvetica,sans-Serif" font-size="10.00">array&#45;001</text>
+<text text-anchor="middle" x="28.25" y="-211.25" font-family="Helvetica,sans-Serif" font-size="10.00">a</text>
+</a>
+</g>
+</g>
+<!-- op&#45;001&#45;&gt;array&#45;001 -->
+<g id="edge1" class="edge">
+<title>op&#45;001&#45;&gt;array&#45;001</title>
+<path fill="none" stroke="black" d="M28.25,-274.45C28.25,-267.16 28.25,-258.48 28.25,-250.29"/>
+<polygon fill="black" stroke="black" points="31.75,-250.37 28.25,-240.37 24.75,-250.37 31.75,-250.37"/>
+</g>
+<!-- op&#45;003 -->
+<g id="node5" class="node">
+<title>op&#45;003</title>
+<g id="a_node5"><a xlink:title="name: op&#45;003&#10;op: blockwise&#10;projected memory: 100.0 MB&#10;tasks: 4&#10;num input blocks: (1, 1)&#10;calls: &lt;module&gt; &#45;&gt; add &#45;&gt; elemwise &#45;&gt; blockwise&#10;line: 1 in &lt;module&gt;">
+<path fill="#dcbeff" stroke="black" d="M80.25,-166.75C80.25,-166.75 50.25,-166.75 50.25,-166.75 44.25,-166.75 38.25,-160.75 38.25,-154.75 38.25,-154.75 38.25,-137 38.25,-137 38.25,-131 44.25,-125 50.25,-125 50.25,-125 80.25,-125 80.25,-125 86.25,-125 92.25,-131 92.25,-137 92.25,-137 92.25,-154.75 92.25,-154.75 92.25,-160.75 86.25,-166.75 80.25,-166.75"/>
+<text text-anchor="middle" x="65.25" y="-153.25" font-family="Helvetica,sans-Serif" font-size="10.00">op&#45;003</text>
+<text text-anchor="middle" x="65.25" y="-142" font-family="Helvetica,sans-Serif" font-size="10.00">add</text>
+<text text-anchor="middle" x="65.25" y="-130.75" font-family="Helvetica,sans-Serif" font-size="10.00">tasks: 4</text>
+</a>
+</g>
+</g>
+<!-- array&#45;001&#45;&gt;op&#45;003 -->
+<g id="edge2" class="edge">
+<title>array&#45;001&#45;&gt;op&#45;003</title>
+<path fill="none" stroke="black" d="M37.02,-202.48C40.85,-194.94 45.45,-185.87 49.82,-177.26"/>
+<polygon fill="black" stroke="black" points="52.88,-178.97 54.28,-168.47 46.64,-175.8 52.88,-178.97"/>
+</g>
+<!-- op&#45;002 -->
+<g id="node3" class="node">
+<title>op&#45;002</title>
+<g id="a_node3"><a xlink:title="name: op&#45;002&#10;op: asarray&#10;calls: &lt;module&gt; &#45;&gt; asarray&#10;line: 1 in &lt;module&gt;">
+<path fill="none" stroke="black" d="M118.25,-310.75C118.25,-310.75 88.25,-310.75 88.25,-310.75 82.25,-310.75 76.25,-304.75 76.25,-298.75 76.25,-298.75 76.25,-286.75 76.25,-286.75 76.25,-280.75 82.25,-274.75 88.25,-274.75 88.25,-274.75 118.25,-274.75 118.25,-274.75 124.25,-274.75 130.25,-280.75 130.25,-286.75 130.25,-286.75 130.25,-298.75 130.25,-298.75 130.25,-304.75 124.25,-310.75 118.25,-310.75"/>
+<text text-anchor="middle" x="103.25" y="-294.5" font-family="Helvetica,sans-Serif" font-size="10.00">op&#45;002</text>
+<text text-anchor="middle" x="103.25" y="-283.25" font-family="Helvetica,sans-Serif" font-size="10.00">asarray</text>
+</a>
+</g>
+</g>
+<!-- array&#45;002 -->
+<g id="node4" class="node">
+<title>array&#45;002</title>
+<g id="a_node4"><a xlink:title="name: array&#45;002&#10;variable: b&#10;shape: (3, 3)&#10;chunks: (2, 2)&#10;dtype: int64&#10;chunk memory: 32 bytes">
+<polygon fill="none" stroke="black" points="131.5,-238.75 75,-238.75 75,-202.75 131.5,-202.75 131.5,-238.75"/>
+<text text-anchor="middle" x="103.25" y="-222.5" font-family="Helvetica,sans-Serif" font-size="10.00">array&#45;002</text>
+<text text-anchor="middle" x="103.25" y="-211.25" font-family="Helvetica,sans-Serif" font-size="10.00">b</text>
+</a>
+</g>
+</g>
+<!-- op&#45;002&#45;&gt;array&#45;002 -->
+<g id="edge3" class="edge">
+<title>op&#45;002&#45;&gt;array&#45;002</title>
+<path fill="none" stroke="black" d="M103.25,-274.45C103.25,-267.16 103.25,-258.48 103.25,-250.29"/>
+<polygon fill="black" stroke="black" points="106.75,-250.37 103.25,-240.37 99.75,-250.37 106.75,-250.37"/>
+</g>
+<!-- array&#45;002&#45;&gt;op&#45;003 -->
+<g id="edge4" class="edge">
+<title>array&#45;002&#45;&gt;op&#45;003</title>
+<path fill="none" stroke="black" d="M94.24,-202.48C90.31,-194.94 85.58,-185.87 81.09,-177.26"/>
+<polygon fill="black" stroke="black" points="84.24,-175.71 76.51,-168.47 78.03,-178.95 84.24,-175.71"/>
+</g>
+<!-- array&#45;003 -->
+<g id="node6" class="node">
+<title>array&#45;003</title>
+<g id="a_node6"><a xlink:title="name: array&#45;003&#10;variable: c&#10;shape: (3, 3)&#10;chunks: (2, 2)&#10;dtype: int64&#10;chunk memory: 32 bytes&#10;nbytes: 72 bytes">
+<polygon fill="#ffd8b1" stroke="black" points="93.5,-89 37,-89 37,-53 93.5,-53 93.5,-89"/>
+<text text-anchor="middle" x="65.25" y="-72.75" font-family="Helvetica,sans-Serif" font-size="10.00">array&#45;003</text>
+<text text-anchor="middle" x="65.25" y="-61.5" font-family="Helvetica,sans-Serif" font-size="10.00">c</text>
+</a>
+</g>
+</g>
+<!-- op&#45;003&#45;&gt;array&#45;003 -->
+<g id="edge5" class="edge">
+<title>op&#45;003&#45;&gt;array&#45;003</title>
+<path fill="none" stroke="black" d="M65.25,-124.58C65.25,-117.19 65.25,-108.7 65.25,-100.73"/>
+<polygon fill="black" stroke="black" points="68.75,-100.74 65.25,-90.74 61.75,-100.74 68.75,-100.74"/>
+</g>
+<!-- create&#45;arrays -->
+<g id="node7" class="node">
+<title>create&#45;arrays</title>
+<g id="a_node7"><a xlink:title="name: create&#45;arrays&#10;op: create&#45;arrays&#10;projected memory: 100.0 MB&#10;tasks: 1">
+<path fill="none" stroke="black" d="M209.75,-310.75C209.75,-310.75 160.75,-310.75 160.75,-310.75 154.75,-310.75 148.75,-304.75 148.75,-298.75 148.75,-298.75 148.75,-286.75 148.75,-286.75 148.75,-280.75 154.75,-274.75 160.75,-274.75 160.75,-274.75 209.75,-274.75 209.75,-274.75 215.75,-274.75 221.75,-280.75 221.75,-286.75 221.75,-286.75 221.75,-298.75 221.75,-298.75 221.75,-304.75 215.75,-310.75 209.75,-310.75"/>
+<text text-anchor="middle" x="185.25" y="-294.5" font-family="Helvetica,sans-Serif" font-size="10.00">create&#45;arrays</text>
+<text text-anchor="middle" x="185.25" y="-283.25" font-family="Helvetica,sans-Serif" font-size="10.00">tasks: 1</text>
+</a>
+</g>
+</g>
+<!-- arrays -->
+<g id="node8" class="node">
+<title>arrays</title>
+<g id="a_node8"><a xlink:title="name: arrays" target="None">
+<polygon fill="none" stroke="black" points="212.25,-238.75 158.25,-238.75 158.25,-202.75 212.25,-202.75 212.25,-238.75"/>
+<text text-anchor="middle" x="185.25" y="-216.88" font-family="Helvetica,sans-Serif" font-size="10.00">arrays</text>
+</a>
+</g>
+</g>
+<!-- create&#45;arrays&#45;&gt;arrays -->
+<g id="edge6" class="edge">
+<title>create&#45;arrays&#45;&gt;arrays</title>
+<path fill="none" stroke="black" d="M185.25,-274.45C185.25,-267.16 185.25,-258.48 185.25,-250.29"/>
+<polygon fill="black" stroke="black" points="188.75,-250.37 185.25,-240.37 181.75,-250.37 188.75,-250.37"/>
+</g>
+</g>
+</svg>
diff --git a/docs/user-guide/diagnostics.md b/docs/user-guide/diagnostics.md
new file mode 100644
index 00000000..3e013323
--- /dev/null
+++ b/docs/user-guide/diagnostics.md
@@ -0,0 +1,66 @@
+# Diagnostics
+
+Cubed provides a variety of tools to understand a computation before running it, to monitor its progress while running, and to view performance statistics after it has completed.
+
+To use these features ensure that the optional dependencies for diagnostics have been installed:
+
+```shell
+python -m pip install "cubed[diagnostics]"
+```
+
+## Visualize the computation plan
+
+Before running a computation, Cubed will create an internal plan that it uses to compute the output arrays.
+
+The plan is a directed acyclic graph (DAG), and it can be useful to visualize it to see the number of steps involved in your computation, the number of tasks in each step (and overall), and the amount of intermediate data written out.
+
+The {py:meth}`Array.visualize() <cubed.Array.visualize()>` method on an array creates an image of the DAG. By default it is saved in a file called *cubed.svg* in the current working directory, but the filename and format can be changed if needed. If running in a Jupyter notebook the image will be rendered in the notebook.
+
+If you are computing multiple arrays at once, then there is a {py:func}`visualize <cubed.visualize>` function that takes multiple array arguments.
+
+This example shows a tiny computation and the resulting plan:
+
+```python
+import cubed.array_api as xp
+import cubed.random
+
+a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2))
+b = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2))
+c = xp.add(a, b)
+
+c.visualize()
+```
+
+![Cubed visualization of a tiny computation](../images/cubed-add.svg)
+
+There are two type of nodes in the plan. Boxes with rounded corners are operations, while boxes with square corners are arrays.
+
+In this case there are three operations (labelled `op-001`, `op-002`, and `op-003`), which produce the three arrays `a`, `b`, and `c`. (There is always an additional operation called `create-arrays`, shown on the right, which Cubed creates automatically.)
+
+Array `c` is coloured orange, which means it is materialized as a Zarr array. Arrays `a` and `b` do not need to be materialized as Zarr arrays since they are small constant arrays that are passed to the workers running the tasks.
+
+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
+
+You can display a progress bar to track your computation by passing callbacks to {py:meth}`compute() <cubed.Array.compute()>`:
+
+```ipython
+>>> from cubed.diagnostics.rich import RichProgressBar
+>>> progress = RichProgressBar()
+>>> c.compute(callbacks=[progress])  # c is the array from above
+  create-arrays 1/1 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:00
+  op-003 add    4/4 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:00
+```
+
+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.
+
+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.)
+
+See the [examples](https://github.com/cubed-dev/cubed/blob/main/examples/README.md) for more information about how to use them.
diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md
index 90e3f90f..22fe6a95 100644
--- a/docs/user-guide/index.md
+++ b/docs/user-guide/index.md
@@ -11,4 +11,5 @@ storage
 memory
 reliability
 scaling
+diagnostics
 ```
diff --git a/docs/user-guide/scaling.md b/docs/user-guide/scaling.md
index e978b7aa..92804d22 100644
--- a/docs/user-guide/scaling.md
+++ b/docs/user-guide/scaling.md
@@ -78,20 +78,7 @@ Different cloud providers' serverless offerings may perform differently. For exa
 
 ## Diagnosing Performance
 
-To understand how your computation could perform better you first need to diagnose the source of any problems.
-
-### Optimized Plan
-
-Use {py:meth}`Plan.visualize() <cubed.Plan.visualize()>` to view the optimized plan. This allows you to see the number of steps involved in your calculation, the number of tasks in each step, and overall.
-
-### History Callback
-
-The history callback function can help determine how much time was spent in worker startup, as well as how much stragglers affected the overall speed.
-
-### Timeline Visualization Callback
-
-A timeline visualization callback can provide a visual representation of the above points. Ideally, we want vertical lines on this plot, which would represent perfect horizontal scaling.
-
+See <project:diagnostics.md>.
 
 ## Tips
 


Test patch:
diff --git a/cubed/tests/test_executor_features.py b/cubed/tests/test_executor_features.py
index f6bb5bd6..6a8dad86 100644
--- a/cubed/tests/test_executor_features.py
+++ b/cubed/tests/test_executor_features.py
@@ -315,3 +315,55 @@ def test_check_runtime_memory_processes(spec, executor):
 
     # OK if we use fewer workers
     c.compute(executor=executor, max_workers=max_workers // 2)
+
+
+COMPILE_FUNCTIONS = [lambda fn: fn]
+
+try:
+    from numba import jit as numba_jit
+    COMPILE_FUNCTIONS.append(numba_jit)
+except ModuleNotFoundError:
+    pass
+
+try:
+    if 'jax' in os.environ.get('CUBED_BACKEND_ARRAY_API_MODULE', ''):
+        from jax import jit as jax_jit
+        COMPILE_FUNCTIONS.append(jax_jit)
+except ModuleNotFoundError:
+    pass
+
+
+@pytest.mark.parametrize("compile_function", COMPILE_FUNCTIONS)
+def test_check_compilation(spec, executor, compile_function):
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
+    c = xp.add(a, b)
+    assert_array_equal(
+        c.compute(executor=executor, compile_function=compile_function), np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]])
+    )
+
+
+def test_compilation_can_fail(spec, executor):
+    def compile_function(func):
+        raise NotImplementedError(f"Cannot compile {func}")
+
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
+    c = xp.add(a, b)
+    with pytest.raises(NotImplementedError) as excinfo:
+        c.compute(executor=executor, compile_function=compile_function)
+    
+    assert "add" in str(excinfo.value), "Compile function was applied to add operation."
+
+
+def test_compilation_with_config_can_fail(spec, executor):
+    def compile_function(func, *, config=None):
+        raise NotImplementedError(f"Cannot compile {func} with {config}")
+
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
+    c = xp.add(a, b)
+    with pytest.raises(NotImplementedError) as excinfo:
+        c.compute(executor=executor, compile_function=compile_function)
+        
+    assert "BlockwiseSpec" in str(excinfo.value), "Compile function was applied with a config argument."