Improve rechunk stats and visualizations (#934)

Author: tomwhite

cubed/array_api/array_object.py

@@ -22,7 +22,7 @@
 from cubed.backend_array_api import namespace as nxp
 from cubed.core.array import CoreArray
 from cubed.core.ops import elemwise
-from cubed.utils import itemsize, memory_repr
+from cubed.utils import memory_repr
 
 ARRAY_SVG_SIZE = (
     120  # cubed doesn't have a config module like dask does so hard-code this for now
@@ -56,7 +56,7 @@ def _repr_html_(self):
 
         if not math.isnan(self.nbytes):
             nbytes = memory_repr(self.nbytes)
-            cbytes = memory_repr(math.prod(self.chunksize) * itemsize(self.dtype))
+            cbytes = memory_repr(self.cbytes)
         else:
             nbytes = "unknown"
             cbytes = "unknown"

cubed/core/array.py

@@ -1,3 +1,4 @@
+import math
 from functools import reduce
 from operator import mul
 from typing import Literal, TypeVar
@@ -103,6 +104,11 @@ def nbytes(self) -> int:
         """Number of bytes in array"""
         return self.size * itemsize(self.dtype)
 
+    @property
+    def cbytes(self) -> int:
+        """Number of bytes in a chunk"""
+        return math.prod(self.chunksize) * itemsize(self.dtype)
+
     @property
     def nchunks(self) -> int:
         """Number of chunks in array"""

cubed/core/rechunk.py

@@ -3,16 +3,15 @@
 import logging
 import warnings
 from collections.abc import Sequence
-from dataclasses import dataclass
-from math import floor, prod
+from dataclasses import dataclass, field
+from math import ceil, floor, prod
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 from cubed.core.plan import FinalizedPlan
 from cubed.types import T_RegularChunks, T_Shape
-from cubed.utils import (
-    normalize_chunks,
-)
+from cubed.utils import normalize_chunks
 from cubed.vendor.rechunker.algorithm import (
     MAX_STAGES,
     ExcessiveIOWarning,
@@ -22,6 +21,9 @@
     consolidate_chunks,
 )
 
+if TYPE_CHECKING:
+    from cubed.array_api.array_object import Array
+
 logger = logging.getLogger(__name__)
 
 
@@ -253,19 +255,46 @@ class RechunkCopy:
     target_chunks: T_RegularChunks
     """The chunks of the target array for this copy operation."""
 
-    source_aligned: bool | None = None
+    num_input_blocks: int = field(init=False)
+    """The number of input blocks read from the source."""
+
+    num_output_blocks: int = field(init=False)
+    """The number of output blocks written to the target."""
+
+    source_aligned: bool = field(init=False)
     """Are copy chunks aligned with source chunks?"""
 
-    target_aligned: bool | None = None
+    target_aligned: bool = field(init=False)
     """
     Are copy chunks aligned with target chunks?
     If not then irregular chunking must be used.
     """
 
+    def __post_init__(self):
+        self.num_input_blocks = prod(
+            ceil(cc / sc) for cc, sc in zip(self.copy_chunks, self.source_chunks)
+        )
+        self.num_output_blocks = prod(
+            ceil(cc / tc) for cc, tc in zip(self.copy_chunks, self.target_chunks)
+        )
+        self.source_aligned = all(
+            (cc == n) or (cc % sc == 0)
+            for n, sc, cc in zip(self.shape, self.source_chunks, self.copy_chunks)
+        )
+        self.target_aligned = all(
+            (cc == n) or (cc % tc == 0)
+            for n, tc, cc in zip(self.shape, self.target_chunks, self.copy_chunks)
+        )
+
 
 @dataclass
 class RechunkPlan:
+    arrays: list["Array"]
     copy_ops: list[RechunkCopy]
+    regular: bool = field(init=False)
+
+    def __post_init__(self):
+        self.regular = not any(not copy_op.target_aligned for copy_op in self.copy_ops)
 
     def _repr_html_(self):
         from cubed.diagnostics.widgets import get_template
@@ -275,20 +304,60 @@ def _repr_html_(self):
         for copy_op in self.copy_ops:
             row = (
                 copy_op,
-                svg(normalize_chunks(copy_op.source_chunks, shape=copy_op.copy_chunks)),
-                svg(normalize_chunks(copy_op.target_chunks, shape=copy_op.copy_chunks)),
+                svg(
+                    normalize_chunks(copy_op.source_chunks, shape=copy_op.copy_chunks),
+                    size=80,
+                ),
+                svg(
+                    normalize_chunks(copy_op.target_chunks, shape=copy_op.copy_chunks),
+                    size=80,
+                ),
             )
             table.append(row)
 
         return get_template("rechunk_plan.j2").render(table=table)
 
+    @property
+    def result(self):
+        return self.arrays[-1]
+
+    @property
+    def array_view(self):
+        return RechunkPlanArrayView(self.arrays)
+
+    @property
+    def stats(self):
+        return RechunkPlanStats.from_plan(self, self.result.plan())
+
+
+@dataclass
+class RechunkPlanArrayView:
+    arrays: list["Array"]
+
+    def _repr_html_(self):
+        from cubed.diagnostics.widgets import get_template
+
+        table = []
+        for array in self.arrays:
+            row = (
+                array,
+                array.to_svg(size=120),
+            )
+            table.append(row)
+
+        return get_template("rechunk_arrays.j2").render(table=table)
+
 
 @dataclass
 class RechunkPlanStats:
     num_copy_ops: int
     num_tasks: int
     total_nbytes_written: int
+    total_task_iops: int
     max_task_iops: int
+    max_task_input_blocks: int
+    max_task_output_blocks: int
+    regular: bool
 
     @classmethod
     def from_plan(cls, rechunk_plan: RechunkPlan, plan: FinalizedPlan):
@@ -302,23 +371,51 @@ def from_plan(cls, rechunk_plan: RechunkPlan, plan: FinalizedPlan):
             + d["pipeline"].config.num_output_blocks[0]
             for _, d in rechunks
         ]
+        total_task_iops = sum(
+            d["primitive_op"].num_tasks
+            * (
+                d["pipeline"].config.num_input_blocks[0]
+                + d["pipeline"].config.num_output_blocks[0]
+            )
+            for _, d in rechunks
+        )
+        num_task_input_blocks = [
+            d["pipeline"].config.num_input_blocks[0] for _, d in rechunks
+        ]
+        num_task_output_blocks = [
+            d["pipeline"].config.num_output_blocks[0] for _, d in rechunks
+        ]
 
         return cls(
             num_copy_ops=len(rechunk_plan.copy_ops),
             num_tasks=plan.num_tasks,
             total_nbytes_written=plan.total_nbytes_written,
+            total_task_iops=total_task_iops,
             max_task_iops=max(num_task_iops),
+            max_task_input_blocks=max(num_task_input_blocks),
+            max_task_output_blocks=max(num_task_output_blocks),
+            regular=rechunk_plan.regular,
         )
 
 
 def rechunk_plan(x, chunks, *, min_mem=None, allow_irregular=True):
     from cubed.core.ops import _rechunk_plan
 
+    out = x
+    arrays = [x]
     copy_ops = []
     source_chunks = x.chunksize
     for copy_chunks, target_chunks in _rechunk_plan(
         x, chunks, min_mem=min_mem, allow_irregular=allow_irregular
     ):
+        from .ops import _rechunk
+
+        out = _rechunk(out, copy_chunks, target_chunks, allow_irregular=allow_irregular)
+        arrays.append(out)
         copy_ops.append(RechunkCopy(x.shape, source_chunks, copy_chunks, target_chunks))
         source_chunks = target_chunks
-    return RechunkPlan(copy_ops)
+
+    # final output must have regular chunking
+    assert copy_ops[-1].target_aligned
+
+    return RechunkPlan(arrays, copy_ops)

cubed/diagnostics/widgets/templates/rechunk_arrays.j2

@@ -0,0 +1,44 @@
+<table>
+    <tbody>
+    {% for row in table %}
+        <tr>
+            <td>
+                <table>
+                    <thead>
+                        <tr>
+                            <td> </td>
+                            <th> Array </th>
+                            <th> Chunk </th>
+                        </tr>
+                    </thead>
+                    <tbody>
+                        <tr>
+                            <th> Bytes </th>
+                            <td> {{ row[0].nbytes | memory_repr }} </td>
+                            <td> {{ row[0].cbytes | memory_repr }} </td>
+                        </tr>
+                        <tr>
+                            <th> Shape </th>
+                            <td> {{ row[0].shape }} </td>
+                            <td> {{ row[0].chunksize }} </td>
+                        </tr>
+                        <tr>
+                            <th> Count </th>
+                            <td> {{ 1 }} </td>
+                            <td> {{ row[0].npartitions }} Chunks </td>
+                        </tr>
+                        <tr>
+                            <th> Type </th>
+                            <td> {{ row[0].dtype }} </td>
+                            <td> {{ arrtype }} </td>
+                        </tr>
+                    </tbody>
+                </table>
+            </td>
+            <td>
+            {{row[1]}}
+            </td>
+        </tr>
+    {% endfor %}
+    </tbody>
+</table>

cubed/diagnostics/widgets/templates/rechunk_plan.j2

@@ -32,6 +32,18 @@
                             <th> Target chunks </th>
                             <td> {{ row[0].target_chunks }} </td>
                         </tr>
+                        <tr>
+                            <th> Num input blocks </th>
+                            <td> {{ row[0].num_input_blocks }} </td>
+                        </tr>
+                        <tr>
+                            <th> Num output blocks </th>
+                            <td> {{ row[0].num_output_blocks }} </td>
+                        </tr>
+                        <tr>
+                            <th> Target aligned </th>
+                            <td> {{ row[0].target_aligned }} </td>
+                        </tr>
                     </tbody>
                 </table>
             </td>

cubed/tests/test_rechunk.py

@@ -6,7 +6,6 @@
 from cubed.backend_array_api import namespace as nxp
 from cubed.core.ops import _store_array, split_chunksizes
 from cubed.core.rechunk import (
-    RechunkPlanStats,
     calculate_regular_stage_chunks,
     multistage_regular_rechunking_plan,
     multspace,
@@ -178,11 +177,10 @@ def test_rechunk_plan_viz():
     assert len(rplan.copy_ops) == 2
     # check generating the repr doesn't raise an exception
     rplan._repr_html_()
+    rplan.array_view._repr_html_()
 
     # check stats (without running the rechunk)
-    b = a.rechunk(target_chunks)
-    plan = b.plan()
-    stats = RechunkPlanStats.from_plan(rplan, plan)
+    stats = rplan.stats
     assert stats.num_copy_ops == 2
     assert stats.max_task_iops == 23