DatasetChunker stops write slices at the 4MB minimum instead of filling toward the max

src/access_moppy/base.py

@@ -219,7 +219,7 @@ class DatasetChunker:
     Rules:
     - Time coordinates: one Dask chunk
     - Time bounds: one Dask chunk
-    - Data variables: target at least 4MB without exceeding 128MB per task
+    - Data variables: batch time steps toward 128MB per task, never below 4MB
     - Spatial dimensions: split when a full spatial slab exceeds 128MB
     """
 
@@ -260,7 +260,6 @@ def calculate_chunk_size_for_variable(self, var: xr.DataArray) -> Dict[str, int]
 
         # Calculate total elements per chunk needed for minimum target size
         element_size = var.dtype.itemsize
-        min_target_elements = self.target_chunk_size_bytes // element_size
         max_target_elements = self.max_chunk_size_bytes // element_size
 
         # For time-dependent variables, start with time dimension
@@ -274,27 +273,19 @@ def calculate_chunk_size_for_variable(self, var: xr.DataArray) -> Dict[str, int]
                     other_elements *= var.sizes[dim]
 
             if other_elements > 0:
-                # How many time steps fit under the max bound -- at least 1,
-                # even if a single step alone already exceeds it (the
+                # Batch as many time steps as fit under the max bound -- at
+                # least 1, even if a single step alone already exceeds it (the
                 # spatial/vertical splitting loop below handles that case).
-                max_time_steps = max(1, max_target_elements // other_elements)
-
-                if other_elements >= min_target_elements:
-                    # A single time step already meets the minimum target.
-                    # Previously this always fell through to 1 step/task
-                    # even when there was headroom left under the max bound
-                    # (e.g. a few-MB/step 3D field, well under 128MB),
-                    # inflating the write-task count for no reason. Batch
-                    # as many steps together as fit under the max instead.
-                    time_chunks = max_time_steps
-                else:
-                    # Multiple steps are needed to reach the minimum
-                    # target; grow toward it, capped at the max bound.
-                    min_time_steps = max(
-                        1,
-                        (min_target_elements + other_elements - 1) // other_elements,
-                    )  # Ceiling division
-                    time_chunks = min(min_time_steps, max_time_steps)
+                #
+                # Growing only as far as ``target_chunk_size_mb`` -- the
+                # *minimum* task size, not a goal -- left every daily variable
+                # writing ~4MB slices whatever the max allowed: 730 slices per
+                # output file for a p19 3-D field, 39 for a surface 2-D one.
+                # Each slice costs a serial graph-cull/submit/gather round trip
+                # in the main process, so the workers idled while that round
+                # trip repeated. The target stays a floor -- any task under the
+                # max bound still clears it.
+                time_chunks = max(1, max_target_elements // other_elements)
 
                 # Don't exceed available time steps
                 time_chunks = min(time_size, time_chunks)
@@ -2248,16 +2239,31 @@ def iter_slices():
             return
 
         pending = deque()
+        # Where the wall time of this loop actually goes. Only the `wait`
+        # segment is work the Dask workers do in parallel; `submit` (graph cull
+        # + serialisation) and `write` (gather + netCDF assignment) run in this
+        # process alone, so a job whose slices are small enough for `submit` to
+        # dominate cannot use more than about one core no matter how many
+        # workers it was given.
+        n_slices = 0
+        submit_s = wait_s = write_s = 0.0
 
         def write_next():
+            nonlocal wait_s, write_s
             slices, future = pending.popleft()
             try:
-                destination[slices] = future.result()
+                t_wait = time.perf_counter()
+                block = future.result()
+                t_write = time.perf_counter()
+                destination[slices] = block
+                wait_s += t_write - t_wait
+                write_s += time.perf_counter() - t_write
             finally:
                 future.release()
 
         try:
             for slices in iter_slices():
+                t_submit = time.perf_counter()
                 indexers = dict(zip(vdat.dims, slices))
                 sliced_data = vdat.isel(indexers).data
                 culled_graph = sliced_data.dask.cull(
@@ -2277,6 +2283,8 @@ def write_next():
                     sliced_data,
                     optimize_graph=False,
                 )
+                submit_s += time.perf_counter() - t_submit
+                n_slices += 1
                 pending.append((slices, future))
                 if len(pending) >= self.write_prefetch:
                     write_next()
@@ -2287,6 +2295,24 @@ def write_next():
             for _, future in pending:
                 future.release()
 
+        slice_mb = (
+            vdat.dtype.itemsize
+            * math.prod(
+                min(int(chunk_sizes[dim]), vdat.sizes[dim]) for dim in vdat.dims
+            )
+            / (1024 * 1024)
+        )
+        logger.info(
+            "Wrote %d slices of %.2fMB with write_prefetch=%d: "
+            "submit %.1fs, wait %.1fs, write %.1fs",
+            n_slices,
+            slice_mb,
+            self.write_prefetch,
+            submit_s,
+            wait_s,
+            write_s,
+        )
+
     def _write_single(self):
         """
         Write the CMORised dataset to an intermediate NetCDF4 file.
@@ -2812,6 +2838,7 @@ def _repack_cmip7_output(self, path: Path):
             "Repacking CMIP7 output with cmip7repack (-d %d): %s", chunk_target, path
         )
 
+        repack_start = time.perf_counter()
         try:
             subprocess.run(  # noqa: S603  # nosec B603
                 cmd,
@@ -2832,6 +2859,14 @@ def _repack_cmip7_output(self, path: Path):
             )
             raise
 
+        # cmip7repack is a single-threaded external process, so its share of
+        # the job's wall time is time no worker can be busy for.
+        logger.info(
+            "cmip7repack finished in %.1fs: %s",
+            time.perf_counter() - repack_start,
+            path,
+        )
+
         # cmip7repack exits 0 even when it rechunks nothing, so the exit
         # status alone would record a silent no-op as a pass. Read the
         # filters back off the file to see what actually happened.

tests/unit/test_base.py

@@ -1855,6 +1855,56 @@ def compute(self, array, *, optimize_graph):
         assert events.count("release") == 3
         np.testing.assert_array_equal(written, vdat.compute().values)
 
+    @pytest.mark.unit
+    def test_write_dask_slices_logs_time_breakdown(
+        self, cmoriser_with_dask_dataset, caplog
+    ):
+        """The write loop reports slice count, slice size and where time went.
+
+        Only the ``wait`` segment is parallel worker time; ``submit`` and
+        ``write`` are serial in this process, so the split is what tells a
+        low-utilisation job apart from a genuinely worker-bound one.
+        """
+        source = da.from_array(np.arange(12).reshape(3, 2, 2), chunks=(1, 2, 2))
+        vdat = xr.DataArray(source, dims=("time", "lat", "lon"))
+        written = np.empty(vdat.shape, dtype=vdat.dtype)
+
+        class Destination:
+            def __setitem__(self, slices, values):
+                written[slices] = values
+
+        class Future:
+            def __init__(self, array):
+                self.array = array
+
+            def result(self):
+                return self.array.compute(scheduler="synchronous")
+
+            def release(self):
+                pass
+
+        class Client:
+            def compute(self, array, *, optimize_graph):
+                return Future(array)
+
+        cmoriser_with_dask_dataset.write_prefetch = 2
+        with caplog.at_level(logging.INFO, logger="access_moppy.base"):
+            with patch("access_moppy.base.get_client", return_value=Client()):
+                cmoriser_with_dask_dataset._write_dask_slices(
+                    Destination(),
+                    vdat,
+                    {"time": 1, "lat": 2, "lon": 2},
+                )
+
+        summaries = [r for r in caplog.records if r.getMessage().startswith("Wrote ")]
+        assert len(summaries) == 1
+        n_slices, slice_mb, prefetch, submit_s, wait_s, write_s = summaries[0].args
+        assert n_slices == 3
+        # int64 x 1x2x2 elements per slice
+        assert slice_mb == pytest.approx(32 / (1024 * 1024))
+        assert prefetch == 2
+        assert min(submit_s, wait_s, write_s) >= 0.0
+
     @pytest.mark.unit
     @pytest.mark.parametrize("write_prefetch", [0, -1])
     def test_write_prefetch_must_be_positive(

tests/unit/test_base_loading_chunker.py

@@ -100,7 +100,9 @@ def test_dataset_chunker_preserves_spatial_slab_when_below_maximum():
 
     chunks = chunker.calculate_chunk_size_for_variable(var)
 
-    assert chunks == {"time": 11, "lev": 10, "j": 100, "i": 100}
+    # 12 steps x 10x100x100 float32 = 4.58MB, still under the 32MB max, so the
+    # whole time axis batches into one task and no spatial dimension is split.
+    assert chunks == {"time": 12, "lev": 10, "j": 100, "i": 100}
 
 
 @pytest.mark.unit
@@ -127,6 +129,33 @@ def test_dataset_chunker_batches_multiple_steps_when_one_step_exceeds_target():
     assert chunks["time"] > 1
 
 
+@pytest.mark.unit
+def test_dataset_chunker_fills_toward_max_when_one_step_is_under_target():
+    """A step smaller than the target must still batch toward the max.
+
+    ``atmos.ta.tavg-p19-hxy-air.day`` is 19 levels on a 145x192 grid --
+    2.02MB/step, under the 4MB target -- so growing only as far as the target
+    gave 2-step/4MB tasks: 730 write slices for one 4-year output file. Each
+    slice is a serial round trip in the main process, which is what held the
+    batch jobs to ~2 busy cores out of 7-18. Filling toward the 128MB max
+    gives 63-step tasks and 24 slices for the same file.
+    """
+    chunker = DatasetChunker(target_chunk_size_mb=4, max_chunk_size_mb=128)
+    var = xr.DataArray(
+        da.empty((1460, 19, 145, 192), dtype=np.float32),
+        dims=("time", "plev", "lat", "lon"),
+    )
+
+    chunks = chunker.calculate_chunk_size_for_variable(var)
+    chunk_bytes = np.dtype(var.dtype).itemsize * np.prod(list(chunks.values()))
+
+    assert chunks["time"] == 63
+    assert chunks["plev"] == 19 and chunks["lat"] == 145 and chunks["lon"] == 192
+    assert chunk_bytes <= 128 * 1024 * 1024
+    # Well past the 4MB target the old calculation stopped at.
+    assert chunk_bytes > 100 * 1024 * 1024
+
+
 @pytest.mark.unit
 def test_dataset_chunker_single_step_over_max_still_clamps_to_one():
     """Complementary case: when a single time step alone already exceeds