Merge pull request #136 from CABLE-LSM/5-spatial-testing

Add payu test suite for spatial configuration

Author: abhaasgoyal

README.md

@@ -7,13 +7,20 @@
 
 - checks out the model versions specified by the user
 - builds the required executables
-- runs each model version across N standard science configurations
+- runs each model version across N standard science configurations for a variety of meteorological forcings
 - performs bitwise comparison checks on model outputs across model versions
 
 The user can then pipe the model outputs into a benchmark analysis via [modelevaluation.org][meorg] to assess model performance.
 
 The full documentation is available at [benchcab.readthedocs.io][docs].
 
+## Supported configurations
+
+`benchcab` currently tests the following model configurations for CABLE:
+
+- **Flux site simulations (offline)** - running CABLE forced with observed eddy covariance data at a single site
+- **Global/regional simulations (offline)** - running CABLE forced with meteorological fields over a region (global or regional)
+
 ## License
 
 `benchcab` is distributed under [an Apache License v2.0][apache-license].

benchcab/benchcab.py

@@ -10,25 +10,21 @@
 from subprocess import CalledProcessError
 from typing import Optional
 
-from benchcab import internal
+from benchcab import fluxsite, internal, spatial
 from benchcab.comparison import run_comparisons, run_comparisons_in_parallel
 from benchcab.config import read_config
 from benchcab.environment_modules import EnvironmentModules, EnvironmentModulesInterface
-from benchcab.fluxsite import (
-    Task,
-    get_fluxsite_comparisons,
-    get_fluxsite_tasks,
-    run_tasks,
-    run_tasks_in_parallel,
-)
 from benchcab.internal import get_met_forcing_file_names
 from benchcab.model import Model
 from benchcab.utils import get_logger
 from benchcab.utils.fs import mkdir, next_path
 from benchcab.utils.pbs import render_job_script
 from benchcab.utils.repo import create_repo
 from benchcab.utils.subprocess import SubprocessWrapper, SubprocessWrapperInterface
-from benchcab.workdir import setup_fluxsite_directory_tree
+from benchcab.workdir import (
+    setup_fluxsite_directory_tree,
+    setup_spatial_directory_tree,
+)
 
 
 class Benchcab:
@@ -57,7 +53,8 @@ def __init__(
 
         self._config: Optional[dict] = None
         self._models: list[Model] = []
-        self.tasks: list[Task] = []  # initialise fluxsite tasks lazily
+        self._fluxsite_tasks: list[fluxsite.FluxsiteTask] = []
+        self._spatial_tasks: list[spatial.SpatialTask] = []
 
         # Get the logger object
         self.logger = get_logger()
@@ -148,16 +145,26 @@ def _get_models(self, config: dict) -> list[Model]:
                 self._models.append(Model(repo=repo, model_id=id, **sub_config))
         return self._models
 
-    def _initialise_tasks(self, config: dict) -> list[Task]:
-        """A helper method that initialises and returns the `tasks` attribute."""
-        self.tasks = get_fluxsite_tasks(
-            models=self._get_models(config),
-            science_configurations=config["science_configurations"],
-            fluxsite_forcing_file_names=get_met_forcing_file_names(
-                config["fluxsite"]["experiment"]
-            ),
-        )
-        return self.tasks
+    def _get_fluxsite_tasks(self, config: dict) -> list[fluxsite.FluxsiteTask]:
+        if not self._fluxsite_tasks:
+            self._fluxsite_tasks = fluxsite.get_fluxsite_tasks(
+                models=self._get_models(config),
+                science_configurations=config["science_configurations"],
+                fluxsite_forcing_file_names=get_met_forcing_file_names(
+                    config["fluxsite"]["experiment"]
+                ),
+            )
+        return self._fluxsite_tasks
+
+    def _get_spatial_tasks(self, config) -> list[spatial.SpatialTask]:
+        if not self._spatial_tasks:
+            self._spatial_tasks = spatial.get_spatial_tasks(
+                models=self._get_models(config),
+                met_forcings=config["spatial"]["met_forcings"],
+                science_configurations=config["science_configurations"],
+                payu_args=config["spatial"]["payu"]["args"],
+            )
+        return self._spatial_tasks
 
     def validate_config(self, config_path: str):
         """Endpoint for `benchcab validate_config`."""
@@ -226,7 +233,7 @@ def checkout(self, config_path: str):
         with rev_number_log_path.open("w", encoding="utf-8") as file:
             file.write(rev_number_log)
 
-    def build(self, config_path: str):
+    def build(self, config_path: str, mpi=False):
         """Endpoint for `benchcab build`."""
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
@@ -239,40 +246,39 @@ def build(self, config_path: str):
                 repo.custom_build(modules=config["modules"])
 
             else:
-                build_mode = "with MPI" if internal.MPI else "serially"
+                build_mode = "with MPI" if mpi else "serially"
                 self.logger.info(
                     f"Compiling CABLE {build_mode} for realisation {repo.name}..."
                 )
-                repo.pre_build()
-                repo.run_build(modules=config["modules"])
-                repo.post_build()
+                repo.pre_build(mpi=mpi)
+                repo.run_build(modules=config["modules"], mpi=mpi)
+                repo.post_build(mpi=mpi)
             self.logger.info(f"Successfully compiled CABLE for realisation {repo.name}")
 
     def fluxsite_setup_work_directory(self, config_path: str):
         """Endpoint for `benchcab fluxsite-setup-work-dir`."""
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
 
-        tasks = self.tasks if self.tasks else self._initialise_tasks(config)
         self.logger.info("Setting up run directory tree for fluxsite tests...")
         setup_fluxsite_directory_tree()
         self.logger.info("Setting up tasks...")
-        for task in tasks:
+        for task in self._get_fluxsite_tasks(config):
             task.setup_task()
         self.logger.info("Successfully setup fluxsite tasks")
 
     def fluxsite_run_tasks(self, config_path: str):
         """Endpoint for `benchcab fluxsite-run-tasks`."""
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
+        tasks = self._get_fluxsite_tasks(config)
 
-        tasks = self.tasks if self.tasks else self._initialise_tasks(config)
         self.logger.info("Running fluxsite tasks...")
         if config["fluxsite"]["multiprocess"]:
             ncpus = config["fluxsite"]["pbs"]["ncpus"]
-            run_tasks_in_parallel(tasks, n_processes=ncpus)
+            fluxsite.run_tasks_in_parallel(tasks, n_processes=ncpus)
         else:
-            run_tasks(tasks)
+            fluxsite.run_tasks(tasks)
         self.logger.info("Successfully ran fluxsite tasks")
 
     def fluxsite_bitwise_cmp(self, config_path: str):
@@ -285,8 +291,9 @@ def fluxsite_bitwise_cmp(self, config_path: str):
                 "nccmp/1.8.5.0"
             )  # use `nccmp -df` for bitwise comparisons
 
-        tasks = self.tasks if self.tasks else self._initialise_tasks(config)
-        comparisons = get_fluxsite_comparisons(tasks)
+        comparisons = fluxsite.get_fluxsite_comparisons(
+            self._get_fluxsite_tasks(config)
+        )
 
         self.logger.info("Running comparison tasks...")
         if config["fluxsite"]["multiprocess"]:
@@ -308,10 +315,44 @@ def fluxsite(self, config_path: str, no_submit: bool, skip: list[str]):
         else:
             self.fluxsite_submit_job(config_path, skip)
 
-    def spatial(self, config_path: str):
+    def spatial_setup_work_directory(self, config_path: str):
+        """Endpoint for `benchcab spatial-setup-work-dir`."""
+        config = self._get_config(config_path)
+        self._validate_environment(project=config["project"], modules=config["modules"])
+
+        self.logger.info("Setting up run directory tree for spatial tests...")
+        setup_spatial_directory_tree()
+        self.logger.info("Setting up tasks...")
+        try:
+            payu_config = config["spatial"]["payu"]["config"]
+        except KeyError:
+            payu_config = None
+        for task in self._get_spatial_tasks(config):
+            task.setup_task(payu_config=payu_config)
+        self.logger.info("Successfully setup spatial tasks")
+
+    def spatial_run_tasks(self, config_path: str):
+        """Endpoint for `benchcab spatial-run-tasks`."""
+        config = self._get_config(config_path)
+        self._validate_environment(project=config["project"], modules=config["modules"])
+
+        self.logger.info("Running spatial tasks...")
+        spatial.run_tasks(tasks=self._get_spatial_tasks(config))
+        self.logger.info("Successfully dispatched payu jobs")
+
+    def spatial(self, config_path: str, skip: list):
         """Endpoint for `benchcab spatial`."""
+        self.checkout(config_path)
+        self.build(config_path, mpi=True)
+        self.spatial_setup_work_directory(config_path)
+        self.spatial_run_tasks(config_path)
 
-    def run(self, config_path: str, no_submit: bool, skip: list[str]):
+    def run(self, config_path: str, skip: list[str]):
         """Endpoint for `benchcab run`."""
-        self.fluxsite(config_path, no_submit, skip)
-        self.spatial(config_path)
+        self.checkout(config_path)
+        self.build(config_path)
+        self.build(config_path, mpi=True)
+        self.fluxsite_setup_work_directory(config_path)
+        self.spatial_setup_work_directory(config_path)
+        self.fluxsite_submit_job(config_path, skip)
+        self.spatial_run_tasks(config_path)

benchcab/cli.py

@@ -38,9 +38,9 @@ def generate_parser(app: Benchcab) -> argparse.ArgumentParser:
         action="store_true",
     )
 
-    # parent parser that contains arguments common to all run specific subcommands
-    args_run_subcommand = argparse.ArgumentParser(add_help=False)
-    args_run_subcommand.add_argument(
+    # parent parser that contains the argument for --no-submit
+    args_no_submit = argparse.ArgumentParser(add_help=False)
+    args_no_submit.add_argument(
         "--no-submit",
         action="store_true",
         help="Force benchcab to execute tasks on the current compute node.",
@@ -80,7 +80,6 @@ def generate_parser(app: Benchcab) -> argparse.ArgumentParser:
         parents=[
             args_help,
             args_subcommand,
-            args_run_subcommand,
             args_composite_subcommand,
         ],
         help="Run all test suites for CABLE.",
@@ -109,7 +108,7 @@ def generate_parser(app: Benchcab) -> argparse.ArgumentParser:
         parents=[
             args_help,
             args_subcommand,
-            args_run_subcommand,
+            args_no_submit,
             args_composite_subcommand,
         ],
         help="Run the fluxsite test suite for CABLE.",
@@ -140,6 +139,11 @@ def generate_parser(app: Benchcab) -> argparse.ArgumentParser:
         config file.""",
         add_help=False,
     )
+    parser_build.add_argument(
+        "--mpi",
+        action="store_true",
+        help="Enable MPI build.",
+    )
     parser_build.set_defaults(func=app.build)
 
     # subcommand: 'benchcab fluxsite-setup-work-dir'
@@ -168,9 +172,9 @@ def generate_parser(app: Benchcab) -> argparse.ArgumentParser:
         "fluxsite-run-tasks",
         parents=[args_help, args_subcommand],
         help="Run the fluxsite tasks of the main fluxsite command.",
-        description="""Runs the fluxsite tasks for the fluxsite test suite. Note, this command should
-        ideally be run inside a PBS job. This command is invoked by the PBS job script generated by
-        `benchcab run`.""",
+        description="""Runs the fluxsite tasks for the fluxsite test suite.
+        Note, this command should ideally be run inside a PBS job. This command
+        is invoked by the PBS job script generated by `benchcab run`.""",
         add_help=False,
     )
     parser_fluxsite_run_tasks.set_defaults(func=app.fluxsite_run_tasks)
@@ -192,11 +196,32 @@ def generate_parser(app: Benchcab) -> argparse.ArgumentParser:
     # subcommand: 'benchcab spatial'
     parser_spatial = subparsers.add_parser(
         "spatial",
-        parents=[args_help, args_subcommand],
+        parents=[args_help, args_subcommand, args_composite_subcommand],
         help="Run the spatial tests only.",
         description="""Runs the default spatial test suite for CABLE.""",
         add_help=False,
     )
     parser_spatial.set_defaults(func=app.spatial)
 
+    # subcommand: 'benchcab spatial-setup-work-dir'
+    parser_spatial_setup_work_dir = subparsers.add_parser(
+        "spatial-setup-work-dir",
+        parents=[args_help, args_subcommand],
+        help="Run the work directory setup step of the spatial command.",
+        description="""Generates the spatial run directory tree in the current working
+        directory so that spatial tasks can be run.""",
+        add_help=False,
+    )
+    parser_spatial_setup_work_dir.set_defaults(func=app.spatial_setup_work_directory)
+
+    # subcommand 'benchcab spatial-run-tasks'
+    parser_spatial_run_tasks = subparsers.add_parser(
+        "spatial-run-tasks",
+        parents=[args_help, args_subcommand],
+        help="Run the spatial tasks of the main spatial command.",
+        description="Runs the spatial tasks for the spatial test suite.",
+        add_help=False,
+    )
+    parser_spatial_run_tasks.set_defaults(func=app.spatial_run_tasks)
+
     return main_parser

benchcab/config.py

@@ -96,6 +96,18 @@ def read_optional_key(config: dict):
         "science_configurations", internal.DEFAULT_SCIENCE_CONFIGURATIONS
     )
 
+    # Default values for spatial
+    config["spatial"] = config.get("spatial", {})
+
+    config["spatial"]["met_forcings"] = config["spatial"].get(
+        "met_forcings", internal.SPATIAL_DEFAULT_MET_FORCINGS
+    )
+
+    config["spatial"]["payu"] = config["spatial"].get("payu", {})
+    config["spatial"]["payu"]["config"] = config["spatial"]["payu"].get("config", {})
+    config["spatial"]["payu"]["args"] = config["spatial"]["payu"].get("args")
+
+    # Default values for fluxsite
     config["fluxsite"] = config.get("fluxsite", {})
 
     config["fluxsite"]["multiprocess"] = config["fluxsite"].get(

benchcab/data/config-schema.yml

@@ -97,4 +97,27 @@ fluxsite:
           schema:
             type: "string"
             required: false
-            
+
+spatial:
+  type: "dict"
+  required: false
+  schema:
+    met_forcings:
+      type: "dict"
+      required: false
+      minlength: 1
+      keysrules:
+        type: "string"
+      valuesrules:
+        type: "string"
+    payu:
+      type: "dict"
+      required: false
+      schema:
+        config:
+          type: "dict"
+          required: false
+        args:
+          nullable: true
+          type: "string"
+          required: false

benchcab/data/test/config-optional.yml

@@ -11,6 +11,14 @@ fluxsite:
     storage:
       - scratch/$PROJECT
 
+spatial:
+  met_forcings:
+    crujra_access: https://github.com/CABLE-LSM/cable_example.git
+  payu:
+    config:
+      walltime: "1:00:00"
+    args: -n 2
+
 science_configurations:
   - cable:
       cable_user: 

benchcab/environment_modules.py

@@ -14,7 +14,7 @@
 try:
     from python import module
 except ImportError:
-    print(
+    get_logger().error(
         "Environment modules error: unable to import "
         "initialization script for python."
     )