cmor_python_script.j2

#!/usr/bin/env python
"""
CMORisation script for variable {{ variable }}
Generated automatically by ACCESS-MOPPy batch processing
"""

import os
import glob
import logging
import sys
from pathlib import Path
import dask.distributed as dd

from access_moppy import ACCESS_ESM_CMORiser
from access_moppy.executors.dask_config import recommend_dask_config
from access_moppy.tracking import TaskTracker


def main():
    import dask
    dask.config.set({
        'distributed.worker.memory.target': 0.6,
        'distributed.worker.memory.spill': 0.7,
        'distributed.worker.memory.pause': 0.8,
    })

    # 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
        variable = os.environ['VARIABLE']
        db_path = os.environ['CMOR_TRACKER_DB']
        experiment_id = os.environ['EXPERIMENT_ID']
        source_id = os.environ['SOURCE_ID']
        variant_label = os.environ['VARIANT_LABEL']
        grid_label = os.environ['GRID_LABEL']
        activity_id = os.environ['ACTIVITY_ID']
        input_folder = os.environ['INPUT_FOLDER']
        output_folder = os.environ['OUTPUT_FOLDER']
        drs_root = os.environ.get('DRS_ROOT') or None
        cmip_version = os.environ.get('CMIP_VERSION', 'CMIP6')

        # File patterns — explicit overrides from the batch config take precedence;
        # if none is supplied the mapping-driven auto-discovery is used instead.
        file_patterns = {{ config.get('file_patterns', {}) | tojson }}
        model_id = os.environ.get('MODEL_ID', 'ACCESS-ESM1-6')

        # Optional year-range filtering for auto-discovery (pi-control / spinup runs).
        # Accepts integers or zero-padded strings (e.g. "0001").
        start_year = {{ config['start_year'] | tojson if 'start_year' in config else 'None' }}
        end_year   = {{ config['end_year']   | tojson if 'end_year'   in config else 'None' }}

        pattern = file_patterns.get(variable)
        if pattern:
            patterns = pattern if isinstance(pattern, list) else [pattern]
            input_files = []
            for p in patterns:
                input_files.extend(glob.glob(os.path.join(input_folder, p.lstrip("/")), recursive=True))
        else:
            # Auto-discover using the model mapping file_discovery config
            from access_moppy.file_discovery import (
                FileDiscoveryError,
                discover_files,
            )
            try:
                input_files = [str(f) for f in discover_files(
                    input_folder, variable, model_id=model_id,
                    start_year=start_year, end_year=end_year,
                )]
            except FileDiscoveryError as exc:
                raise ValueError(
                    f"Could not auto-discover files for '{variable}': {exc}\n"
                    "Add a 'file_patterns' entry in the batch config or a "
                    "'file_pattern' field to the variable's mapping entry."
                ) 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,
            )
            raise ValueError(f"No files found for variable '{variable}'. {diag}")

        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))

        if tracker.is_done(variable, experiment_id):
            print(f'Skipped: {variable} (already done)')
            return

        tracker.mark_running(variable, experiment_id)

        # Run CMORisation
        cmoriser = ACCESS_ESM_CMORiser(
            input_paths=input_files,
            compound_name=variable,
            experiment_id=experiment_id,
            source_id=source_id,
            variant_label=variant_label,
            grid_label=grid_label,
            activity_id=activity_id,
            output_path=output_folder,
            drs_root=drs_root,
            cmip_version=cmip_version,
            enable_resampling={{ 'True' if (config['enable_resampling'] | default(true)) else 'False' }},
        )

        cmoriser.run()
        cmoriser.write()

        tracker.mark_done(variable, experiment_id)
        print(f'Completed: {variable}')

    except Exception as e:
        print(f'Error processing {variable}: {e}', file=sys.stderr)
        try:
            tracker.mark_failed(variable, experiment_id, str(e))
        except:
            pass
        sys.exit(1)

    finally:
        # 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.
        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__':
    main()