Add setup_env input to IconTask

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."

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):

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)

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:

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