diff --git a/.conda/benchcab-dev.yaml b/.conda/benchcab-dev.yaml index 0228564c..5cd0cf84 100644 --- a/.conda/benchcab-dev.yaml +++ b/.conda/benchcab-dev.yaml @@ -9,3 +9,4 @@ dependencies: - pytest-cov - pyyaml - flatdict + - gitpython diff --git a/benchcab/benchcab.py b/benchcab/benchcab.py index 878dcdb0..dc3adab4 100644 --- a/benchcab/benchcab.py +++ b/benchcab/benchcab.py @@ -1,5 +1,6 @@ """Contains the main program entry point for `benchcab`.""" +import functools import grp import os import shutil @@ -8,24 +9,20 @@ from subprocess import CalledProcessError from typing import Optional -from benchcab import internal +from benchcab import fluxsite, internal, spatial from benchcab.cli import generate_parser 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.repository import CableRepository from benchcab.utils.fs import mkdir, next_path from benchcab.utils.pbs import render_job_script 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: @@ -48,7 +45,6 @@ def __init__( CableRepository(**config, repo_id=id) for id, config in enumerate(self.config["realisations"]) ] - self.tasks: list[Task] = [] # initialise fluxsite tasks lazily self.benchcab_exe_path = benchcab_exe_path if validate_env: @@ -102,9 +98,10 @@ def _validate_environment(self, project: str, modules: list): ) sys.exit(1) - def _initialise_tasks(self) -> list[Task]: - """A helper method that initialises and returns the `tasks` attribute.""" - self.tasks = get_fluxsite_tasks( + @functools.cache + def _fluxsite_tasks(self) -> list[fluxsite.FluxsiteTask]: + """Generate the list of fluxsite tasks and cache the result.""" + return fluxsite.get_fluxsite_tasks( repos=self.repos, science_configurations=self.config.get( "science_configurations", internal.DEFAULT_SCIENCE_CONFIGURATIONS @@ -113,7 +110,18 @@ def _initialise_tasks(self) -> list[Task]: self.config["experiment"] ), ) - return self.tasks + + @functools.cache + def _spatial_tasks(self) -> list[spatial.SpatialTask]: + """Generate the list of spatial tasks and cache the result.""" + return spatial.get_spatial_tasks( + repos=self.repos, + met_forcings=internal.SPATIAL_DEFAULT_FORCINGS, + science_configurations=self.config.get( + "science_configurations", internal.DEFAULT_SCIENCE_CONFIGURATIONS + ), + payu_args=self.config.get("spatial", {}).get("payu", {}).get("args"), + ) def fluxsite_submit_job(self) -> None: """Submits the PBS job script step in the fluxsite test workflow.""" @@ -187,7 +195,7 @@ def checkout(self): print("") - def build(self): + def build(self, mpi=False): """Endpoint for `benchcab build`.""" for repo in self.repos: if repo.build_script: @@ -199,30 +207,28 @@ def build(self): modules=self.config["modules"], verbose=self.args.verbose ) else: - build_mode = "with MPI" if internal.MPI else "serially" + build_mode = "with MPI" if mpi else "serially" print(f"Compiling CABLE {build_mode} for realisation {repo.name}...") - repo.pre_build(verbose=self.args.verbose) + repo.pre_build(mpi=mpi, verbose=self.args.verbose) repo.run_build( - modules=self.config["modules"], verbose=self.args.verbose + modules=self.config["modules"], mpi=mpi, verbose=self.args.verbose ) - repo.post_build(verbose=self.args.verbose) + repo.post_build(mpi=mpi, verbose=self.args.verbose) 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() print("Setting up run directory tree for fluxsite tests...") setup_fluxsite_directory_tree(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() print("Running fluxsite tasks...") try: multiprocess = self.config["fluxsite"]["multiprocess"] @@ -232,9 +238,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("") @@ -245,8 +253,7 @@ 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) + comparisons = fluxsite.get_fluxsite_comparisons(self._fluxsite_tasks()) print("Running comparison tasks...") try: @@ -277,13 +284,42 @@ 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...") + try: + payu_config = self.config["spatial"]["payu"]["config"] + except KeyError: + payu_config = None + for task in self._spatial_tasks(): + task.setup_task(payu_config=payu_config, 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) + 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`.""" @@ -294,7 +330,7 @@ def main(self): self.checkout() if self.args.subcommand == "build": - self.build() + self.build(mpi=self.args.mpi) if self.args.subcommand == "fluxsite": self.fluxsite() @@ -314,6 +350,12 @@ def main(self): if self.args.subcommand == "spatial": self.spatial() + if self.args.subcommand == "spatial-setup-work-dir": + self.spatial_setup_work_directory() + + if self.args.subcommand == "spatial-run-tasks": + self.spatial_run_tasks() + def main(): """Main program entry point for `benchcab`. diff --git a/benchcab/cli.py b/benchcab/cli.py index 48bf9ef9..b3b6827e 100644 --- a/benchcab/cli.py +++ b/benchcab/cli.py @@ -30,9 +30,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.", @@ -72,7 +72,6 @@ def generate_parser() -> argparse.ArgumentParser: parents=[ args_help, args_subcommand, - args_run_subcommand, args_composite_subcommand, ], help="Run all test suites for CABLE.", @@ -87,7 +86,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.", @@ -108,7 +107,7 @@ 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.", @@ -116,6 +115,11 @@ def generate_parser() -> argparse.ArgumentParser: 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( @@ -141,9 +145,9 @@ def generate_parser() -> 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, ) @@ -163,10 +167,29 @@ def generate_parser() -> argparse.ArgumentParser: # subcommand: 'benchcab 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, ) + # subcommand: 'benchcab 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, + ) + + # subcommand 'benchcab 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, + ) + return main_parser diff --git a/benchcab/config.py b/benchcab/config.py index f6aaabf4..df027e9c 100644 --- a/benchcab/config.py +++ b/benchcab/config.py @@ -92,6 +92,34 @@ def check_config(config: dict): msg = "The 'multiprocessing' key must be a boolean." raise TypeError(msg) + # the "spatial" key is optional + if "spatial" in config: + if not isinstance(config["spatial"], dict): + msg = "The 'spatial' key must be a dictionary." + raise TypeError(msg) + # the "payu" key is optional + if "payu" in config["spatial"]: + if not isinstance(config["spatial"]["payu"], dict): + msg = "The 'payu' key must be a dictionary." + raise TypeError(msg) + # the "config" key is optional + if "config" in config["spatial"]["payu"]: + if not isinstance(config["spatial"]["payu"]["config"], dict): + msg = "The 'config' key must be a dictionary." + raise TypeError(msg) + # the "args" key is optional + if "args" in config["spatial"]["payu"]: + if not isinstance(config["spatial"]["payu"]["args"], str): + msg = "The 'args' key must be a string." + raise TypeError(msg) + # the "met_forcings" key is optional + if "met_forcings" in config["spatial"]: + if not isinstance(config["spatial"]["met_forcings"], list) or any( + not isinstance(val, str) for val in config["spatial"]["met_forcings"] + ): + msg = "The 'met_forcings' key must be a list of strings." + raise TypeError(msg) + valid_experiments = ( list(internal.MEORG_EXPERIMENTS) + internal.MEORG_EXPERIMENTS["five-site-test"] ) diff --git a/benchcab/fluxsite.py b/benchcab/fluxsite.py index 6bca611e..08c02a7b 100644 --- a/benchcab/fluxsite.py +++ b/benchcab/fluxsite.py @@ -6,7 +6,6 @@ import sys from pathlib import Path from subprocess import CalledProcessError -from typing import Any, Dict, TypeVar import f90nml import flatdict @@ -16,70 +15,9 @@ from benchcab.comparison import ComparisonTask from benchcab.repository import CableRepository from benchcab.utils.fs import chdir, mkdir +from benchcab.utils.namelist import patch_namelist, patch_remove_namelist from benchcab.utils.subprocess import SubprocessWrapper, SubprocessWrapperInterface -# fmt: off -# ====================================================== -# Copyright (c) 2017 - 2022 Samuel Colvin and other contributors -# from https://github.com/pydantic/pydantic/blob/fd2991fe6a73819b48c906e3c3274e8e47d0f761/pydantic/utils.py#L200 - -KeyType = TypeVar('KeyType') - - -def deep_update(mapping: Dict[KeyType, Any], *updating_mappings: Dict[KeyType, Any]) -> Dict[KeyType, Any]: - updated_mapping = mapping.copy() - for updating_mapping in updating_mappings: - for k, v in updating_mapping.items(): - if k in updated_mapping and isinstance(updated_mapping[k], dict) and isinstance(v, dict): - updated_mapping[k] = deep_update(updated_mapping[k], v) - else: - updated_mapping[k] = v - return updated_mapping - -# ====================================================== -# fmt: on - - -def deep_del( - mapping: Dict[KeyType, Any], *updating_mappings: Dict[KeyType, Any] -) -> Dict[KeyType, Any]: - """Deletes all key-value 'leaf nodes' in `mapping` specified by `updating_mappings`.""" - updated_mapping = mapping.copy() - for updating_mapping in updating_mappings: - for key, value in updating_mapping.items(): - if isinstance(updated_mapping[key], dict) and isinstance(value, dict): - updated_mapping[key] = deep_del(updated_mapping[key], value) - else: - del updated_mapping[key] - return updated_mapping - - -def patch_namelist(nml_path: Path, patch: dict): - """Writes a namelist patch specified by `patch` to `nml_path`. - - The `patch` dictionary must comply with the `f90nml` api. - """ - if not nml_path.exists(): - f90nml.write(patch, nml_path) - return - - nml = f90nml.read(nml_path) - f90nml.write(deep_update(nml, patch), nml_path, force=True) - - -def patch_remove_namelist(nml_path: Path, patch_remove: dict): - """Removes a subset of namelist parameters specified by `patch_remove` from `nml_path`. - - The `patch_remove` dictionary must comply with the `f90nml` api. - """ - nml = f90nml.read(nml_path) - try: - f90nml.write(deep_del(nml, patch_remove), nml_path, force=True) - except KeyError as exc: - msg = f"Namelist parameters specified in `patch_remove` do not exist in {nml_path.name}." - raise KeyError(msg) from exc - - f90_logical_repr = {True: ".true.", False: ".false."} @@ -87,7 +25,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 @@ -361,10 +299,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, @@ -377,14 +315,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.""" run_task = operator.methodcaller("run", verbose=verbose) @@ -393,7 +333,7 @@ def run_tasks_in_parallel( 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. diff --git a/benchcab/internal.py b/benchcab/internal.py index c6a54d2a..1af50a14 100644 --- a/benchcab/internal.py +++ b/benchcab/internal.py @@ -16,7 +16,6 @@ "walltime": "6:00:00", "storage": [], } -MPI = False FLUXSITE_DEFAULT_MULTIPROCESS = True # DIRECTORY PATHS/STRUCTURE: @@ -59,6 +58,7 @@ # Fluxsite directory tree FLUXSITE_DIRS = {} + # Relative path to root directory for CABLE fluxsite runs FLUXSITE_DIRS["RUN"] = RUN_DIR / "fluxsite" @@ -75,18 +75,44 @@ FLUXSITE_DIRS["ANALYSIS"] = FLUXSITE_DIRS["RUN"] / "analysis" # Relative path to directory that stores bitwise comparison results -FLUXSITE_DIRS["BITWISE_CMP"] = ( - FLUXSITE_DIRS["ANALYSIS"] / "bitwise-comparisons" -) +FLUXSITE_DIRS["BITWISE_CMP"] = FLUXSITE_DIRS["ANALYSIS"] / "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" -# Path to met files: +# Path to PLUMBER2 site forcing data directory (doi: 10.25914/5fdb0902607e1): MET_DIR = Path("/g/data/ks32/CLEX_Data/PLUMBER2/v1-0/Met/") +# Path to spatial meteorological forcing data set collection: +SPATIAL_MET_DIR = Path("/g/data/tm70/sb8430/cable/MetForcing") + +# Default data set names of meteorolgical forcings to use for spatial runs: +SPATIAL_DEFAULT_FORCINGS = ["crujra_access"] + # CABLE SVN root url: CABLE_SVN_ROOT = "https://trac.nci.org.au/svn/cable" +# Relative path to temporary build directory (serial) +TMP_BUILD_DIR = Path("offline", ".tmp") + +# Relative path to temporary build directory (MPI) +TMP_BUILD_DIR_MPI = Path("offline", ".mpitmp") + # 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" diff --git a/benchcab/repository.py b/benchcab/repository.py index 009a8790..7a3a0e02 100644 --- a/benchcab/repository.py +++ b/benchcab/repository.py @@ -115,10 +115,12 @@ def custom_build(self, modules: list[str], verbose=False): verbose=verbose, ) - def pre_build(self, verbose=False): + def pre_build(self, mpi=False, verbose=False): """Runs CABLE pre-build steps.""" path_to_repo = self.root_dir / internal.SRC_DIR / self.name - tmp_dir = path_to_repo / "offline" / ".tmp" + tmp_dir = path_to_repo / ( + internal.TMP_BUILD_DIR_MPI if mpi else internal.TMP_BUILD_DIR + ) if not tmp_dir.exists(): if verbose: print(f"mkdir {tmp_dir.relative_to(self.root_dir)}") @@ -152,10 +154,12 @@ def pre_build(self, verbose=False): verbose=verbose, ) - def run_build(self, modules: list[str], verbose=False): + def run_build(self, modules: list[str], mpi=False, verbose=False): """Runs CABLE build scripts.""" path_to_repo = self.root_dir / internal.SRC_DIR / self.name - tmp_dir = path_to_repo / "offline" / ".tmp" + tmp_dir = path_to_repo / ( + internal.TMP_BUILD_DIR_MPI if mpi else internal.TMP_BUILD_DIR + ) with chdir(tmp_dir), self.modules_handler.load(modules, verbose=verbose): env = os.environ.copy() @@ -164,26 +168,29 @@ def run_build(self, modules: list[str], verbose=False): env["CFLAGS"] = "-O2 -fp-model precise" env["LDFLAGS"] = f"-L{env['NETCDF_ROOT']}/lib/Intel -O0" env["LD"] = "-lnetcdf -lnetcdff" - env["FC"] = "mpif90" if internal.MPI else "ifort" + env["FC"] = "mpif90" if mpi else "ifort" self.subprocess_handler.run_cmd( "make -f Makefile", env=env, verbose=verbose ) self.subprocess_handler.run_cmd( - f"./{'parallel_cable' if internal.MPI else 'serial_cable'} \"{env['FC']}\" " + f"./{'parallel_cable' if mpi else 'serial_cable'} \"{env['FC']}\" " f"\"{env['CFLAGS']}\" \"{env['LDFLAGS']}\" \"{env['LD']}\" \"{env['NCMOD']}\"", env=env, verbose=verbose, ) - def post_build(self, verbose=False): + def post_build(self, mpi=False, verbose=False): """Runs CABLE post-build steps.""" path_to_repo = self.root_dir / internal.SRC_DIR / self.name - tmp_dir = path_to_repo / "offline" / ".tmp" + tmp_dir = path_to_repo / ( + internal.TMP_BUILD_DIR_MPI if mpi else internal.TMP_BUILD_DIR + ) + exe = internal.CABLE_MPI_EXE if mpi else internal.CABLE_EXE rename( - (tmp_dir / internal.CABLE_EXE).relative_to(self.root_dir), - (path_to_repo / "offline" / internal.CABLE_EXE).relative_to(self.root_dir), + (tmp_dir / exe).relative_to(self.root_dir), + (path_to_repo / "offline" / exe).relative_to(self.root_dir), verbose=verbose, ) diff --git a/benchcab/spatial.py b/benchcab/spatial.py new file mode 100644 index 00000000..cc30c279 --- /dev/null +++ b/benchcab/spatial.py @@ -0,0 +1,227 @@ +"""A module containing functions and data structures for running spatial tasks.""" + +import shutil +from pathlib import Path +from typing import Optional + +import yaml + +from benchcab import internal +from benchcab.repository import CableRepository +from benchcab.utils.dict import deep_update +from benchcab.utils.fs import chdir +from benchcab.utils.git import GitRepo, GitRepoInterface +from benchcab.utils.namelist import patch_namelist, patch_remove_namelist +from benchcab.utils.subprocess import SubprocessWrapper, SubprocessWrapperInterface + + +class SpatialTask: + """A class used to represent a single spatial task.""" + + root_dir: Path = internal.CWD + subprocess_handler: SubprocessWrapperInterface = SubprocessWrapper() + git_repo_handler: GitRepoInterface = GitRepo() + + def __init__( + self, + repo: CableRepository, + met_forcing: str, + sci_conf_id: int, + sci_config: dict, + payu_args: Optional[str] = None, + ) -> None: + self.repo = repo + self.met_forcing = met_forcing + self.sci_conf_id = sci_conf_id + self.sci_config = sci_config + self.payu_args = payu_args + + def get_task_name(self) -> str: + """Returns the file name convention used for this task.""" + return f"{self.met_forcing}_R{self.repo.repo_id}_S{self.sci_conf_id}" + + def setup_task(self, payu_config: dict = None, verbose=False): + """Does all file manipulations to run cable with payu for this task.""" + + if verbose: + print(f"Setting up task: {self.get_task_name()}") + + task_dir = self.root_dir / internal.SPATIAL_TASKS_DIR / self.get_task_name() + + self.git_repo_handler.clone_from( + internal.EXPERIMENT_TEMPLATE_SPATIAL, + task_dir.relative_to(self.root_dir), + verbose=verbose, + ) + + if verbose: + print( + f" Copying namelist files from {self.root_dir / internal.NAMELIST_DIR} " + f"to {task_dir}" + ) + + shutil.copytree( + self.root_dir / internal.NAMELIST_DIR, task_dir, dirs_exist_ok=True + ) + + exp_config_path = task_dir / "config.yaml" + with exp_config_path.open("r", encoding="utf-8") as file: + config = yaml.safe_load(file) + if config is None: + config = {} + + if verbose: + print( + " Updating experiment config parameters in", + internal.SPATIAL_TASKS_DIR / self.get_task_name() / "config.yaml", + ) + + if payu_config: + config = deep_update(config, payu_config) + + config["exe"] = str( + self.root_dir + / internal.SRC_DIR + / self.repo.name + / "offline" + / internal.CABLE_MPI_EXE + ) + + config["input"] = [ + str(internal.SPATIAL_MET_DIR / self.met_forcing / "LWdown"), + str(internal.SPATIAL_MET_DIR / self.met_forcing / "PSurf"), + str(internal.SPATIAL_MET_DIR / self.met_forcing / "Qair"), + str(internal.SPATIAL_MET_DIR / self.met_forcing / "Rainf"), + str(internal.SPATIAL_MET_DIR / self.met_forcing / "Snowf"), + str(internal.SPATIAL_MET_DIR / self.met_forcing / "SWdown"), + str(internal.SPATIAL_MET_DIR / self.met_forcing / "Tair"), + str(internal.SPATIAL_MET_DIR / self.met_forcing / "Wind"), + str( + internal.SPATIAL_MET_DIR + / self.met_forcing + / f"{self.met_forcing}_landmask.nc" + ), + str( + internal.SPATIAL_MET_DIR + / self.met_forcing + / f"{self.met_forcing}_gridinfo.nc" + ), + str( + self.root_dir + / internal.CABLE_AUX_DIR + / "core" + / "biogeophys" + / "def_veg_params_zr_clitt_albedo_fix.txt" + ), + str( + self.root_dir + / internal.CABLE_AUX_DIR + / "core" + / "biogeophys" + / "def_soil_params.txt" + ), + ] + + config["laboratory"] = str(self.root_dir / internal.PAYU_LABORATORY_DIR) + + with exp_config_path.open("w", encoding="utf-8") as file: + yaml.dump(config, file) + + nml_path = ( + self.root_dir + / internal.SPATIAL_TASKS_DIR + / self.get_task_name() + / internal.CABLE_NML + ) + + if verbose: + print(f" Adding base configurations to CABLE namelist file {nml_path}") + patch_namelist( + nml_path, + { + "cable": { + "filename": { + "out": "cable_out.nc", + "log": "cable_log.nc", + "restart_in": "RESTART/restart.nc", + "restart_out": "RESTART/restart.nc", + "type": f"INPUT/{self.met_forcing}_gridinfo.nc", + "veg": "INPUT/def_veg_params_zr_clitt_albedo_fix.txt", + "soil": "INPUT/def_soil_params.txt", + }, + "spinup": False, + "cable_user": { + "gswp3": True, + "mettype": "gswp3", + }, + "gswpfile": { + "mask": f"INPUT/{self.met_forcing}_landmask.nc", + "LWdown": f"INPUT/{self.met_forcing}_LWdown_$year.nc", + "PSurf": f"INPUT/{self.met_forcing}_PSurf_$year.nc", + "Qair": f"INPUT/{self.met_forcing}_Qair_$year.nc", + "Rainf": f"INPUT/{self.met_forcing}_Rainf_$year.nc", + "Snowf": f"INPUT/{self.met_forcing}_Snowf_$year.nc", + "SWdown": f"INPUT/{self.met_forcing}_SWdown_$year.nc", + "Tair": f"INPUT/{self.met_forcing}_Tair_$year.nc", + "Wind": f"INPUT/{self.met_forcing}_Wind_$year.nc", + }, + }, + }, + ) + + if verbose: + print(f" Adding science configurations to CABLE namelist file {nml_path}") + patch_namelist(nml_path, self.sci_config) + + if self.repo.patch: + if verbose: + print( + f" Adding branch specific configurations to CABLE namelist file {nml_path}" + ) + patch_namelist(nml_path, self.repo.patch) + + if self.repo.patch_remove: + if verbose: + print( + f" Removing branch specific configurations from CABLE namelist file {nml_path}" + ) + patch_remove_namelist(nml_path, self.repo.patch_remove) + + def run(self, verbose=False) -> None: + """Runs a single spatial task.""" + + task_dir = self.root_dir / internal.SPATIAL_TASKS_DIR / self.get_task_name() + with chdir(task_dir): + self.subprocess_handler.run_cmd( + f"payu run {self.payu_args}" if self.payu_args else "payu run", + verbose=verbose, + ) + + +def run_tasks(tasks: list[SpatialTask], verbose=False): + """Runs tasks in `tasks` sequentially.""" + + for task in tasks: + task.run(verbose=verbose) + + +def get_spatial_tasks( + repos: list[CableRepository], + met_forcings: list[str], + science_configurations: list[dict], + payu_args: Optional[str] = None, +): + """Returns a list of spatial tasks to run.""" + tasks = [ + SpatialTask( + repo=repo, + met_forcing=met_forcing, + sci_conf_id=sci_conf_id, + sci_config=sci_config, + payu_args=payu_args, + ) + for repo in repos + for met_forcing in met_forcings + for sci_conf_id, sci_config in enumerate(science_configurations) + ] + return tasks diff --git a/benchcab/utils/dict.py b/benchcab/utils/dict.py new file mode 100644 index 00000000..f6fbc75e --- /dev/null +++ b/benchcab/utils/dict.py @@ -0,0 +1,38 @@ +"""Utility functions for manipulating nested dictionaries.""" + +from typing import Any, Dict, TypeVar + +# fmt: off +# ====================================================== +# Copyright (c) 2017 - 2022 Samuel Colvin and other contributors +# from https://github.com/pydantic/pydantic/blob/fd2991fe6a73819b48c906e3c3274e8e47d0f761/pydantic/utils.py#L200 + +KeyType = TypeVar('KeyType') + + +def deep_update(mapping: Dict[KeyType, Any], *updating_mappings: Dict[KeyType, Any]) -> Dict[KeyType, Any]: # noqa + updated_mapping = mapping.copy() + for updating_mapping in updating_mappings: + for k, v in updating_mapping.items(): + if k in updated_mapping and isinstance(updated_mapping[k], dict) and isinstance(v, dict): + updated_mapping[k] = deep_update(updated_mapping[k], v) + else: + updated_mapping[k] = v + return updated_mapping + +# ====================================================== +# fmt: on + + +def deep_del( + mapping: Dict[KeyType, Any], *updating_mappings: Dict[KeyType, Any] +) -> Dict[KeyType, Any]: + """Deletes all key-value 'leaf nodes' in `mapping` specified by `updating_mappings`.""" + updated_mapping = mapping.copy() + for updating_mapping in updating_mappings: + for key, value in updating_mapping.items(): + if isinstance(updated_mapping[key], dict) and isinstance(value, dict): + updated_mapping[key] = deep_del(updated_mapping[key], value) + else: + del updated_mapping[key] + return updated_mapping diff --git a/benchcab/utils/git.py b/benchcab/utils/git.py new file mode 100644 index 00000000..d80fbddf --- /dev/null +++ b/benchcab/utils/git.py @@ -0,0 +1,27 @@ +"""Contains a wrapper around the git API.""" + +from abc import ABC as AbstractBaseClass # noqa: N811 +from abc import abstractmethod + +import git + + +class GitRepoInterface(AbstractBaseClass): + """An abstract class (interface) for interacting with the git API. + + An interface is defined so that we can easily mock git API requests in our + unit tests. + """ + + @abstractmethod + def clone_from(self, url: str, to_path: str, verbose=False) -> None: + """A wrapper around `git clone `.""" + + +class GitRepo(GitRepoInterface): + """A concrete implementation of the `GitRepoInterface` abstract class.""" + + def clone_from(self, url: str, to_path: str, verbose=False) -> None: + if verbose: + print(f"git clone {url} {to_path}") + git.Repo.clone_from(url, to_path) diff --git a/benchcab/utils/namelist.py b/benchcab/utils/namelist.py new file mode 100644 index 00000000..2f5173fa --- /dev/null +++ b/benchcab/utils/namelist.py @@ -0,0 +1,33 @@ +"""Contains utility functions for manipulating Fortran namelist files.""" + +from pathlib import Path + +import f90nml + +from benchcab.utils.dict import deep_del, deep_update + + +def patch_namelist(nml_path: Path, patch: dict): + """Writes a namelist patch specified by `patch` to `nml_path`. + + The `patch` dictionary must comply with the `f90nml` api. + """ + if not nml_path.exists(): + f90nml.write(patch, nml_path) + return + + nml = f90nml.read(nml_path) + f90nml.write(deep_update(nml, patch), nml_path, force=True) + + +def patch_remove_namelist(nml_path: Path, patch_remove: dict): + """Removes a subset of namelist parameters specified by `patch_remove` from `nml_path`. + + The `patch_remove` dictionary must comply with the `f90nml` api. + """ + nml = f90nml.read(nml_path) + try: + f90nml.write(deep_del(nml, patch_remove), nml_path, force=True) + except KeyError as exc: + msg = f"Namelist parameters specified in `patch_remove` do not exist in {nml_path.name}." + raise KeyError(msg) from exc diff --git a/benchcab/workdir.py b/benchcab/workdir.py index 084c5bfd..50d8e796 100644 --- a/benchcab/workdir.py +++ b/benchcab/workdir.py @@ -25,3 +25,14 @@ def setup_fluxsite_directory_tree(verbose=False): """ for path in internal.FLUXSITE_DIRS.values(): mkdir(path, verbose=verbose, parents=True, exist_ok=True) + + +def setup_spatial_directory_tree(verbose=False): + """Generate the directory structure for running spatial tests.""" + + for path in [ + internal.SPATIAL_RUN_DIR, + internal.SPATIAL_TASKS_DIR, + internal.PAYU_LABORATORY_DIR, + ]: + mkdir(path, verbose=verbose, parents=True, exist_ok=True) diff --git a/tests/conftest.py b/tests/conftest.py index 807e0203..7f54c560 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,6 +10,7 @@ import pytest from benchcab.environment_modules import EnvironmentModulesInterface +from benchcab.utils.git import GitRepoInterface from benchcab.utils.subprocess import SubprocessWrapperInterface @@ -90,6 +91,17 @@ def config(): }, "multiprocessing": True, }, + "spatial": { + "met_forcings": ["crujra_access", "gswp3"], + "payu": { + "config": { + "ncpus": 16, + "walltime": "1:00:00", + "mem": "64GB", + }, + "args": "-n 2", + }, + }, } @@ -155,3 +167,25 @@ def module_unload(self, *args: str) -> None: self.commands.append("module unload " + " ".join(args)) return MockEnvironmentModules() + + +@pytest.fixture() +def mock_git_repo(): + """Returns a mock implementation of `GitRepoInterface`.""" + + class MockGitRepo(GitRepoInterface): + """A mock implementation of `GitRepoInterface` used for testing.""" + + def __init__(self) -> None: + self.files: list[Path] = [] + self.dirs: list[Path] = [] + + def clone_from(self, url: str, to_path: str, verbose=False) -> None: + repo_path = Path(to_path) + repo_path.mkdir() + for path in self.files: + (repo_path / path).touch() + for path in self.dirs: + (repo_path / path).mkdir() + + return MockGitRepo() diff --git a/tests/test_cli.py b/tests/test_cli.py index 8e77967d..3ec7fe89 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -14,7 +14,6 @@ def test_cli_parser(): assert res == { "subcommand": "run", "config": "config.yaml", - "no_submit": False, "verbose": False, "skip": [], } @@ -25,7 +24,12 @@ def test_cli_parser(): # Success case: default build command res = vars(parser.parse_args(["build"])) - assert res == {"subcommand": "build", "config": "config.yaml", "verbose": False} + assert res == { + "subcommand": "build", + "config": "config.yaml", + "verbose": False, + "mpi": False, + } # Success case: default fluxsite command res = vars(parser.parse_args(["fluxsite"])) @@ -76,6 +80,23 @@ def test_cli_parser(): "subcommand": "spatial", "config": "config.yaml", "verbose": False, + "skip": [], + } + + # Success case: default spatial-setup-work-dir command + res = vars(parser.parse_args(["spatial-setup-work-dir"])) + assert res == { + "subcommand": "spatial-setup-work-dir", + "config": "config.yaml", + "verbose": False, + } + + # Success case: default spatial-run-tasks command + res = vars(parser.parse_args(["spatial-run-tasks"])) + assert res == { + "subcommand": "spatial-run-tasks", + "config": "config.yaml", + "verbose": False, } # Failure case: pass --no-submit to a non 'run' command diff --git a/tests/test_fluxsite.py b/tests/test_fluxsite.py index 7c16ab3f..d3936ba0 100644 --- a/tests/test_fluxsite.py +++ b/tests/test_fluxsite.py @@ -8,7 +8,6 @@ import contextlib import io import math -from pathlib import Path import f90nml import netCDF4 @@ -17,12 +16,10 @@ from benchcab import __version__, internal from benchcab.fluxsite import ( CableError, - Task, + FluxsiteTask, get_comparison_name, get_fluxsite_comparisons, get_fluxsite_tasks, - patch_namelist, - patch_remove_namelist, ) from benchcab.repository import CableRepository @@ -42,8 +39,8 @@ def repo(mock_cwd, mock_subprocess_handler): @pytest.fixture() def task(repo, mock_cwd, mock_subprocess_handler): - """Returns a mock `Task` instance.""" - _task = Task( + """Returns a mock `FluxsiteTask` instance.""" + _task = FluxsiteTask( repo=repo, met_forcing_file="forcing-file.nc", sci_conf_id=0, @@ -55,7 +52,7 @@ def task(repo, mock_cwd, mock_subprocess_handler): class TestGetTaskName: - """tests for `Task.get_task_name()`.""" + """tests for `FluxsiteTask.get_task_name()`.""" def test_task_name_convention(self, task): """Success case: check task name convention.""" @@ -63,7 +60,7 @@ def test_task_name_convention(self, task): class TestGetLogFilename: - """Tests for `Task.get_log_filename()`.""" + """Tests for `FluxsiteTask.get_log_filename()`.""" def test_log_filename_convention(self, task): """Success case: check log file name convention.""" @@ -71,7 +68,7 @@ def test_log_filename_convention(self, task): class TestGetOutputFilename: - """Tests for `Task.get_output_filename()`.""" + """Tests for `FluxsiteTask.get_output_filename()`.""" def test_output_filename_convention(self, task): """Success case: check output file name convention.""" @@ -79,11 +76,11 @@ def test_output_filename_convention(self, task): class TestFetchFiles: - """Tests for `Task.fetch_files()`.""" + """Tests for `FluxsiteTask.fetch_files()`.""" @pytest.fixture(autouse=True) def _setup(self, task): - """Setup precondition for `Task.fetch_files()`.""" + """Setup precondition for `FluxsiteTask.fetch_files()`.""" internal.NAMELIST_DIR.mkdir() (internal.NAMELIST_DIR / internal.CABLE_NML).touch() (internal.NAMELIST_DIR / internal.CABLE_SOIL_NML).touch() @@ -109,11 +106,11 @@ def test_required_files_are_copied_to_task_dir(self, task): class TestCleanTask: - """Tests for `Task.clean_task()`.""" + """Tests for `FluxsiteTask.clean_task()`.""" @pytest.fixture(autouse=True) def _setup(self, task): - """Setup precondition for `Task.clean_task()`.""" + """Setup precondition for `FluxsiteTask.clean_task()`.""" task_dir = internal.FLUXSITE_DIRS["TASKS"] / task.get_task_name() task_dir.mkdir(parents=True) (task_dir / internal.CABLE_NML).touch() @@ -141,91 +138,12 @@ def test_clean_files(self, task): assert not (internal.FLUXSITE_DIRS["LOG"] / task.get_log_filename()).exists() -class TestPatchNamelist: - """Tests for `patch_namelist()`.""" - - @pytest.fixture() - def nml_path(self): - """Return a path to a namelist file used for testing.""" - return Path("test.nml") - - def test_patch_on_non_existing_namelist_file(self, nml_path): - """Success case: patch non-existing namelist file.""" - patch = {"cable": {"file": "/path/to/file", "bar": 123}} - patch_namelist(nml_path, patch) - assert f90nml.read(nml_path) == patch - - def test_patch_on_non_empty_namelist_file(self, nml_path): - """Success case: patch non-empty namelist file.""" - f90nml.write({"cable": {"file": "/path/to/file", "bar": 123}}, nml_path) - patch_namelist(nml_path, {"cable": {"some": {"parameter": True}, "bar": 456}}) - assert f90nml.read(nml_path) == { - "cable": { - "file": "/path/to/file", - "bar": 456, - "some": {"parameter": True}, - } - } - - def test_empty_patch_does_nothing(self, nml_path): - """Success case: empty patch does nothing.""" - f90nml.write({"cable": {"file": "/path/to/file", "bar": 123}}, nml_path) - prev = f90nml.read(nml_path) - patch_namelist(nml_path, {}) - assert f90nml.read(nml_path) == prev - - -class TestPatchRemoveNamelist: - """Tests for `patch_remove_namelist()`.""" - - @pytest.fixture() - def nml(self): - """Return a namelist dictionary used for testing.""" - return { - "cable": { - "cable_user": { - "some_parameter": True, - "new_feature": True, - }, - }, - } - - @pytest.fixture() - def nml_path(self, nml): - """Create a namelist file and return its path.""" - _nml_path = Path("test.nml") - f90nml.write(nml, _nml_path) - return _nml_path - - def test_remove_namelist_parameter_from_derived_type(self, nml_path): - """Success case: remove a namelist parameter from derrived type.""" - patch_remove_namelist( - nml_path, {"cable": {"cable_user": {"new_feature": True}}} - ) - assert f90nml.read(nml_path) == { - "cable": {"cable_user": {"some_parameter": True}} - } - - def test_empty_patch_remove_does_nothing(self, nml_path, nml): - """Success case: empty patch_remove does nothing.""" - patch_remove_namelist(nml_path, {}) - assert f90nml.read(nml_path) == nml - - def test_key_error_raised_for_non_existent_namelist_parameter(self, nml_path): - """Failure case: test patch_remove KeyError exeption.""" - with pytest.raises( - KeyError, - match=f"Namelist parameters specified in `patch_remove` do not exist in {nml_path.name}.", - ): - patch_remove_namelist(nml_path, {"cable": {"foo": {"bar": True}}}) - - class TestSetupTask: - """Tests for `Task.setup_task()`.""" + """Tests for `FluxsiteTask.setup_task()`.""" @pytest.fixture(autouse=True) def _setup(self, task): - """Setup precondition for `Task.setup_task()`.""" + """Setup precondition for `FluxsiteTask.setup_task()`.""" (internal.NAMELIST_DIR).mkdir() (internal.NAMELIST_DIR / internal.CABLE_NML).touch() (internal.NAMELIST_DIR / internal.CABLE_SOIL_NML).touch() @@ -308,11 +226,11 @@ def test_standard_output(self, task, verbosity, expected): class TestRunCable: - """Tests for `Task.run_cable()`.""" + """Tests for `FluxsiteTask.run_cable()`.""" @pytest.fixture(autouse=True) def _setup(self, task): - """Setup precondition for `Task.run_cable()`.""" + """Setup precondition for `FluxsiteTask.run_cable()`.""" task_dir = internal.FLUXSITE_DIRS["TASKS"] / task.get_task_name() task_dir.mkdir(parents=True) @@ -341,7 +259,7 @@ def test_cable_error_exception(self, task, mock_subprocess_handler): class TestAddProvenanceInfo: - """Tests for `Task.add_provenance_info()`.""" + """Tests for `FluxsiteTask.add_provenance_info()`.""" @pytest.fixture() def nml(self): @@ -365,7 +283,7 @@ def nc_output_path(self, task): @pytest.fixture(autouse=True) def _setup(self, task, nml): - """Setup precondition for `Task.add_provenance_info()`.""" + """Setup precondition for `FluxsiteTask.add_provenance_info()`.""" task_dir = internal.FLUXSITE_DIRS["TASKS"] / task.get_task_name() task_dir.mkdir(parents=True) fluxsite_output_dir = internal.FLUXSITE_DIRS["OUTPUT"] @@ -464,7 +382,7 @@ class TestGetFluxsiteComparisons: def test_comparisons_for_two_branches_with_two_tasks(self, mock_cwd): """Success case: comparisons for two branches with two tasks.""" tasks = [ - Task( + FluxsiteTask( repo=CableRepository("path/to/repo", repo_id=repo_id), met_forcing_file="foo.nc", sci_config={"foo": "bar"}, @@ -491,7 +409,7 @@ def test_comparisons_for_two_branches_with_two_tasks(self, mock_cwd): def test_comparisons_for_three_branches_with_three_tasks(self, mock_cwd): """Success case: comparisons for three branches with three tasks.""" tasks = [ - Task( + FluxsiteTask( repo=CableRepository("path/to/repo", repo_id=repo_id), met_forcing_file="foo.nc", sci_config={"foo": "bar"}, diff --git a/tests/test_namelist.py b/tests/test_namelist.py new file mode 100644 index 00000000..87d629cf --- /dev/null +++ b/tests/test_namelist.py @@ -0,0 +1,87 @@ +"""`pytest` tests for namelist.py.""" + +from pathlib import Path + +import f90nml +import pytest + +from benchcab.utils.namelist import patch_namelist, patch_remove_namelist + + +class TestPatchNamelist: + """Tests for `patch_namelist()`.""" + + @pytest.fixture() + def nml_path(self): + """Return a path to a namelist file used for testing.""" + return Path("test.nml") + + def test_patch_on_non_existing_namelist_file(self, nml_path): + """Success case: patch non-existing namelist file.""" + patch = {"cable": {"file": "/path/to/file", "bar": 123}} + patch_namelist(nml_path, patch) + assert f90nml.read(nml_path) == patch + + def test_patch_on_non_empty_namelist_file(self, nml_path): + """Success case: patch non-empty namelist file.""" + f90nml.write({"cable": {"file": "/path/to/file", "bar": 123}}, nml_path) + patch_namelist(nml_path, {"cable": {"some": {"parameter": True}, "bar": 456}}) + assert f90nml.read(nml_path) == { + "cable": { + "file": "/path/to/file", + "bar": 456, + "some": {"parameter": True}, + } + } + + def test_empty_patch_does_nothing(self, nml_path): + """Success case: empty patch does nothing.""" + f90nml.write({"cable": {"file": "/path/to/file", "bar": 123}}, nml_path) + prev = f90nml.read(nml_path) + patch_namelist(nml_path, {}) + assert f90nml.read(nml_path) == prev + + +class TestPatchRemoveNamelist: + """Tests for `patch_remove_namelist()`.""" + + @pytest.fixture() + def nml(self): + """Return a namelist dictionary used for testing.""" + return { + "cable": { + "cable_user": { + "some_parameter": True, + "new_feature": True, + }, + }, + } + + @pytest.fixture() + def nml_path(self, nml): + """Create a namelist file and return its path.""" + _nml_path = Path("test.nml") + f90nml.write(nml, _nml_path) + return _nml_path + + def test_remove_namelist_parameter_from_derived_type(self, nml_path): + """Success case: remove a namelist parameter from derrived type.""" + patch_remove_namelist( + nml_path, {"cable": {"cable_user": {"new_feature": True}}} + ) + assert f90nml.read(nml_path) == { + "cable": {"cable_user": {"some_parameter": True}} + } + + def test_empty_patch_remove_does_nothing(self, nml_path, nml): + """Success case: empty patch_remove does nothing.""" + patch_remove_namelist(nml_path, {}) + assert f90nml.read(nml_path) == nml + + def test_key_error_raised_for_non_existent_namelist_parameter(self, nml_path): + """Failure case: test patch_remove KeyError exeption.""" + with pytest.raises( + KeyError, + match=f"Namelist parameters specified in `patch_remove` do not exist in {nml_path.name}.", + ): + patch_remove_namelist(nml_path, {"cable": {"foo": {"bar": True}}}) diff --git a/tests/test_repository.py b/tests/test_repository.py index 06bef617..97f1cb25 100644 --- a/tests/test_repository.py +++ b/tests/test_repository.py @@ -28,6 +28,12 @@ def repo(mock_cwd, mock_subprocess_handler, mock_environment_modules_handler): return _repo +@pytest.fixture(params=[False, True]) +def mpi(request): + """Return a parametrized mpi flag for testing.""" + return request.param + + class TestRepoID: """Tests for `CableRepository.repo_id`.""" @@ -118,10 +124,18 @@ def _setup(self, repo): (internal.SRC_DIR / repo.name / "offline" / "serial_cable").touch() (internal.SRC_DIR / repo.name / "offline" / "foo.f90").touch() - def test_source_files_and_scripts_are_copied_to_tmp_dir(self, repo): + @pytest.fixture() + def tmp_dir(self, repo, mpi): + """Return the relative path to the temporary build directory.""" + return ( + internal.SRC_DIR + / repo.name + / (internal.TMP_BUILD_DIR_MPI if mpi else internal.TMP_BUILD_DIR) + ) + + def test_source_files_and_scripts_are_copied_to_tmp_dir(self, repo, mpi, tmp_dir): """Success case: test source files and scripts are copied to .tmp.""" - repo.pre_build() - tmp_dir = internal.SRC_DIR / repo.name / "offline" / ".tmp" + repo.pre_build(mpi=mpi) assert (tmp_dir / "Makefile").exists() assert (tmp_dir / "parallel_cable").exists() assert (tmp_dir / "serial_cable").exists() @@ -165,37 +179,52 @@ def modules(self): return ["foo", "bar"] @pytest.fixture() - def env(self, netcdf_root): - """Return a dictionary containing the required environment variables.""" + def expected_env(self, netcdf_root, mpi): + """Return a dictionary of expected environment variables to be defined.""" return { "NCDIR": f"{netcdf_root}/lib/Intel", "NCMOD": f"{netcdf_root}/include/Intel", "CFLAGS": "-O2 -fp-model precise", "LDFLAGS": f"-L{netcdf_root}/lib/Intel -O0", "LD": "-lnetcdf -lnetcdff", - "FC": "ifort", + "FC": "mpif90" if mpi else "ifort", } @pytest.fixture(autouse=True) def _setup(self, repo, netcdf_root): """Setup precondition for `CableRepository.run_build()`.""" - (internal.SRC_DIR / repo.name / "offline" / ".tmp").mkdir(parents=True) + (internal.SRC_DIR / repo.name / internal.TMP_BUILD_DIR).mkdir(parents=True) + (internal.SRC_DIR / repo.name / internal.TMP_BUILD_DIR_MPI).mkdir(parents=True) # This is required so that we can use the NETCDF_ROOT environment variable # when running `make`, and `serial_cable` and `parallel_cable` scripts: os.environ["NETCDF_ROOT"] = netcdf_root + @pytest.fixture() + def expected_commands(self, netcdf_root, mpi): + """Return a list of expected commands to be executed.""" + return ( + [ + "make -f Makefile", + './parallel_cable "mpif90" "-O2 -fp-model precise"' + f' "-L{netcdf_root}/lib/Intel -O0" "-lnetcdf -lnetcdff" ' + f'"{netcdf_root}/include/Intel"', + ] + if mpi + else [ + "make -f Makefile", + './serial_cable "ifort" "-O2 -fp-model precise"' + f' "-L{netcdf_root}/lib/Intel -O0" "-lnetcdf -lnetcdff" ' + f'"{netcdf_root}/include/Intel"', + ] + ) + def test_build_command_execution( - self, repo, mock_subprocess_handler, modules, netcdf_root + self, repo, mock_subprocess_handler, modules, mpi, expected_commands ): """Success case: test build commands are run.""" - repo.run_build(modules) - assert mock_subprocess_handler.commands == [ - "make -f Makefile", - './serial_cable "ifort" "-O2 -fp-model precise"' - f' "-L{netcdf_root}/lib/Intel -O0" "-lnetcdf -lnetcdff" ' - f'"{netcdf_root}/include/Intel"', - ] + repo.run_build(modules, mpi=mpi) + assert mock_subprocess_handler.commands == expected_commands def test_modules_loaded_at_runtime( self, repo, mock_environment_modules_handler, modules @@ -210,11 +239,11 @@ def test_modules_loaded_at_runtime( ) in mock_environment_modules_handler.commands def test_commands_are_run_with_environment_variables( - self, repo, mock_subprocess_handler, modules, env + self, repo, mock_subprocess_handler, modules, mpi, expected_env ): """Success case: test commands are run with the correct environment variables.""" - repo.run_build(modules) - for kv in env.items(): + repo.run_build(modules, mpi=mpi) + for kv in expected_env.items(): assert kv in mock_subprocess_handler.env.items() @pytest.mark.parametrize( @@ -237,16 +266,26 @@ class TestPostBuild: @pytest.fixture(autouse=True) def _setup(self, repo): """Setup precondition for `CableRepository.post_build()`.""" - (internal.SRC_DIR / repo.name / "offline" / ".tmp").mkdir(parents=True) - (internal.SRC_DIR / repo.name / "offline" / ".tmp" / internal.CABLE_EXE).touch() + tmp_build_dir = internal.SRC_DIR / repo.name / internal.TMP_BUILD_DIR + tmp_build_dir.mkdir(parents=True) + (tmp_build_dir / internal.CABLE_EXE).touch() + + tmp_build_dir_mpi = internal.SRC_DIR / repo.name / internal.TMP_BUILD_DIR_MPI + tmp_build_dir_mpi.mkdir(parents=True) + (tmp_build_dir_mpi / internal.CABLE_MPI_EXE).touch() - def test_exe_moved_to_offline_dir(self, repo): + def test_exe_moved_to_offline_dir(self, repo, mpi): """Success case: test executable is moved to offline directory.""" - repo.post_build() - tmp_dir = internal.SRC_DIR / repo.name / "offline" / ".tmp" - assert not (tmp_dir / internal.CABLE_EXE).exists() + repo.post_build(mpi=mpi) + tmp_dir = ( + internal.SRC_DIR + / repo.name + / (internal.TMP_BUILD_DIR_MPI if mpi else internal.TMP_BUILD_DIR) + ) + exe = internal.CABLE_MPI_EXE if mpi else internal.CABLE_EXE + assert not (tmp_dir / exe).exists() offline_dir = internal.SRC_DIR / repo.name / "offline" - assert (offline_dir / internal.CABLE_EXE).exists() + assert (offline_dir / exe).exists() @pytest.mark.parametrize( ("verbosity", "expected"), diff --git a/tests/test_spatial.py b/tests/test_spatial.py new file mode 100644 index 00000000..509bc696 --- /dev/null +++ b/tests/test_spatial.py @@ -0,0 +1,224 @@ +"""`pytest` tests for spatial.py. + +Note: explicit teardown for generated files and directories are not required as +the working directory used for testing is cleaned up in the `_run_around_tests` +pytest autouse fixture. +""" + +import contextlib +import io + +import f90nml +import pytest +import yaml + +from benchcab import internal +from benchcab.repository import CableRepository +from benchcab.spatial import SpatialTask, get_spatial_tasks + + +@pytest.fixture() +def repo(mock_cwd, mock_subprocess_handler): + """Returns a `CableRepository` instance.""" + _repo = CableRepository( + repo_id=1, + path="path/to/test-branch", + patch={"cable": {"some_branch_specific_setting": True}}, + ) + _repo.subprocess_handler = mock_subprocess_handler + _repo.root_dir = mock_cwd + return _repo + + +@pytest.fixture() +def task(repo, mock_cwd, mock_subprocess_handler, mock_git_repo): + """Returns a mock `SpatialTask` instance.""" + _task = SpatialTask( + repo=repo, + met_forcing="crujra_access", + sci_conf_id=0, + sci_config={"cable": {"some_setting": True}}, + ) + _task.subprocess_handler = mock_subprocess_handler + _task.git_repo_handler = mock_git_repo + _task.root_dir = mock_cwd + return _task + + +class TestGetTaskName: + """Tests for `SpatialTask.get_task_name()`.""" + + def test_task_name_convention(self, task): + """Success case: check task name convention.""" + assert task.get_task_name() == "crujra_access_R1_S0" + + +class TestSetupTask: + """Tests for `SpatialTask.setup_task()`.""" + + @pytest.fixture(autouse=True) + def _setup(self, mock_git_repo): + internal.NAMELIST_DIR.mkdir() + (internal.NAMELIST_DIR / internal.CABLE_NML).touch() + (internal.NAMELIST_DIR / internal.CABLE_SOIL_NML).touch() + (internal.NAMELIST_DIR / internal.CABLE_VEGETATION_NML).touch() + + internal.SPATIAL_TASKS_DIR.mkdir(parents=True) + + mock_git_repo.files = ["config.yaml", "cable.nml"] + + def test_namelist_parameters_are_patched(self, task): + """Success case: test namelist parameters are patched.""" + task.setup_task() + res_nml = f90nml.read( + str(internal.SPATIAL_TASKS_DIR / task.get_task_name() / internal.CABLE_NML) + ) + assert res_nml["cable"] == { + "filename": { + "out": "cable_out.nc", + "log": "cable_log.nc", + "restart_in": "RESTART/restart.nc", + "restart_out": "RESTART/restart.nc", + "type": "INPUT/crujra_access_gridinfo.nc", + "veg": "INPUT/def_veg_params_zr_clitt_albedo_fix.txt", + "soil": "INPUT/def_soil_params.txt", + }, + "spinup": False, + "cable_user": { + "gswp3": True, + "mettype": "gswp3", + }, + "gswpfile": { + "mask": "INPUT/crujra_access_landmask.nc", + "lwdown": "INPUT/crujra_access_LWdown_$year.nc", + "psurf": "INPUT/crujra_access_PSurf_$year.nc", + "qair": "INPUT/crujra_access_Qair_$year.nc", + "rainf": "INPUT/crujra_access_Rainf_$year.nc", + "snowf": "INPUT/crujra_access_Snowf_$year.nc", + "swdown": "INPUT/crujra_access_SWdown_$year.nc", + "tair": "INPUT/crujra_access_Tair_$year.nc", + "wind": "INPUT/crujra_access_Wind_$year.nc", + }, + "some_setting": True, + "some_branch_specific_setting": True, + } + + def test_payu_config_parameters(self, task, mock_cwd): + """Success case: check config.yaml parameters.""" + task.setup_task(payu_config={"some_parameter": "foo"}) + config_path = internal.SPATIAL_TASKS_DIR / task.get_task_name() / "config.yaml" + with config_path.open("r", encoding="utf-8") as file: + config = yaml.safe_load(file) + assert ( + config["exe"] + == f"{mock_cwd}/src/test-branch/offline/{internal.CABLE_MPI_EXE}" + ) + assert config["input"] == [ + f"{internal.SPATIAL_MET_DIR}/crujra_access/LWdown", + f"{internal.SPATIAL_MET_DIR}/crujra_access/PSurf", + f"{internal.SPATIAL_MET_DIR}/crujra_access/Qair", + f"{internal.SPATIAL_MET_DIR}/crujra_access/Rainf", + f"{internal.SPATIAL_MET_DIR}/crujra_access/Snowf", + f"{internal.SPATIAL_MET_DIR}/crujra_access/SWdown", + f"{internal.SPATIAL_MET_DIR}/crujra_access/Tair", + f"{internal.SPATIAL_MET_DIR}/crujra_access/Wind", + f"{internal.SPATIAL_MET_DIR}/crujra_access/crujra_access_landmask.nc", + f"{internal.SPATIAL_MET_DIR}/crujra_access/crujra_access_gridinfo.nc", + f"{mock_cwd}/src/CABLE-AUX/core/biogeophys/def_veg_params_zr_clitt_albedo_fix.txt", + f"{mock_cwd}/src/CABLE-AUX/core/biogeophys/def_soil_params.txt", + ] + assert config["laboratory"] == str(mock_cwd / internal.PAYU_LABORATORY_DIR) + assert config["some_parameter"] == "foo" + + # TODO(Sean) fix for issue https://github.com/CABLE-LSM/benchcab/issues/162 + @pytest.mark.skip( + reason="""This will always fail since `parametrize()` parameters are + dependent on the `mock_cwd` fixture.""" + ) + @pytest.mark.parametrize( + ("verbosity", "expected"), + [ + (False, ""), + ( + True, + "Setting up task: crujra_access_R1_S0\n" + " Copying namelist files from namelists to " + "runs/spatial/tasks/crujra_access_R1_S0\n" + " Updating experiment config parameters in " + "config.yaml\n" + " Adding base configurations to CABLE namelist file " + "runs/spatial/tasks/crujra_access_R1_S0/cable.nml\n" + " Adding science configurations to CABLE namelist file " + "runs/spatial/tasks/crujra_access_R1_S0/cable.nml\n" + " Adding branch specific configurations to CABLE namelist file " + "runs/spatial/tasks/crujra_access_R1_S0/cable.nml\n", + ), + ], + ) + def test_standard_output(self, task, verbosity, expected): + """Success case: test standard output.""" + with contextlib.redirect_stdout(io.StringIO()) as buf: + task.setup_task(verbose=verbosity) + assert buf.getvalue() == expected + + +class TestRun: + """Tests for `SpatialTask.run()`.""" + + @pytest.fixture(autouse=True) + def _setup(self, task): + task_dir = internal.SPATIAL_TASKS_DIR / task.get_task_name() + task_dir.mkdir(parents=True, exist_ok=True) + + def test_payu_run_command(self, task, mock_subprocess_handler): + """Success case: test payu run command.""" + task.run() + assert "payu run" in mock_subprocess_handler.commands + + def test_payu_run_with_optional_arguments(self, task, mock_subprocess_handler): + """Success case: test payu run command with optional arguments.""" + task.payu_args = "--some-flag" + task.run() + assert "payu run --some-flag" in mock_subprocess_handler.commands + + +class TestGetSpatialTasks: + """Tests for `get_spatial_tasks()`.""" + + @pytest.fixture() + def repos(self, config): + """Return a list of `CableRepository` instances used for testing.""" + return [ + CableRepository(**branch_config, repo_id=id) + for id, branch_config in enumerate(config["realisations"]) + ] + + @pytest.fixture() + def met_forcings(self): + """Return a list of forcing file names used for testing.""" + return ["foo", "bar"] + + @pytest.fixture() + def science_configurations(self, config): + """Return a list of science configurations used for testing.""" + return config["science_configurations"] + + def test_task_product_across_branches_forcings_and_configurations( + self, repos, met_forcings, science_configurations + ): + """Success case: test task product across branches, forcings and configurations.""" + tasks = get_spatial_tasks( + repos=repos, + met_forcings=met_forcings, + science_configurations=science_configurations, + ) + assert [(task.repo, task.met_forcing, task.sci_config) for task in tasks] == [ + (repos[0], met_forcings[0], science_configurations[0]), + (repos[0], met_forcings[0], science_configurations[1]), + (repos[0], met_forcings[1], science_configurations[0]), + (repos[0], met_forcings[1], science_configurations[1]), + (repos[1], met_forcings[0], science_configurations[0]), + (repos[1], met_forcings[0], science_configurations[1]), + (repos[1], met_forcings[1], science_configurations[0]), + (repos[1], met_forcings[1], science_configurations[1]), + ] diff --git a/tests/test_workdir.py b/tests/test_workdir.py index 62dd710c..ae67676e 100644 --- a/tests/test_workdir.py +++ b/tests/test_workdir.py @@ -12,6 +12,7 @@ from benchcab.workdir import ( clean_directory_tree, setup_fluxsite_directory_tree, + setup_spatial_directory_tree, ) @@ -37,6 +38,25 @@ def test_directory_structure_generated(self, fluxsite_directory_list): assert path.exists() +class TestSetupSpatialDirectoryTree: + """Tests for `setup_spatial_directory_tree()`.""" + + @pytest.fixture() + def spatial_directory_list(self): + """Return the list of work directories we want benchcab to create.""" + return [ + Path("runs", "spatial"), + Path("runs", "spatial", "tasks"), + Path("runs", "payu-laboratory"), + ] + + def test_directory_structure_generated(self, spatial_directory_list): + """Success case: generate spatial directory structure.""" + setup_spatial_directory_tree() + for path in spatial_directory_list: + assert path.exists() + + class TestCleanDirectoryTree: """Tests for `clean_directory_tree()`."""