Add option to specify the computer for a task

src/sirocco/parsing/_yaml_data_models.py

@@ -244,6 +244,7 @@ def check_period_is_not_negative_or_zero(self) -> ConfigCycle:
 
 @dataclass
 class ConfigBaseTaskSpecs:
+    computer: str | None = None
     host: str | None = None
     account: str | None = None
     uenv: dict | None = None

src/sirocco/workgraph.py

@@ -6,6 +6,7 @@
 import aiida.common
 import aiida.orm
 import aiida_workgraph.engine.utils  # type: ignore[import-untyped]
+from aiida.common.exceptions import NotExistent
 from aiida_workgraph import WorkGraph
 
 from sirocco.core._tasks.icon_task import IconTask
@@ -179,14 +180,25 @@ def _create_task_node(self, task: graph_items.Task):
             command_full_path = task.command if command_path.is_absolute() else task.config_rootdir / command_path
             command = str(command_full_path)
 
-            # Source file
+            # metadata
+            metadata = {}
+            ## Source file
             env_source_paths = [
                 env_source_path
                 if (env_source_path := Path(env_source_file)).is_absolute()
                 else (task.config_rootdir / env_source_path)
                 for env_source_file in task.env_source_files
             ]
             prepend_text = "\n".join([f"source {env_source_path}" for env_source_path in env_source_paths])
+            metadata["options"] = {"prepend_text": prepend_text}
+
+            ## computer
+            if task.computer is not None:
+                try:
+                    metadata["computer"] = aiida.orm.load_computer(task.computer)
+                except NotExistent as err:
+                    msg = f"Could not find computer {task.computer} for task {task}."
+                    raise ValueError(msg) from err
 
             # NOTE: We don't pass the `nodes` dictionary here, as then we would need to have the sockets available when
             # we create the task. Instead, they are being updated via the WG internals when linking inputs/outputs to
@@ -197,7 +209,7 @@ def _create_task_node(self, task: graph_items.Task):
                 command=command,
                 arguments=[],
                 outputs=[],
-                metadata={"options": {"prepend_text": prepend_text}},
+                metadata=metadata,
             )
 
             self._aiida_task_nodes[label] = workgraph_task

tests/cases/small/config/test_config_small.yml

@@ -26,11 +26,13 @@ cycles:
                   date: 2026-05-01T00:00
 tasks:
   - icon:
+      computer: localhost
       plugin: shell
       command: scripts/icon.py
       cli_arguments: "{--restart icon_restart} {--init initial_conditions}"
       env_source_files: data/dummy_source_file.sh
   - cleanup:
+      computer: localhost
       plugin: shell
       command: scripts/cleanup.py
 data:

tests/cases/small/data/test_config_small.txt

@@ -10,6 +10,7 @@ cycles:
               - icon_restart [date: 2026-01-01 00:00:00]
             name: 'icon'
             coordinates: {'date': datetime.datetime(2026, 1, 1, 0, 0)}
+            computer: 'localhost'
             start date: 2026-01-01 00:00:00
             end date: 2026-06-01 00:00:00
             plugin: 'shell'
@@ -27,6 +28,7 @@ cycles:
               - icon_restart [date: 2026-03-01 00:00:00]
             name: 'icon'
             coordinates: {'date': datetime.datetime(2026, 3, 1, 0, 0)}
+            computer: 'localhost'
             start date: 2026-01-01 00:00:00
             end date: 2026-06-01 00:00:00
             plugin: 'shell'
@@ -44,6 +46,7 @@ cycles:
               - icon_restart [date: 2026-05-01 00:00:00]
             name: 'icon'
             coordinates: {'date': datetime.datetime(2026, 5, 1, 0, 0)}
+            computer: 'localhost'
             start date: 2026-01-01 00:00:00
             end date: 2026-06-01 00:00:00
             plugin: 'shell'
@@ -57,6 +60,7 @@ cycles:
               - icon [date: 2026-05-01 00:00:00]
             name: 'cleanup'
             coordinates: {}
+            computer: 'localhost'
             plugin: 'shell'
             command: 'scripts/cleanup.py'
             cli arguments: []

tests/test_wc_workflow.py

@@ -76,12 +76,15 @@ def test_vizgraph(config_paths):
         "tests/cases/parameters/config/test_config_parameters.yml",
     ],
 )
-def test_run_workgraph(config_path):
+def test_run_workgraph(config_path, aiida_computer):
     """Tests end-to-end the parsing from file up to running the workgraph.
 
     Automatically uses the aiida_profile fixture to create a new profile. Note to debug the test with your profile
     please run this in a separate file as the profile is deleted after test finishes.
     """
+    # some configs reference computer "localhost" which we need to create beforehand
+    aiida_computer("localhost").store()
+
     core_workflow = Workflow.from_yaml(config_path)
     aiida_workflow = AiidaWorkGraph(core_workflow)
     out = aiida_workflow.run()