Skip to content

Commit

Permalink
Arbitrary CLI args for shell tasks (#66)
Browse files Browse the repository at this point in the history
This PR adds the `_CliArgsBaseModel` class (inspired from `_WhenBaseModel`), which replaces the `command_option` and `input_arg_options` of the `ConfigShellTaskSpecs`. Validation is applied on the correctness of the keyword and positional arguments to ensure they start with `-` or `--`.

The three test YAML files are updated accordingly, leading to changes in the pretty-print test text files, which is also part of the PR.

Actually making something useful out of these arguments only happens when the WG is created, so will require #45 to be merged.
  • Loading branch information
GeigerJ2 authored Dec 17, 2024
1 parent 7ee6861 commit 3520519
Show file tree
Hide file tree
Showing 8 changed files with 199 additions and 341 deletions.
2 changes: 1 addition & 1 deletion src/sirocco/core/workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@


class Workflow:
"""Internal reprensentation of a workflow"""
"""Internal representation of a workflow"""

def __init__(self, workflow_config: ConfigWorkflow) -> None:
self.name = workflow_config.name
Expand Down
41 changes: 38 additions & 3 deletions src/sirocco/parsing/_yaml_data_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,43 @@ def convert_datetime(cls, value) -> datetime:
return datetime.fromisoformat(value)


class _CliArgsBaseModel(BaseModel):
"""Base class for cli_arguments specifications"""

# TODO: Even allow for `str`, or always require list?
positional: str | list[str] | None = None
# Field needed for child class doing pydantic parsing
keyword: dict[str, str] | None = Field(default_factory=dict)
flags: str | list[str] | None = None
source_file: str | list[str] | None = None

# TODO: Should we allow users to pass it without the hyphen(s), and prepend them automatically?
# TODO: While convenient, it could be a bad idea, if users put in wrong things. Better to be explicit.
@field_validator("keyword", mode="before")
@classmethod
def validate_keyword_args(cls, value):
"""Ensure keyword arguments start with '-' or '--'."""
if value is not None:
invalid_keys = [key for key in value if not key.startswith(("-", "--"))]
if invalid_keys:
invalid_kwarg_exc = f"Invalid keyword arguments: {', '.join(invalid_keys)}"
raise ValueError(invalid_kwarg_exc)
return value

@field_validator("flags", mode="before")
@classmethod
def validate_flag_args(cls, value):
"""Ensure positional arguments start with '-' or '--'."""
if value is not None:
if isinstance(value, str):
value = [value]
invalid_flags = [arg for arg in value if not arg.startswith(("-", "--"))]
if invalid_flags:
invalid_flags_exc = f"Invalid positional arguments: {', '.join(invalid_flags)}"
raise ValueError(invalid_flags_exc)
return value


class TargetNodesBaseModel(_NamedBaseModel):
"""class for targeting other task or data nodes in the graph
Expand Down Expand Up @@ -273,9 +310,7 @@ class ConfigRootTask(ConfigBaseTask):
class ConfigShellTaskSpecs:
plugin: ClassVar[Literal["shell"]] = "shell"
command: str = ""
command_option: str = ""
input_arg_options: dict[str, str] = Field(default_factory=dict) # noqa: RUF009 Field needed
# for child class doing pydantic parsing
cli_arguments: _CliArgsBaseModel | None = None
src: str | None = None


Expand Down
57 changes: 36 additions & 21 deletions tests/files/configs/test_config_large.yml
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,6 @@ cycles:
inputs:
- stream_2:
lag: ['P0M', 'P2M', 'P4M', 'P6M', 'P8M', 'P10M']
arg_option: --input
outputs: [postout_2]
- store_and_clean_2:
inputs:
Expand All @@ -63,9 +62,11 @@ tasks:
- extpar:
plugin: shell # no extpar plugin available yet
command: $PWD/examples/files/scripts/extpar
command_option: '--verbose' # todo implement support
input_arg_options:
obs_data: '--input'
cli_arguments:
keyword:
--input: obs_data
flags:
- --verbose
uenv:
squashfs: path/to/squashfs
mount_point: runtime/mount/point
Expand All @@ -74,10 +75,13 @@ tasks:
- preproc:
plugin: shell
command: $PWD/examples/files/scripts/cleanup.sh
input_arg_options:
grid_file: '-g'
extpar_file : '-p'
ERA5: '-e'
cli_arguments:
positional:
- grid_file
keyword:
-p: extpar_file
-e: ERA5
source_file: dummy_source_file
nodes: 4
walltime: 00:02:00
uenv:
Expand All @@ -86,9 +90,10 @@ tasks:
- icon:
plugin: icon
command: $PWD/examples/files/scripts/icon
input_arg_options:
grid_file: '-g'
icon_input: '--input'
cli_arguments:
keyword:
-g: grid_file
--input: icon_input
nodes: 40
walltime: 23:59:59
namelists:
Expand All @@ -100,8 +105,9 @@ tasks:
- postproc_1:
plugin: shell
command: $PWD/examples/files/scripts/main_script_ocn.sh
input_arg_options:
stream_1: '--input'
cli_arguments:
keyword:
--input: stream_1
nodes: 2
walltime: 00:05:00
uenv:
Expand All @@ -110,8 +116,12 @@ tasks:
- postproc_2:
plugin: shell
command: $PWD/examples/files/scripts/main_script_atm.sh
input_arg_options:
stream_2: '--input'
cli_arguments:
keyword:
--input: stream_2
# `arg_option` should be in `tasks` section instead
# How to implement this? Even needed with keyword-arguments?
# arg_option: --input
nodes: 2
walltime: 00:05:00
src: path/to/src/dir
Expand All @@ -121,17 +131,19 @@ tasks:
- store_and_clean_1:
plugin: shell
command: $PWD/examples/files/scripts/post_clean.sh
input_arg_options:
postout_1: '--input'
stream_1: '--stream'
icon_input: '--icon_input'
cli_arguments:
keyword:
--input: postout_1
--stream: stream_1
--icon_input: icon_input
nodes: 1
walltime: 00:01:00
- store_and_clean_2:
plugin: shell
command: $PWD/examples/files/scripts/post_clean.sh
input_arg_options:
postout_2: '--input'
cli_arguments:
keyword:
--input: postout_2
nodes: 1
walltime: 00:01:00
data:
Expand All @@ -145,6 +157,9 @@ data:
- ERA5:
type: file
src: $PWD/examples/files/data/era5
- dummy_source_file:
type: file
src: $PWD/examples/files/data/dummy_source_file.sh
generated:
- extpar_file:
type: file
Expand Down
5 changes: 3 additions & 2 deletions tests/files/configs/test_config_parameters.yml
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,9 @@ tasks:
- icon:
plugin: shell
command: $PWD/tests/files/scripts/icon.py
input_arg_options:
icon_restart: '--restart'
cli_arguments:
keyword:
--restart: icon_restart
parameters: [foo, bar]
- statistics_foo:
plugin: shell
Expand Down
5 changes: 3 additions & 2 deletions tests/files/configs/test_config_small.yml
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,9 @@ tasks:
- icon:
plugin: shell
command: $PWD/tests/files/scripts/icon.py
input_arg_options:
icon_restart: '--restart'
cli_arguments:
keyword:
--restart: icon_restart
- cleanup:
plugin: shell
command: $PWD/tests/files/scripts/cleanup.py
Expand Down
Loading

0 comments on commit 3520519

Please sign in to comment.