diff --git a/docs/shared/ungrib.yaml b/docs/shared/ungrib.yaml new file mode 100644 index 000000000..59a7cbd9d --- /dev/null +++ b/docs/shared/ungrib.yaml @@ -0,0 +1,10 @@ +ungrib: + execution: + mpicmd: srun + batchargs: + cores: 1 + walltime: 00:01:00 + executable: "/path/to/ungrib.exe" + gfs_file: "/path/to/dir/gfs.t12z.pgrb2.0p25.f000" + run_dir: "/path/to/dir/run" + vtable: "/path/to/Vtable.GFS" diff --git a/src/uwtools/api/ungrib.py b/src/uwtools/api/ungrib.py new file mode 100644 index 000000000..7f998df68 --- /dev/null +++ b/src/uwtools/api/ungrib.py @@ -0,0 +1,53 @@ +""" +API access to the ``uwtools`` ``ungrib`` driver. +""" +import datetime as dt +from pathlib import Path +from typing import Dict, Optional + +import uwtools.drivers.support as _support +from uwtools.drivers.ungrib import Ungrib as _Ungrib + + +def execute( + task: str, + config_file: Path, + cycle: dt.datetime, + batch: bool = False, + dry_run: bool = False, + graph_file: Optional[Path] = None, +) -> bool: + """ + Execute an ``ungrib`` task. + + If ``batch`` is specified, a runscript will be written and submitted to the batch system. + Otherwise, the executable will be run directly on the current system. + + :param task: The task to execute + :param config_file: Path to YAML config file + :param cycle: The cycle to run + :param batch: Submit run to the batch system + :param dry_run: Do not run the executable, just report what would have been done + :param graph_file: Write Graphviz DOT output here + :return: ``True`` if task completes without raising an exception + """ + obj = _Ungrib(config_file=config_file, cycle=cycle, batch=batch, dry_run=dry_run) + getattr(obj, task)() + if graph_file: + with open(graph_file, "w", encoding="utf-8") as f: + print(graph(), file=f) + return True + + +def graph() -> str: + """ + Returns Graphviz DOT code for the most recently executed task. + """ + return _support.graph() + + +def tasks() -> Dict[str, str]: + """ + Returns a mapping from task names to their one-line descriptions. + """ + return _support.tasks(_Ungrib) diff --git a/src/uwtools/cli.py b/src/uwtools/cli.py index e450fcc86..27bdf1a13 100644 --- a/src/uwtools/cli.py +++ b/src/uwtools/cli.py @@ -18,6 +18,7 @@ import uwtools.api.rocoto import uwtools.api.sfc_climo_gen import uwtools.api.template +import uwtools.api.ungrib import uwtools.config.jinja2 import uwtools.rocoto from uwtools.exceptions import UWConfigRealizeError, UWError, UWTemplateRenderError @@ -61,6 +62,7 @@ def main() -> None: STR.rocoto: _dispatch_rocoto, STR.sfcclimogen: _dispatch_sfc_climo_gen, STR.template: _dispatch_template, + STR.ungrib: _dispatch_ungrib, } sys.exit(0 if modes[args[STR.mode]](args) else 1) except UWError as e: @@ -658,6 +660,60 @@ def _dispatch_template_translate(args: Args) -> bool: ) +# Mode ungrib + + +def _add_subparser_ungrib(subparsers: Subparsers) -> ModeChecks: + """ + Subparser for mode: ungrib + + :param subparsers: Parent parser's subparsers, to add this subparser to. + """ + parser = _add_subparser(subparsers, STR.ungrib, "Execute Ungrib tasks") + _basic_setup(parser) + subparsers = _add_subparsers(parser, STR.action, STR.task.upper()) + return { + task: _add_subparser_ungrib_task(subparsers, task, helpmsg) + for task, helpmsg in uwtools.api.ungrib.tasks().items() + } + + +def _add_subparser_ungrib_task(subparsers: Subparsers, task: str, helpmsg: str) -> ActionChecks: + """ + Subparser for mode: ungrib + + :param subparsers: Parent parser's subparsers, to add this subparser to. + :param task: The task to add a subparser for. + :param helpmsg: Help message for task. + """ + parser = _add_subparser(subparsers, task, helpmsg.rstrip(".")) + required = parser.add_argument_group(TITLE_REQ_ARG) + _add_arg_config_file(group=required, required=True) + _add_arg_cycle(required) + optional = _basic_setup(parser) + _add_arg_batch(optional) + _add_arg_dry_run(optional) + _add_arg_graph_file(optional) + checks = _add_args_verbosity(optional) + return checks + + +def _dispatch_ungrib(args: Args) -> bool: + """ + Dispatch logic for ungrib mode. + + :param args: Parsed command-line args. + """ + return uwtools.api.ungrib.execute( + task=args[STR.action], + config_file=args[STR.cfgfile], + cycle=args[STR.cycle], + batch=args[STR.batch], + dry_run=args[STR.dryrun], + graph_file=args[STR.graphfile], + ) + + # Arguments # pylint: disable=missing-function-docstring @@ -1031,6 +1087,7 @@ def _parse_args(raw_args: List[str]) -> Tuple[Args, Checks]: STR.rocoto: _add_subparser_rocoto(subparsers), STR.sfcclimogen: _add_subparser_sfc_climo_gen(subparsers), STR.template: _add_subparser_template(subparsers), + STR.ungrib: _add_subparser_ungrib(subparsers), } return vars(parser.parse_args(raw_args)), checks diff --git a/src/uwtools/drivers/ungrib.py b/src/uwtools/drivers/ungrib.py new file mode 100644 index 000000000..6080331f8 --- /dev/null +++ b/src/uwtools/drivers/ungrib.py @@ -0,0 +1,148 @@ +""" +A driver for the ungrib component. +""" + +from datetime import datetime +from pathlib import Path +from typing import Any, Dict + +from iotaa import asset, dryrun, task, tasks + +from uwtools.config.formats.nml import NMLConfig +from uwtools.drivers.driver import Driver +from uwtools.strings import STR +from uwtools.utils.tasks import file + + +class Ungrib(Driver): + """ + A driver for ungrib. + """ + + def __init__( + self, config_file: Path, cycle: datetime, dry_run: bool = False, batch: bool = False + ): + """ + The driver. + + :param config_file: Path to config file. + :param cycle: The forecast cycle. + :param dry_run: Run in dry-run mode? + :param batch: Run component via the batch system? + """ + super().__init__(config_file=config_file, dry_run=dry_run, batch=batch) + self._config.dereference(context={"cycle": cycle}) + if self._dry_run: + dryrun() + self._cycle = cycle + + # Workflow tasks + + @task + def gribfile_aaa(self): + """ + The gribfile. + """ + path = self._rundir / "GRIBFILE.AAA" + yield self._taskname(str(path)) + yield asset(path, path.is_symlink) + infile = Path(self._driver_config["gfs_file"]) + yield file(path=infile) + path.parent.mkdir(parents=True, exist_ok=True) + path.symlink_to(infile) + + @task + def namelist_file(self): + """ + The namelist file. + """ + d = { + "update_values": { + "share": { + "end_date": self._cycle.strftime("%Y-%m-%d_%H:00:00"), + "interval_seconds": 1, + "max_dom": 1, + "start_date": self._cycle.strftime("%Y-%m-%d_%H:00:00"), + "wrf_core": "ARW", + }, + "ungrib": { + "out_format": "WPS", + "prefix": "FILE", + }, + } + } + path = self._rundir / "namelist.wps" + yield self._taskname(str(path)) + yield asset(path, path.is_file) + yield None + path.parent.mkdir(parents=True, exist_ok=True) + self._create_user_updated_config( + config_class=NMLConfig, + config_values=d, + path=path, + ) + + @tasks + def provisioned_run_directory(self): + """ + Run directory provisioned with all required content. + """ + yield self._taskname("provisioned run directory") + yield [ + self.gribfile_aaa(), + self.namelist_file(), + self.runscript(), + self.vtable(), + ] + + @task + def runscript(self): + """ + The runscript. + """ + path = self._runscript_path + yield self._taskname(path.name) + yield asset(path, path.is_file) + yield None + self._write_runscript(path=path, envvars={}) + + @task + def vtable(self): + """ + The Vtable. + """ + path = self._rundir / "Vtable" + yield self._taskname(str(path)) + yield asset(path, path.is_symlink) + yield None + path.parent.mkdir(parents=True, exist_ok=True) + path.symlink_to(Path(self._driver_config["vtable"])) + + # Private helper methods + + @property + def _driver_name(self) -> str: + """ + Returns the name of this driver. + """ + return STR.ungrib + + @property + def _resources(self) -> Dict[str, Any]: + """ + Returns configuration data for the runscript. + """ + return { + "account": self._config["platform"]["account"], + "rundir": self._rundir, + "scheduler": self._config["platform"]["scheduler"], + **self._driver_config.get("execution", {}).get("batchargs", {}), + } + + def _taskname(self, suffix: str) -> str: + """ + Returns a common tag for graph-task log messages. + + :param suffix: Log-string suffix. + """ + return "%s ungrib %s" % (self._cycle.strftime("%Y%m%d %HZ"), suffix) diff --git a/src/uwtools/resources/jsonschema/ungrib.jsonschema b/src/uwtools/resources/jsonschema/ungrib.jsonschema new file mode 100644 index 000000000..d3bbc9208 --- /dev/null +++ b/src/uwtools/resources/jsonschema/ungrib.jsonschema @@ -0,0 +1,29 @@ +{ + "properties": { + "ungrib": { + "additionalProperties": false, + "properties": { + "execution": { + "$ref": "urn:uwtools:execution" + }, + "gfs_file": { + "type": "string" + }, + "run_dir": { + "type": "string" + }, + "vtable": { + "type": "string" + } + }, + "required": [ + "execution", + "gfs_file", + "run_dir", + "vtable" + ], + "type": "object" + } + }, + "type": "object" +} diff --git a/src/uwtools/strings.py b/src/uwtools/strings.py index cd1ee9d3f..ebb8f85a2 100644 --- a/src/uwtools/strings.py +++ b/src/uwtools/strings.py @@ -104,6 +104,7 @@ class STR: template: str = "template" total: str = "total" translate: str = "translate" + ungrib: str = "ungrib" validate: str = "validate" valsfile: str = "values_file" valsfmt: str = "values_format" diff --git a/src/uwtools/tests/api/test_ungrib.py b/src/uwtools/tests/api/test_ungrib.py new file mode 100644 index 000000000..52fe1a63d --- /dev/null +++ b/src/uwtools/tests/api/test_ungrib.py @@ -0,0 +1,34 @@ +# pylint: disable=missing-function-docstring,protected-access + +import datetime as dt +from unittest.mock import patch + +from uwtools.api import ungrib + + +def test_execute(tmp_path): + dot = tmp_path / "graph.dot" + args: dict = { + "config_file": "config.yaml", + "cycle": dt.datetime.utcnow(), + "batch": False, + "dry_run": True, + "graph_file": dot, + } + with patch.object(ungrib, "_Ungrib") as Ungrib: + assert ungrib.execute(**args, task="foo") is True + del args["graph_file"] + Ungrib.assert_called_once_with(**args) + Ungrib().foo.assert_called_once_with() + + +def test_graph(): + with patch.object(ungrib._support, "graph") as graph: + ungrib.graph() + graph.assert_called_once_with() + + +def test_tasks(): + with patch.object(ungrib._support, "tasks") as _tasks: + ungrib.tasks() + _tasks.assert_called_once_with(ungrib._Ungrib) diff --git a/src/uwtools/tests/drivers/test_ungrib.py b/src/uwtools/tests/drivers/test_ungrib.py new file mode 100644 index 000000000..213731c7f --- /dev/null +++ b/src/uwtools/tests/drivers/test_ungrib.py @@ -0,0 +1,182 @@ +# pylint: disable=missing-function-docstring,protected-access,redefined-outer-name +""" +Ungrib driver tests. +""" +import datetime as dt +from unittest.mock import DEFAULT as D +from unittest.mock import patch + +import f90nml # type: ignore +import yaml +from pytest import fixture + +from uwtools.drivers import ungrib + +# Fixtures + + +@fixture +def cycle(): + return dt.datetime(2024, 2, 1, 18) + + +# Driver fixtures + + +@fixture +def config(tmp_path): + return { + "ungrib": { + "execution": { + "mpicmd": "srun", + "batchargs": { + "cores": 1, + "walltime": "00:01:00", + }, + "executable": str(tmp_path / "ungrib.exe"), + }, + "gfs_file": str(tmp_path / "gfs.t12z.pgrb2.0p25.f000"), + "run_dir": str(tmp_path), + "vtable": str(tmp_path / "Vtable.GFS"), + }, + "platform": { + "account": "wrfruc", + "scheduler": "slurm", + }, + } + + +@fixture +def config_file(config, tmp_path): + path = tmp_path / "config.yaml" + with open(path, "w", encoding="utf-8") as f: + yaml.dump(config, f) + return path + + +@fixture +def driverobj(config_file, cycle): + return ungrib.Ungrib(config_file=config_file, cycle=cycle, batch=True) + + +# Driver tests + + +def test_Ungrib(driverobj): + assert isinstance(driverobj, ungrib.Ungrib) + + +def test_Ungrib_dry_run(config_file, cycle): + with patch.object(ungrib, "dryrun") as dryrun: + driverobj = ungrib.Ungrib(config_file=config_file, cycle=cycle, batch=True, dry_run=True) + assert driverobj._dry_run is True + dryrun.assert_called_once_with() + + +def test_Ungrib_gribfile_aaa(driverobj): + src = driverobj._rundir / "GRIBFILE.AAA.in" + src.touch() + driverobj._driver_config["gfs_file"] = src + dst = driverobj._rundir / "GRIBFILE.AAA" + assert not dst.is_symlink() + driverobj.gribfile_aaa() + assert dst.is_symlink() + + +def test_Ungrib_namelist_file(driverobj): + dst = driverobj._rundir / "namelist.wps" + assert not dst.is_file() + driverobj.namelist_file() + assert dst.is_file() + assert isinstance(f90nml.read(dst), f90nml.Namelist) + + +def test_Ungrib_provisioned_run_directory(driverobj): + with patch.multiple( + driverobj, + gribfile_aaa=D, + namelist_file=D, + runscript=D, + vtable=D, + ) as mocks: + driverobj.provisioned_run_directory() + for m in mocks: + mocks[m].assert_called_once_with() + + +def test_Ungrib_run_batch(driverobj): + with patch.object(driverobj, "_run_via_batch_submission") as func: + driverobj.run() + func.assert_called_once_with() + + +def test_Ungrib_run_local(driverobj): + driverobj._batch = False + with patch.object(driverobj, "_run_via_local_execution") as func: + driverobj.run() + func.assert_called_once_with() + + +def test_Ungrib_runscript(driverobj): + dst = driverobj._rundir / "runscript.ungrib" + assert not dst.is_file() + driverobj._driver_config["execution"].update( + { + "batchargs": {"walltime": "00:01:00"}, + "envcmds": ["cmd1", "cmd2"], + "mpicmd": "runit", + "threads": 8, + } + ) + driverobj._config["platform"] = {"account": "me", "scheduler": "slurm"} + driverobj.runscript() + with open(dst, "r", encoding="utf-8") as f: + lines = f.read().split("\n") + # Check directives: + assert "#SBATCH --account=me" in lines + assert "#SBATCH --time=00:01:00" in lines + # Check environment commands: + assert "cmd1" in lines + assert "cmd2" in lines + # Check execution: + assert "test $? -eq 0 && touch %s/done" % driverobj._rundir + + +def test_Ungrib_vtable(driverobj): + src = driverobj._rundir / "Vtable.GFS.in" + src.touch() + driverobj._driver_config["vtable"] = src + dst = driverobj._rundir / "Vtable" + assert not dst.is_symlink() + driverobj.vtable() + assert dst.is_symlink() + + +def test_Ungrib__driver_config(driverobj): + assert driverobj._driver_config == driverobj._config["ungrib"] + + +def test_Ungrib__resources(driverobj): + account = "me" + scheduler = "slurm" + walltime = "00:01:00" + driverobj._driver_config["execution"].update({"batchargs": {"walltime": walltime}}) + driverobj._config["platform"] = {"account": account, "scheduler": scheduler} + assert driverobj._resources == { + "account": account, + "rundir": driverobj._rundir, + "scheduler": scheduler, + "walltime": walltime, + } + + +def test_Ungrib__runscript_path(driverobj): + assert driverobj._runscript_path == driverobj._rundir / "runscript.ungrib" + + +def test_Ungrib__taskanme(driverobj): + assert driverobj._taskname("foo") == "20240201 18Z ungrib foo" + + +def test_Ungrib__validate(driverobj): + driverobj._validate() diff --git a/src/uwtools/tests/test_cli.py b/src/uwtools/tests/test_cli.py index dc3f0b8e7..3e4a018a4 100644 --- a/src/uwtools/tests/test_cli.py +++ b/src/uwtools/tests/test_cli.py @@ -16,8 +16,10 @@ import uwtools.api.rocoto import uwtools.api.sfc_climo_gen import uwtools.api.template +import uwtools.api.ungrib import uwtools.drivers.fv3 import uwtools.drivers.sfc_climo_gen +import uwtools.drivers.ungrib from uwtools import cli from uwtools.cli import STR from uwtools.exceptions import UWConfigRealizeError, UWError, UWTemplateRenderError @@ -167,6 +169,18 @@ def test__add_subparser_template_translate(subparsers): assert subparsers.choices[STR.translate] +def test__add_subparser_ungrib(subparsers): + cli._add_subparser_ungrib(subparsers) + assert actions(subparsers.choices[STR.ungrib]) == [ + "gribfile_aaa", + "namelist_file", + "provisioned_run_directory", + "run", + "runscript", + "vtable", + ] + + @pytest.mark.parametrize( "vals", [ @@ -604,6 +618,19 @@ def test__dispatch_template_translate_no_optional(): ) +def test__dispatch_ungrib(): + args: dict = { + "batch": True, + "config_file": "config.yaml", + "cycle": dt.datetime.now(), + "dry_run": False, + "graph_file": None, + } + with patch.object(uwtools.api.ungrib, "execute") as execute: + cli._dispatch_ungrib({**args, "action": "foo"}) + execute.assert_called_once_with(**{**args, "task": "foo"}) + + @pytest.mark.parametrize("quiet", [False, True]) @pytest.mark.parametrize("verbose", [False, True]) def test_main_fail_checks(capsys, quiet, verbose):