|
| 1 | +"""Adaptive Dask client sizing for batch CMORisation. |
| 2 | +
|
| 3 | +Each variable is CMORised in its own PBS job. The job's cost is dominated by |
| 4 | +reading thousands of netCDF4 files, and the netCDF4/HDF5 library serialises |
| 5 | +every read through a single global lock, so a threaded client (processes=False) |
| 6 | +collapses to ~1 core. Multi-process workers each hold their own HDF5 lock, so |
| 7 | +reads parallelise across workers. |
| 8 | +
|
| 9 | +Worker *count* is a memory trade-off, and the right count depends on the |
| 10 | +variable: |
| 11 | +
|
| 12 | +* Light single-field variables (e.g. ``Amon.tas``) read a fraction of a MB per |
| 13 | + file, so the allocation can be split into many small workers -> maximum read |
| 14 | + parallelism. |
| 15 | +* Carbon-pool variables (``cVeg``/``cSoil``/``nbp``) read several tiled model |
| 16 | + variables per file and need a moderate per-worker budget. |
| 17 | +* Daily-input variables that resample to monthly (``tasmax``/``tasmin``) push |
| 18 | + large time-axis intermediates plus unmanaged HDF5 buffers through a worker and |
| 19 | + need a generous per-worker budget, or a worker exceeds its limit and is killed. |
| 20 | +
|
| 21 | +Rather than hard-code a worker count, :func:`recommend_dask_config` probes one |
| 22 | +input file for the variable's actual read footprint and picks the largest worker |
| 23 | +count whose per-worker budget covers that footprint. |
| 24 | +""" |
| 25 | + |
| 26 | +from __future__ import annotations |
| 27 | + |
| 28 | +import os |
| 29 | +import sys |
| 30 | + |
| 31 | +# Per-worker memory floors (GB) for the three intensity tiers. These are the |
| 32 | +# minimum RAM a single worker needs to process a variable of that class without |
| 33 | +# being killed -- a property of the *workload*, not the node -- so they are |
| 34 | +# absolute GB, not fractions of the allocation. The defaults are calibrated from |
| 35 | +# real runs (daily variables were killed at ~11GB/worker; carbon pools paused |
| 36 | +# hard at ~11GB; light variables were fine), and each can be overridden per job |
| 37 | +# via an environment variable set in the batch config's worker_init, so nothing |
| 38 | +# is locked in code. |
| 39 | +_FLOOR_ENV = { |
| 40 | + "light": ("MOPPY_WORKER_GB_LIGHT", 12), |
| 41 | + "medium": ("MOPPY_WORKER_GB_MEDIUM", 20), |
| 42 | + "heavy": ("MOPPY_WORKER_GB_HEAVY", 28), |
| 43 | +} |
| 44 | + |
| 45 | +# A variable reading at least this many MB of model-variable data per file is |
| 46 | +# treated as at least "medium" intensity (carbon pools read ~8-10 MB/file). |
| 47 | +# Overridable for unusual mappings via MOPPY_MEDIUM_MB_PER_FILE. |
| 48 | +_MEDIUM_MB_PER_FILE_DEFAULT = 4.0 |
| 49 | + |
| 50 | + |
| 51 | +def _floor_gb(tier): |
| 52 | + """Per-worker memory floor (GB) for a tier, honouring an env override.""" |
| 53 | + env_name, default = _FLOOR_ENV[tier] |
| 54 | + try: |
| 55 | + return max(1, int(os.environ.get(env_name, default))) |
| 56 | + except (TypeError, ValueError): |
| 57 | + return default |
| 58 | + |
| 59 | + |
| 60 | +def _estimate_worker_memory_gb(variable, input_files, model_id): |
| 61 | + """Estimate the per-worker memory floor (GB) for ``variable``. |
| 62 | +
|
| 63 | + Probes the first input file for (a) the number of time steps per file -- |
| 64 | + sub-monthly input implies a resample-to-monthly with large time-axis |
| 65 | + intermediates -- and (b) the total size of the variable's model variables, |
| 66 | + which captures tiled / multi-pool variables. Falls back to the heavy tier if |
| 67 | + anything about the probe is uncertain, so a bad estimate never under-sizes |
| 68 | + memory. |
| 69 | + """ |
| 70 | + if not input_files: |
| 71 | + return _floor_gb("heavy") |
| 72 | + try: |
| 73 | + import xarray as xr |
| 74 | + |
| 75 | + from access_moppy.utilities import load_model_mappings |
| 76 | + |
| 77 | + _, cmor_name = variable.split(".", 1) |
| 78 | + mapping = load_model_mappings(variable, model_id) |
| 79 | + model_vars = (mapping.get(cmor_name, {}) or {}).get("model_variables", []) or [] |
| 80 | + |
| 81 | + with xr.open_dataset(input_files[0], decode_cf=False) as ds0: |
| 82 | + steps_per_file = int(ds0.sizes.get("time", 1)) |
| 83 | + present = [v for v in model_vars if v in ds0.variables] |
| 84 | + per_file_mb = ( |
| 85 | + sum(ds0[v].size * ds0[v].dtype.itemsize for v in present) / 1e6 |
| 86 | + ) |
| 87 | + |
| 88 | + try: |
| 89 | + medium_mb = float( |
| 90 | + os.environ.get("MOPPY_MEDIUM_MB_PER_FILE", _MEDIUM_MB_PER_FILE_DEFAULT) |
| 91 | + ) |
| 92 | + except (TypeError, ValueError): |
| 93 | + medium_mb = _MEDIUM_MB_PER_FILE_DEFAULT |
| 94 | + |
| 95 | + # Sub-monthly input -> monthly resampling: fat workers. |
| 96 | + if steps_per_file > 1: |
| 97 | + return _floor_gb("heavy") |
| 98 | + # Multiple tiled model variables (carbon pools): moderate workers. |
| 99 | + if per_file_mb >= medium_mb: |
| 100 | + return _floor_gb("medium") |
| 101 | + return _floor_gb("light") |
| 102 | + except Exception as exc: # noqa: BLE001 - never let sizing crash the job |
| 103 | + print( |
| 104 | + f"Worker-memory estimate failed ({exc}); assuming heavy tier.", |
| 105 | + file=sys.stderr, |
| 106 | + ) |
| 107 | + return _floor_gb("heavy") |
| 108 | + |
| 109 | + |
| 110 | +def recommend_dask_config(variable, input_files, model_id, n_cpus=None, mem_gb=None): |
| 111 | + """Return ``dask.distributed.Client`` kwargs sized to ``variable``. |
| 112 | +
|
| 113 | + Parameters mirror what the generated batch script already has to hand. The |
| 114 | + result is a dict with ``n_workers``, ``threads_per_worker`` and |
| 115 | + ``memory_limit`` (per worker), for use as ``Client(processes=True, **cfg)``. |
| 116 | + """ |
| 117 | + import psutil |
| 118 | + |
| 119 | + if n_cpus is None: |
| 120 | + n_cpus = int(os.environ.get("PBS_NCPUS", psutil.cpu_count() or 1)) |
| 121 | + if mem_gb is None: |
| 122 | + mem_gb = int(os.environ.get("PBS_MEM_GB", "16")) |
| 123 | + |
| 124 | + system_memory_gb = psutil.virtual_memory().total / (1024**3) |
| 125 | + effective_memory = min(mem_gb, system_memory_gb * 0.9) |
| 126 | + |
| 127 | + floor_gb = _estimate_worker_memory_gb(variable, input_files, model_id) |
| 128 | + |
| 129 | + # Largest worker count whose per-worker share still meets the floor, capped |
| 130 | + # by CPU count. All parallelism comes from *processes*: netCDF4/HDF5 reads |
| 131 | + # serialise on a global lock (and the GIL), so extra threads per worker add |
| 132 | + # no read throughput while doubling the in-flight data resident per worker -- |
| 133 | + # the opposite of what the memory-limited heavy tier needs. Hence one thread |
| 134 | + # per worker; memory-limited variables simply use fewer cores. |
| 135 | + n_workers = max(1, min(n_cpus, int(effective_memory // floor_gb))) |
| 136 | + threads_per_worker = 1 |
| 137 | + per_worker_gb = max(1, int(effective_memory / n_workers)) |
| 138 | + |
| 139 | + idle_cores = max(0, n_cpus - n_workers) |
| 140 | + print( |
| 141 | + f"Dask sizing for {variable}: {n_cpus} CPUs, {effective_memory:.0f}GB usable, " |
| 142 | + f"per-worker floor {floor_gb}GB -> {n_workers} workers x {threads_per_worker} " |
| 143 | + f"thread, {per_worker_gb}GB each" |
| 144 | + + (f" ({idle_cores} cores idle: memory-limited)" if idle_cores else "") |
| 145 | + ) |
| 146 | + |
| 147 | + return { |
| 148 | + "n_workers": n_workers, |
| 149 | + "threads_per_worker": threads_per_worker, |
| 150 | + "memory_limit": f"{per_worker_gb}GB", |
| 151 | + } |
0 commit comments