From 497e3961252cab074a483eebe61975f7deb68c67 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Fri, 29 Aug 2025 17:24:01 +0200 Subject: [PATCH 1/4] First implementation --- src/sirocco/core/_tasks/icon_task.py | 13 ++++++++++--- src/sirocco/parsing/yaml_data_models.py | 5 +++++ src/sirocco/workgraph.py | 5 ++++- tests/cases/small-icon/config/config.yml | 1 + tests/cases/small-icon/config/scripts/setup_env.sh | 0 tests/unit_tests/test_workgraph.py | 1 + 6 files changed, 21 insertions(+), 4 deletions(-) create mode 100644 tests/cases/small-icon/config/scripts/setup_env.sh diff --git a/src/sirocco/core/_tasks/icon_task.py b/src/sirocco/core/_tasks/icon_task.py index 21c96db0..bec6391d 100644 --- a/src/sirocco/core/_tasks/icon_task.py +++ b/src/sirocco/core/_tasks/icon_task.py @@ -53,7 +53,14 @@ def __post_init__(self): self._model_namelist = model_namelist if self.wrapper_script is not None: - self.wrapper_script = self._validate_wrapper_script(self.wrapper_script, self.config_rootdir) + self.wrapper_script = self._validate_local_script_path( + script_path=self.wrapper_script, config_rootdir=self.config_rootdir + ) + + if self.setup_env is not None: + self.setup_env = self._validate_local_script_path( + script_path=self.setup_env, config_rootdir=self.config_rootdir + ) @property def master_namelist(self) -> NamelistFile: @@ -120,9 +127,9 @@ def build_from_config(cls: type[Self], config: models.ConfigTask, config_rootdir self.update_icon_namelists_from_workflow() return self - def _validate_wrapper_script(self, wrapper_script: Path, config_rootdir: Path) -> Path: + def _validate_local_script_path(self, script_path: Path, config_rootdir: Path) -> Path: """Validate and resolve wrapper script path""" - resolved_path = wrapper_script if wrapper_script.is_absolute() else config_rootdir / wrapper_script + resolved_path = script_path if script_path.is_absolute() else config_rootdir / script_path if not resolved_path.exists(): msg = f"Wrapper script in path {resolved_path} does not exist." diff --git a/src/sirocco/parsing/yaml_data_models.py b/src/sirocco/parsing/yaml_data_models.py index aaa23b18..4eeb0446 100644 --- a/src/sirocco/parsing/yaml_data_models.py +++ b/src/sirocco/parsing/yaml_data_models.py @@ -491,6 +491,11 @@ class ConfigIconTaskSpecs: repr=False, metadata={"description": "Path to wrapper script file relative to the config directory or absolute."}, ) + setup_env: Path | None = field( + default=None, + repr=False, + metadata={"description": ("A file that is sourced before the execution of ICON to set environment variables.")}, + ) class ConfigIconTask(ConfigBaseTask, ConfigIconTaskSpecs): diff --git a/src/sirocco/workgraph.py b/src/sirocco/workgraph.py index 56f083c7..4285da96 100644 --- a/src/sirocco/workgraph.py +++ b/src/sirocco/workgraph.py @@ -354,6 +354,9 @@ def _create_icon_task_node(self, task: core.IconTask): if wrapper_script_data is not None: builder.wrapper_script = wrapper_script_data + if task.setup_env is not None: + builder.setup_env = aiida.orm.SinglefileData(str(task.setup_env)) + # Set runtime information options = {} options.update(self._from_task_get_scheduler_options(task)) @@ -576,7 +579,7 @@ def _get_default_wrapper_script() -> aiida.orm.SinglefileData | None: """Get default wrapper script based on task type""" # Import the script directory from aiida-icon - from aiida_icon.site_support.cscs.todi import SCRIPT_DIR + from aiida_icon.site_support.cscs.alps import SCRIPT_DIR default_script_path = SCRIPT_DIR / "todi_cpu.sh" return aiida.orm.SinglefileData(file=default_script_path) diff --git a/tests/cases/small-icon/config/config.yml b/tests/cases/small-icon/config/config.yml index 08d99923..77937e50 100644 --- a/tests/cases/small-icon/config/config.yml +++ b/tests/cases/small-icon/config/config.yml @@ -54,6 +54,7 @@ tasks: ntasks_per_node: 1 cpus_per_task: 1 wrapper_script: scripts/dummy_wrapper.sh + setup_env: scripts/setup_env.sh uenv: "icon-wcp/v1:rc4" view: "icon" - cleanup: diff --git a/tests/cases/small-icon/config/scripts/setup_env.sh b/tests/cases/small-icon/config/scripts/setup_env.sh new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/test_workgraph.py b/tests/unit_tests/test_workgraph.py index 9f6aa94d..7f5a29c8 100644 --- a/tests/unit_tests/test_workgraph.py +++ b/tests/unit_tests/test_workgraph.py @@ -148,6 +148,7 @@ def test_aiida_icon_task_metadata(config_paths): for aiida_icon_task in [task for task in aiida_workflow._workgraph.tasks if task.identifier == "IconCalculation"]: # noqa: SLF001 # testing wrapper script assert aiida_icon_task.inputs.wrapper_script.value.filename == "dummy_wrapper.sh" + assert aiida_icon_task.inputs.setup_env.value.filename == "setup_env.sh" # testing uenv assert ( "#SBATCH --uenv=icon-wcp/v1:rc4" in aiida_icon_task.inputs.metadata.options.custom_scheduler_commands.value From 83bd4fc60de871bfbb5e45f05c544954585e15ef Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Fri, 29 Aug 2025 17:49:29 +0200 Subject: [PATCH 2/4] wip --- src/sirocco/core/_tasks/icon_task.py | 13 +++++-------- src/sirocco/parsing/yaml_data_models.py | 2 +- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/src/sirocco/core/_tasks/icon_task.py b/src/sirocco/core/_tasks/icon_task.py index bec6391d..f2350ec1 100644 --- a/src/sirocco/core/_tasks/icon_task.py +++ b/src/sirocco/core/_tasks/icon_task.py @@ -52,15 +52,12 @@ def __post_init__(self): raise ValueError(msg) self._model_namelist = model_namelist + # TODO: Could use a different method here, e.g., if self.wrapper_script is not None: - self.wrapper_script = self._validate_local_script_path( - script_path=self.wrapper_script, config_rootdir=self.config_rootdir - ) + self.wrapper_script = self._resolve_local_script_path(rel_path=self.wrapper_script) if self.setup_env is not None: - self.setup_env = self._validate_local_script_path( - script_path=self.setup_env, config_rootdir=self.config_rootdir - ) + self.setup_env = self._resolve_local_script_path(rel_path=self.setup_env) @property def master_namelist(self) -> NamelistFile: @@ -127,9 +124,9 @@ def build_from_config(cls: type[Self], config: models.ConfigTask, config_rootdir self.update_icon_namelists_from_workflow() return self - def _validate_local_script_path(self, script_path: Path, config_rootdir: Path) -> Path: + def _resolve_local_script_path(self, rel_path: Path) -> Path: """Validate and resolve wrapper script path""" - resolved_path = script_path if script_path.is_absolute() else config_rootdir / script_path + resolved_path = rel_path if rel_path.is_absolute() else self.config_rootdir / rel_path if not resolved_path.exists(): msg = f"Wrapper script in path {resolved_path} does not exist." diff --git a/src/sirocco/parsing/yaml_data_models.py b/src/sirocco/parsing/yaml_data_models.py index 4eeb0446..aebb5bea 100644 --- a/src/sirocco/parsing/yaml_data_models.py +++ b/src/sirocco/parsing/yaml_data_models.py @@ -494,7 +494,7 @@ class ConfigIconTaskSpecs: setup_env: Path | None = field( default=None, repr=False, - metadata={"description": ("A file that is sourced before the execution of ICON to set environment variables.")}, + metadata={"description": "A file that is sourced before the execution of ICON to set environment variables."}, ) From 02b0a047b11b580ae864f813bfefcc3db64ead75 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Mon, 1 Sep 2025 15:10:24 +0200 Subject: [PATCH 3/4] wip modify test setup --- src/sirocco/core/_tasks/icon_task.py | 1 - tests/cases/small-icon/config/ICON/icon_master.namelist | 2 +- .../config/ICON/{model.namelist => models/atm.namelist} | 0 tests/cases/small-icon/config/config.yml | 2 +- 4 files changed, 2 insertions(+), 3 deletions(-) rename tests/cases/small-icon/config/ICON/{model.namelist => models/atm.namelist} (100%) diff --git a/src/sirocco/core/_tasks/icon_task.py b/src/sirocco/core/_tasks/icon_task.py index f2350ec1..b9b1e74a 100644 --- a/src/sirocco/core/_tasks/icon_task.py +++ b/src/sirocco/core/_tasks/icon_task.py @@ -52,7 +52,6 @@ def __post_init__(self): raise ValueError(msg) self._model_namelist = model_namelist - # TODO: Could use a different method here, e.g., if self.wrapper_script is not None: self.wrapper_script = self._resolve_local_script_path(rel_path=self.wrapper_script) diff --git a/tests/cases/small-icon/config/ICON/icon_master.namelist b/tests/cases/small-icon/config/ICON/icon_master.namelist index 37910d0e..c31397d6 100644 --- a/tests/cases/small-icon/config/ICON/icon_master.namelist +++ b/tests/cases/small-icon/config/ICON/icon_master.namelist @@ -11,7 +11,7 @@ / &master_model_nml model_name="atm" - model_namelist_filename="model.namelist" + model_namelist_filename="atm.namelist" model_type=1 model_min_rank=0 model_max_rank=65535 diff --git a/tests/cases/small-icon/config/ICON/model.namelist b/tests/cases/small-icon/config/ICON/models/atm.namelist similarity index 100% rename from tests/cases/small-icon/config/ICON/model.namelist rename to tests/cases/small-icon/config/ICON/models/atm.namelist diff --git a/tests/cases/small-icon/config/config.yml b/tests/cases/small-icon/config/config.yml index 77937e50..85dd2ce7 100644 --- a/tests/cases/small-icon/config/config.yml +++ b/tests/cases/small-icon/config/config.yml @@ -47,7 +47,7 @@ tasks: bin: /TESTS_ROOTDIR/tests/cases/small-icon/config/ICON/bin/icon namelists: - ./ICON/icon_master.namelist - - ./ICON/model.namelist + - ./ICON/models/atm.namelist mpi_cmd: "mpirun -np {MPI_TOTAL_PROCS}" walltime: 00:05:00 nodes: 1 From a912102b2a9bdf27a8fb74a6d5f81bec01533d45 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Mon, 1 Sep 2025 15:10:32 +0200 Subject: [PATCH 4/4] Revert "wip modify test setup" This reverts commit 02b0a047b11b580ae864f813bfefcc3db64ead75. --- src/sirocco/core/_tasks/icon_task.py | 1 + tests/cases/small-icon/config/ICON/icon_master.namelist | 2 +- .../config/ICON/{models/atm.namelist => model.namelist} | 0 tests/cases/small-icon/config/config.yml | 2 +- 4 files changed, 3 insertions(+), 2 deletions(-) rename tests/cases/small-icon/config/ICON/{models/atm.namelist => model.namelist} (100%) diff --git a/src/sirocco/core/_tasks/icon_task.py b/src/sirocco/core/_tasks/icon_task.py index b9b1e74a..f2350ec1 100644 --- a/src/sirocco/core/_tasks/icon_task.py +++ b/src/sirocco/core/_tasks/icon_task.py @@ -52,6 +52,7 @@ def __post_init__(self): raise ValueError(msg) self._model_namelist = model_namelist + # TODO: Could use a different method here, e.g., if self.wrapper_script is not None: self.wrapper_script = self._resolve_local_script_path(rel_path=self.wrapper_script) diff --git a/tests/cases/small-icon/config/ICON/icon_master.namelist b/tests/cases/small-icon/config/ICON/icon_master.namelist index c31397d6..37910d0e 100644 --- a/tests/cases/small-icon/config/ICON/icon_master.namelist +++ b/tests/cases/small-icon/config/ICON/icon_master.namelist @@ -11,7 +11,7 @@ / &master_model_nml model_name="atm" - model_namelist_filename="atm.namelist" + model_namelist_filename="model.namelist" model_type=1 model_min_rank=0 model_max_rank=65535 diff --git a/tests/cases/small-icon/config/ICON/models/atm.namelist b/tests/cases/small-icon/config/ICON/model.namelist similarity index 100% rename from tests/cases/small-icon/config/ICON/models/atm.namelist rename to tests/cases/small-icon/config/ICON/model.namelist diff --git a/tests/cases/small-icon/config/config.yml b/tests/cases/small-icon/config/config.yml index 85dd2ce7..77937e50 100644 --- a/tests/cases/small-icon/config/config.yml +++ b/tests/cases/small-icon/config/config.yml @@ -47,7 +47,7 @@ tasks: bin: /TESTS_ROOTDIR/tests/cases/small-icon/config/ICON/bin/icon namelists: - ./ICON/icon_master.namelist - - ./ICON/models/atm.namelist + - ./ICON/model.namelist mpi_cmd: "mpirun -np {MPI_TOTAL_PROCS}" walltime: 00:05:00 nodes: 1