Optimise memory allocation in batchprocessing (#501)

* update batch_processing dask config

* update dask defination in batch processing script

* pre-commit fix

* Batch Processing Dask Optimisation — Version 2

Author: rhaegar325

src/access_moppy/base.py

@@ -486,10 +486,19 @@ def _preprocess(ds):
                 # combine when multiple required variables are requested.
                 prefer_by_coords = bool(required_vars and len(required_vars) > 1)
 
+                # One dask chunk per file along time. Sub-monthly inputs (e.g.
+                # daily tasmax/tasmin) are stored with per-timestep on-disk HDF5
+                # chunking; the default chunks={} inherits that, so a 31-day file
+                # becomes 31 chunks and the task graph explodes (~650k tasks over
+                # a multi-decade run), which is what drives the distributed
+                # workers out of memory on those variables. Collapsing to one
+                # chunk per file cuts the graph ~26x and the compute memory ~5x
+                # with bit-identical results; monthly inputs (time size 1) are
+                # unaffected.
                 common_kwargs = {
                     "engine": "netcdf4",
                     "decode_cf": False,
-                    "chunks": {},
+                    "chunks": {"time": -1},
                     "data_vars": "minimal",
                     "coords": "minimal",
                     "compat": "override",

src/access_moppy/derivations/calc_utils.py

@@ -233,7 +233,11 @@ def _iter_markers(value):
         if np.isfinite(marker):
             # Match both exact values and float32-rounded encodings (e.g. 1e20).
             atol = max(1e-12, abs(float(np.spacing(np.float32(marker)))))
-            condition = np.isclose(da, marker, rtol=0.0, atol=atol)
+            # Use abs(da - marker) <= atol, not np.isclose: np.isclose is not a
+            # ufunc and eagerly computes a Dask-backed array to NumPy (loading
+            # the whole series into memory). This form is Dask-native/lazy and,
+            # with rtol=0, is equivalent to np.isclose(da, marker, atol=atol).
+            condition = abs(da - marker) <= atol
         else:
             condition = da == marker
         mask = condition if mask is None else (mask | condition)

src/access_moppy/executors/dask_config.py

@@ -31,15 +31,21 @@
 # Per-worker memory floors (GB) for the three intensity tiers. These are the
 # minimum RAM a single worker needs to process a variable of that class without
 # being killed -- a property of the *workload*, not the node -- so they are
-# absolute GB, not fractions of the allocation. The defaults are calibrated from
-# real runs (daily variables were killed at ~11GB/worker; carbon pools paused
-# hard at ~11GB; light variables were fine), and each can be overridden per job
+# absolute GB, not fractions of the allocation. Each can be overridden per job
 # via an environment variable set in the batch config's worker_init, so nothing
 # is locked in code.
+#
+# Calibrated from full-scale (5352-file / ~446-year) runs once the pipeline was
+# made fully streaming (lazy missing-value masks + chunked reads/writes). At that
+# point measured per-worker peaks were ~12.5GB (daily), ~10GB (carbon pools) and
+# ~6GB (light single-field), all with zero memory-pressure warnings, so the
+# floors sit ~1.3-1.5x above the observed peaks. Before streaming, daily
+# variables materialised the whole ~18GB series into one worker and needed a
+# 28GB floor; that is no longer the case.
 _FLOOR_ENV = {
     "light": ("MOPPY_WORKER_GB_LIGHT", 12),
-    "medium": ("MOPPY_WORKER_GB_MEDIUM", 20),
-    "heavy": ("MOPPY_WORKER_GB_HEAVY", 28),
+    "medium": ("MOPPY_WORKER_GB_MEDIUM", 14),
+    "heavy": ("MOPPY_WORKER_GB_HEAVY", 16),
 }
 
 # A variable reading at least this many MB of model-variable data per file is

src/access_moppy/templates/cmor_python_script.j2

@@ -123,6 +123,13 @@ def main():
             drs_root=drs_root,
             cmip_version=cmip_version,
             enable_resampling={{ 'True' if (config['enable_resampling'] | default(true)) else 'False' }},
+            # Stream the write in time-chunks instead of materialising the whole
+            # variable in one .compute(). Without this the write is a single
+            # monolithic gather ('finalize') over the entire time series, which
+            # exceeds a worker's memory on long (multi-century) daily runs and
+            # kills the job. Chunked writing bounds per-worker memory regardless
+            # of dataset length.
+            enable_chunking={{ 'True' if (config['enable_chunking'] | default(true)) else 'False' }},
         )
 
         cmoriser.run()

src/access_moppy/vocabulary_processors.py

@@ -633,17 +633,22 @@ def normalize_dataset_missing_values(dataset):
             current_missing = var.attrs.get("missing_value")
             current_fill = var.attrs.get("_FillValue")
 
-            # Build conditions for values that should become NaN
+            # Build conditions for values that should become NaN.
+            # NOTE: use ``abs(var - m) <= atol`` rather than np.isclose(var, m).
+            # np.isclose is not a ufunc, so on a Dask-backed array it does not
+            # dispatch lazily — it eagerly materialises the whole array to NumPy.
+            # For long daily variables that means loading the entire (multi-GB)
+            # series into one worker and OOM-killing it. ``abs`` / ``<=`` are
+            # Dask-native and stay lazy. With rtol=0, this is exactly np.isclose.
             nan_conditions = []
+            atol = 1e-3
 
             # Check for current missing_value
             if current_missing is not None:
                 try:
                     current_missing = float(current_missing)
                     if not np.isnan(current_missing):  # Don't double-convert NaN
-                        nan_conditions.append(
-                            np.isclose(var, current_missing, rtol=0, atol=1e-3)
-                        )
+                        nan_conditions.append(abs(var - current_missing) <= atol)
                 except (ValueError, TypeError):
                     pass
 
@@ -652,9 +657,7 @@ def normalize_dataset_missing_values(dataset):
                 try:
                     current_fill = float(current_fill)
                     if not np.isnan(current_fill):  # Don't double-convert NaN
-                        nan_conditions.append(
-                            np.isclose(var, current_fill, rtol=0, atol=1e-3)
-                        )
+                        nan_conditions.append(abs(var - current_fill) <= atol)
                 except (ValueError, TypeError):
                     pass