Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion src/access_moppy/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
6 changes: 5 additions & 1 deletion src/access_moppy/derivations/calc_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
16 changes: 11 additions & 5 deletions src/access_moppy/executors/dask_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions src/access_moppy/templates/cmor_python_script.j2
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
17 changes: 10 additions & 7 deletions src/access_moppy/vocabulary_processors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand Down