From a385008286c66ef14a8f2749c5784b5efda3b28f Mon Sep 17 00:00:00 2001 From: Alexander Goscinski Date: Fri, 3 Jan 2025 09:26:51 +0100 Subject: [PATCH] Add option to specify the computer for a task The user can specify a computer for a task by the aiida label. A computer be initialized through the verdi CLI or programatically. This is for now the only feasible solution to support the whole scope of connections ssh and firecrest without implementing our own CLI resolution logic. --- src/sirocco/parsing/_yaml_data_models.py | 1 + src/sirocco/workgraph.py | 16 ++++++++++++++-- tests/cases/small/config/test_config_small.yml | 2 ++ tests/cases/small/data/test_config_small.txt | 4 ++++ tests/test_wc_workflow.py | 5 ++++- 5 files changed, 25 insertions(+), 3 deletions(-) diff --git a/src/sirocco/parsing/_yaml_data_models.py b/src/sirocco/parsing/_yaml_data_models.py index e48bf1bf..d5920d02 100644 --- a/src/sirocco/parsing/_yaml_data_models.py +++ b/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 diff --git a/src/sirocco/workgraph.py b/src/sirocco/workgraph.py index bd374b5e..5eb9127d 100644 --- a/src/sirocco/workgraph.py +++ b/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,7 +180,9 @@ 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() @@ -187,6 +190,15 @@ def _create_task_node(self, task: graph_items.Task): 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 diff --git a/tests/cases/small/config/test_config_small.yml b/tests/cases/small/config/test_config_small.yml index 40394a42..3dd6fbeb 100644 --- a/tests/cases/small/config/test_config_small.yml +++ b/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: diff --git a/tests/cases/small/data/test_config_small.txt b/tests/cases/small/data/test_config_small.txt index d774c79f..3e799e13 100644 --- a/tests/cases/small/data/test_config_small.txt +++ b/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: [] diff --git a/tests/test_wc_workflow.py b/tests/test_wc_workflow.py index badf577d..0b3c3b9d 100644 --- a/tests/test_wc_workflow.py +++ b/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()