From 3b7b877ec8b1f027fd44671dd68d20b5d4514d70 Mon Sep 17 00:00:00 2001 From: rhaegar325 Date: Wed, 8 Jul 2026 16:10:10 +1000 Subject: [PATCH 1/3] update batch_processing dask config --- .../templates/cmor_python_script.j2 | 32 +++++++++++++------ 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/src/access_moppy/templates/cmor_python_script.j2 b/src/access_moppy/templates/cmor_python_script.j2 index 565a5960..cfcaf6ab 100644 --- a/src/access_moppy/templates/cmor_python_script.j2 +++ b/src/access_moppy/templates/cmor_python_script.j2 @@ -6,6 +6,7 @@ Generated automatically by ACCESS-MOPPy batch processing import os import glob +import logging import sys from pathlib import Path import dask.distributed as dd @@ -15,11 +16,17 @@ from access_moppy import ACCESS_ESM_CMORiser from access_moppy.tracking import TaskTracker def detect_optimal_dask_config(): - """Detect Dask configuration for single-node processes=False mode. + """Detect Dask configuration for single-node processes=True mode. - All workers share one Python process, so memory_limit must reflect - the full process budget. n_workers=1 ensures the single worker's - limit equals the PBS allocation rather than a fractional share. + CMORisation is dominated by reading thousands of netCDF4 files. The + netCDF4/HDF5 library serialises every read through a single global lock, + so a threaded scheduler (processes=False) collapses to ~1 core regardless + of thread count. Multi-process workers each hold their own HDF5 lock, so + file reads parallelise across cores. + + With one worker per CPU the PBS memory allocation is shared, so + memory_limit here is the *per-worker* budget (total / n_workers) to avoid + over-committing the node. """ n_cpus = int(os.environ.get('PBS_NCPUS', psutil.cpu_count())) allocated_memory_gb = int(os.environ.get('PBS_MEM_GB', '16')) @@ -27,12 +34,16 @@ def detect_optimal_dask_config(): system_memory_gb = psutil.virtual_memory().total / (1024**3) effective_memory = min(allocated_memory_gb, system_memory_gb * 0.9) + # One process per core maximises parallel netCDF4 reads (I/O-bound). + n_workers = max(1, n_cpus) + per_worker_gb = max(1, int(effective_memory * 0.9 / n_workers)) + print(f"Detected: {n_cpus} CPUs, {allocated_memory_gb}GB allocated, {system_memory_gb:.1f}GB system") return { - 'n_workers': 1, - 'threads_per_worker': n_cpus, - 'memory_per_worker': f"{int(effective_memory * 0.9)}GB", + 'n_workers': n_workers, + 'threads_per_worker': 1, + 'memory_per_worker': f"{per_worker_gb}GB", } @@ -52,12 +63,15 @@ def main(): print(f"Optimal Dask Config: {n_workers} workers, {threads_per_worker} threads each, {memory_per_worker} per worker") + # processes=True: each worker is a separate process with its own HDF5 + # lock, so netCDF4 reads run in parallel across cores (a threaded client + # serialises them on the global HDF5 lock and collapses to ~1 core). client = dd.Client( - processes=False, + processes=True, threads_per_worker=threads_per_worker, n_workers=n_workers, memory_limit=memory_per_worker, - silence_logs=False + silence_logs=logging.WARNING, ) From da5dd0fc9310525cbcd5e2575242491880a240bf Mon Sep 17 00:00:00 2001 From: rhaegar325 Date: Thu, 9 Jul 2026 20:04:44 +1000 Subject: [PATCH 2/3] update dask defination in batch processing script --- src/access_moppy/executors/dask_config.py | 150 ++++++++++++++++++ .../templates/cmor_python_script.j2 | 87 ++++------ src/access_moppy/utilities.py | 17 +- 3 files changed, 193 insertions(+), 61 deletions(-) create mode 100644 src/access_moppy/executors/dask_config.py diff --git a/src/access_moppy/executors/dask_config.py b/src/access_moppy/executors/dask_config.py new file mode 100644 index 00000000..3512f0fc --- /dev/null +++ b/src/access_moppy/executors/dask_config.py @@ -0,0 +1,150 @@ +"""Adaptive Dask client sizing for batch CMORisation. + +Each variable is CMORised in its own PBS job. The job's cost is dominated by +reading thousands of netCDF4 files, and the netCDF4/HDF5 library serialises +every read through a single global lock, so a threaded client (processes=False) +collapses to ~1 core. Multi-process workers each hold their own HDF5 lock, so +reads parallelise across workers. + +Worker *count* is a memory trade-off, and the right count depends on the +variable: + +* Light single-field variables (e.g. ``Amon.tas``) read a fraction of a MB per + file, so the allocation can be split into many small workers -> maximum read + parallelism. +* Carbon-pool variables (``cVeg``/``cSoil``/``nbp``) read several tiled model + variables per file and need a moderate per-worker budget. +* Daily-input variables that resample to monthly (``tasmax``/``tasmin``) push + large time-axis intermediates plus unmanaged HDF5 buffers through a worker and + need a generous per-worker budget, or a worker exceeds its limit and is killed. + +Rather than hard-code a worker count, :func:`recommend_dask_config` probes one +input file for the variable's actual read footprint and picks the largest worker +count whose per-worker budget covers that footprint. +""" +from __future__ import annotations + +import os +import sys + +# 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 +# via an environment variable set in the batch config's worker_init, so nothing +# is locked in code. +_FLOOR_ENV = { + "light": ("MOPPY_WORKER_GB_LIGHT", 12), + "medium": ("MOPPY_WORKER_GB_MEDIUM", 20), + "heavy": ("MOPPY_WORKER_GB_HEAVY", 28), +} + +# A variable reading at least this many MB of model-variable data per file is +# treated as at least "medium" intensity (carbon pools read ~8-10 MB/file). +# Overridable for unusual mappings via MOPPY_MEDIUM_MB_PER_FILE. +_MEDIUM_MB_PER_FILE_DEFAULT = 4.0 + + +def _floor_gb(tier): + """Per-worker memory floor (GB) for a tier, honouring an env override.""" + env_name, default = _FLOOR_ENV[tier] + try: + return max(1, int(os.environ.get(env_name, default))) + except (TypeError, ValueError): + return default + + +def _estimate_worker_memory_gb(variable, input_files, model_id): + """Estimate the per-worker memory floor (GB) for ``variable``. + + Probes the first input file for (a) the number of time steps per file -- + sub-monthly input implies a resample-to-monthly with large time-axis + intermediates -- and (b) the total size of the variable's model variables, + which captures tiled / multi-pool variables. Falls back to the heavy tier if + anything about the probe is uncertain, so a bad estimate never under-sizes + memory. + """ + if not input_files: + return _floor_gb("heavy") + try: + import xarray as xr + + from access_moppy.utilities import load_model_mappings + + _, cmor_name = variable.split(".", 1) + mapping = load_model_mappings(variable, model_id) + model_vars = (mapping.get(cmor_name, {}) or {}).get("model_variables", []) or [] + + with xr.open_dataset(input_files[0], decode_cf=False) as ds0: + steps_per_file = int(ds0.sizes.get("time", 1)) + present = [v for v in model_vars if v in ds0.variables] + per_file_mb = sum( + ds0[v].size * ds0[v].dtype.itemsize for v in present + ) / 1e6 + + try: + medium_mb = float( + os.environ.get("MOPPY_MEDIUM_MB_PER_FILE", _MEDIUM_MB_PER_FILE_DEFAULT) + ) + except (TypeError, ValueError): + medium_mb = _MEDIUM_MB_PER_FILE_DEFAULT + + # Sub-monthly input -> monthly resampling: fat workers. + if steps_per_file > 1: + return _floor_gb("heavy") + # Multiple tiled model variables (carbon pools): moderate workers. + if per_file_mb >= medium_mb: + return _floor_gb("medium") + return _floor_gb("light") + except Exception as exc: # noqa: BLE001 - never let sizing crash the job + print( + f"Worker-memory estimate failed ({exc}); assuming heavy tier.", + file=sys.stderr, + ) + return _floor_gb("heavy") + + +def recommend_dask_config(variable, input_files, model_id, n_cpus=None, mem_gb=None): + """Return ``dask.distributed.Client`` kwargs sized to ``variable``. + + Parameters mirror what the generated batch script already has to hand. The + result is a dict with ``n_workers``, ``threads_per_worker`` and + ``memory_limit`` (per worker), for use as ``Client(processes=True, **cfg)``. + """ + import psutil + + if n_cpus is None: + n_cpus = int(os.environ.get("PBS_NCPUS", psutil.cpu_count() or 1)) + if mem_gb is None: + mem_gb = int(os.environ.get("PBS_MEM_GB", "16")) + + system_memory_gb = psutil.virtual_memory().total / (1024**3) + effective_memory = min(mem_gb, system_memory_gb * 0.9) + + floor_gb = _estimate_worker_memory_gb(variable, input_files, model_id) + + # Largest worker count whose per-worker share still meets the floor, capped + # by CPU count. All parallelism comes from *processes*: netCDF4/HDF5 reads + # serialise on a global lock (and the GIL), so extra threads per worker add + # no read throughput while doubling the in-flight data resident per worker -- + # the opposite of what the memory-limited heavy tier needs. Hence one thread + # per worker; memory-limited variables simply use fewer cores. + n_workers = max(1, min(n_cpus, int(effective_memory // floor_gb))) + threads_per_worker = 1 + per_worker_gb = max(1, int(effective_memory / n_workers)) + + idle_cores = max(0, n_cpus - n_workers) + print( + f"Dask sizing for {variable}: {n_cpus} CPUs, {effective_memory:.0f}GB usable, " + f"per-worker floor {floor_gb}GB -> {n_workers} workers x {threads_per_worker} " + f"thread, {per_worker_gb}GB each" + + (f" ({idle_cores} cores idle: memory-limited)" if idle_cores else "") + ) + + return { + "n_workers": n_workers, + "threads_per_worker": threads_per_worker, + "memory_limit": f"{per_worker_gb}GB", + } diff --git a/src/access_moppy/templates/cmor_python_script.j2 b/src/access_moppy/templates/cmor_python_script.j2 index cfcaf6ab..bcd8e90f 100644 --- a/src/access_moppy/templates/cmor_python_script.j2 +++ b/src/access_moppy/templates/cmor_python_script.j2 @@ -10,42 +10,11 @@ import logging import sys from pathlib import Path import dask.distributed as dd -import psutil from access_moppy import ACCESS_ESM_CMORiser +from access_moppy.executors.dask_config import recommend_dask_config from access_moppy.tracking import TaskTracker -def detect_optimal_dask_config(): - """Detect Dask configuration for single-node processes=True mode. - - CMORisation is dominated by reading thousands of netCDF4 files. The - netCDF4/HDF5 library serialises every read through a single global lock, - so a threaded scheduler (processes=False) collapses to ~1 core regardless - of thread count. Multi-process workers each hold their own HDF5 lock, so - file reads parallelise across cores. - - With one worker per CPU the PBS memory allocation is shared, so - memory_limit here is the *per-worker* budget (total / n_workers) to avoid - over-committing the node. - """ - n_cpus = int(os.environ.get('PBS_NCPUS', psutil.cpu_count())) - allocated_memory_gb = int(os.environ.get('PBS_MEM_GB', '16')) - - system_memory_gb = psutil.virtual_memory().total / (1024**3) - effective_memory = min(allocated_memory_gb, system_memory_gb * 0.9) - - # One process per core maximises parallel netCDF4 reads (I/O-bound). - n_workers = max(1, n_cpus) - per_worker_gb = max(1, int(effective_memory * 0.9 / n_workers)) - - print(f"Detected: {n_cpus} CPUs, {allocated_memory_gb}GB allocated, {system_memory_gb:.1f}GB system") - - return { - 'n_workers': n_workers, - 'threads_per_worker': 1, - 'memory_per_worker': f"{per_worker_gb}GB", - } - def main(): import dask @@ -55,27 +24,11 @@ def main(): 'distributed.worker.memory.pause': 0.8, }) - dask_config = detect_optimal_dask_config() - - n_workers = dask_config['n_workers'] - threads_per_worker = dask_config['threads_per_worker'] - memory_per_worker = dask_config['memory_per_worker'] - - print(f"Optimal Dask Config: {n_workers} workers, {threads_per_worker} threads each, {memory_per_worker} per worker") - - # processes=True: each worker is a separate process with its own HDF5 - # lock, so netCDF4 reads run in parallel across cores (a threaded client - # serialises them on the global HDF5 lock and collapses to ~1 core). - client = dd.Client( - processes=True, - threads_per_worker=threads_per_worker, - n_workers=n_workers, - memory_limit=memory_per_worker, - silence_logs=logging.WARNING, - ) - - - print(f'Dask dashboard: {client.dashboard_link}') + # The Dask client is created *after* input files are discovered, so it can + # be sized to this variable's actual read footprint (recommend_dask_config + # reads PBS_NCPUS / PBS_MEM_GB, so it adapts to whatever the job requested). + client = None + tracker = None try: # Get values from environment @@ -111,7 +64,6 @@ def main(): # Auto-discover using the model mapping file_discovery config from access_moppy.file_discovery import ( FileDiscoveryError, - _diagnose_no_files, discover_files, ) try: @@ -127,6 +79,7 @@ def main(): ) from exc if not input_files: + from access_moppy.file_discovery import _diagnose_no_files diag = _diagnose_no_files( Path(input_folder), variable, model_id, start_year=start_year, end_year=end_year, @@ -135,6 +88,19 @@ def main(): print(f'Processing {variable} with {len(input_files)} files') + # Size the Dask client to this variable's footprint, then start it. + # processes=True: each worker is a separate process with its own HDF5 + # lock, so netCDF4 reads run in parallel across workers (a threaded + # client serialises them on the global HDF5 lock and collapses to ~1 + # core). Worker count / memory adapt to PBS_NCPUS / PBS_MEM_GB. + dask_config = recommend_dask_config(variable, input_files, model_id) + client = dd.Client( + processes=True, + silence_logs=logging.WARNING, + **dask_config, + ) + print(f'Dask dashboard: {client.dashboard_link}') + # Initialize tracker tracker = TaskTracker(Path(db_path)) @@ -156,7 +122,7 @@ def main(): output_path=output_folder, drs_root=drs_root, cmip_version=cmip_version, - enable_resampling=True, + enable_resampling={{ 'True' if (config['enable_resampling'] | default(true)) else 'False' }}, ) cmoriser.run() @@ -177,11 +143,12 @@ def main(): # Release the sqlite handle before the process exits so the journal # file (DELETE mode) is unlinked promptly. On Lustre, lingering # handles can leave .db-journal behind and confuse downstream tools. - try: - tracker.close() - except Exception: - pass - if client.status != 'closed': + if tracker is not None: + try: + tracker.close() + except Exception: + pass + if client is not None and client.status != 'closed': client.close() if __name__ == '__main__': diff --git a/src/access_moppy/utilities.py b/src/access_moppy/utilities.py index eac0488a..b4dba5e3 100644 --- a/src/access_moppy/utilities.py +++ b/src/access_moppy/utilities.py @@ -953,7 +953,22 @@ def _validate_monthly_compatibility( monthly_max = 35 * 86400 # 35 days in seconds if not (monthly_min <= freq_seconds <= monthly_max): - # If concatenated detection doesn't give monthly range, validate individual files + # Sub-monthly input (e.g. daily) is expected and valid for variables + # that aggregate to monthly (tasmax/tasmin). The sampled concatenated + # detection already confirms the frequency, so accept it directly + # rather than re-opening every input file one-by-one -- that + # exhaustive per-file loop can take tens of minutes for a multi-year + # daily run while verifying something already known. + if allow_submonthly and freq_seconds < monthly_min: + logger.debug( + "Sub-monthly frequency (%s) confirmed from sample - valid for " + "monthly aggregation, skipping per-file validation", + detected_freq, + ) + return detected_freq + + # Otherwise the frequency is unexpected: fall back to opening each + # file individually so any inconsistent/rogue file can be pinpointed. logger.debug( "Concatenated frequency (%s) not in monthly range, validating individual files...", detected_freq, From cd80472134b33d120e6749854a6941480ea9ba35 Mon Sep 17 00:00:00 2001 From: rhaegar325 Date: Thu, 9 Jul 2026 20:20:28 +1000 Subject: [PATCH 3/3] pre-commit fix --- src/access_moppy/executors/dask_config.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/access_moppy/executors/dask_config.py b/src/access_moppy/executors/dask_config.py index 3512f0fc..134826eb 100644 --- a/src/access_moppy/executors/dask_config.py +++ b/src/access_moppy/executors/dask_config.py @@ -22,6 +22,7 @@ input file for the variable's actual read footprint and picks the largest worker count whose per-worker budget covers that footprint. """ + from __future__ import annotations import os @@ -80,9 +81,9 @@ def _estimate_worker_memory_gb(variable, input_files, model_id): with xr.open_dataset(input_files[0], decode_cf=False) as ds0: steps_per_file = int(ds0.sizes.get("time", 1)) present = [v for v in model_vars if v in ds0.variables] - per_file_mb = sum( - ds0[v].size * ds0[v].dtype.itemsize for v in present - ) / 1e6 + per_file_mb = ( + sum(ds0[v].size * ds0[v].dtype.itemsize for v in present) / 1e6 + ) try: medium_mb = float(