Speed up batch CMORisation: multi-process Dask, per-variable worker sizing, and skip redundant daily-file validation

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",
+    }

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__':

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,