Skip to content

Commit

Permalink
Add payu test suite for spatial configuration
Browse files Browse the repository at this point in the history
The spatial test suite runs CABLE with CRUJRA forcing at ACCESS
resolution (see [here][cable_example] for more details) over all science
configurations and model versions.

Spatial tests use the [payu framework][payu]. The payu framework was
chosen so that we:
- Encourage uptake of payu amongst users of CABLE
- Have the foundations in place for running coupled models (atmosphere +
  land) with payu
- Can easily test longer running simulations (payu makes it easy to run
  a model multiple times and have state persist in the model via restart
  files)

The run directory structure is organised as follows:

runs/
├── spatial
│   └── tasks
│	├── <spatial-task-name> (a payu control / experiment directory)
│	└── ...
├── payu-laboratory
│   └── ...
└── fluxsite
    └── ...

Note we have a separate payu-laboratory directory. This is so we keep
all CABLE outputs produced by benchcab under the bench_example work
directory.

Add the ability to build the CABLE executable with MPI at runtime so
that we run the spatial configurations with MPI.

Add the --mpi flag to benchcab build command so that the user can
run the MPI build step independently.

Add utility functions for git API requests and manipulating namelist
files.

Add subcommands to run each step of the spatial workflow in isolation.

Add payu key in the benchcab config file so that users can easily
configure payu experiments and pass in command line arguments to the
`payu run` command.

Fixes #5

[payu]: https://github.com/payu-org/payu
[cable_example]: https://github.com/CABLE-LSM/cable_example
  • Loading branch information
SeanBryan51 committed Oct 2, 2023
1 parent b5fe7fd commit 9e2df50
Show file tree
Hide file tree
Showing 20 changed files with 1,042 additions and 275 deletions.
1 change: 1 addition & 0 deletions .conda/benchcab-dev.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,4 @@ dependencies:
- pytest-cov
- pyyaml
- flatdict
- gitpython
117 changes: 86 additions & 31 deletions benchcab/benchcab.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -48,7 +47,18 @@ 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,
met_forcings=internal.SPATIAL_DEFAULT_FORCINGS,
science_configurations=self.science_configurations,
payu_args=self.config.get("spatial", {}).get("payu", {}).get("args"),
)
self.benchcab_exe_path = benchcab_exe_path

if validate_env:
Expand Down Expand Up @@ -103,18 +113,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."""
Expand Down Expand Up @@ -189,7 +197,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:
Expand All @@ -201,30 +209,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"]
Expand All @@ -234,9 +250,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("")

Expand All @@ -248,8 +266,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:
Expand Down Expand Up @@ -280,13 +300,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`."""
Expand All @@ -298,7 +347,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()
Expand All @@ -318,6 +367,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`.
Expand Down
43 changes: 33 additions & 10 deletions benchcab/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand Down Expand Up @@ -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.",
Expand All @@ -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.",
Expand All @@ -110,14 +109,19 @@ def generate_parser() -> argparse.ArgumentParser:
)

# subcommand: 'benchcab build'
subparsers.add_parser(
build_parser = subparsers.add_parser(
"build",
parents=[args_help, args_subcommand],
help="Run the build step in the benchmarking workflow.",
description="""Build the CABLE offline executable for each repository specified in the
config file.""",
add_help=False,
)
build_parser.add_argument(
"--mpi",
action="store_true",
help="Enable MPI build.",
)

# subcommand: 'benchcab fluxsite-setup-work-dir'
subparsers.add_parser(
Expand All @@ -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,
)

Expand All @@ -165,10 +169,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
23 changes: 23 additions & 0 deletions benchcab/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,29 @@ def check_config(config: dict):
):
raise TypeError("The 'multiprocessing' key must be a boolean.")

# the "spatial" key is optional
if "spatial" in config:
if not isinstance(config["spatial"], dict):
raise TypeError("The 'spatial' key must be a dictionary.")
# the "payu" key is optional
if "payu" in config["spatial"]:
if not isinstance(config["spatial"]["payu"], dict):
raise TypeError("The 'payu' key must be a dictionary.")
# the "config" key is optional
if "config" in config["spatial"]["payu"]:
if not isinstance(config["spatial"]["payu"]["config"], dict):
raise TypeError("The 'config' key must be a dictionary.")
# the "args" key is optional
if "args" in config["spatial"]["payu"]:
if not isinstance(config["spatial"]["payu"]["args"], str):
raise TypeError("The 'args' key must be a string.")
# 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"]
):
raise TypeError("The 'met_forcings' key must be a list of strings.")

valid_experiments = (
list(internal.MEORG_EXPERIMENTS) + internal.MEORG_EXPERIMENTS["five-site-test"]
)
Expand Down
Loading

0 comments on commit 9e2df50

Please sign in to comment.