diff --git a/src/access_moppy/executors/dask_config.py b/src/access_moppy/executors/dask_config.py new file mode 100644 index 00000000..134826eb --- /dev/null +++ b/src/access_moppy/executors/dask_config.py @@ -0,0 +1,151 @@ +"""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 565a5960..bcd8e90f 100644 --- a/src/access_moppy/templates/cmor_python_script.j2 +++ b/src/access_moppy/templates/cmor_python_script.j2 @@ -6,35 +6,15 @@ Generated automatically by ACCESS-MOPPy batch processing import os import glob +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=False 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. - """ - 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) - - 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", - } - def main(): import dask @@ -44,24 +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") - - client = dd.Client( - processes=False, - threads_per_worker=threads_per_worker, - n_workers=n_workers, - memory_limit=memory_per_worker, - silence_logs=False - ) - - - 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 @@ -97,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: @@ -113,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, @@ -121,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)) @@ -142,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() @@ -163,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,