spatial wip

Author: SeanBryan51

.conda/benchcab-dev.yaml

@@ -10,3 +10,4 @@ dependencies:
   - pytest-cov
   - pyyaml
   - flatdict
+  - gitpython

benchcab/benchcab.py

@@ -11,15 +11,14 @@
 from benchcab import internal
 from benchcab.internal import get_met_forcing_file_names
 from benchcab.config import read_config
-from benchcab.workdir import setup_fluxsite_directory_tree, setup_src_dir
-from benchcab.repository import CableRepository
-from benchcab.fluxsite import (
-    get_fluxsite_tasks,
-    get_fluxsite_comparisons,
-    run_tasks,
-    run_tasks_in_parallel,
-    Task,
+from benchcab.workdir import (
+    setup_src_dir,
+    setup_fluxsite_directory_tree,
+    setup_spatial_directory_tree,
 )
+from benchcab.repository import CableRepository
+from benchcab import fluxsite
+from benchcab import spatial
 from benchcab.comparison import run_comparisons, run_comparisons_in_parallel
 from benchcab.cli import generate_parser
 from benchcab.environment_modules import EnvironmentModules, EnvironmentModulesInterface
@@ -48,7 +47,15 @@ def __init__(
             CableRepository(**config, repo_id=id)
             for id, config in enumerate(self.config["realisations"])
         ]
-        self.tasks: list[Task] = []  # initialise fluxsite tasks lazily
+        self.science_configurations = self.config.get(
+            "science_configurations", internal.DEFAULT_SCIENCE_CONFIGURATIONS
+        )
+        self.fluxsite_tasks: list[
+            fluxsite.FluxsiteTask
+        ] = []  # initialise fluxsite tasks lazily
+        self.spatial_tasks = spatial.get_spatial_tasks(
+            repos=self.repos, science_configurations=self.science_configurations
+        )
         self.benchcab_exe_path = benchcab_exe_path
 
         if validate_env:
@@ -103,18 +110,16 @@ def _validate_environment(self, project: str, modules: list):
                 )
                 sys.exit(1)
 
-    def _initialise_tasks(self) -> list[Task]:
+    def _initialise_tasks(self) -> list[fluxsite.FluxsiteTask]:
         """A helper method that initialises and returns the `tasks` attribute."""
-        self.tasks = get_fluxsite_tasks(
+        self.fluxsite_tasks = fluxsite.get_fluxsite_tasks(
             repos=self.repos,
-            science_configurations=self.config.get(
-                "science_configurations", internal.DEFAULT_SCIENCE_CONFIGURATIONS
-            ),
+            science_configurations=self.science_configurations,
             fluxsite_forcing_file_names=get_met_forcing_file_names(
                 self.config["experiment"]
             ),
         )
-        return self.tasks
+        return self.fluxsite_tasks
 
     def fluxsite_submit_job(self) -> None:
         """Submits the PBS job script step in the fluxsite test workflow."""
@@ -189,27 +194,40 @@ def checkout(self):
 
         print("")
 
-    def build(self):
+    def build(self, mpi=None, clean=False):
         """Endpoint for `benchcab build`."""
         for repo in self.repos:
-            repo.build(modules=self.config["modules"], verbose=self.args.verbose)
+            repo.build(
+                modules=self.config["modules"],
+                verbose=self.args.verbose,
+                mpi=mpi if mpi else self.args.mpi,
+                clean=clean,
+            )
             print(f"Successfully compiled CABLE for realisation {repo.name}")
         print("")
 
     def fluxsite_setup_work_directory(self):
         """Endpoint for `benchcab fluxsite-setup-work-dir`."""
-        tasks = self.tasks if self.tasks else self._initialise_tasks()
+
+        if not self.fluxsite_tasks:
+            self._initialise_tasks()
+
         print("Setting up run directory tree for fluxsite tests...")
-        setup_fluxsite_directory_tree(fluxsite_tasks=tasks, verbose=self.args.verbose)
+        setup_fluxsite_directory_tree(
+            fluxsite_tasks=self.fluxsite_tasks, verbose=self.args.verbose
+        )
         print("Setting up tasks...")
-        for task in tasks:
+        for task in self.fluxsite_tasks:
             task.setup_task(verbose=self.args.verbose)
         print("Successfully setup fluxsite tasks")
         print("")
 
     def fluxsite_run_tasks(self):
         """Endpoint for `benchcab fluxsite-run-tasks`."""
-        tasks = self.tasks if self.tasks else self._initialise_tasks()
+
+        if not self.fluxsite_tasks:
+            self._initialise_tasks()
+
         print("Running fluxsite tasks...")
         try:
             multiprocess = self.config["fluxsite"]["multiprocess"]
@@ -219,9 +237,11 @@ def fluxsite_run_tasks(self):
             ncpus = self.config.get("pbs", {}).get(
                 "ncpus", internal.FLUXSITE_DEFAULT_PBS["ncpus"]
             )
-            run_tasks_in_parallel(tasks, n_processes=ncpus, verbose=self.args.verbose)
+            fluxsite.run_tasks_in_parallel(
+                self.fluxsite_tasks, n_processes=ncpus, verbose=self.args.verbose
+            )
         else:
-            run_tasks(tasks, verbose=self.args.verbose)
+            fluxsite.run_tasks(self.fluxsite_tasks, verbose=self.args.verbose)
         print("Successfully ran fluxsite tasks")
         print("")
 
@@ -233,8 +253,10 @@ def fluxsite_bitwise_cmp(self):
                 "nccmp/1.8.5.0"
             )  # use `nccmp -df` for bitwise comparisons
 
-        tasks = self.tasks if self.tasks else self._initialise_tasks()
-        comparisons = get_fluxsite_comparisons(tasks)
+        if not self.fluxsite_tasks:
+            self._initialise_tasks()
+
+        comparisons = fluxsite.get_fluxsite_comparisons(self.fluxsite_tasks)
 
         print("Running comparison tasks...")
         try:
@@ -265,13 +287,38 @@ def fluxsite(self):
         else:
             self.fluxsite_submit_job()
 
+    def spatial_setup_work_directory(self):
+        """Endpoint for `benchcab spatial-setup-work-dir`."""
+        print("Setting up run directory tree for spatial tests...")
+        setup_spatial_directory_tree()
+        print("Setting up tasks...")
+        for task in self.spatial_tasks:
+            task.setup_task(verbose=self.args.verbose)
+        print("Successfully setup spatial tasks")
+        print("")
+
+    def spatial_run_tasks(self):
+        """Endpoint for `benchcab spatial-run-tasks`."""
+        print("Running spatial tasks...")
+        spatial.run_tasks(tasks=self.spatial_tasks, verbose=self.args.verbose)
+        print("")
+
     def spatial(self):
         """Endpoint for `benchcab spatial`."""
+        self.checkout()
+        self.build(mpi=True)
+        self.spatial_setup_work_directory()
+        self.spatial_run_tasks()
 
     def run(self):
         """Endpoint for `benchcab run`."""
-        self.fluxsite()
-        self.spatial()
+        self.checkout()
+        self.build()
+        self.build(mpi=True, clean=True)
+        self.fluxsite_setup_work_directory()
+        self.spatial_setup_work_directory()
+        self.fluxsite_submit_job()
+        self.spatial_run_tasks()
 
     def main(self):
         """Main function for `benchcab`."""

benchcab/cli.py

@@ -32,9 +32,9 @@ def generate_parser() -> 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.",
@@ -74,7 +74,6 @@ def generate_parser() -> argparse.ArgumentParser:
         parents=[
             args_help,
             args_subcommand,
-            args_run_subcommand,
             args_composite_subcommand,
         ],
         help="Run all test suites for CABLE.",
@@ -89,7 +88,7 @@ def generate_parser() -> argparse.ArgumentParser:
         parents=[
             args_help,
             args_subcommand,
-            args_run_subcommand,
+            args_no_submit,
             args_composite_subcommand,
         ],
         help="Run the fluxsite test suite for CABLE.",
@@ -110,14 +109,19 @@ def generate_parser() -> argparse.ArgumentParser:
     )
 
     # subcommand: 'benchcab build'
-    subparsers.add_parser(
+    build_parser = subparsers.add_parser(
         "build",
         parents=[args_help, args_subcommand],
         help="Run the build step in the benchmarking workflow.",
         description="""Build the CABLE offline executable for each repository specified in the
         config file.""",
         add_help=False,
     )
+    build_parser.add_argument(
+        "--mpi",
+        action="store_true",
+        help="Enable MPI build.",
+    )
 
     # subcommand: 'benchcab fluxsite-setup-work-dir'
     subparsers.add_parser(

benchcab/fluxsite.py

@@ -66,7 +66,7 @@ class CableError(Exception):
     """Custom exception class for CABLE errors."""
 
 
-class Task:
+class FluxsiteTask:
     """A class used to represent a single fluxsite task."""
 
     root_dir: Path = internal.CWD
@@ -312,10 +312,10 @@ def get_fluxsite_tasks(
     repos: list[CableRepository],
     science_configurations: list[dict],
     fluxsite_forcing_file_names: list[str],
-) -> list[Task]:
+) -> list[FluxsiteTask]:
     """Returns a list of fluxsite tasks to run."""
     tasks = [
-        Task(
+        FluxsiteTask(
             repo=repo,
             met_forcing_file=file_name,
             sci_conf_id=sci_conf_id,
@@ -328,14 +328,16 @@ def get_fluxsite_tasks(
     return tasks
 
 
-def run_tasks(tasks: list[Task], verbose=False):
+def run_tasks(tasks: list[FluxsiteTask], verbose=False):
     """Runs tasks in `tasks` serially."""
     for task in tasks:
         task.run(verbose=verbose)
 
 
 def run_tasks_in_parallel(
-    tasks: list[Task], n_processes=internal.FLUXSITE_DEFAULT_PBS["ncpus"], verbose=False
+    tasks: list[FluxsiteTask],
+    n_processes=internal.FLUXSITE_DEFAULT_PBS["ncpus"],
+    verbose=False,
 ):
     """Runs tasks in `tasks` in parallel across multiple processes."""
 
@@ -364,7 +366,7 @@ def worker_run(task_queue: multiprocessing.Queue, verbose=False):
 
 
 def get_fluxsite_comparisons(
-    tasks: list[Task], root_dir=internal.CWD
+    tasks: list[FluxsiteTask], root_dir=internal.CWD
 ) -> list[ComparisonTask]:
     """Returns a list of `ComparisonTask` objects to run comparisons with.
 

benchcab/internal.py

@@ -17,7 +17,6 @@
     "walltime": "6:00:00",
     "storage": [],
 }
-MPI = False
 FLUXSITE_DEFAULT_MULTIPROCESS = True
 
 # DIRECTORY PATHS/STRUCTURE:
@@ -76,14 +75,38 @@
 # Relative path to directory that stores bitwise comparison results
 FLUXSITE_BITWISE_CMP_DIR = FLUXSITE_ANALYSIS_DIR / "bitwise-comparisons"
 
+# Relative path to root directory for CABLE spatial runs
+SPATIAL_RUN_DIR = RUN_DIR / "spatial"
+
+# Relative path to tasks directory (contains payu control directories configured
+# for each spatial task)
+SPATIAL_TASKS_DIR = SPATIAL_RUN_DIR / "tasks"
+
+# A custom payu laboratory directory for payu runs
+PAYU_LABORATORY_DIR = RUN_DIR / "payu-laboratory"
+
+# URL to a payu experiment template for offline spatial runs:
+EXPERIMENT_TEMPLATE_SPATIAL = "https://github.com/CABLE-LSM/cable_example.git"
+
+# TODO(Sean) add more info here e.g. site forcing data, what is PLUMBER2?
+# Observations from where? What has been done to the data? Is it available?:
 # Path to met files:
 MET_DIR = Path("/g/data/ks32/CLEX_Data/PLUMBER2/v1-0/Met/")
 
+# TODO(Sean) add more info here
+SPATIAL_MET_FORCING_CRUJRA = {
+    "name": "CRUJRA-ACCESS",
+    "path": "/g/data/tm70/ccc561/CABLE/CABLE-as-ACCESS_Yingping/cruncep10",
+}
+
 # CABLE SVN root url:
 CABLE_SVN_ROOT = "https://trac.nci.org.au/svn/cable"
 
 # CABLE executable file name:
-CABLE_EXE = "cable-mpi" if MPI else "cable"
+CABLE_EXE = "cable"
+
+# CABLE MPI executable file name:
+CABLE_MPI_EXE = "cable-mpi"
 
 # CABLE namelist file name:
 CABLE_NML = "cable.nml"

benchcab/repository.py

@@ -72,7 +72,7 @@ def svn_info_show_item(self, item: str) -> str:
         )
         return proc.stdout.strip()
 
-    def build(self, modules: list[str], verbose=False) -> None:
+    def build(self, modules: list[str], mpi=False, clean=False, verbose=False) -> None:
         """Build CABLE using the default build script or a custom build script."""
 
         if self.build_script:
@@ -82,7 +82,7 @@ def build(self, modules: list[str], verbose=False) -> None:
             )
         else:
             print(
-                f"Compiling CABLE {'with MPI' if internal.MPI else 'serially'} for "
+                f"Compiling CABLE {'with MPI' if mpi else 'serially'} for "
                 f"realisation {self.name}..."
             )
 
@@ -100,6 +100,16 @@ def build(self, modules: list[str], verbose=False) -> None:
                 "'build_script' option in config.yaml?",
             )
 
+        if clean:
+            # TODO(Sean) ideally we should not have to do this. We do this
+            # because the build script is designed to build the serial and MPI
+            # executable using the same .tmp directory. This forces us to delete
+            # the .tmp directory everytime we switch from building serial to
+            # building with MPI and visa versa. A simple solution to this would
+            # be to have a separate temporary directory for serial and MPI
+            # builds.
+            shutil.rmtree(build_script_path.parent / ".tmp")
+
         tmp_script_path = build_script_path.parent / "tmp-build.sh"
 
         if verbose:
@@ -118,7 +128,7 @@ def build(self, modules: list[str], verbose=False) -> None:
         remove_module_lines(tmp_script_path)
 
         args: list[str] = []
-        if internal.MPI and self.build_script is None:
+        if mpi and self.build_script is None:
             args.append("mpi")
 
         with chdir(build_script_path.parent), self.modules_handler.load(