Enhance job script and Python template for variable-specific resource allocation and optimal Dask configuration

Author: rbeucher

src/access_mopper/batch_cmoriser.py

@@ -42,13 +42,22 @@ def create_job_script(variable, config, db_path, script_dir):
     pbs_template = Template(pbs_template_content)
     python_template = Template(python_template_content)
 
+    # Get variable-specific resources if available
+    variable_config = config.copy()
+    if "variable_resources" in config and variable in config["variable_resources"]:
+        # Override with variable-specific settings
+        variable_config.update(config["variable_resources"][variable])
+        print(
+            f"Using custom resources for {variable}: {config['variable_resources'][variable]}"
+        )
+
     # Get the package path for sys.path.insert
     package_path = Path(__file__).parent.parent
 
     # Create Python script
     python_script_content = python_template.render(
         variable=variable,
-        config=config,
+        config=variable_config,  # Use variable-specific config
         db_path=db_path,
         package_path=package_path,
     )
@@ -60,7 +69,7 @@ def create_job_script(variable, config, db_path, script_dir):
     # Create PBS script
     pbs_script_content = pbs_template.render(
         variable=variable,
-        config=config,
+        config=variable_config,  # Use variable-specific config
         script_dir=script_dir,
         python_script_path=python_script_path,
         db_path=db_path,

src/access_mopper/examples/batch_config.yml

@@ -53,6 +53,26 @@ walltime: "02:00:00"
 scheduler_options: "#PBS -P tm70"
 storage: "gdata/p73+gdata/tm70+scratch/tm70+gdata/xp65"
 
+# Variable-specific resource overrides
+variable_resources:
+  # 3D ocean variables need more resources
+  Omon.thetao:
+    cpus_per_node: 28
+    mem: 128GB
+    walltime: "06:00:00"
+
+  # Daily data needs more memory
+  day.pr:
+    cpus_per_node: 14
+    mem: 64GB
+    walltime: "04:00:00"
+
+  # Simple 2D monthly variables
+  Amon.pr:
+    cpus_per_node: 8
+    mem: 32GB
+    walltime: "02:00:00"
+
 # Environment setup for each job
 worker_init: |
   source /g/data/tm70/rb5533/miniforge3/bin/activate

src/access_mopper/templates/cmor_job_script.j2

@@ -15,6 +15,10 @@
 {{ config.get('worker_init', 'module load netcdf-python') }}
 
 export OMP_NUM_THREADS=1
+# Extract memory allocation for Python script
+{% set mem_value = config.get('mem', '16GB') %}
+{% set mem_gb = mem_value.replace('GB', '').replace('gb', '') | int %}
+export PBS_MEM_GB="{{ mem_gb }}"
 
 # Set environment variables for this job
 export CMOR_TRACKER_DB="{{ db_path }}"
@@ -35,6 +39,7 @@ mkdir -p $TMPDIR/cmor_temp
 {% endif %}
 
 # Debug information
+echo "PBS Resources: $PBS_NCPUS CPUs, ${PBS_MEM_GB}GB memory"
 echo "Database path: $CMOR_TRACKER_DB"
 echo "Database exists: $(test -f "$CMOR_TRACKER_DB" && echo "YES" || echo "NO")"
 echo "Variable: $VARIABLE"

src/access_mopper/templates/cmor_python_script.j2

@@ -9,17 +9,93 @@ import glob
 import sys
 from pathlib import Path
 import dask.distributed as dd
+import psutil
 
 from access_mopper import ACCESS_ESM_CMORiser
 from access_mopper.tracking import TaskTracker
 
+def detect_optimal_dask_config():
+    """Automatically detect optimal Dask configuration based on available resources."""
+
+    # Get PBS allocation
+    n_cpus = int(os.environ.get('PBS_NCPUS', psutil.cpu_count()))
+    allocated_memory_gb = int(os.environ.get('PBS_MEM_GB', '16'))
+
+    # Get system memory info as fallback
+    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")
+
+    # Heuristic for optimal configuration
+    if effective_memory >= 128:
+        # Very high memory jobs
+        return {
+            'n_workers': max(2, n_cpus // 8),
+            'threads_per_worker': 8,
+            'memory_per_worker': f"{int(effective_memory // max(2, n_cpus // 8) - 4)}GB",
+            'chunk_size': '1GB'
+        }
+    elif effective_memory >= 64:
+        # High memory jobs
+        return {
+            'n_workers': max(2, n_cpus // 4),
+            'threads_per_worker': 4,
+            'memory_per_worker': f"{int(effective_memory // max(2, n_cpus // 4) - 2)}GB",
+            'chunk_size': '512MB'
+        }
+    elif effective_memory >= 32:
+        # Medium memory jobs
+        return {
+            'n_workers': max(2, n_cpus // 2),
+            'threads_per_worker': 2,
+            'memory_per_worker': f"{int(effective_memory // max(2, n_cpus // 2) - 1)}GB",
+            'chunk_size': '256MB'
+        }
+    else:
+        # Low memory jobs
+        return {
+            'n_workers': 1,
+            'threads_per_worker': min(n_cpus, 4),
+            'memory_per_worker': f"{int(effective_memory - 1)}GB",
+            'chunk_size': '128MB'
+        }
+
 def main():
-    # Start a local Dask client for this job
+    # Get PBS resource allocation
+    n_cpus = int(os.environ.get('PBS_NCPUS', '4'))
+
+    # Get memory allocation from PBS (parse from PBS_JOBID or estimate)
+    # PBS doesn't directly expose memory allocation, so we'll use a heuristic
+    # or pass it via environment variable
+    allocated_memory_gb = int(os.environ.get('PBS_MEM_GB', '16'))
+
+    # Configure Dask based on available resources
+    import dask
+    dask.config.set({
+        'array.chunk-size': '512MB',
+        'distributed.worker.memory.target': 0.8,
+        'distributed.worker.memory.spill': 0.9,
+        'distributed.worker.memory.pause': 0.95,
+    })
+
+    dask_config = detect_optimal_dask_config()
+
+    # Dynamic worker configuration based on memory
+    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=1,
-        n_workers=int(os.environ.get('PBS_NCPUS', '1'))
+        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}')
 
     try: