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

[rhaegar325]

Summary

Batch CMORisation runs one PBS job per variable. Each job requested
ncpus=12 mem=190GB but effectively used ~1 core, so the heavy variables
(cVeg/cSoil/nbp) took 1.5–2.7 h and the daily-input variables
(tasmax/tasmin) blew past walltime. This PR fixes the root causes:

1. The generated script ran Dask as a single process with 12 threads
   (processes=False), so every netCDF4 read serialised on the global HDF5
   lock and collapsed to ~1 core.
2. A first fix (one process per core) then split the 190 GB allocation into
   pools too small for the heavy/daily variables, causing constant memory
   pausing and killing tasmax/tasmin.
3. Independently, tasmax/tasmin spent ~32 min of single-core time in a
   frequency-validation step that re-opened all 2064 daily files just to
   re-confirm they are daily.

The result is a multi-process Dask client whose worker count adapts to each
variable's memory footprint and to the PBS resources, plus removal of the
redundant serial validation for daily variables. Numerical output is
unchanged (verified bit-identical).

Root cause

1. Threaded Dask client + HDF5 global lock → ~1 core

Reads go through xr.open_mfdataset(..., engine="netcdf4"). netCDF4/HDF5 is not
thread-safe, so xarray serialises every read on a single global HDF5_LOCK.
Under a threaded client (processes=False), all 2064 file reads run on one
core regardless of thread count.

Evidence (real cVeg read+sum on a 12-core node, 1032 cold files each):

client	wall	cpu/wall
threads (before)	273.9 s	1.2 (~1 core)
processes (after)	63.1 s	8.4 (~8 cores)

Results were bit-identical (sum=1.19124e+49 in both).

2. Why only some variables were slow

The lock caps every variable at ~1 core; the heavy ones simply move far more
data along that single serial lane:

variable	read/file	vs tas	extra work
tas	0.11 MB	1x	direct read
cVeg	9.58 MB	~86x	4 carbon pools x 17 tile levels
nbp	7.91 MB	~71x	5 pools/fluxes x 17 levels
tasmax/tasmin	3.45 MB	~31x	daily input (31 steps/file) + resample

3. Naive one-worker-per-core over-divided memory

Splitting 190 GB into 12 x ~11 GB workers caused:

• constant memory pausing (cSoil paused 105x, dragging it to 2:15) while
  ~95 GB of the node sat idle;
• tasmax/tasmin workers exceeding their ~11 GB limit and being killed by
  the nanny, failing the final gather.

4. Redundant serial validation of all daily files

For tasmax/tasmin, validate_cmip6_frequency_compatibility sampled 10 files
and correctly detected "daily", but because daily is outside the 20–35 day
"monthly" band it fell back to _validate_monthly_files_individually, which
opened all 2064 files one at a time in the main process (no Dask), just to
re-confirm sub-monthly frequency. Measured ~941 ms/file -> ~32 min/job.

Changes

src/access_moppy/executors/dask_config.py (new)

recommend_dask_config(variable, input_files, model_id) sizes a
processes=True client:

• Resources are read from the environment (PBS_NCPUS, PBS_MEM_GB), so
  worker count and per-worker memory adapt to whatever the PBS job requested —
  nothing about node size is hard-coded.
• Per-variable intensity is measured by probing one input file:
  sub-monthly input (needs resample) -> heavy; multi-MB tiled model variables
  (carbon pools) -> medium; single small field -> light.
• All parallelism comes from processes; threads_per_worker=1. Extra
  threads add no read throughput (HDF5 lock + GIL) and would double a worker's
  resident memory — the opposite of what the memory-limited tiers need.
• Per-worker memory floors (light 12 / medium 20 / heavy 28 GB) are the
  minimum RAM a worker of that class needs; they are env-overridable
  (MOPPY_WORKER_GB_LIGHT|MEDIUM|HEAVY, MOPPY_MEDIUM_MB_PER_FILE).

Resulting sizing on a 12-core / 190 GB job:

class	example	workers x threads	per-worker
light	tas, pr	12 x 1	15 GB
carbon	cVeg, cSoil, nbp	9 x 1	21 GB
daily	tasmax, tasmin	6 x 1	31 GB

src/access_moppy/templates/cmor_python_script.j2

• Switched the generated Dask client to processes=True.
• Input files are now discovered before the client is created, so the
  client can be sized to the variable via recommend_dask_config.
• Guarded client/tracker teardown for the new ordering; dropped the unused
  psutil import; quieter worker logs.

src/access_moppy/utilities.py

• In _validate_monthly_compatibility, when a variable allows sub-monthly input
  (tasmax/tasmin) and the sampled detection already confirms a sub-monthly
  frequency, return it directly instead of re-opening every file. The exhaustive
  per-file loop is retained as the fallback for genuinely unexpected
  frequencies; monthly variables are unaffected.

Validation

• Parallelism: 12-core node, cVeg read+sum, cpu/wall 1.2 -> 8.4, wall
  273.9 s -> 63.1 s, bit-identical results, 70.8/190 GB memory.
• Adaptive sizing: correct tier for all 7 target variables; adapts across
  PBS shapes (4–48 cpus); env override confirmed (heavy 28 -> 40 => 6 -> 4
  workers).
• Validation short-circuit: tasmax frequency validation ~451 s -> 14.7 s
  for 480 files (near-constant regardless of file count; ~32 min saved on a full
  2064-file job). Monthly cVeg unchanged (freq=31 days, normal path).
• Correctness of tasmax: CMORiser output verified bit-exact equal to the
  independently computed monthly maximum of the daily data (and clearly not
  the mean), with the production setting enable_resampling=False.

Risk

• Output values are unchanged (scheduler-only change for the read/compute path;
  validation short-circuit only skips redundant re-confirmation).
• _preprocess is dispatched to workers under processes=True; verified to
  round-trip correctly (bit-identical results at full scale).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.1%. Comparing base (965399e) to head (cd80472).
⚠️ Report is 13 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##            main    #499     +/-   ##
=======================================
- Coverage   78.5%   78.1%   -0.4%     
=======================================
  Files         33      33             
  Lines       6520    6626    +106     
  Branches    1233    1260     +27     
=======================================
+ Hits        5117    5173     +56     
- Misses      1129    1180     +51     
+ Partials     274     273      -1     

Flag	Coverage Δ
unit	78.1% <100.0%> (-0.4%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.