Add support of icon task in workgraph module (#143)

Adds a new test case `small-icon` and renamed test case `small` to
`small-shell`. The `small-icon` test case is a copy of `small` replacing the
usage of the dummy icon scripts with real icon. Because the `small` test case
increases coverage of  `ShellTask` it is kept.

- Add dispatching for `IconTask` in multiple functions in responsible for
  creating the WorkGraph.
- Pass `computer` argument to `core.Task` since it is required for the creation
  of a `IconCaclculation`. In the `small-shell` config the usage of `computer`
  has been removed temporarly until PR #136 fixes the usage.
- Changes in `core.AvailableData`:
  - Use `config_rootdir` to resolve location of data if relative.
  - The `src` member is now compulsory and validated. This change required to
    implement the `from_config` constructor for `AvailableData` and
    `GeneratedData` as the validation does not happen for `GeneratedData`.
- Fixing how `is_restart` is determined: It was not considering the data
  structure correctly data items. It was only checking the existance of the
  restart key in the input items, it was, however, not considering that the data
  structure still lists the input item with an empty list when the `when`
  keyword is used. Now it validates correctly to `False`.
- In tests we use now `pytest.fixture.usefixtures` when a fixture is not
  directly used but is required to be executed
- Exclude `test/cases/*` in type check as it contains dummy python scripts
- Update the default options of `hatch test` to run without icon
- Rename pytest fixture `icon_grid_simple_path` to `icon_grid_path`

Author: agoscinski

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
   "pydantic",
   "ruamel.yaml",
   "aiida-core>=2.5",
+  "aiida-icon>=0.4.0",
   "aiida-workgraph==0.5.2",
   "termcolor",
   "pygraphviz",
@@ -95,7 +96,7 @@ extra-dependencies = [
     "ipdb"
 ]
 default-args = []
-extra-args = ["--doctest-modules", '-m not slow', '-m not requires_icon']
+extra-args = ["--doctest-modules", "-m", "not slow and not requires_icon"]
 
 [[tool.hatch.envs.hatch-test.matrix]]
 python = ["3.12"]
@@ -135,7 +136,7 @@ extra-dependencies = [
 ]
 
 [tool.hatch.envs.types.scripts]
-check = "mypy --no-incremental {args:.}"
+check = "mypy --exclude 'tests/cases/*' --no-incremental {args:.}"
 
 [[tool.mypy.overrides]]
 module = ["isoduration", "isoduration.*"]

src/sirocco/core/_tasks/icon_task.py

@@ -22,7 +22,6 @@ class IconTask(models.ConfigIconTaskSpecs, Task):
 
     def __post_init__(self):
         super().__post_init__()
-
         # detect master namelist
         master_namelist = None
         for namelist in self.namelists:
@@ -33,6 +32,7 @@ def __post_init__(self):
             msg = f"Failed to read master namelists. Could not find {self._MASTER_NAMELIST_NAME!r} in namelists {self.namelists}"
             raise ValueError(msg)
         self._master_namelist = master_namelist
+        self.src = self._validate_src(self.src, self.config_rootdir)
 
         # retrieve model namelist name from master namelist
         if (master_model_nml := self._master_namelist.namelist.get(self._MASTER_MODEL_NML_SECTION, None)) is None:
@@ -64,7 +64,8 @@ def model_namelist(self) -> NamelistFile:
     @property
     def is_restart(self) -> bool:
         """Check if the icon task starts from the restart file."""
-        return self._AIIDA_ICON_RESTART_FILE_PORT_NAME in self.inputs
+        # restart port must be present and nonempty
+        return bool(self.inputs.get(self._AIIDA_ICON_RESTART_FILE_PORT_NAME, False))
 
     def update_icon_namelists_from_workflow(self):
         if not isinstance(self.cycle_point, DateCyclePoint):
@@ -115,3 +116,18 @@ def build_from_config(cls: type[Self], config: models.ConfigTask, **kwargs: Any)
         )
         self.update_icon_namelists_from_workflow()
         return self
+
+    @staticmethod
+    def _validate_src(config_src: Path, config_rootdir: Path | None = None) -> Path:
+        if config_rootdir is None and not config_src.is_absolute():
+            msg = f"Cannot specify relative path {config_src} for namelist while the rootdir is None"
+            raise ValueError(msg)
+
+        src = config_src if config_rootdir is None else (config_rootdir / config_src)
+        if not src.exists():
+            msg = f"Icon executable in path {src} does not exist."
+            raise FileNotFoundError(msg)
+        if not src.is_file():
+            msg = f"Icon executable in path {src} is not a file."
+            raise OSError(msg)
+        return src

src/sirocco/core/graph_items.py

@@ -45,23 +45,53 @@ class Data(ConfigBaseDataSpecs, GraphItem):
     color: ClassVar[str] = field(default="light_blue", repr=False)
 
     @classmethod
-    def from_config(cls, config: ConfigBaseData, coordinates: dict) -> AvailableData | GeneratedData:
+    def from_config(
+        cls, config: ConfigBaseData, config_rootdir: Path, coordinates: dict
+    ) -> AvailableData | GeneratedData:
         data_class = AvailableData if isinstance(config, ConfigAvailableData) else GeneratedData
-        return data_class(
+
+        return data_class.from_config(
+            config=config,
+            config_rootdir=config_rootdir,
+            coordinates=coordinates,
+        )
+
+
+class AvailableData(Data):
+    src: Path
+
+    @classmethod
+    def from_config(cls, config: ConfigBaseData, config_rootdir: Path, coordinates: dict) -> Self:
+        src = cls._validate_src(config.src, config_rootdir)
+        return cls(
             name=config.name,
             computer=config.computer,
             type=config.type,
-            src=config.src,
+            src=src,
             coordinates=coordinates,
         )
 
+    @staticmethod
+    def _validate_src(config_src: Path | None, config_rootdir: Path | None = None) -> Path | None:
+        if config_src is None:
+            return None
+        if config_rootdir is None and not config_src.is_absolute():
+            msg = f"Cannot specify relative path {config_src} for namelist while the rootdir is None"
+            raise ValueError(msg)
 
-class AvailableData(Data):
-    pass
+        return config_src if config_rootdir is None else (config_rootdir / config_src)
 
 
 class GeneratedData(Data):
-    pass
+    @classmethod
+    def from_config(cls, config: ConfigBaseData, config_rootdir: Path, coordinates: dict) -> Self:  # noqa: ARG003 # we need to keep same signature as for AvailableData
+        return cls(
+            name=config.name,
+            computer=config.computer,
+            type=config.type,
+            src=config.src,
+            coordinates=coordinates,
+        )
 
 
 @dataclass(kw_only=True)
@@ -92,9 +122,15 @@ def __init_subclass__(cls, **kwargs):
     def input_data_nodes(self) -> Iterator[Data]:
         yield from chain(*self.inputs.values())
 
+    def input_data_items(self) -> Iterator[tuple[str, Data]]:
+        yield from ((key, value) for key, values in self.inputs.items() for value in values)
+
     def output_data_nodes(self) -> Iterator[Data]:
         yield from chain(*self.outputs.values())
 
+    def output_data_items(self) -> Iterator[tuple[str | None, Data]]:
+        yield from ((key, value) for key, values in self.outputs.items() for value in values)
+
     @classmethod
     def from_config(
         cls: type[Self],

src/sirocco/core/workflow.py

@@ -56,7 +56,11 @@ def iter_coordinates(cycle_point: CyclePoint, param_refs: list[str]) -> Iterator
         # 1 - create availalbe data nodes
         for available_data_config in config_data.available:
             for coordinates in iter_coordinates(OneOffPoint(), available_data_config.parameters):
-                self.data.add(Data.from_config(config=available_data_config, coordinates=coordinates))
+                self.data.add(
+                    Data.from_config(
+                        config=available_data_config, config_rootdir=config_rootdir, coordinates=coordinates
+                    )
+                )
 
         # 2 - create output data nodes
         for cycle_config in config_cycles:
@@ -65,7 +69,11 @@ def iter_coordinates(cycle_point: CyclePoint, param_refs: list[str]) -> Iterator
                     for data_ref in task_ref.outputs:
                         data_config = config_data_dict[data_ref.name]
                         for coordinates in iter_coordinates(cycle_point, data_config.parameters):
-                            self.data.add(Data.from_config(config=data_config, coordinates=coordinates))
+                            self.data.add(
+                                Data.from_config(
+                                    config=data_config, config_rootdir=config_rootdir, coordinates=coordinates
+                                )
+                            )
 
         # 3 - create cycles and tasks
         for cycle_config in config_cycles:

src/sirocco/parsing/yaml_data_models.py

@@ -441,9 +441,13 @@ def merge_path_key(cls, data: Any) -> dict[str, Any]:
 @dataclass(kw_only=True)
 class ConfigIconTaskSpecs:
     plugin: ClassVar[Literal["icon"]] = "icon"
+    # PRCOMMENT this is later resolved to an absolute path so we cannot use it for serialization
+    #           as the tests would fail. Since the main purpose of the repr is for regression tests
+    #           I disable its prnting here
+    src: Path = field(repr=False)
 
 
-class ConfigIconTask(ConfigBaseTask):
+class ConfigIconTask(ConfigBaseTask, ConfigIconTaskSpecs):
     """Class representing an ICON task configuration from a workflow file
 
     Examples:
@@ -460,6 +464,7 @@ class ConfigIconTask(ConfigBaseTask):
         ...           - path/to/case_nml:
         ...               block_1:
         ...                 param_name: param_value
+        ...         src: path/to/icon
         ...     '''
         ... )
         >>> icon_task_cfg = validate_yaml_content(ConfigIconTask, snippet)
@@ -487,7 +492,7 @@ class DataType(enum.StrEnum):
 @dataclass(kw_only=True)
 class ConfigBaseDataSpecs:
     type: DataType
-    src: Path
+    src: Path | None = None
     format: str | None = None
     computer: str | None = None
 
@@ -522,7 +527,7 @@ class ConfigBaseData(_NamedBaseModel, ConfigBaseDataSpecs):
 
 
 class ConfigAvailableData(ConfigBaseData):
-    pass
+    src: Path
 
 
 class ConfigGeneratedData(ConfigBaseData):

src/sirocco/workgraph.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import functools
+import io
+import uuid
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, TypeAlias
 
@@ -9,6 +11,7 @@
 import aiida_workgraph  # type: ignore[import-untyped] # does not have proper typing and stubs
 import aiida_workgraph.tasks.factory.shelljob_task  # type: ignore[import-untyped]  # is only for a workaround
 from aiida.common.exceptions import NotExistent
+from aiida_icon.calculations import IconCalculation
 
 from sirocco import core
 
@@ -93,13 +96,13 @@ def __init__(self, core_workflow: core.Workflow):
         for task in self._core_workflow.tasks:
             self.create_task_node(task)
             # Create and link corresponding output sockets
-            for output in task.output_data_nodes():
-                self._link_output_node_to_task(task, output)
+            for port, output in task.output_data_items():
+                self._link_output_node_to_task(task, port, output)
 
         # link input nodes to workgraph tasks
         for task in self._core_workflow.tasks:
-            for input_ in task.input_data_nodes():
-                self._link_input_node_to_task(task, input_)
+            for port, input_ in task.input_data_items():
+                self._link_input_node_to_task(task, port, input_)
 
         # set shelljob arguments
         for task in self._core_workflow.tasks:
@@ -180,7 +183,7 @@ def _add_available_data(self):
             if isinstance(data, core.AvailableData):
                 self._add_aiida_input_data_node(data)
 
-    def _add_aiida_input_data_node(self, data: core.Data):
+    def _add_aiida_input_data_node(self, data: core.AvailableData):
         """
         Create an `aiida.orm.Data` instance from the provided graph item.
         """
@@ -267,23 +270,76 @@ def _create_shell_task_node(self, task: core.ShellTask):
 
         self._aiida_task_nodes[label] = workgraph_task
 
-    def _link_output_node_to_task(self, task: core.Task, output: core.Data):
+    @create_task_node.register
+    def _create_icon_task_node(self, task: core.IconTask):
+        task_label = self.get_aiida_label_from_graph_item(task)
+
+        try:
+            # PRCOMMENT move to parsing? But then it has aiida logic
+            computer = aiida.orm.Computer.collection.get(label=task.computer)
+        except NotExistent as err:
+            msg = f"Could not find computer {task.computer!r} in AiiDA database. One needs to create and configure the computer before running a workflow."
+            raise ValueError(msg) from err
+
+        label_uuid = str(uuid.uuid4())
+        icon_code = aiida.orm.InstalledCode(
+            label=f"icon-{label_uuid}",
+            description="aiida_icon",
+            default_calc_job_plugin="icon.icon",
+            computer=computer,
+            filepath_executable=str(task.src),
+            with_mpi=False,
+        ).store()
+
+        builder = IconCalculation.get_builder()
+        builder.code = icon_code
+
+        task.update_icon_namelists_from_workflow()
+
+        with io.StringIO() as buffer:
+            task.master_namelist.namelist.write(buffer)
+            buffer.seek(0)
+            builder.master_namelist = aiida.orm.SinglefileData(buffer, task.master_namelist.name)
+
+        with io.StringIO() as buffer:
+            task.model_namelist.namelist.write(buffer)
+            buffer.seek(0)
+            builder.model_namelist = aiida.orm.SinglefileData(buffer, task.model_namelist.name)
+
+        self._aiida_task_nodes[task_label] = self._workgraph.add_task(builder)
+
+    @functools.singledispatchmethod
+    def _link_output_node_to_task(self, task: core.Task, port: str, output: core.Data):  # noqa: ARG002
+        """Dispatch linking input to task based on task type."""
+
+        msg = f"method not implemented for task type {type(task)}"
+        raise NotImplementedError(msg)
+
+    @_link_output_node_to_task.register
+    def _link_output_node_to_shell_task(self, task: core.ShellTask, _: str, output: core.Data):
         """Links the output to the workgraph task."""
 
         workgraph_task = self.task_from_core(task)
         output_label = self.get_aiida_label_from_graph_item(output)
         output_socket = workgraph_task.add_output("workgraph.any", str(output.src))
         self._aiida_socket_nodes[output_label] = output_socket
 
+    @_link_output_node_to_task.register
+    def _link_output_node_to_icon_task(self, task: core.IconTask, port: str, output: core.Data):
+        workgraph_task = self.task_from_core(task)
+        output_label = self.get_aiida_label_from_graph_item(output)
+        output_socket = workgraph_task.outputs._sockets.get(port)  # noqa SLF001 # there so public accessor
+        self._aiida_socket_nodes[output_label] = output_socket
+
     @functools.singledispatchmethod
-    def _link_input_node_to_task(self, task: core.Task, input_: core.Data):  # noqa: ARG002
+    def _link_input_node_to_task(self, task: core.Task, port: str, input_: core.Data):  # noqa: ARG002
         """ "Dispatch linking input to task based on task type"""
 
         msg = f"method not implemented for task type {type(task)}"
         raise NotImplementedError(msg)
 
     @_link_input_node_to_task.register
-    def _link_input_node_to_shelltask(self, task: core.ShellTask, input_: core.Data):
+    def _link_input_node_to_shell_task(self, task: core.ShellTask, _: str, input_: core.Data):
         """Links the input to the workgraph shell task."""
 
         workgraph_task = self.task_from_core(task)
@@ -305,6 +361,20 @@ def _link_input_node_to_shelltask(self, task: core.ShellTask, input_: core.Data)
         else:
             raise TypeError
 
+    @_link_input_node_to_task.register
+    def _link_input_node_to_icon_task(self, task: core.IconTask, port: str, input_: core.Data):
+        """Links the input to the workgraph shell task."""
+
+        workgraph_task = self.task_from_core(task)
+
+        # resolve data
+        if isinstance(input_, core.AvailableData):
+            setattr(workgraph_task.inputs, f"{port}", self.data_from_core(input_))
+        elif isinstance(input_, core.GeneratedData):
+            setattr(workgraph_task.inputs, f"{port}", self.socket_from_core(input_))
+        else:
+            raise TypeError
+
     def _link_wait_on_to_task(self, task: core.Task):
         """link wait on tasks to workgraph task"""
 
@@ -349,7 +419,7 @@ def _set_shelljob_filenames(self, task: core.ShellTask):
 
             if task.computer and input_.computer and isinstance(input_, core.AvailableData):
                 # For RemoteData on the same computer, use just the filename
-                filename = Path(input_.src).name
+                filename = input_.src.name
                 filenames[input_.name] = filename
             else:
                 # For other cases (including GeneratedData), we need to handle parameterized data
@@ -368,7 +438,7 @@ def _set_shelljob_filenames(self, task: core.ShellTask):
                     filename = input_label
                 else:
                     # Single data node with this name - can use simple filename
-                    filename = Path(input_.src).name if hasattr(input_, "src") else input_.name
+                    filename = input_.src.name if input_.src is not None else input_.name
 
                 # The key in filenames dict should be the input label (what's used in nodes dict)
                 filenames[input_label] = filename

tests/cases/small/config/data/initial_conditions renamed to tests/cases/large/config/ICON/bin/icon

@@ -114,6 +114,7 @@ tasks:
         mount_point: runtime/mount/point
   - icon:
       plugin: icon
+      src: ./ICON/bin/icon
       nodes: 40
       walltime: 23:59:59
       namelists:

tests/cases/large/config/config.yml

@@ -1,6 +1,6 @@
 &master_nml
-    lrestart = .true.
-    read_restart_namelists = .true.
+    lrestart = .false.
+    read_restart_namelists = .false.
 /
 
 &master_time_control_nml