diff --git a/.conda/benchcab-dev.yaml b/.conda/benchcab-dev.yaml index 6be15d1e..dc121c01 100644 --- a/.conda/benchcab-dev.yaml +++ b/.conda/benchcab-dev.yaml @@ -10,3 +10,4 @@ dependencies: - pytest-cov - pyyaml - flatdict + - gitpython diff --git a/benchcab/benchcab.py b/benchcab/benchcab.py index c6fc1fcd..53807144 100644 --- a/benchcab/benchcab.py +++ b/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,7 +194,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: @@ -201,30 +206,38 @@ 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() + + 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"] @@ -234,9 +247,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("") @@ -248,8 +263,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: @@ -280,13 +297,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) + 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`.""" @@ -298,7 +340,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() diff --git a/benchcab/cli.py b/benchcab/cli.py index 2dc266d9..8a82e7d7 100644 --- a/benchcab/cli.py +++ b/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,7 +109,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.", @@ -118,6 +117,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( @@ -143,9 +147,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, ) @@ -165,7 +169,7 @@ 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, diff --git a/benchcab/fluxsite.py b/benchcab/fluxsite.py index 921a76a1..33fae914 100644 --- a/benchcab/fluxsite.py +++ b/benchcab/fluxsite.py @@ -5,7 +5,6 @@ import multiprocessing import queue from pathlib import Path -from typing import TypeVar, Dict, Any from subprocess import CalledProcessError import flatdict @@ -17,77 +16,7 @@ from benchcab.comparison import ComparisonTask from benchcab.utils.subprocess import SubprocessWrapperInterface, SubprocessWrapper from benchcab.utils.fs import chdir - - -# fmt: off -# pylint: disable=invalid-name,missing-function-docstring,line-too-long -# ====================================================== -# 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 - -# ====================================================== -# pylint: enable=invalid-name,missing-function-docstring,line-too-long -# 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) - # remove namelist file as f90nml cannot write to an existing file - nml_path.unlink() - f90nml.write(deep_update(nml, patch), nml_path) - - -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) - # remove namelist file as f90nml cannot write to an existing file - nml_path.unlink() - try: - f90nml.write(deep_del(nml, patch_remove), nml_path) - except KeyError as exc: - raise KeyError( - f"Namelist parameters specified in `patch_remove` do not exist in {nml_path.name}." - ) from exc +from benchcab.utils.namelist import patch_namelist, patch_remove_namelist f90_logical_repr = {True: ".true.", False: ".false."} @@ -97,7 +26,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 @@ -350,10 +279,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, @@ -366,14 +295,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.""" @@ -402,7 +333,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. diff --git a/benchcab/internal.py b/benchcab/internal.py index bb08b7e3..c53c8671 100644 --- a/benchcab/internal.py +++ b/benchcab/internal.py @@ -17,7 +17,6 @@ "walltime": "6:00:00", "storage": [], } -MPI = False FLUXSITE_DEFAULT_MULTIPROCESS = True # DIRECTORY PATHS/STRUCTURE: @@ -76,14 +75,36 @@ # Relative path to directory that stores bitwise comparison results FLUXSITE_BITWISE_CMP_DIR = FLUXSITE_ANALYSIS_DIR / "bitwise-comparisons" -# Path to met files: +# 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 PLUMBER2 site forcing data directory (doi: 10.25914/5fdb0902607e1): MET_DIR = Path("/g/data/ks32/CLEX_Data/PLUMBER2/v1-0/Met/") +# Path to CRUJRA (ACCESS resolution) spatial forcing data: +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" diff --git a/benchcab/repository.py b/benchcab/repository.py index 61e62a5e..4b5c57a7 100644 --- a/benchcab/repository.py +++ b/benchcab/repository.py @@ -114,11 +114,11 @@ 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 / "offline" / (".mpitmp" if mpi else ".tmp") if not tmp_dir.exists(): if verbose: print(f"mkdir {tmp_dir.relative_to(self.root_dir)}") @@ -152,11 +152,11 @@ 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 / "offline" / (".mpitmp" if mpi else ".tmp") with chdir(tmp_dir), self.modules_handler.load(modules, verbose=verbose): env = os.environ.copy() @@ -165,27 +165,28 @@ 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 / "offline" / (".mpitmp" if mpi else ".tmp") + 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..0e1ed4d5 --- /dev/null +++ b/benchcab/spatial.py @@ -0,0 +1,148 @@ +"""A module containing functions and data structures for running spatial tasks.""" + +from pathlib import Path +import yaml + +from benchcab import internal +from benchcab.repository import CableRepository +from benchcab.utils.fs import chdir +from benchcab.utils.subprocess import SubprocessWrapper, SubprocessWrapperInterface +from benchcab.utils.git import GitRepo, GitRepoInterface +from benchcab.utils.namelist import patch_namelist, patch_remove_namelist + + +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, + sci_conf_id: int, + sci_config: dict, + ) -> None: + self.repo = repo + self.sci_conf_id = sci_conf_id + self.sci_config = sci_config + + def get_task_name(self) -> str: + """Returns the file name convention used for this task.""" + met_name = internal.SPATIAL_MET_FORCING_CRUJRA["name"] + return f"{met_name}_R{self.repo.repo_id}_S{self.sci_conf_id}" + + def setup_task(self, 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, + ) + + 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", + ) + config["exe"] = str( + self.root_dir + / internal.SRC_DIR + / self.repo.name + / "offline" + / internal.CABLE_MPI_EXE + ) + + config["input"] = [ + "/g/data/tm70/sb8430/cable_inputs/landsea-access-esm.nc", + "/g/data/tm70/ccc561/CABLE/CABLE-as-ACCESS_Yingping/input/gridinfo-access-esm.nc", + "/g/data/tm70/ccc561/CABLE/CABLE-as-ACCESS_Yingping/cruncep10", + 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 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("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], + science_configurations: list[dict], +): + """Returns a list of spatial tasks to run.""" + tasks = [ + SpatialTask( + repo=repo, + sci_conf_id=sci_conf_id, + sci_config=sci_config, + ) + for repo in repos + for sci_conf_id, sci_config in enumerate(science_configurations) + ] + return tasks diff --git a/benchcab/utils/git.py b/benchcab/utils/git.py new file mode 100644 index 00000000..3e0cee2f --- /dev/null +++ b/benchcab/utils/git.py @@ -0,0 +1,25 @@ +"""Contains a wrapper around the git API.""" + +from abc import ABC as AbstractBaseClass, 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..7fb34528 --- /dev/null +++ b/benchcab/utils/namelist.py @@ -0,0 +1,75 @@ +"""Contains utility functions for manipulating Fortran namelist files.""" + +from pathlib import Path +from typing import TypeVar, Dict, Any +import f90nml + +# fmt: off +# pylint: disable=invalid-name,missing-function-docstring,line-too-long +# ====================================================== +# 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 + +# ====================================================== +# pylint: enable=invalid-name,missing-function-docstring,line-too-long +# 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) + # remove namelist file as f90nml cannot write to an existing file + nml_path.unlink() + f90nml.write(deep_update(nml, patch), nml_path) + + +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) + # remove namelist file as f90nml cannot write to an existing file + nml_path.unlink() + try: + f90nml.write(deep_del(nml, patch_remove), nml_path) + except KeyError as exc: + raise KeyError( + f"Namelist parameters specified in `patch_remove` do not exist in {nml_path.name}." + ) from exc diff --git a/benchcab/workdir.py b/benchcab/workdir.py index fa7fa32d..c8d68cc9 100644 --- a/benchcab/workdir.py +++ b/benchcab/workdir.py @@ -5,7 +5,7 @@ import shutil from benchcab import internal -from benchcab.fluxsite import Task +from benchcab.fluxsite import FluxsiteTask def clean_directory_tree(root_dir=internal.CWD): @@ -29,9 +29,9 @@ def setup_src_dir(root_dir=internal.CWD): def setup_fluxsite_directory_tree( - fluxsite_tasks: list[Task], root_dir=internal.CWD, verbose=False + fluxsite_tasks: list[FluxsiteTask], root_dir=internal.CWD, verbose=False ): - """Generate the directory structure used of `benchcab`.""" + """Generate the directory structure for running fluxsite tests.""" run_dir = Path(root_dir, internal.RUN_DIR) if not run_dir.exists(): @@ -85,3 +85,15 @@ def setup_fluxsite_directory_tree( if verbose: print(f"Creating {task_dir.relative_to(root_dir)}: " f"{task_dir}") os.makedirs(task_dir) + + +def setup_spatial_directory_tree(root_dir=internal.CWD): + """Generate the directory structure for running spatial tests.""" + + spatial_tasks_dir = root_dir / internal.SPATIAL_TASKS_DIR + if not spatial_tasks_dir.exists(): + spatial_tasks_dir.mkdir(parents=True) + + payu_laboratory_dir = root_dir / internal.PAYU_LABORATORY_DIR + if not payu_laboratory_dir.exists(): + payu_laboratory_dir.mkdir(parents=True) diff --git a/tests/common.py b/tests/common.py index 335315f2..4ef8e298 100644 --- a/tests/common.py +++ b/tests/common.py @@ -6,6 +6,7 @@ from typing import Optional from benchcab.utils.subprocess import SubprocessWrapperInterface +from benchcab.utils.git import GitRepoInterface from benchcab.environment_modules import EnvironmentModulesInterface MOCK_CWD = TMP_DIR = Path(tempfile.mkdtemp(prefix="benchcab_tests")) @@ -116,3 +117,19 @@ def module_load(self, *args: str) -> None: def module_unload(self, *args: str) -> None: self.commands.append("module unload " + " ".join(args)) + + +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() diff --git a/tests/test_cli.py b/tests/test_cli.py index df1d8ba9..af2856d8 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,7 @@ def test_cli_parser(): "subcommand": "spatial", "config": "config.yaml", "verbose": False, + "skip": [], } # Failure case: pass --no-submit to a non 'run' command diff --git a/tests/test_fluxsite.py b/tests/test_fluxsite.py index 53fa5bf7..516158dd 100644 --- a/tests/test_fluxsite.py +++ b/tests/test_fluxsite.py @@ -8,12 +8,10 @@ import netCDF4 from benchcab.fluxsite import ( - patch_namelist, - patch_remove_namelist, get_fluxsite_tasks, get_fluxsite_comparisons, get_comparison_name, - Task, + FluxsiteTask, CableError, ) from benchcab import internal @@ -24,8 +22,8 @@ def get_mock_task( subprocess_handler: SubprocessWrapperInterface = MockSubprocessWrapper(), -) -> Task: - """Returns a mock `Task` instance.""" +) -> FluxsiteTask: + """Returns a mock `FluxsiteTask` instance.""" repo = CableRepository( repo_id=1, @@ -35,7 +33,7 @@ def get_mock_task( repo.subprocess_handler = subprocess_handler repo.root_dir = MOCK_CWD - task = Task( + task = FluxsiteTask( repo=repo, met_forcing_file="forcing-file.nc", sci_conf_id=0, @@ -66,7 +64,7 @@ def setup_mock_namelists_directory(): assert cable_vegetation_nml_path.exists() -def setup_mock_run_directory(task: Task): +def setup_mock_run_directory(task: FluxsiteTask): """Setup mock run directory for a single task.""" task_dir = MOCK_CWD / internal.FLUXSITE_TASKS_DIR / task.get_task_name() task_dir.mkdir(parents=True) @@ -87,7 +85,7 @@ def do_mock_checkout_and_build(): assert cable_exe_path.exists() -def do_mock_run(task: Task): +def do_mock_run(task: FluxsiteTask): """Make mock log files and output files as if benchcab has just been run.""" output_path = Path( MOCK_CWD, internal.FLUXSITE_OUTPUT_DIR, task.get_output_filename() @@ -195,74 +193,6 @@ def test_clean_task(): ).exists() -def test_patch_namelist(): - """Tests for `patch_namelist()`.""" - - nml_path = MOCK_CWD / "test.nml" - - # Success case: patch non-existing namelist file - assert not nml_path.exists() - patch = {"cable": {"file": "/path/to/file", "bar": 123}} - patch_namelist(nml_path, patch) - assert f90nml.read(nml_path) == patch - - # Success case: patch non-empty namelist file - 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}, - } - } - - # Success case: empty patch does nothing - prev = f90nml.read(nml_path) - patch_namelist(nml_path, {}) - assert f90nml.read(nml_path) == prev - - -def test_patch_remove_namelist(): - """Tests for `patch_remove_namelist()`.""" - - nml_path = MOCK_CWD / "test.nml" - - # Success case: remove a namelist parameter from derrived type - nml = {"cable": {"cable_user": {"some_parameter": True}}} - f90nml.write(nml, nml_path) - patch_remove_namelist(nml_path, nml) - assert not f90nml.read(nml_path)["cable"] - nml_path.unlink() - - # Success case: test existing namelist parameters are preserved - # when removing a namelist parameter - to_remove = {"cable": {"cable_user": {"new_feature": True}}} - nml = {"cable": {"cable_user": {"some_parameter": True, "new_feature": True}}} - f90nml.write(nml, nml_path) - patch_remove_namelist(nml_path, to_remove) - assert f90nml.read(nml_path) == {"cable": {"cable_user": {"some_parameter": True}}} - nml_path.unlink() - - # Success case: empty patch_remove does nothing - nml = {"cable": {"cable_user": {"some_parameter": True}}} - f90nml.write(nml, nml_path) - patch_remove_namelist(nml_path, {}) - assert f90nml.read(nml_path) == nml - nml_path.unlink() - - # Failure case: patch_remove should raise KeyError when namelist parameters don't exist in - # nml_path - to_remove = {"cable": {"foo": {"bar": True}}} - nml = {"cable": {"cable_user": {"some_parameter": True, "new_feature": True}}} - f90nml.write(nml, nml_path) - with pytest.raises( - KeyError, - match=f"Namelist parameters specified in `patch_remove` do not exist in {nml_path.name}.", - ): - patch_remove_namelist(nml_path, to_remove) - nml_path.unlink(missing_ok=True) - - def test_setup_task(): """Tests for `setup_task()`.""" @@ -431,13 +361,13 @@ def test_get_fluxsite_comparisons(): # Success case: comparisons for two branches with two tasks # met0_S0_R0 met0_S0_R1 - task_a = Task( + task_a = FluxsiteTask( repo=CableRepository("path/to/repo_a", repo_id=0), met_forcing_file="foo.nc", sci_config={"foo": "bar"}, sci_conf_id=0, ) - task_b = Task( + task_b = FluxsiteTask( repo=CableRepository("path/to/repo_b", repo_id=1), met_forcing_file="foo.nc", sci_config={"foo": "bar"}, @@ -453,19 +383,19 @@ def test_get_fluxsite_comparisons(): # Success case: comparisons for three branches with three tasks # met0_S0_R0 met0_S0_R1 met0_S0_R2 - task_a = Task( + task_a = FluxsiteTask( repo=CableRepository("path/to/repo_a", repo_id=0), met_forcing_file="foo.nc", sci_config={"foo": "bar"}, sci_conf_id=0, ) - task_b = Task( + task_b = FluxsiteTask( repo=CableRepository("path/to/repo_b", repo_id=1), met_forcing_file="foo.nc", sci_config={"foo": "bar"}, sci_conf_id=0, ) - task_c = Task( + task_c = FluxsiteTask( repo=CableRepository("path/to/repo_b", repo_id=2), 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..b6ca0cd0 --- /dev/null +++ b/tests/test_namelist.py @@ -0,0 +1,75 @@ +"""`pytest` tests for namelist.py""" + +import pytest +import f90nml + +from benchcab.utils.namelist import patch_namelist, patch_remove_namelist +from .common import MOCK_CWD + + +def test_patch_namelist(): + """Tests for `patch_namelist()`.""" + + nml_path = MOCK_CWD / "test.nml" + + # Success case: patch non-existing namelist file + assert not nml_path.exists() + patch = {"cable": {"file": "/path/to/file", "bar": 123}} + patch_namelist(nml_path, patch) + assert f90nml.read(nml_path) == patch + + # Success case: patch non-empty namelist file + 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}, + } + } + + # Success case: empty patch does nothing + prev = f90nml.read(nml_path) + patch_namelist(nml_path, {}) + assert f90nml.read(nml_path) == prev + + +def test_patch_remove_namelist(): + """Tests for `patch_remove_namelist()`.""" + + nml_path = MOCK_CWD / "test.nml" + + # Success case: remove a namelist parameter from derrived type + nml = {"cable": {"cable_user": {"some_parameter": True}}} + f90nml.write(nml, nml_path) + patch_remove_namelist(nml_path, nml) + assert not f90nml.read(nml_path)["cable"] + nml_path.unlink() + + # Success case: test existing namelist parameters are preserved + # when removing a namelist parameter + to_remove = {"cable": {"cable_user": {"new_feature": True}}} + nml = {"cable": {"cable_user": {"some_parameter": True, "new_feature": True}}} + f90nml.write(nml, nml_path) + patch_remove_namelist(nml_path, to_remove) + assert f90nml.read(nml_path) == {"cable": {"cable_user": {"some_parameter": True}}} + nml_path.unlink() + + # Success case: empty patch_remove does nothing + nml = {"cable": {"cable_user": {"some_parameter": True}}} + f90nml.write(nml, nml_path) + patch_remove_namelist(nml_path, {}) + assert f90nml.read(nml_path) == nml + nml_path.unlink() + + # Failure case: patch_remove should raise KeyError when namelist parameters don't exist in + # nml_path + to_remove = {"cable": {"foo": {"bar": True}}} + nml = {"cable": {"cable_user": {"some_parameter": True, "new_feature": True}}} + f90nml.write(nml, nml_path) + with pytest.raises( + KeyError, + match=f"Namelist parameters specified in `patch_remove` do not exist in {nml_path.name}.", + ): + patch_remove_namelist(nml_path, to_remove) + nml_path.unlink(missing_ok=True) diff --git a/tests/test_repository.py b/tests/test_repository.py index a578665e..75a5b519 100644 --- a/tests/test_repository.py +++ b/tests/test_repository.py @@ -130,6 +130,16 @@ def test_pre_build(): assert (offline_dir / ".tmp" / "foo.f90").exists() shutil.rmtree(offline_dir / ".tmp") + # Success case: test source files and scripts are copied to .mpitmp when MPI + # is enabled + repo = get_mock_repo() + repo.pre_build(mpi=True) + assert (offline_dir / ".mpitmp" / "Makefile").exists() + assert (offline_dir / ".mpitmp" / "parallel_cable").exists() + assert (offline_dir / ".mpitmp" / "serial_cable").exists() + assert (offline_dir / ".mpitmp" / "foo.f90").exists() + shutil.rmtree(offline_dir / ".mpitmp") + # Success case: test non-verbose standard output repo = get_mock_repo() with contextlib.redirect_stdout(io.StringIO()) as buf: @@ -157,6 +167,7 @@ def test_run_build(): mock_netcdf_root = "/mock/path/to/root" mock_modules = ["foo", "bar"] (MOCK_CWD / internal.SRC_DIR / "trunk" / "offline" / ".tmp").mkdir(parents=True) + (MOCK_CWD / internal.SRC_DIR / "trunk" / "offline" / ".mpitmp").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: @@ -173,6 +184,17 @@ def test_run_build(): f'"{mock_netcdf_root}/include/Intel"', ] + # Success case: test MPI build commands are run + mock_subprocess = MockSubprocessWrapper() + repo = get_mock_repo(subprocess_handler=mock_subprocess) + repo.run_build(mock_modules, mpi=True) + assert mock_subprocess.commands == [ + "make -f Makefile", + './parallel_cable "mpif90" "-O2 -fp-model precise"' + f' "-L{mock_netcdf_root}/lib/Intel -O0" "-lnetcdf -lnetcdff" ' + f'"{mock_netcdf_root}/include/Intel"', + ] + # Success case: test modules are loaded at runtime mock_environment_modules = MockEnvironmentModules() repo = get_mock_repo(modules_handler=mock_environment_modules) @@ -200,6 +222,23 @@ def test_run_build(): }.items() ) + # Success case: test MPI commands are run with the correct environment + # variables + mock_subprocess = MockSubprocessWrapper() + repo = get_mock_repo(subprocess_handler=mock_subprocess) + repo.run_build(mock_modules, mpi=True) + assert all( + kv in mock_subprocess.env.items() + for kv in { + "NCDIR": f"{mock_netcdf_root}/lib/Intel", + "NCMOD": f"{mock_netcdf_root}/include/Intel", + "CFLAGS": "-O2 -fp-model precise", + "LDFLAGS": f"-L{mock_netcdf_root}/lib/Intel -O0", + "LD": "-lnetcdf -lnetcdff", + "FC": "mpif90", + }.items() + ) + # Success case: test non-verbose standard output repo = get_mock_repo() with contextlib.redirect_stdout(io.StringIO()) as buf: @@ -222,15 +261,24 @@ def test_post_build(): repo_dir = MOCK_CWD / internal.SRC_DIR / "trunk" offline_dir = repo_dir / "offline" tmp_dir = offline_dir / ".tmp" + mpitmp_dir = offline_dir / ".mpitmp" # Success case: test executable is moved to offline directory tmp_dir.mkdir(parents=True) (tmp_dir / internal.CABLE_EXE).touch() repo = get_mock_repo() repo.post_build() - assert not (offline_dir / ".tmp" / internal.CABLE_EXE).exists() + assert not (tmp_dir / internal.CABLE_EXE).exists() assert (offline_dir / internal.CABLE_EXE).exists() + # Success case: test MPI executable is moved to offline directory + mpitmp_dir.mkdir(parents=True) + (mpitmp_dir / internal.CABLE_MPI_EXE).touch() + repo = get_mock_repo() + repo.post_build(mpi=True) + assert not (mpitmp_dir / internal.CABLE_MPI_EXE).exists() + assert (offline_dir / internal.CABLE_MPI_EXE).exists() + # Success case: test non-verbose standard output (tmp_dir / internal.CABLE_EXE).touch() repo = get_mock_repo() diff --git a/tests/test_spatial.py b/tests/test_spatial.py new file mode 100644 index 00000000..4aadedf8 --- /dev/null +++ b/tests/test_spatial.py @@ -0,0 +1,149 @@ +"""`pytest` tests for spatial.py""" + +import shutil +import contextlib +import io +import yaml +import f90nml + +from benchcab import internal +from benchcab.spatial import SpatialTask, get_spatial_tasks +from benchcab.repository import CableRepository +from benchcab.utils.subprocess import SubprocessWrapperInterface +from benchcab.utils.git import GitRepoInterface +from .common import MOCK_CWD, get_mock_config, MockSubprocessWrapper, MockGitRepo + + +def get_mock_task( + subprocess_handler: SubprocessWrapperInterface = MockSubprocessWrapper(), + git_repo_handler: GitRepoInterface = MockGitRepo(), +) -> SpatialTask: + """Returns a mock `SpatialTask` instance.""" + repo = CableRepository( + repo_id=1, + path="path/to/test-branch", + patch={"cable": {"some_branch_specific_setting": True}}, + ) + repo.subprocess_handler = subprocess_handler + repo.root_dir = MOCK_CWD + + task = SpatialTask( + repo=repo, sci_conf_id=0, sci_config={"cable": {"some_setting": True}} + ) + task.subprocess_handler = subprocess_handler + task.git_repo_handler = git_repo_handler + task.root_dir = MOCK_CWD + + return task + + +def test_get_task_name(): + """Tests for `SpatialTask.get_task_name()`.""" + # Success case: check task name convention + task = get_mock_task() + assert task.get_task_name() == "CRUJRA-ACCESS_R1_S0" + + +def test_setup_task(): + """Tests for `SpatialTask.setup_task()`.""" + + (MOCK_CWD / internal.NAMELIST_DIR).mkdir() + (MOCK_CWD / internal.NAMELIST_DIR / internal.CABLE_NML).touch() + (MOCK_CWD / internal.NAMELIST_DIR / internal.CABLE_SOIL_NML).touch() + (MOCK_CWD / internal.NAMELIST_DIR / internal.CABLE_VEGETATION_NML).touch() + + # Success case: test namelist parameters are patched + (MOCK_CWD / internal.SPATIAL_TASKS_DIR).mkdir(parents=True) + git_repo_handler = MockGitRepo() + git_repo_handler.files = ["config.yaml", "cable.nml"] + task = get_mock_task(git_repo_handler=git_repo_handler) + task.setup_task() + task_dir = MOCK_CWD / internal.SPATIAL_TASKS_DIR / task.get_task_name() + res_nml = f90nml.read(str(task_dir / internal.CABLE_NML)) + assert res_nml["cable"] == { + "some_setting": True, + "some_branch_specific_setting": True, + } + shutil.rmtree(MOCK_CWD / internal.SPATIAL_TASKS_DIR) + + # Success case: check config.yaml parameters + (MOCK_CWD / internal.SPATIAL_TASKS_DIR).mkdir(parents=True) + git_repo_handler = MockGitRepo() + git_repo_handler.files = ["config.yaml"] + task = get_mock_task(git_repo_handler=git_repo_handler) + task.setup_task() + task_dir = MOCK_CWD / internal.SPATIAL_TASKS_DIR / task.get_task_name() + with (task_dir / "config.yaml").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"] == [ + "/g/data/tm70/sb8430/cable_inputs/landsea-access-esm.nc", + "/g/data/tm70/ccc561/CABLE/CABLE-as-ACCESS_Yingping/input/gridinfo-access-esm.nc", + "/g/data/tm70/ccc561/CABLE/CABLE-as-ACCESS_Yingping/cruncep10", + 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) + shutil.rmtree(MOCK_CWD / internal.SPATIAL_TASKS_DIR) + + # Success case: test non-verbose standard output + (MOCK_CWD / internal.SPATIAL_TASKS_DIR).mkdir(parents=True) + git_repo_handler = MockGitRepo() + git_repo_handler.files = ["config.yaml"] + task = get_mock_task(git_repo_handler=git_repo_handler) + with contextlib.redirect_stdout(io.StringIO()) as buf: + task.setup_task() + assert not buf.getvalue() + shutil.rmtree(MOCK_CWD / internal.SPATIAL_TASKS_DIR) + + # Success case: test verbose standard output + (MOCK_CWD / internal.SPATIAL_TASKS_DIR).mkdir(parents=True) + git_repo_handler = MockGitRepo() + git_repo_handler.files = ["config.yaml"] + task = get_mock_task(git_repo_handler=git_repo_handler) + with contextlib.redirect_stdout(io.StringIO()) as buf: + task.setup_task(verbose=True) + assert buf.getvalue() == ( + "Setting up task: CRUJRA-ACCESS_R1_S0\n" + " Updating experiment config parameters in " + f"{internal.SPATIAL_TASKS_DIR / task.get_task_name() / 'config.yaml'}\n" + " Adding science configurations to CABLE namelist file " + f"{MOCK_CWD}/runs/spatial/tasks/CRUJRA-ACCESS_R1_S0/cable.nml\n" + " Adding branch specific configurations to CABLE namelist file " + f"{MOCK_CWD}/runs/spatial/tasks/CRUJRA-ACCESS_R1_S0/cable.nml\n" + ) + shutil.rmtree(MOCK_CWD / internal.SPATIAL_TASKS_DIR) + + +def test_run(): + """Tests for `SpatialTask.run()`.""" + + # Success case: test payu run command + mock_subprocess = MockSubprocessWrapper() + task = get_mock_task(subprocess_handler=mock_subprocess) + task_dir = MOCK_CWD / internal.SPATIAL_TASKS_DIR / task.get_task_name() + task_dir.mkdir(parents=True) + task.run() + assert "payu run" in mock_subprocess.commands + + +def test_get_spatial_tasks(): + """Tests for `get_spatial_tasks()`.""" + + # Success case: get task list for two branches and two science + # configurations + config = get_mock_config() + repos = [ + CableRepository(**branch_config, repo_id=id) + for id, branch_config in enumerate(config["realisations"]) + ] + sci_a, sci_b = config["science_configurations"] + tasks = get_spatial_tasks(repos, config["science_configurations"]) + assert [(task.repo, task.sci_config) for task in tasks] == [ + (repos[0], sci_a), + (repos[0], sci_b), + (repos[1], sci_a), + (repos[1], sci_b), + ] diff --git a/tests/test_workdir.py b/tests/test_workdir.py index 2df618a0..d244431b 100644 --- a/tests/test_workdir.py +++ b/tests/test_workdir.py @@ -7,48 +7,43 @@ from tests.common import MOCK_CWD -from tests.common import get_mock_config -from benchcab.fluxsite import Task +from benchcab.fluxsite import FluxsiteTask from benchcab.repository import CableRepository from benchcab.workdir import ( setup_fluxsite_directory_tree, + setup_spatial_directory_tree, clean_directory_tree, setup_src_dir, ) -def setup_mock_tasks() -> list[Task]: +def get_mock_fluxsite_tasks() -> list[FluxsiteTask]: """Return a mock list of fluxsite tasks.""" - - config = get_mock_config() - repo_a = CableRepository("trunk", repo_id=0) - repo_b = CableRepository("path/to/my-branch", repo_id=1) - met_forcing_file_a, met_forcing_file_b = "site_foo", "site_bar" - (sci_id_a, sci_config_a), (sci_id_b, sci_config_b) = enumerate( - config["science_configurations"] - ) - + trunk = CableRepository("trunk", repo_id=0) tasks = [ - Task(repo_a, met_forcing_file_a, sci_id_a, sci_config_a), - Task(repo_a, met_forcing_file_a, sci_id_b, sci_config_b), - Task(repo_a, met_forcing_file_b, sci_id_a, sci_config_a), - Task(repo_a, met_forcing_file_b, sci_id_b, sci_config_b), - Task(repo_b, met_forcing_file_a, sci_id_a, sci_config_a), - Task(repo_b, met_forcing_file_a, sci_id_b, sci_config_b), - Task(repo_b, met_forcing_file_b, sci_id_a, sci_config_a), - Task(repo_b, met_forcing_file_b, sci_id_b, sci_config_b), + FluxsiteTask( + trunk, + met_forcing_file="site_foo", + sci_config={"cable": {"some_option": True}}, + sci_conf_id=0, + ), + FluxsiteTask( + trunk, + met_forcing_file="site_foo", + sci_config={"cable": {"some_option": False}}, + sci_conf_id=1, + ), ] - return tasks def test_setup_directory_tree(): """Tests for `setup_fluxsite_directory_tree()`.""" + tasks = get_mock_fluxsite_tasks() + # Success case: generate fluxsite directory structure - tasks = setup_mock_tasks() setup_fluxsite_directory_tree(fluxsite_tasks=tasks, root_dir=MOCK_CWD) - assert len(list(MOCK_CWD.glob("*"))) == 1 assert Path(MOCK_CWD, "runs").exists() assert Path(MOCK_CWD, "runs", "fluxsite").exists() @@ -58,16 +53,8 @@ def test_setup_directory_tree(): assert Path( MOCK_CWD, "runs", "fluxsite", "analysis", "bitwise-comparisons" ).exists() - assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R0_S0").exists() assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R0_S1").exists() - assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R0_S0").exists() - assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R0_S1").exists() - assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R1_S0").exists() - assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_foo_R1_S1").exists() - assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R1_S0").exists() - assert Path(MOCK_CWD, "runs", "fluxsite", "tasks", "site_bar_R1_S1").exists() - shutil.rmtree(MOCK_CWD / "runs") # Success case: test non-verbose output @@ -82,7 +69,6 @@ def test_setup_directory_tree(): "/runs/fluxsite/analysis/bitwise-comparisons\n" f"Creating task directories...\n" ) - shutil.rmtree(MOCK_CWD / "runs") # Success case: test verbose output @@ -102,20 +88,7 @@ def test_setup_directory_tree(): f"{MOCK_CWD}/runs/fluxsite/tasks/site_foo_R0_S0\n" f"Creating runs/fluxsite/tasks/site_foo_R0_S1: " f"{MOCK_CWD}/runs/fluxsite/tasks/site_foo_R0_S1\n" - f"Creating runs/fluxsite/tasks/site_bar_R0_S0: " - f"{MOCK_CWD}/runs/fluxsite/tasks/site_bar_R0_S0\n" - f"Creating runs/fluxsite/tasks/site_bar_R0_S1: " - f"{MOCK_CWD}/runs/fluxsite/tasks/site_bar_R0_S1\n" - f"Creating runs/fluxsite/tasks/site_foo_R1_S0: " - f"{MOCK_CWD}/runs/fluxsite/tasks/site_foo_R1_S0\n" - f"Creating runs/fluxsite/tasks/site_foo_R1_S1: " - f"{MOCK_CWD}/runs/fluxsite/tasks/site_foo_R1_S1\n" - f"Creating runs/fluxsite/tasks/site_bar_R1_S0: " - f"{MOCK_CWD}/runs/fluxsite/tasks/site_bar_R1_S0\n" - f"Creating runs/fluxsite/tasks/site_bar_R1_S1: " - f"{MOCK_CWD}/runs/fluxsite/tasks/site_bar_R1_S1\n" ) - shutil.rmtree(MOCK_CWD / "runs") @@ -123,7 +96,7 @@ def test_clean_directory_tree(): """Tests for `clean_directory_tree()`.""" # Success case: directory tree does not exist after clean - tasks = setup_mock_tasks() + tasks = get_mock_fluxsite_tasks() setup_fluxsite_directory_tree(fluxsite_tasks=tasks, root_dir=MOCK_CWD) clean_directory_tree(root_dir=MOCK_CWD) @@ -140,3 +113,12 @@ def test_setup_src_dir(): # Success case: make src directory setup_src_dir(root_dir=MOCK_CWD) assert Path(MOCK_CWD, "src").exists() + + +def test_setup_spatial_directory_tree(): + """Tests for `setup_spatial_directory_tree()`.""" + + # Success case: generate spatial directory structure + setup_spatial_directory_tree(root_dir=MOCK_CWD) + assert (MOCK_CWD / "runs" / "spatial" / "tasks").exists() + assert (MOCK_CWD / "runs" / "payu-laboratory").exists()