diff --git a/src/sirocco/core/_tasks/icon_task.py b/src/sirocco/core/_tasks/icon_task.py index 21c96db0..f2350ec1 100644 --- a/src/sirocco/core/_tasks/icon_task.py +++ b/src/sirocco/core/_tasks/icon_task.py @@ -52,8 +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_wrapper_script(self.wrapper_script, 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._resolve_local_script_path(rel_path=self.setup_env) @property def master_namelist(self) -> NamelistFile: @@ -120,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_wrapper_script(self, wrapper_script: Path, config_rootdir: Path) -> Path: + def _resolve_local_script_path(self, rel_path: Path) -> Path: """Validate and resolve wrapper script path""" - resolved_path = wrapper_script if wrapper_script.is_absolute() else config_rootdir / wrapper_script + 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 aaa23b18..aebb5bea 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