diff --git a/src/access_moppy/base.py b/src/access_moppy/base.py index 5a51f752..07a31fd1 100644 --- a/src/access_moppy/base.py +++ b/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. diff --git a/tests/unit/test_base.py b/tests/unit/test_base.py index 7fa5bd2d..8eca4f1c 100644 --- a/tests/unit/test_base.py +++ b/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( diff --git a/tests/unit/test_base_loading_chunker.py b/tests/unit/test_base_loading_chunker.py index 859f62cb..c3f984fd 100644 --- a/tests/unit/test_base_loading_chunker.py +++ b/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