diff --git a/src/access_moppy/base.py b/src/access_moppy/base.py index 83a13731..70de2683 100644 --- a/src/access_moppy/base.py +++ b/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", diff --git a/src/access_moppy/derivations/calc_utils.py b/src/access_moppy/derivations/calc_utils.py index 72c278ba..c26c31f9 100644 --- a/src/access_moppy/derivations/calc_utils.py +++ b/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) diff --git a/src/access_moppy/executors/dask_config.py b/src/access_moppy/executors/dask_config.py index 134826eb..f0c602d7 100644 --- a/src/access_moppy/executors/dask_config.py +++ b/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 diff --git a/src/access_moppy/templates/cmor_python_script.j2 b/src/access_moppy/templates/cmor_python_script.j2 index bcd8e90f..61a782ed 100644 --- a/src/access_moppy/templates/cmor_python_script.j2 +++ b/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() diff --git a/src/access_moppy/vocabulary_processors.py b/src/access_moppy/vocabulary_processors.py index 86c6af86..b3fa3109 100644 --- a/src/access_moppy/vocabulary_processors.py +++ b/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