From 1063865caf6f0a4bcbb6c1a4f705b3cf37de8e5e Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Mon, 20 Jul 2026 21:33:09 +0200 Subject: [PATCH 01/19] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20`WorkChain`:=20make?= =?UTF-8?q?=20the=20stepper=20pluggable?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `WorkChain` hard-binds its execution strategy to the outline declared on the spec: `run` calls `spec().get_outline().create_stepper(self)`, and `load_instance_state` calls the matching `recreate_stepper`. Since `run`, `on_run`, `to_context`, `on_exiting` and `on_wait` are all `@Protect.final`, a subclass cannot substitute a different strategy, and the only way to get one is to bypass `WorkChain` entirely and reimplement the parts that have nothing to do with stepping: awaitables, context, checkpointing and node lifecycle. Nothing needs inventing to fix this, because plumpy already defines the strategy interface. `plumpy.workchains.Stepper` is `step() -> (finished, result)` plus its own `save_instance_state`/`load_instance_state`. Add two overridable hooks, `_create_stepper` and `_recreate_stepper`, both defaulting to exactly the previous outline behaviour, and route the two call sites through them. They are deliberately not `@Protect.final`: they are the extension point. The restore hook matters as much as the create one, since without it a process using a custom stepper could not be reconstructed from a checkpoint. Default behaviour is unchanged: with neither hook overridden, a work chain still steps through its outline exactly as before. This is what lets a dependency-graph scheduler exist as a strategy over `WorkChain` rather than as a fork of it (issue #6754). Tests cover both hooks: one work chain whose outline raises if it is ever stepped, proving the custom stepper drove execution; and a bundle and unbundle round trip, asserting the reloaded process resumes at the saved position rather than repeating completed steps. --- .../engine/processes/workchains/workchain.py | 26 ++++- tests/engine/test_work_chain.py | 108 ++++++++++++++++++ 2 files changed, 132 insertions(+), 2 deletions(-) diff --git a/src/aiida/engine/processes/workchains/workchain.py b/src/aiida/engine/processes/workchains/workchain.py index 500e69cff4..e66aaedd87 100644 --- a/src/aiida/engine/processes/workchains/workchain.py +++ b/src/aiida/engine/processes/workchains/workchain.py @@ -174,13 +174,35 @@ def load_instance_state(self, saved_state, load_context): self._stepper = None stepper_state = saved_state.get(self._STEPPER_STATE, None) if stepper_state is not None: - self._stepper = self.spec().get_outline().recreate_stepper(stepper_state, self) # type: ignore[arg-type] + self._stepper = self._recreate_stepper(stepper_state) self.set_logger(self.node.logger) if self._awaitables: self._action_awaitables() + def _create_stepper(self) -> Stepper: + """Return the stepper that drives this work chain. + + This is the seam for supplying a different execution strategy. The default steps through the outline declared + on the spec, but a subclass may return any :class:`plumpy.workchains.Stepper`, for example one that derives the + order of execution from a graph of data dependencies instead of a static outline. + + A subclass that overrides this should also override :meth:`_recreate_stepper`, otherwise its processes cannot + be restored from a checkpoint. + """ + return self.spec().get_outline().create_stepper(self) # type: ignore[arg-type] + + def _recreate_stepper(self, saved_state: t.Any) -> Stepper: + """Restore the stepper from the state it wrote to the checkpoint. + + The counterpart of :meth:`_create_stepper`, called when a process is loaded from a checkpoint rather than + started fresh. + + :param saved_state: the state previously returned by ``Stepper.save()`` + """ + return self.spec().get_outline().recreate_stepper(saved_state, self) # type: ignore[arg-type] + @Protect.final def on_run(self): super().on_run() @@ -299,7 +321,7 @@ def _update_process_status(self) -> None: @override @Protect.final async def run(self) -> t.Any: - self._stepper = self.spec().get_outline().create_stepper(self) # type: ignore[arg-type] + self._stepper = self._create_stepper() return await run_with_portal(self._do_step) def _do_step(self) -> t.Any: diff --git a/tests/engine/test_work_chain.py b/tests/engine/test_work_chain.py index 43ac897270..9c299ca37a 100644 --- a/tests/engine/test_work_chain.py +++ b/tests/engine/test_work_chain.py @@ -1769,3 +1769,111 @@ def define(cls, spec): async def run(self): pass + + +class SequenceStepper(plumpy.workchains.Stepper): + """A minimal alternative execution strategy, running a work chain's `STEPS` in order. + + Stands in for a stepper that derives its order from something other than the outline, such as a graph of data + dependencies. + """ + + POSITION = 'position' + + def __init__(self, workchain, position=0): + super().__init__(workchain) + self._position = position + + def save_instance_state(self, out_state, save_context): + super().save_instance_state(out_state, save_context) + out_state[self.POSITION] = self._position + + def load_instance_state(self, saved_state, load_context): + super().load_instance_state(saved_state, load_context) + self._position = saved_state[self.POSITION] + + def step(self): + steps = self._workchain.STEPS + getattr(self._workchain, steps[self._position])() + self._position += 1 + return self._position >= len(steps), None + + def __str__(self): + return f'{self._position}/{len(self._workchain.STEPS)}' + + +class CustomStepperWorkChain(WorkChain): + """A work chain whose execution is driven by its own stepper instead of the outline.""" + + STEPS = ('step_a', 'step_b', 'step_c') + + @classmethod + def define(cls, spec): + super().define(spec) + # Deliberately an outline that fails if it is ever stepped through. + spec.outline(cls.outline_must_not_run) + spec.output('result', valid_type=Int) + + def _create_stepper(self): + return SequenceStepper(self) + + def _recreate_stepper(self, saved_state): + return SequenceStepper(self, position=saved_state[SequenceStepper.POSITION]) + + def outline_must_not_run(self): + raise AssertionError('the outline drove execution instead of the custom stepper') + + def step_a(self): + self.ctx.trail = 'a' + + def step_b(self): + self.ctx.trail += 'b' + + def step_c(self): + self.ctx.trail += 'c' + self.out('result', Int(len(self.ctx.trail)).store()) + + +class PausingStepperWorkChain(CustomStepperWorkChain): + """As above, but pauses after the first step so the checkpoint path can be exercised.""" + + def step_a(self): + super().step_a() + self.pause() + + +class TestCustomStepper: + """Test that a `WorkChain` subclass can supply its own stepping strategy.""" + + def test_custom_stepper_drives_execution(self): + """The stepper returned by `_create_stepper` runs, and the outline does not.""" + result, node = launch.run_get_node(CustomStepperWorkChain) + assert node.is_finished_ok, node.exit_status + assert result['result'] == 3 + + def test_custom_stepper_survives_checkpoint(self): + """`_recreate_stepper` restores position, so a resumed process does not redo completed steps.""" + runner = get_manager().get_runner() + workchain = PausingStepperWorkChain() + runner.schedule(workchain) + + async def run_async(wc): + await run_until_paused(wc) + assert wc.ctx.trail == 'a' + + bundle = plumpy.Bundle(wc) + wc.close() + + reloaded = bundle.unbundle() + assert reloaded.ctx.trail == 'a' + + runner.schedule(reloaded) + reloaded.play() + await reloaded.future() + + # 'aabc' would mean the stepper restarted rather than resumed + assert reloaded.ctx.trail == 'abc' + + wc.future().set_result(None) + + runner.loop.run_until_complete(run_async(workchain)) From cafb11cacf14239709ae09632eb9681c63aba0ec Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Tue, 21 Jul 2026 13:42:13 +0200 Subject: [PATCH 02/19] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20`WorkChain`:=20make?= =?UTF-8?q?=20the=20awaitable=20barrier=20a=20policy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `WorkChain` hardcodes one execution model: `_do_step` clears `self._awaitables` at the start of every step, and a finished child resumes the process only once every awaitable is done. That is the outline model, where a step waits for everything it launched before the next begins. A stepper that schedules by data dependencies wants the opposite (keep the awaitables, resume as each child finishes, so independent branches stay in flight), and today the only way to get it is to override `_do_step`, `_on_awaitable_finished` and `_action_awaitables` wholesale, i.e. to fork the awaitable machinery. Make the barrier a property of the stepping strategy instead. A stepper declares `awaitable_barrier = False` to opt into the streaming model; `WorkChain._awaitable_barrier` reads it and defaults to `True`. Three call sites consult it: `_do_step` clears the awaitables only under the barrier, and `_on_awaitable_finished` resumes on any awaitable under streaming rather than only when none remain. `_action_awaitables` now skips an awaitable whose callback is already registered, which a streaming stepper needs because the same awaitable is seen on every pass through the waiting state; under the barrier the awaitables are cleared each step so it never triggers. A new `_on_awaitable_resolved` hook (default no-op) lets a subclass run per-child bookkeeping before the resume decision without reimplementing the callback. Default behaviour is unchanged: with no stepper declaring the flag, every outline work chain clears, waits and resumes exactly as before. This is what lets a dependency-graph stepper stream by setting one flag rather than forking `WorkChain` (issue #6754), and it is a general capability: any fan-out or dependency-aware stepper can use it. Tests cover the clearing policy (barrier clears, streaming keeps), the default, and the registration guard. --- .../engine/processes/workchains/workchain.py | 54 ++++++++++-- tests/engine/test_work_chain.py | 82 +++++++++++++++++++ 2 files changed, 130 insertions(+), 6 deletions(-) diff --git a/src/aiida/engine/processes/workchains/workchain.py b/src/aiida/engine/processes/workchains/workchain.py index e66aaedd87..6acaba102d 100644 --- a/src/aiida/engine/processes/workchains/workchain.py +++ b/src/aiida/engine/processes/workchains/workchain.py @@ -131,6 +131,9 @@ def __init__( self._stepper: Stepper | None = None self._awaitables: list[Awaitable] = [] + # The pks of awaitables whose completion callback is already registered. This is runtime state, callbacks + # do not survive a checkpoint, so it is not persisted and is reset in `load_instance_state`. + self._registered_awaitable_pks: set[int] = set() self._context = AttributeDict() @classmethod @@ -178,6 +181,8 @@ def load_instance_state(self, saved_state, load_context): self.set_logger(self.node.logger) + # Callbacks do not survive the checkpoint, so nothing is registered yet on the reloaded process. + self._registered_awaitable_pks = set() if self._awaitables: self._action_awaitables() @@ -203,6 +208,20 @@ def _recreate_stepper(self, saved_state: t.Any) -> Stepper: """ return self.spec().get_outline().recreate_stepper(saved_state, self) # type: ignore[arg-type] + @property + def _awaitable_barrier(self) -> bool: + """Whether each step waits for everything it launched before the next one begins. + + This is the difference between the two execution models, and it is a property of the stepping strategy, so + the value is taken from the stepper. ``True``, the default, is the outline model: :meth:`_do_step` clears + the awaitables at the start of every step, so a step forms a barrier over the children it launched and the + process only resumes once all of them have finished. A stepper that schedules by data dependencies wants + ``False``: the awaitables persist across steps and the process resumes as each child finishes, so + independent branches stay in flight together. A stepper opts into the streaming model by defining + ``awaitable_barrier = False`` on itself. + """ + return getattr(self._stepper, 'awaitable_barrier', True) + @Protect.final def on_run(self): super().on_run() @@ -334,7 +353,11 @@ def _do_step(self) -> t.Any: """ from .context import ToContext - self._awaitables = [] + # Under the barrier model the awaitables belong to a single step and are cleared before the next one, which + # is what forces every step to wait for all the children it launched. A streaming stepper keeps them, so + # children launched in earlier steps stay in flight while later steps run. + if self._awaitable_barrier: + self._awaitables = [] result: t.Any = None try: @@ -402,24 +425,41 @@ def on_wait(self, awaitables: t.Sequence[t.Awaitable]): self.call_soon(self.resume) def _action_awaitables(self) -> None: - """Handle the awaitables that are currently registered with the work chain. + """Register the completion callback for each awaitable that does not already have one. Depending on the class type of the awaitable's target a different callback function will be bound with the awaitable and the runner will be asked to - call it when the target is completed + call it when the target is completed. + + The registration is guarded against duplicates: under the barrier model the awaitables are cleared each + step so the same one is never seen twice, but a streaming stepper keeps its awaitables across steps and + would otherwise register a further callback for the same awaitable on every pass through the waiting state. """ for awaitable in self._awaitables: + if awaitable.pk in self._registered_awaitable_pks: + continue if awaitable.target == AwaitableTarget.PROCESS: callback = functools.partial(self.call_soon, self._on_awaitable_finished, awaitable) self.runner.call_on_process_finish(awaitable.pk, callback) + self._registered_awaitable_pks.add(awaitable.pk) else: raise AssertionError(f"invalid awaitable target '{awaitable.target}'") + def _on_awaitable_resolved(self, awaitable: Awaitable) -> None: + """Hook called once a finished awaitable has been resolved onto the context, before the resume decision. + + Defaults to doing nothing. A subclass can use it to run bookkeeping that must see the resolved value and + must happen before the process is resumed, without having to reimplement :meth:`_on_awaitable_finished`. + + :param awaitable: the awaitable that has just been resolved + """ + def _on_awaitable_finished(self, awaitable: Awaitable) -> None: """Callback function, for when an awaitable process instance is completed. - The awaitable will be effectuated on the context of the work chain and removed from the internal list. If all - awaitables have been dealt with, the work chain process is resumed. + The awaitable will be effectuated on the context of the work chain and removed from the internal list. The + process is then resumed: under the barrier model only once every awaitable has finished, and under the + streaming model as soon as this one does, so a finished child can unblock its dependents while others run. :param awaitable: an Awaitable instance """ @@ -436,6 +476,8 @@ def _on_awaitable_finished(self, awaitable: Awaitable) -> None: value = node # type: ignore[assignment] self._resolve_awaitable(awaitable, value) + self._registered_awaitable_pks.discard(awaitable.pk) + self._on_awaitable_resolved(awaitable) - if self.state == ProcessState.WAITING and not self._awaitables: + if self.state == ProcessState.WAITING and (not self._awaitable_barrier or not self._awaitables): self.resume() diff --git a/tests/engine/test_work_chain.py b/tests/engine/test_work_chain.py index 9c299ca37a..75df26a27e 100644 --- a/tests/engine/test_work_chain.py +++ b/tests/engine/test_work_chain.py @@ -1877,3 +1877,85 @@ async def run_async(wc): wc.future().set_result(None) runner.loop.run_until_complete(run_async(workchain)) + + +class OneShotStepper(plumpy.workchains.Stepper): + """A stepper that finishes in a single step, used to drive `_do_step` once in isolation.""" + + def step(self): + return True, None + + +class StreamingStepper(OneShotStepper): + """A stepper in the streaming (data-dependency) model, where awaitables persist across steps.""" + + awaitable_barrier = False + + +class TestAwaitableBarrier: + """The awaitable barrier is a property of the stepping strategy. + + Under the barrier model (the outline default) each step waits for every child it launched before the next + begins; a streaming stepper keeps its awaitables so independent children stay in flight. The switch lives on + the stepper (``awaitable_barrier``), and the work chain reads it. + """ + + @staticmethod + def _work_chain(): + class _WorkChain(WorkChain): + @classmethod + def define(cls, spec): + super().define(spec) + spec.outline(cls._noop) + + def _noop(self): + pass + + return _WorkChain() + + @pytest.mark.parametrize( + 'stepper_cls, expected', + [(OneShotStepper, []), (StreamingStepper, ['keep'])], + ids=['barrier', 'streaming'], + ) + def test_do_step_clears_awaitables_only_under_barrier(self, stepper_cls, expected): + """`_do_step` clears the awaitables at the start of a step under the barrier model, and keeps them under + the streaming model. The clearing is what forces an outline step to wait for all its children.""" + work_chain = self._work_chain() + work_chain._stepper = stepper_cls(work_chain) + work_chain._awaitables = ['keep'] + + work_chain._do_step() + + assert work_chain._awaitables == expected + + def test_barrier_is_the_default(self): + """A stepper that does not declare the flag (every outline stepper) gets the barrier.""" + work_chain = self._work_chain() + work_chain._stepper = OneShotStepper(work_chain) + assert work_chain._awaitable_barrier is True + + def test_action_awaitables_registers_each_awaitable_once(self, monkeypatch): + """A streaming stepper sees the same awaitable on every pass through the waiting state, so the callback + must be registered once, not once per pass, and again only after the awaitable is resolved.""" + from types import SimpleNamespace + + from aiida.engine.processes.workchains.awaitable import AwaitableTarget + + work_chain = self._work_chain() + work_chain._stepper = StreamingStepper(work_chain) + + registered: list[int] = [] + monkeypatch.setattr(work_chain.runner, 'call_on_process_finish', lambda pk, callback: registered.append(pk)) + + awaitable = SimpleNamespace(pk=123, target=AwaitableTarget.PROCESS) + work_chain._awaitables = [awaitable] + + work_chain._action_awaitables() + work_chain._action_awaitables() + assert registered == [123], 'the persisting awaitable was registered more than once' + + # Once resolved, its pk is forgotten and a fresh awaitable reusing it would register again. + work_chain._registered_awaitable_pks.discard(123) + work_chain._action_awaitables() + assert registered == [123, 123] From 1b930796992d7df2cc312807397875a4c622e556 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Tue, 21 Jul 2026 17:23:20 +0200 Subject: [PATCH 03/19] =?UTF-8?q?=E2=9C=A8=20`orm`:=20add=20`WorkGraphNode?= =?UTF-8?q?`=20for=20work=20graphs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `WorkGraphNode` is the process node for a running WorkGraph, storing the per-task runtime state (state, process, action, execution count, map info) on top of what `WorkChainNode` records. It lived in the aiida-workgraph package; this moves it into core as the first relocation in bringing the WorkGraph runtime into aiida-core. It is a clean subclass of `WorkChainNode` with no node-graph or plugin dependency, so unlike the rest of that runtime it can live in core unconditionally: a database written by aiida-workgraph stays loadable on plain aiida-core even without the eventual workgraph extra installed. The `aiida.node` entry point keeps its name so the stored `node_type` (`process.workflow.workgraph.WorkGraphNode.`) is unchanged and existing nodes load against the moved class. The accompanying aiida-workgraph change drops its own copy and registration and imports the node from `aiida.orm`. Tests cover the task accessors, the accessor/bulk-attribute consistency, and pin the `node_type` string; the field-coverage regression gains its generated entry for the new node type. --- pyproject.toml | 1 + src/aiida/orm/__init__.py | 1 + src/aiida/orm/nodes/__init__.py | 1 + src/aiida/orm/nodes/process/__init__.py | 1 + .../orm/nodes/process/workflow/__init__.py | 2 + .../orm/nodes/process/workflow/workgraph.py | 143 ++++++++++++++++++ tests/orm/nodes/process/test_workgraph.py | 45 ++++++ ...ocess.workflow.workgraph.WorkGraphNode.yml | 36 +++++ 8 files changed, 230 insertions(+) create mode 100644 src/aiida/orm/nodes/process/workflow/workgraph.py create mode 100644 tests/orm/nodes/process/test_workgraph.py create mode 100644 tests/orm/test_fields/fields_aiida.node.process.workflow.workgraph.WorkGraphNode.yml diff --git a/pyproject.toml b/pyproject.toml index 3ace4824d2..14885e2889 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -152,6 +152,7 @@ requires-python = '>=3.10' 'process.workflow' = 'aiida.orm.nodes.process.workflow.workflow:WorkflowNode' 'process.workflow.workchain' = 'aiida.orm.nodes.process.workflow.workchain:WorkChainNode' 'process.workflow.workfunction' = 'aiida.orm.nodes.process.workflow.workfunction:WorkFunctionNode' +'process.workflow.workgraph' = 'aiida.orm.nodes.process.workflow.workgraph:WorkGraphNode' [project.entry-points.'aiida.orm'] 'core.auth_info' = 'aiida.orm.authinfos:AuthInfo' diff --git a/src/aiida/orm/__init__.py b/src/aiida/orm/__init__.py index 8ef057ef12..b65cb0b611 100644 --- a/src/aiida/orm/__init__.py +++ b/src/aiida/orm/__init__.py @@ -104,6 +104,7 @@ 'User', 'WorkChainNode', 'WorkFunctionNode', + 'WorkGraphNode', 'WorkflowNode', 'XyData', 'cif_from_ase', diff --git a/src/aiida/orm/nodes/__init__.py b/src/aiida/orm/nodes/__init__.py index e8ef1da246..c4696f0fec 100644 --- a/src/aiida/orm/nodes/__init__.py +++ b/src/aiida/orm/nodes/__init__.py @@ -62,6 +62,7 @@ 'UpfData', 'WorkChainNode', 'WorkFunctionNode', + 'WorkGraphNode', 'WorkflowNode', 'XyData', 'cif_from_ase', diff --git a/src/aiida/orm/nodes/process/__init__.py b/src/aiida/orm/nodes/process/__init__.py index e6c672a526..a7844c7ea1 100644 --- a/src/aiida/orm/nodes/process/__init__.py +++ b/src/aiida/orm/nodes/process/__init__.py @@ -23,6 +23,7 @@ 'ProcessNode', 'WorkChainNode', 'WorkFunctionNode', + 'WorkGraphNode', 'WorkflowNode', ) diff --git a/src/aiida/orm/nodes/process/workflow/__init__.py b/src/aiida/orm/nodes/process/workflow/__init__.py index 72a2dcadfe..373afebba3 100644 --- a/src/aiida/orm/nodes/process/workflow/__init__.py +++ b/src/aiida/orm/nodes/process/workflow/__init__.py @@ -15,10 +15,12 @@ from .workchain import * from .workflow import * from .workfunction import * +from .workgraph import * __all__ = ( 'WorkChainNode', 'WorkFunctionNode', + 'WorkGraphNode', 'WorkflowNode', ) diff --git a/src/aiida/orm/nodes/process/workflow/workgraph.py b/src/aiida/orm/nodes/process/workflow/workgraph.py new file mode 100644 index 0000000000..5a85e1b528 --- /dev/null +++ b/src/aiida/orm/nodes/process/workflow/workgraph.py @@ -0,0 +1,143 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Module with `Node` sub class for work graph processes.""" + +import logging + +from aiida.common.lang import classproperty +from aiida.orm.nodes.process.workflow.workchain import WorkChainNode + +__all__ = ('WorkGraphNode',) + + +def make_dict_property(attribute_key: str, default=None): + """ + Return a property object that gets/sets a dict attribute from `self.base.attributes`. + + :param attribute_key: the key in `self.base.attributes` for this dict + :param default: default value to return if nothing is set + """ + + def getter(self): + return self.base.attributes.get(attribute_key, default) + + def setter(self, value): + self.base.attributes.set(attribute_key, value) + + return property(getter, setter) + + +def get_item_from_dict(base, attribute_key: str, item_key: str, default=None): + """ + Get one value from a dict attribute (by item_key). + """ + dct = base.attributes.get(attribute_key, {}) + return dct.get(item_key, default) + + +def set_item_in_dict(base, attribute_key: str, item_key: str, value): + """ + Set one value in a dict attribute (by item_key). + """ + dct = base.attributes.get(attribute_key, {}) + dct[item_key] = value + base.attributes.set(attribute_key, dct) + + +class WorkGraphNode(WorkChainNode): + """ORM class for all nodes representing the execution of a WorkGraph.""" + + TASK_STATES_KEY = 'task_states' + TASK_PROCESSES_KEY = 'task_processes' + TASK_ACTIONS_KEY = 'task_actions' + TASK_EXECUTORS_KEY = 'task_executors' + TASK_ERROR_HANDLERS_KEY = 'task_error_handlers' + TASK_EXECUTION_COUNTS_KEY = 'task_execution_counts' + TASK_MAP_INFO_KEY = 'task_map_info' + TASK_INPUTS_KEY = 'task_inputs' + WORKGRAPH_DATA_KEY = 'workgraph_data' + WORKGRAPH_DATA_SHORT_KEY = 'workgraph_data_short' + WORKGRAPH_ERROR_HANDLERS_KEY = 'workgraph_error_handlers' + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Use the same logger as WorkChainNode, this ensures the log level is set correctly. `getLogger` is + # statically typed as returning a `Logger`, but at runtime aiida's global logger class makes it the + # `AiidaLoggerType` that `_logger` expects. + self._logger = logging.getLogger( # type: ignore[assignment] + 'aiida.orm.nodes.process.workflow.workchain.WorkChainNode' + ) + + @classproperty + def _updatable_attributes(cls) -> tuple[str, ...]: # noqa: N805 + return super()._updatable_attributes + ( + cls.WORKGRAPH_DATA_KEY, + cls.TASK_INPUTS_KEY, + cls.WORKGRAPH_DATA_SHORT_KEY, + cls.WORKGRAPH_ERROR_HANDLERS_KEY, + cls.TASK_STATES_KEY, + cls.TASK_PROCESSES_KEY, + cls.TASK_ACTIONS_KEY, + cls.TASK_EXECUTORS_KEY, + cls.TASK_ERROR_HANDLERS_KEY, + cls.TASK_EXECUTION_COUNTS_KEY, + cls.TASK_MAP_INFO_KEY, + ) + + task_states = make_dict_property(TASK_STATES_KEY, default={}) + task_processes = make_dict_property(TASK_PROCESSES_KEY, default={}) + task_actions = make_dict_property(TASK_ACTIONS_KEY, default={}) + task_executors = make_dict_property(TASK_EXECUTORS_KEY, default={}) + task_error_handlers = make_dict_property(TASK_ERROR_HANDLERS_KEY, default={}) + task_execution_counts = make_dict_property(TASK_EXECUTION_COUNTS_KEY, default={}) + task_map_info = make_dict_property(TASK_MAP_INFO_KEY, default={}) + workgraph_data = make_dict_property(WORKGRAPH_DATA_KEY, default=None) + task_inputs = make_dict_property(TASK_INPUTS_KEY, default=None) + workgraph_data_short = make_dict_property(WORKGRAPH_DATA_SHORT_KEY, default=None) + workgraph_error_handlers = make_dict_property(WORKGRAPH_ERROR_HANDLERS_KEY, default=None) + + def get_task_state(self, task_name: str) -> str | None: + """Return the state of a single task.""" + return get_item_from_dict(self.base, self.TASK_STATES_KEY, task_name, default='') + + def set_task_state(self, task_name: str, task_state: str) -> None: + """Set the state of a single task.""" + set_item_in_dict(self.base, self.TASK_STATES_KEY, task_name, task_state) + + def get_task_process(self, task_name: str) -> str | None: + """Return the process info of a single task.""" + return get_item_from_dict(self.base, self.TASK_PROCESSES_KEY, task_name, default=None) + + def set_task_process(self, task_name: str, task_process: str) -> None: + """Set the process info of a single task.""" + set_item_in_dict(self.base, self.TASK_PROCESSES_KEY, task_name, task_process) + + def get_task_action(self, task_name: str) -> str | None: + """Return the action info of a single task.""" + return get_item_from_dict(self.base, self.TASK_ACTIONS_KEY, task_name, default='') + + def set_task_action(self, task_name: str, task_action: str) -> None: + """Set the action info of a single task.""" + set_item_in_dict(self.base, self.TASK_ACTIONS_KEY, task_name, task_action) + + def get_task_execution_count(self, task_name: str) -> int: + """Return the execution count of a single task.""" + return get_item_from_dict(self.base, self.TASK_EXECUTION_COUNTS_KEY, task_name, default=0) + + def set_task_execution_count(self, task_name: str, count: int) -> None: + """Set the execution count of a single task.""" + set_item_in_dict(self.base, self.TASK_EXECUTION_COUNTS_KEY, task_name, count) + + def get_task_map_info(self, task_name: str) -> str | None: + """Return the map info of a single task.""" + return get_item_from_dict(self.base, self.TASK_MAP_INFO_KEY, task_name, default='') + + def set_task_map_info(self, task_name: str, task_map_info: str) -> None: + """Set the map info of a single task.""" + set_item_in_dict(self.base, self.TASK_MAP_INFO_KEY, task_name, task_map_info) diff --git a/tests/orm/nodes/process/test_workgraph.py b/tests/orm/nodes/process/test_workgraph.py new file mode 100644 index 0000000000..78cdca354c --- /dev/null +++ b/tests/orm/nodes/process/test_workgraph.py @@ -0,0 +1,45 @@ +"""Tests for :mod:`aiida.orm.nodes.process.workflow.workgraph`.""" + +import pytest + +from aiida.orm import WorkChainNode, WorkGraphNode + + +def test_is_workchain_node_subclass(): + """A work graph is executed as a work chain, so its node specialises `WorkChainNode`.""" + assert issubclass(WorkGraphNode, WorkChainNode) + + +def test_node_type_is_stable(): + """The entry point keeps the `node_type` that existing work graph nodes were stored with. + + The class moved from aiida-workgraph into core under the same `aiida.node` entry point name, so nodes written + by the standalone package stay loadable. This pins that string against an accidental rename. + """ + assert WorkGraphNode().node_type == 'process.workflow.workgraph.WorkGraphNode.' + + +@pytest.mark.parametrize( + 'setter, getter, value', + [ + ('set_task_state', 'get_task_state', 'RUNNING'), + ('set_task_process', 'get_task_process', 12345), + ('set_task_action', 'get_task_action', 'pause'), + ('set_task_execution_count', 'get_task_execution_count', 3), + ('set_task_map_info', 'get_task_map_info', 'parent'), + ], + ids=['state', 'process', 'action', 'execution_count', 'map_info'], +) +def test_task_accessor_roundtrip(setter, getter, value): + """Each per-task accessor stores a value under a task name and reads the same value back.""" + node = WorkGraphNode() + getattr(node, setter)('task_a', value) + assert getattr(node, getter)('task_a') == value + + +def test_per_task_accessor_writes_the_bulk_attribute(): + """The per-task setter and the bulk dict property are two views of the same stored attribute.""" + node = WorkGraphNode() + node.set_task_state('task_a', 'RUNNING') + node.set_task_state('task_b', 'FINISHED') + assert node.task_states == {'task_a': 'RUNNING', 'task_b': 'FINISHED'} diff --git a/tests/orm/test_fields/fields_aiida.node.process.workflow.workgraph.WorkGraphNode.yml b/tests/orm/test_fields/fields_aiida.node.process.workflow.workgraph.WorkGraphNode.yml new file mode 100644 index 0000000000..87012cbe29 --- /dev/null +++ b/tests/orm/test_fields/fields_aiida.node.process.workflow.workgraph.WorkGraphNode.yml @@ -0,0 +1,36 @@ +attributes: QbAttributesField('attributes', dtype=, + doc='The node attributes') +computer: QbNumericField('computer', dtype=int | None, doc='The PK of the computer') +ctime: QbNumericField('ctime', dtype=, doc='The creation + time of the node') +description: QbStrField('description', dtype=, doc='The node description') +exception: QbStrField('attributes.exception', dtype=str | None, doc='The process exception + message') +exit_message: QbStrField('attributes.exit_message', dtype=str | None, doc='The process + exit message') +exit_status: QbNumericField('attributes.exit_status', dtype=int | None, doc='The process + exit status') +extras: QbDictField('extras', dtype=dict[str, typing.Any], doc='The node extras') +label: QbStrField('label', dtype=, doc='The node label') +mtime: QbNumericField('mtime', dtype=, doc='The modification + time of the node') +node_type: QbStrField('node_type', dtype=typing.Literal['process.workflow.workgraph.WorkGraphNode.'], + doc='The type of the node.') +paused: QbAnyField('attributes.paused', dtype=bool | None, doc='Whether the process + is paused') +pk: QbNumericField('pk', dtype=, doc='The primary key of the entity') +process_label: QbStrField('attributes.process_label', dtype=str | None, doc='The process + label') +process_state: QbStrField('attributes.process_state', dtype=str | None, doc='The process + state enum') +process_status: QbStrField('attributes.process_status', dtype=str | None, doc='The + process status is a generic status message') +process_type: QbStrField('process_type', dtype=str | None, doc='The process type of + the node') +repository_metadata: QbDictField('repository_metadata', dtype=dict[str, typing.Any], + doc='Virtual hierarchy of the file repository') +sealed: QbAnyField('attributes.sealed', dtype=, doc='Whether the node + is sealed') +user: QbNumericField('user', dtype=, doc='The PK of the user who owns + the node') +uuid: QbAnyField('uuid', dtype=, doc='The UUID of the node') From 49e2bc8f583ae004ad79a96b19a146492f1b2b78 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Tue, 21 Jul 2026 17:38:00 +0200 Subject: [PATCH 04/19] =?UTF-8?q?=E2=9C=A8=20`workgraph`:=20add=20the=20`e?= =?UTF-8?q?nums`=20module?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Start the `aiida.workgraph` subpackage, the home for the AiiDA WorkGraph language and runtime as they move into core, and add its first module: the task enums (`TaskState`, `TaskAction`, `TERMINAL_TASK_STATES`, `RuntimeInfoKey`, `TaskActionMessage`), moved verbatim from aiida-workgraph. The enums are pure stdlib with no node-graph dependency, so they import without the eventual workgraph extra. The package `__init__` is deliberately minimal and imports nothing node-graph-dependent, so a plain `import aiida` stays free of that dependency; the node-graph-bound parts of the subpackage will be imported lazily. The accompanying aiida-workgraph change drops its own copy and imports from `aiida.workgraph.enums`. The unit tests move across with the code. --- src/aiida/workgraph/__init__.py | 30 ++++++++++++ src/aiida/workgraph/enums.py | 82 +++++++++++++++++++++++++++++++++ tests/workgraph/__init__.py | 8 ++++ tests/workgraph/test_enums.py | 71 ++++++++++++++++++++++++++++ 4 files changed, 191 insertions(+) create mode 100644 src/aiida/workgraph/__init__.py create mode 100644 src/aiida/workgraph/enums.py create mode 100644 tests/workgraph/__init__.py create mode 100644 tests/workgraph/test_enums.py diff --git a/src/aiida/workgraph/__init__.py b/src/aiida/workgraph/__init__.py new file mode 100644 index 0000000000..02514aeea4 --- /dev/null +++ b/src/aiida/workgraph/__init__.py @@ -0,0 +1,30 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""The AiiDA WorkGraph: a data-dependency workflow language and runtime. + +This subpackage is the AiiDA-specific layer over the generic node-graph SDK. It is imported lazily and depends on +node-graph, so it must not be imported from aiida-core's own import path; a plain ``import aiida`` stays free of the +node-graph dependency, which the (optional) workgraph install provides. +""" + +# AUTO-GENERATED + +# fmt: off + +from .enums import * + +__all__ = ( + 'TERMINAL_TASK_STATES', + 'RuntimeInfoKey', + 'TaskAction', + 'TaskActionMessage', + 'TaskState', +) + +# fmt: on diff --git a/src/aiida/workgraph/enums.py b/src/aiida/workgraph/enums.py new file mode 100644 index 0000000000..b15b7b25af --- /dev/null +++ b/src/aiida/workgraph/enums.py @@ -0,0 +1,82 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Enums used by WorkGraph.""" + +from __future__ import annotations + +from enum import Enum, unique +from typing import Final, Literal, TypedDict + +__all__ = ( + 'TERMINAL_TASK_STATES', + 'RuntimeInfoKey', + 'TaskAction', + 'TaskActionMessage', + 'TaskState', +) + + +@unique +class TaskState(str, Enum): + """Lifecycle state of a single task within a running WorkGraph.""" + + PLANNED = 'PLANNED' + READY = 'READY' + CREATED = 'CREATED' + RUNNING = 'RUNNING' + FINISHED = 'FINISHED' + FAILED = 'FAILED' + SKIPPED = 'SKIPPED' + MAPPED = 'MAPPED' + + def __str__(self) -> str: + # Return the bare value ('RUNNING'), not 'TaskState.RUNNING', uniformly + # across Python versions so reports and logs read naturally. + return self.value + + @property + def is_terminal(self) -> bool: + """Whether the task has settled and will not transition further.""" + return self in TERMINAL_TASK_STATES + + +#: States a task does not transition out of; a task in any of these is "done" for +#: readiness and finished checks. +TERMINAL_TASK_STATES: Final[frozenset[TaskState]] = frozenset({TaskState.FINISHED, TaskState.SKIPPED, TaskState.FAILED}) + + +@unique +class TaskAction(str, Enum): + """Externally-triggered action requested on a task.""" + + PAUSE = 'PAUSE' + PLAY = 'PLAY' + KILL = 'KILL' + SKIP = 'SKIP' + RESET = 'RESET' + + def __str__(self) -> str: + return self.value + + +#: Keys addressing a task's runtime info on the process node, used as the dispatch +#: key in ``get_task_runtime_info`` / ``set_task_runtime_info``. +RuntimeInfoKey = Literal['process', 'state', 'action', 'execution_count', 'map_info'] + + +class TaskActionMessage(TypedDict): + """RPC payload sent to a running WorkGraph to act on its tasks. + + Built by ``create_task_action`` and consumed by ``apply_task_actions``. + """ + + intent: str + catalog: str + action: str + tasks: list[str] diff --git a/tests/workgraph/__init__.py b/tests/workgraph/__init__.py new file mode 100644 index 0000000000..c56ff0a1f8 --- /dev/null +++ b/tests/workgraph/__init__.py @@ -0,0 +1,8 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### diff --git a/tests/workgraph/test_enums.py b/tests/workgraph/test_enums.py new file mode 100644 index 0000000000..f29459dd36 --- /dev/null +++ b/tests/workgraph/test_enums.py @@ -0,0 +1,71 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Unit tests for the task-state and task-action enums. + +These cover the behaviour the engine relies on: the human-facing string form +drops the ``TaskState.``/``TaskAction.`` prefix, only canonical values are +accepted, and the terminal-state set matches both enum members and bare strings. +""" + +import pytest + +from aiida.workgraph.enums import TERMINAL_TASK_STATES, TaskAction, TaskState + + +@pytest.mark.parametrize('enum_cls', [TaskState, TaskAction]) +def test_str_and_format_drop_class_prefix(enum_cls): + """``__str__``/f-strings must yield the bare value on every supported Python + version, so report messages stay e.g. ``Action: RESET`` (test_workgraph).""" + member = next(iter(enum_cls)) + assert str(member) == member.value + assert f'{member}' == member.value + + +@pytest.mark.parametrize('enum_cls', [TaskState, TaskAction]) +@pytest.mark.parametrize( + 'bad', + [ + pytest.param('running', id='lowercased-state'), + pytest.param('reset', id='lowercased-action'), + pytest.param('nope', id='gibberish'), + pytest.param('', id='empty'), + ], +) +def test_construction_from_noncanonical_value_raises(enum_cls, bad): + """Only the canonical uppercase values are valid; a typo or wrong case fails + loud instead of silently never matching.""" + with pytest.raises(ValueError): + enum_cls(bad) + + +@pytest.mark.parametrize( + 'state, terminal', + [ + (TaskState.FINISHED, True), + (TaskState.SKIPPED, True), + (TaskState.FAILED, True), + (TaskState.PLANNED, False), + (TaskState.RUNNING, False), + (TaskState.CREATED, False), + (TaskState.READY, False), + (TaskState.MAPPED, False), + ], +) +def test_is_terminal(state, terminal): + assert state.is_terminal is terminal + + +def test_terminal_set_matches_bare_strings_and_not_unset_default(): + """The engine checks ``stored_state in TERMINAL_TASK_STATES`` where the + stored value is a bare string and the unset default is ``''``.""" + assert TERMINAL_TASK_STATES == frozenset({TaskState.FINISHED, TaskState.SKIPPED, TaskState.FAILED}) + for state in TERMINAL_TASK_STATES: + assert state.value in TERMINAL_TASK_STATES # hash-equality with bare strings + assert '' not in TERMINAL_TASK_STATES + assert TaskState.PLANNED not in TERMINAL_TASK_STATES From 8023feda252bff6eeb795daaafdd2e3309ccb9d2 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Tue, 21 Jul 2026 20:31:58 +0200 Subject: [PATCH 05/19] =?UTF-8?q?=E2=9C=A8=20`workgraph`:=20add=20`utils`?= =?UTF-8?q?=20nested-dict=20helpers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `aiida/workgraph/utils.py` with the generic helpers the WorkGraph runtime relies on: dotted-key access into nested dictionaries (`get_nested_dict`, `update_nested_dict`, `update_nested_dict_with_special_keys`) and resolving AiiDA `NodeLinksManager` structures into plain dictionaries (`resolve_node_link_managers`), moved from aiida-workgraph. They depend only on aiida-core (`NodeLinksManager`), with no node-graph or plugin import, so `aiida.workgraph` still imports without the eventual workgraph extra. The node-graph-coupled workgraph-data serialization stays downstream for now and moves later with the engine. The accompanying aiida-workgraph change deletes its copies and imports these from `aiida.workgraph.utils`. Tests cover the dict helpers. --- src/aiida/workgraph/__init__.py | 5 ++ src/aiida/workgraph/utils.py | 147 ++++++++++++++++++++++++++++++++ tests/workgraph/test_utils.py | 59 +++++++++++++ 3 files changed, 211 insertions(+) create mode 100644 src/aiida/workgraph/utils.py create mode 100644 tests/workgraph/test_utils.py diff --git a/src/aiida/workgraph/__init__.py b/src/aiida/workgraph/__init__.py index 02514aeea4..cac16f23e8 100644 --- a/src/aiida/workgraph/__init__.py +++ b/src/aiida/workgraph/__init__.py @@ -18,6 +18,7 @@ # fmt: off from .enums import * +from .utils import * __all__ = ( 'TERMINAL_TASK_STATES', @@ -25,6 +26,10 @@ 'TaskAction', 'TaskActionMessage', 'TaskState', + 'get_nested_dict', + 'resolve_node_link_managers', + 'update_nested_dict', + 'update_nested_dict_with_special_keys', ) # fmt: on diff --git a/src/aiida/workgraph/utils.py b/src/aiida/workgraph/utils.py new file mode 100644 index 0000000000..e166ad8e28 --- /dev/null +++ b/src/aiida/workgraph/utils.py @@ -0,0 +1,147 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Generic helpers used by the WorkGraph runtime. + +These are the node-graph-free, plugin-free utilities the engine relies on: dotted-key access into nested +dictionaries, and resolving AiiDA ``NodeLinksManager`` structures into plain dictionaries. +""" + +from __future__ import annotations + +from typing import Any + +from aiida.orm.utils.managers import NodeLinksManager + +__all__ = ( + 'get_nested_dict', + 'resolve_node_link_managers', + 'update_nested_dict', + 'update_nested_dict_with_special_keys', +) + + +def get_nested_dict(d: Any, name: str, **kwargs: Any) -> Any: + """Get the value from a nested dictionary. + + ``d`` is deliberately ``Any``: the traversal descends through both plain dicts and AiiDA + ``NodeLinksManager`` containers, whose values are heterogeneous. + + If default is provided, return the default value if the key is not found. + Otherwise, raise ValueError. + For example: + d = {"base": {"pw": {"parameters": 2}}} + name = "base.pw.parameters" + """ + keys = name.split('.') + current = d + for key in keys: + if key not in current: + if 'default' in kwargs: + return kwargs.get('default') + if isinstance(current, dict): + avaiable_keys = list(current.keys()) + elif isinstance(current, NodeLinksManager): + avaiable_keys = list(current._get_keys()) + else: + avaiable_keys = [] + raise ValueError(f'{name} not exist. Available keys: {avaiable_keys}') + current = current[key] + return current + + +def merge_dicts(dict1: Any, dict2: Any) -> Any: + """Recursively merges two dictionaries.""" + for key, value in dict2.items(): + if key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict): + # Recursively merge dictionaries + dict1[key] = merge_dicts(dict1[key], value) + else: + # Overwrite or add the key + dict1[key] = value + return dict1 + + +def update_nested_dict(base: dict[str, Any] | None, key_path: str, value: Any) -> dict[str, Any]: + """ + Update or create a nested dictionary structure based on a dotted key path. + + This function allows updating a nested dictionary or creating one if `d` is `None`. + Given a dictionary and a key path (e.g., "base.pw.parameters"), it will traverse + or create the necessary nested structure to set the provided value at the specified + key location. If intermediate dictionaries do not exist, they will be created. + If the resulting dictionary is empty, it is set to `None`. + + Args: + base (Dict[str, Any] | None): The dictionary to update, which can be `None`. + If `None`, an empty dictionary will be created. + key (str): A dotted key path string representing the nested structure. + value (Any): The value to set at the specified key. + + Example: + base = None + key = "scf.pw.parameters" + value = 2 + After running: + update_nested_dict(d, key, value) + The result will be: + base = {"scf": {"pw": {"parameters": 2}}} + + Edge Case: + If the resulting dictionary is empty after the update, it will be set to `None`. + + """ + if base is None: + base = {} + keys = key_path.split('.') + current_key = keys[0] + if len(keys) == 1: + # Base case: Merge dictionaries or set the value directly. + if isinstance(base.get(current_key), dict) and isinstance(value, dict): + base[current_key] = merge_dicts(base[current_key], value) + else: + base[current_key] = value + else: + # Recursive case: Ensure the key exists and is a dictionary, then recurse. + if current_key not in base or not isinstance(base[current_key], dict): + base[current_key] = {} + base[current_key] = update_nested_dict(base[current_key], '.'.join(keys[1:]), value) + + return base + + +def update_nested_dict_with_special_keys(data: dict[str, Any]) -> dict[str, Any]: + """Update the nested dictionary with special keys like "base.pw.parameters".""" + # Remove None + data = {k: v for k, v in data.items() if v is not None} + special_keys = [k for k in data.keys() if '.' in k] + for key in special_keys: + value = data.pop(key) + update_nested_dict(data, key, value) + return data + + +def resolve_node_link_managers(data: Any) -> Any: + """Recursively resolve all NodeLinksManagers either in a dictionary or a NodeLinksManager.""" + if isinstance(data, dict): + return {key: resolve_node_link_managers(value) for key, value in data.items()} + if isinstance(data, NodeLinksManager): + return convert_node_link_manager_to_dict(data) + return data + + +def convert_node_link_manager_to_dict(node_link_manager: NodeLinksManager) -> dict[str, Any]: + """Recursively convert a NodeLinksManager to a dictionary representation.""" + data = {} + for name in node_link_manager._get_keys(): + item = node_link_manager._get_node_by_link_label(name) + if isinstance(item, NodeLinksManager): + data[name] = convert_node_link_manager_to_dict(item) + else: + data[name] = item + return data diff --git a/tests/workgraph/test_utils.py b/tests/workgraph/test_utils.py new file mode 100644 index 0000000000..46c31d105a --- /dev/null +++ b/tests/workgraph/test_utils.py @@ -0,0 +1,59 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Unit tests for :mod:`aiida.workgraph.utils`.""" + +import pytest + +from aiida.workgraph.utils import ( + get_nested_dict, + resolve_node_link_managers, + update_nested_dict, + update_nested_dict_with_special_keys, +) + + +def test_get_nested_dict_returns_leaf(): + assert get_nested_dict({'base': {'pw': {'parameters': 2}}}, 'base.pw.parameters') == 2 + + +def test_get_nested_dict_default_when_missing(): + assert get_nested_dict({'base': {'pw': {}}}, 'base.pw.parameters', default=None) is None + + +def test_get_nested_dict_raises_without_default(): + with pytest.raises(ValueError, match='not exist'): + get_nested_dict({'base': {'pw': {}}}, 'base.pw.parameters') + + +def test_update_nested_dict_creates_path_from_none(): + assert update_nested_dict(None, 'scf.pw.parameters', 2) == {'scf': {'pw': {'parameters': 2}}} + + +def test_update_nested_dict_merges_into_existing_branch(): + base = {'scf': {'pw': {'a': 1}}} + assert update_nested_dict(base, 'scf.pw', {'b': 2}) == {'scf': {'pw': {'a': 1, 'b': 2}}} + + +def test_update_nested_dict_overwrites_non_dict_leaf(): + assert update_nested_dict({'x': 1}, 'x', 2) == {'x': 2} + + +def test_update_nested_dict_with_special_keys_expands_and_drops_none(): + data = {'base.pw.parameters': 2, 'plain': 'keep', 'gone': None} + assert update_nested_dict_with_special_keys(data) == { + 'base': {'pw': {'parameters': 2}}, + 'plain': 'keep', + } + + +def test_resolve_node_link_managers_passes_through_plain_data(): + """Values that are not ``NodeLinksManager`` are returned unchanged, recursing into dicts. The manager + conversion itself is exercised end to end by every work graph run.""" + data = {'a': 1, 'nested': {'b': [1, 2], 'c': 'str'}} + assert resolve_node_link_managers(data) == data From 6b6f71b09c004e9aa3534bbf645efe216fd9d5ed Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Tue, 21 Jul 2026 21:21:13 +0200 Subject: [PATCH 06/19] =?UTF-8?q?=E2=9C=A8=20`NoneData`:=20node=20type=20f?= =?UTF-8?q?or=20Python=20`None`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `aiida/orm/nodes/data/none.py`, a `Data` subclass that explicitly represents a Python `None`. It has no repository content and no attributes, so every instance shares one content hash: `None` is a single value. A dedicated node is needed because `None` cannot be a simple `BaseType` (`Int`, `Bool`, ...), yet a serialized value must always map to a node. Register it with `to_aiida_type` for `type(None)` and as the `core.none` `aiida.data` entry point, so `None` serializes with no special case in the caller. This is the first increment of the serializer reconcile (step 5a) that moves the aiida-pythonjob serialization stack into aiida-core: it lands the one moved data type the WorkGraph engine references by name, built on core's existing `to_aiida_type` rather than duplicating a plugin mapping. --- pyproject.toml | 1 + src/aiida/orm/__init__.py | 1 + src/aiida/orm/nodes/__init__.py | 1 + src/aiida/orm/nodes/data/__init__.py | 2 + src/aiida/orm/nodes/data/none.py | 46 +++++++++++++++++ tests/orm/nodes/data/test_none.py | 51 +++++++++++++++++++ .../fields_aiida.data.core.none.NoneData.yml | 21 ++++++++ 7 files changed, 123 insertions(+) create mode 100644 src/aiida/orm/nodes/data/none.py create mode 100644 tests/orm/nodes/data/test_none.py create mode 100644 tests/orm/test_fields/fields_aiida.data.core.none.NoneData.yml diff --git a/pyproject.toml b/pyproject.toml index 14885e2889..515895fcb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -125,6 +125,7 @@ requires-python = '>=3.10' 'core.int' = 'aiida.orm.nodes.data.int:Int' 'core.jsonable' = 'aiida.orm.nodes.data.jsonable:JsonableData' 'core.list' = 'aiida.orm.nodes.data.list:List' +'core.none' = 'aiida.orm.nodes.data.none:NoneData' 'core.numeric' = 'aiida.orm.nodes.data.numeric:NumericType' 'core.orbital' = 'aiida.orm.nodes.data.orbital:OrbitalData' 'core.remote' = 'aiida.orm.nodes.data.remote.base:RemoteData' diff --git a/src/aiida/orm/__init__.py b/src/aiida/orm/__init__.py index b65cb0b611..780fad23db 100644 --- a/src/aiida/orm/__init__.py +++ b/src/aiida/orm/__init__.py @@ -77,6 +77,7 @@ 'NodeEntityLoader', 'NodeLinksManager', 'NodeRepository', + 'NoneData', 'NumericType', 'OrbitalData', 'OrderSpecifier', diff --git a/src/aiida/orm/nodes/__init__.py b/src/aiida/orm/nodes/__init__.py index c4696f0fec..f4ff37f640 100644 --- a/src/aiida/orm/nodes/__init__.py +++ b/src/aiida/orm/nodes/__init__.py @@ -44,6 +44,7 @@ 'Node', 'NodeAttributes', 'NodeRepository', + 'NoneData', 'NumericType', 'OrbitalData', 'PortableCode', diff --git a/src/aiida/orm/nodes/data/__init__.py b/src/aiida/orm/nodes/data/__init__.py index 983cf1519e..3ad38aa1a9 100644 --- a/src/aiida/orm/nodes/data/__init__.py +++ b/src/aiida/orm/nodes/data/__init__.py @@ -25,6 +25,7 @@ from .int import * from .jsonable import * from .list import * +from .none import * from .numeric import * from .orbital import * from .remote import * @@ -53,6 +54,7 @@ 'Kind', 'KpointsData', 'List', + 'NoneData', 'NumericType', 'OrbitalData', 'PortableCode', diff --git a/src/aiida/orm/nodes/data/none.py b/src/aiida/orm/nodes/data/none.py new file mode 100644 index 0000000000..5ceb9c5bf3 --- /dev/null +++ b/src/aiida/orm/nodes/data/none.py @@ -0,0 +1,46 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""`Data` sub class to explicitly represent a Python ``None``.""" + +from __future__ import annotations + +from .base import to_aiida_type +from .data import Data + +__all__ = ('NoneData',) + + +@to_aiida_type.register(type(None)) +def _(value): + return NoneData() + + +class NoneData(Data): + """A ``Data`` node that explicitly represents a Python ``None``. + + It carries no repository content and no attributes, so every instance has an identical content hash. That is the + intended behaviour: ``None`` is a single value. A dedicated node type is needed because ``None`` cannot be stored as + one of the simple ``BaseType`` nodes (``Int``, ``Bool``, ...), yet a serialized value must always map to a node. + """ + + @property + def value(self) -> None: + """Return the represented value, which is always ``None``.""" + return None + + @property + def obj(self) -> None: + """Alias of :attr:`value`, mirroring the ``.obj`` accessor of other wrapping data types.""" + return None + + def __repr__(self) -> str: + return 'NoneData()' + + def __str__(self) -> str: + return 'NoneData()' diff --git a/tests/orm/nodes/data/test_none.py b/tests/orm/nodes/data/test_none.py new file mode 100644 index 0000000000..084b5ecdf9 --- /dev/null +++ b/tests/orm/nodes/data/test_none.py @@ -0,0 +1,51 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Tests for the :class:`aiida.orm.nodes.data.none.NoneData` data type.""" + +from aiida import orm +from aiida.orm import NoneData, load_node + + +def test_value_and_obj_are_none(): + """Both accessors report the represented value.""" + node = NoneData() + assert node.value is None + assert node.obj is None + + +def test_repr_and_str(): + assert repr(NoneData()) == 'NoneData()' + assert str(NoneData()) == 'NoneData()' + + +def test_to_aiida_type_dispatches_none(): + """``to_aiida_type(None)`` returns an (unstored) ``NoneData``.""" + node = orm.to_aiida_type(None) + assert isinstance(node, NoneData) + assert node.is_stored is False + + +def test_stores_without_attributes(): + """A ``None`` carries no state, so the stored node has no attributes.""" + node = NoneData().store() + assert node.is_stored is True + assert node.base.attributes.all == {} + + +def test_identical_content_hash(): + """Every ``None`` is the same value, so all instances hash identically.""" + assert NoneData().store().base.caching.compute_hash() == NoneData().store().base.caching.compute_hash() + + +def test_roundtrip_through_entry_point(): + """A stored node reloads as ``NoneData`` via its registered ``node_type``.""" + pk = NoneData().store().pk + loaded = load_node(pk) + assert isinstance(loaded, NoneData) + assert loaded.value is None diff --git a/tests/orm/test_fields/fields_aiida.data.core.none.NoneData.yml b/tests/orm/test_fields/fields_aiida.data.core.none.NoneData.yml new file mode 100644 index 0000000000..9ede304ffb --- /dev/null +++ b/tests/orm/test_fields/fields_aiida.data.core.none.NoneData.yml @@ -0,0 +1,21 @@ +attributes: QbAttributesField('attributes', dtype=, + doc='The node attributes') +computer: QbNumericField('computer', dtype=int | None, doc='The PK of the computer') +ctime: QbNumericField('ctime', dtype=, doc='The creation + time of the node') +description: QbStrField('description', dtype=, doc='The node description') +extras: QbDictField('extras', dtype=dict[str, typing.Any], doc='The node extras') +label: QbStrField('label', dtype=, doc='The node label') +mtime: QbNumericField('mtime', dtype=, doc='The modification + time of the node') +node_type: QbStrField('node_type', dtype=typing.Literal['data.core.none.NoneData.'], + doc='The type of the node.') +pk: QbNumericField('pk', dtype=, doc='The primary key of the entity') +process_type: QbStrField('process_type', dtype=str | None, doc='The process type of + the node') +repository_metadata: QbDictField('repository_metadata', dtype=dict[str, typing.Any], + doc='Virtual hierarchy of the file repository') +source: QbDictField('attributes.source', dtype=dict | None, doc='Source of the data') +user: QbNumericField('user', dtype=, doc='The PK of the user who owns + the node') +uuid: QbAnyField('uuid', dtype=, doc='The UUID of the node') From dd3307e4f2560b442c95eb424a379ab0b0064e1e Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Tue, 21 Jul 2026 22:10:04 +0200 Subject: [PATCH 07/19] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20`JsonableData`:=20ac?= =?UTF-8?q?cept=20more=20serializable=20inputs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Generalise the wrapper so it can act as the fallback for arbitrary Python value serialization. Besides the existing `as_dict` / `from_dict` contract it now also accepts objects exposing `to_dict` / `todict` / `asdict`, `dataclasses.dataclass` instances, and `pydantic.BaseModel` instances, coerces numpy scalars and arrays in the produced dictionary to their JSON-native form, and reconstructs through `model_validate` (pydantic), the constructor (dataclass), or `from_dict` / `fromdict`. The change is storage-compatible (same `@class` / `@module` + attribute layout) and strictly more permissive: the `as_dict` / `from_dict` path, the inf/-inf/NaN round-trip, and MSONable support are unchanged. The one visible behaviour change is the error raised when an object implements none of the supported methods; its message now names the wider set. This is increment 2 of the serializer reconcile (step 5a): rather than port aiida-pythonjob's separate, more permissive `JsonableData`, the one core `JsonableData` absorbs its flexibility, so the moved serializer can fall back to a single JSON wrapper instead of duplicating one. --- src/aiida/orm/nodes/data/jsonable.py | 111 +++++++++++++++++++++++--- tests/orm/nodes/data/test_jsonable.py | 66 ++++++++++++++- 2 files changed, 165 insertions(+), 12 deletions(-) diff --git a/src/aiida/orm/nodes/data/jsonable.py b/src/aiida/orm/nodes/data/jsonable.py index 4b34060202..f559a547bd 100644 --- a/src/aiida/orm/nodes/data/jsonable.py +++ b/src/aiida/orm/nodes/data/jsonable.py @@ -2,11 +2,13 @@ from __future__ import annotations +import dataclasses import importlib import json import typing -from pydantic import ConfigDict, WithJsonSchema +import numpy as np +from pydantic import BaseModel, ConfigDict, WithJsonSchema from aiida.orm.pydantic import OrmFieldsAsModelDump, OrmMetadataField, OrmModel @@ -50,6 +52,13 @@ class JsonableData(Data): Of course, this requires that the class of the originally wrapped instance can be imported in the current environment, or an ``ImportError`` will be raised. + + Besides the ``as_dict`` / ``from_dict`` contract above, the wrapper also accepts objects whose dictionary is + produced by ``to_dict`` / ``todict`` / ``asdict``, as well as :func:`dataclasses.dataclass` instances and + ``pydantic.BaseModel`` instances (reconstructed with the constructor and ``model_validate`` respectively). Numpy + scalars and arrays returned in the dictionary are coerced to their JSON-native counterparts. This broader support + is what lets it act as the generic fallback for arbitrary-Python-value serialization; the ``as_dict`` / + ``from_dict`` path is unchanged. """ class AttributesModel(OrmFieldsAsModelDump, Data.AttributesModel): @@ -89,24 +98,29 @@ class ConstructorArgsModel(OrmModel): ), ] + #: Method names, tried in order, that an object may implement to produce its serializable dictionary. ``as_dict`` + #: (the historical, MSONable-style contract) is tried first so existing behaviour is unchanged. + _DICT_METHODS = ('as_dict', 'to_dict', 'todict', 'asdict') + + #: Class-method names, tried in order, that a class may implement to rebuild an instance from its dictionary. + _FROM_DICT_METHODS = ('from_dict', 'fromdict') + def __init__(self, obj: JsonSerializableProtocol, *args, **kwargs): """Construct the node for the to be wrapped object.""" if obj is None: raise TypeError('the `obj` argument cannot be `None`.') - if not hasattr(obj, 'as_dict') or not callable(getattr(obj, 'as_dict')): - raise TypeError('the `obj` argument does not have the required `as_dict` method.') + dictionary = self._extract_dict(obj) super().__init__(*args, **kwargs) self._obj = obj - dictionary = obj.as_dict() - if '@class' not in dictionary: - dictionary['@class'] = obj.__class__.__name__ + dictionary.setdefault('@class', obj.__class__.__name__) + dictionary.setdefault('@module', obj.__class__.__module__) - if '@module' not in dictionary: - dictionary['@module'] = obj.__class__.__module__ + # Coerce numpy scalars and arrays that ``as_dict`` may return into JSON-native types before the round-trip. + dictionary = self._make_jsonable(dictionary) # Even though the dictionary returned by ``as_dict`` should be JSON-serializable and therefore this should be # sufficient to be able to generate a JSON representation and thus store it in the database, there is a @@ -125,6 +139,83 @@ def __init__(self, obj: JsonSerializableProtocol, *args, **kwargs): self.base.attributes.set_many(serialized) + def _extract_dict(self, obj: typing.Any) -> dict: + """Return the serializable dictionary for ``obj``. + + Pydantic models (``model_dump``) and dataclasses (``dataclasses.asdict``) are supported natively; any other + object must implement one of :attr:`_DICT_METHODS`. Raises ``TypeError`` if none applies. + """ + if self._is_pydantic_instance(obj): + return obj.model_dump(exclude_none=False) + if self._is_dataclass_instance(obj): + return dataclasses.asdict(obj) + for method_name in self._DICT_METHODS: + method = getattr(obj, method_name, None) + if callable(method): + return method() + raise TypeError( + f'the `obj` argument does not implement any of the supported serialization methods ' + f'({", ".join(self._DICT_METHODS)}), nor is it a dataclass or pydantic model.' + ) + + @classmethod + def _make_jsonable(cls, data: typing.Any) -> typing.Any: + """Recursively coerce numpy scalars and arrays in ``data`` into JSON-native types (``ndarray`` to ``list``, + ``numpy.generic`` to its Python scalar), leaving everything else untouched.""" + if isinstance(data, dict): + return {key: cls._make_jsonable(value) for key, value in data.items()} + if isinstance(data, list): + return [cls._make_jsonable(value) for value in data] + if isinstance(data, tuple): + return tuple(cls._make_jsonable(value) for value in data) + if isinstance(data, np.ndarray): + return data.tolist() + if isinstance(data, np.generic): + return data.item() + return data + + def _rebuild_object(self, cls_: typing.Any, attributes: dict) -> typing.Any: + """Reconstruct an instance of ``cls_`` from its stored ``attributes``. + + Pydantic types (``model_validate`` / ``parse_obj``) and dataclasses (constructor) are handled natively; any + other class is rebuilt through one of :attr:`_FROM_DICT_METHODS`, falling back to the plain constructor. + """ + if self._is_pydantic_type(cls_): + if hasattr(cls_, 'model_validate'): + return cls_.model_validate(attributes) + if hasattr(cls_, 'parse_obj'): + return cls_.parse_obj(attributes) + return cls_(**attributes) + if self._is_dataclass_type(cls_): + return cls_(**attributes) + for method_name in self._FROM_DICT_METHODS: + from_dict = getattr(cls_, method_name, None) + if callable(from_dict): + return from_dict(attributes) + try: + return cls_(**attributes) + except TypeError as exc: + raise TypeError( + f'cannot rebuild an object of type `{cls_}`: it implements none of the from-dict methods ' + f'({", ".join(self._FROM_DICT_METHODS)}) and its constructor does not accept the stored attributes.' + ) from exc + + @staticmethod + def _is_pydantic_instance(obj: typing.Any) -> bool: + return isinstance(obj, BaseModel) + + @staticmethod + def _is_pydantic_type(cls_: typing.Any) -> bool: + return isinstance(cls_, type) and issubclass(cls_, BaseModel) + + @staticmethod + def _is_dataclass_instance(obj: typing.Any) -> bool: + return dataclasses.is_dataclass(obj) and not isinstance(obj, type) + + @staticmethod + def _is_dataclass_type(cls_: typing.Any) -> bool: + return isinstance(cls_, type) and dataclasses.is_dataclass(cls_) + @property def the_module(self) -> str: """Return the module name of the wrapped object.""" @@ -198,7 +289,7 @@ def _get_object(self) -> JsonSerializableProtocol: ) from exc deserialized = self._deserialize_float_constants(attributes) - self._obj = cls.from_dict(deserialized) + self._obj = self._rebuild_object(cls, deserialized) return self._obj @@ -215,5 +306,5 @@ def to_model_field_values( schema=schema, ) if schema and issubclass(schema, self.WritableFields): - fields['attributes'] |= self.obj.as_dict() + fields['attributes'] |= self._extract_dict(self.obj) return fields diff --git a/tests/orm/nodes/data/test_jsonable.py b/tests/orm/nodes/data/test_jsonable.py index 1c235ed853..e63adfd70e 100644 --- a/tests/orm/nodes/data/test_jsonable.py +++ b/tests/orm/nodes/data/test_jsonable.py @@ -1,9 +1,12 @@ """Tests for the :class:`aiida.orm.nodes.data.jsonable.JsonableData` data type.""" +import dataclasses import datetime import math +import numpy import pytest +from pydantic import BaseModel from pymatgen.core.structure import Molecule from aiida.orm import load_node @@ -45,12 +48,12 @@ def test_construct(): def test_invalid_class_no_as_dict(): - """Test the ``JsonableData`` constructor raises if object does not implement ``as_dict``.""" + """Test the ``JsonableData`` constructor raises if object implements no supported serialization method.""" class InvalidClass: pass - with pytest.raises(TypeError, match=r'the `obj` argument does not have the required `as_dict` method.'): + with pytest.raises(TypeError, match=r'does not implement any of the supported serialization methods'): JsonableData(InvalidClass()) @@ -147,3 +150,62 @@ def test_msonable(): loaded = load_node(node.pk) assert loaded is not node assert loaded.obj == obj + + +class ToDictClass: + """Object that follows the ``to_dict`` / ``from_dict`` convention instead of ``as_dict``.""" + + def __init__(self, value): + self.value = value + + def to_dict(self): + return {'value': self.value} + + @classmethod + def from_dict(cls, dictionary): + return cls(dictionary['value']) + + +@dataclasses.dataclass +class DataclassObj: + """A plain dataclass, wrapped via ``dataclasses.asdict`` and rebuilt through its constructor.""" + + x: int + y: str + + +class PydanticObj(BaseModel): + """A pydantic model, wrapped via ``model_dump`` and rebuilt through ``model_validate``.""" + + a: int + b: str + + +def test_wrap_object_with_to_dict(): + """An object exposing ``to_dict`` (not ``as_dict``) round-trips.""" + node = JsonableData(ToDictClass(7)).store() + loaded = load_node(node.pk) + assert isinstance(loaded.obj, ToDictClass) + assert loaded.obj.value == 7 + + +def test_wrap_dataclass(): + """A dataclass instance round-trips through ``asdict`` and its constructor.""" + node = JsonableData(DataclassObj(x=1, y='a')).store() + loaded = load_node(node.pk) + assert loaded.obj == DataclassObj(x=1, y='a') + + +def test_wrap_pydantic_model(): + """A pydantic model round-trips through ``model_dump`` and ``model_validate``.""" + node = JsonableData(PydanticObj(a=2, b='b')).store() + loaded = load_node(node.pk) + assert loaded.obj == PydanticObj(a=2, b='b') + + +def test_numpy_values_are_coerced(): + """Numpy scalars and arrays in the serialized dictionary are stored as JSON-native types.""" + obj = JsonableClass({'arr': numpy.array([1, 2, 3]), 'scalar': numpy.int64(5)}) + node = JsonableData(obj).store() + loaded = load_node(node.pk) + assert loaded.obj.data == {'arr': [1, 2, 3], 'scalar': 5} From bd6d34819acc60a4bc1dc0eeb087081de58c0572 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Tue, 21 Jul 2026 22:12:05 +0200 Subject: [PATCH 08/19] =?UTF-8?q?=E2=9C=A8=20`serializer`:=20serialize=20v?= =?UTF-8?q?alues=20into=20data=20nodes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `aiida/orm/nodes/data/serializer.py`, the generic service that turns an arbitrary Python value into an AiiDA data node. `general_serializer` dispatches in three layers: existing nodes and `AttributeDict` namespaces pass through unchanged; core-owned value types (the scalars, list, dict, numpy, enum and `None`) go through `to_aiida_type`; foreign types are resolved through an `aiida.data` entry-point registry keyed by `module.ClassName`; and anything JSON-able falls back to `JsonableData`, raising an actionable `ValueError` if none applies. `serialize_to_aiida_nodes` maps it over a dict. Both are exported from `aiida.orm`. The registry (`get_serializers`) is built lazily and cached, so importing `aiida.orm` triggers no entry-point scan, and dropping the value type's built-in mappings avoids duplicating `to_aiida_type` (which already covers them). Custom serializers are supplied through the `serializers` argument rather than a config file. Increment 3 of the serializer reconcile (step 5a): this is the moved aiida-pythonjob serializer, rebuilt on core's existing `to_aiida_type` and `JsonableData` instead of carrying its own copies, and it is what the WorkGraph engine and function-based calculations will serialize with. --- src/aiida/orm/__init__.py | 2 + src/aiida/orm/nodes/__init__.py | 2 + src/aiida/orm/nodes/data/__init__.py | 3 + src/aiida/orm/nodes/data/serializer.py | 149 ++++++++++++++++++++++++ tests/orm/nodes/data/test_serializer.py | 107 +++++++++++++++++ 5 files changed, 263 insertions(+) create mode 100644 src/aiida/orm/nodes/data/serializer.py create mode 100644 tests/orm/nodes/data/test_serializer.py diff --git a/src/aiida/orm/__init__.py b/src/aiida/orm/__init__.py index 780fad23db..32e97fd428 100644 --- a/src/aiida/orm/__init__.py +++ b/src/aiida/orm/__init__.py @@ -110,6 +110,7 @@ 'XyData', 'cif_from_ase', 'find_bandgap', + 'general_serializer', 'get_loader', 'get_query_type_from_type_string', 'get_type_string_from_class', @@ -121,6 +122,7 @@ 'load_node', 'load_node_class', 'pycifrw_from_cif', + 'serialize_to_aiida_nodes', 'to_aiida_type', 'validate_link', ) diff --git a/src/aiida/orm/nodes/__init__.py b/src/aiida/orm/nodes/__init__.py index f4ff37f640..8619c5fb69 100644 --- a/src/aiida/orm/nodes/__init__.py +++ b/src/aiida/orm/nodes/__init__.py @@ -68,8 +68,10 @@ 'XyData', 'cif_from_ase', 'find_bandgap', + 'general_serializer', 'has_pycifrw', 'pycifrw_from_cif', + 'serialize_to_aiida_nodes', 'to_aiida_type', ) diff --git a/src/aiida/orm/nodes/data/__init__.py b/src/aiida/orm/nodes/data/__init__.py index 3ad38aa1a9..505682eb5d 100644 --- a/src/aiida/orm/nodes/data/__init__.py +++ b/src/aiida/orm/nodes/data/__init__.py @@ -29,6 +29,7 @@ from .numeric import * from .orbital import * from .remote import * +from .serializer import * from .singlefile import * from .str import * from .structure import * @@ -73,8 +74,10 @@ 'XyData', 'cif_from_ase', 'find_bandgap', + 'general_serializer', 'has_pycifrw', 'pycifrw_from_cif', + 'serialize_to_aiida_nodes', 'to_aiida_type', ) diff --git a/src/aiida/orm/nodes/data/serializer.py b/src/aiida/orm/nodes/data/serializer.py new file mode 100644 index 0000000000..adee903ee3 --- /dev/null +++ b/src/aiida/orm/nodes/data/serializer.py @@ -0,0 +1,149 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Serialize arbitrary Python values into AiiDA data nodes. + +This is the generic value-to-node serialization service shared by function-based calculations and the WorkGraph +engine. It layers three strategies: + +1. :func:`~aiida.orm.nodes.data.base.to_aiida_type` for the value types core owns (the ``BaseType`` scalars, lists, + dicts, numpy scalars and arrays, enums and ``None``), dispatched by type. +2. an entry-point-driven registry for types core does not own, so a plugin can serialize a type it does not import. +3. :class:`~aiida.orm.nodes.data.jsonable.JsonableData` as a JSON-able last resort. +""" + +from __future__ import annotations + +import functools +import typing as t +from importlib import import_module + +from .base import to_aiida_type +from .jsonable import JsonableData + +if t.TYPE_CHECKING: + from aiida.orm import User + +__all__ = ('general_serializer', 'serialize_to_aiida_nodes') + + +def import_from_path(path: str) -> t.Any: + """Import and return the object referenced by a ``module.qualname`` path. + + :param path: dotted path whose last segment is the attribute name, e.g. ``aiida.orm.nodes.data.list.List``. + :raises AttributeError: if the attribute does not exist in the resolved module. + """ + module_name, object_name = path.rsplit('.', 1) + module = import_module(module_name) + try: + return getattr(module, object_name) + except AttributeError as exc: + msg = f'`{object_name}` not found in module `{module_name}`.' + raise AttributeError(msg) from exc + + +@functools.cache +def get_serializers() -> dict[str, str]: + """Return the ``{type_key: import_path}`` serializer registry built from the ``aiida.data`` entry points. + + The entry-point *name* encodes the value type it serializes: the first dotted segment is a namespace and the rest + is the value's ``f'{type.__module__}.{type.__name__}'`` key. For example ``core.none`` registered as + ``pythonjob.builtins.NoneType`` yields the key ``builtins.NoneType``. Names without a dot after the namespace are + plain class registrations (``core.dict``, ``core.array``), not type serializers, and are skipped; hierarchical + core names (``core.array.bands``) yield keys that match no real value type and are harmless. + + Built-in and numpy value types are intentionally absent here: they are served by :func:`to_aiida_type`. The result + is cached; on the rare competing registration for one key the lexicographically first import path wins so the + choice is deterministic. Pass an explicit ``serializers`` mapping to :func:`general_serializer` to override. + """ + from aiida.plugins.entry_point import get_entry_points + + grouped: dict[str, list[str]] = {} + for entry_point in get_entry_points('aiida.data'): + key = entry_point.name.split('.', 1)[-1] + if '.' not in key: + continue + grouped.setdefault(key, []).append(entry_point.value.replace(':', '.')) + + return {key: sorted(paths)[0] for key, paths in grouped.items()} + + +def general_serializer( + data: t.Any, + serializers: dict[str, str] | None = None, + store: bool = True, + user: User | None = None, +) -> t.Any: + """Serialize a single Python value to an AiiDA data node. + + Existing nodes and :class:`~aiida.common.extendeddicts.AttributeDict` namespaces are returned unchanged. Otherwise + the value is converted through :func:`to_aiida_type` (core-owned types), then the entry-point registry (foreign + types), then :class:`JsonableData` (JSON-able objects), raising :class:`ValueError` with guidance if none applies. + + :param data: the value to serialize. + :param serializers: optional ``{type_key: import_path}`` override; defaults to :func:`get_serializers`. + :param store: whether to store the created node (an already-existing node is never re-stored). + :param user: the user to assign to a newly created node. + :raises ValueError: if the value cannot be serialized by any strategy. + """ + from aiida.common.extendeddicts import AttributeDict + from aiida.orm import Node + + if serializers is None: + serializers = get_serializers() + + if isinstance(data, Node): + return data + if isinstance(data, AttributeDict): + return data + + # Core-owned value types (scalars, list, dict, numpy, enum, None) dispatch by type through ``to_aiida_type``. + try: + node = to_aiida_type(data) + except TypeError: + node = None + if node is not None: + if store: + node.store() + return node + + # Foreign / plugin-owned types resolved by their ``aiida.data`` entry-point registration. + type_key = f'{type(data).__module__}.{type(data).__name__}' + if type_key in serializers: + try: + serializer = import_from_path(serializers[type_key]) + new_node = serializer(data, user=user) + except Exception as exc: + # A plugin-provided serializer may raise anything; wrap it with context about which registration failed. + msg = f'error serializing `{type_key}` with `{serializers[type_key]}`: {exc}' + raise ValueError(msg) from exc + if store: + new_node.store() + return new_node + + # Last resort: wrap any JSON-representable object. + try: + node = JsonableData(data, user=user) + except (TypeError, ValueError) as exc: + msg = ( + f'cannot serialize the object of type `{type_key}`.\n' + 'To fix this, either:\n' + ' 1. register a type-specific `aiida.orm.Data` subclass as an `aiida.data` entry point, or\n' + ' 2. make the class JSON-able for `JsonableData` (an `as_dict`/`to_dict` method plus a `from_dict`\n' + ' class method, or make it a dataclass or pydantic model), or\n' + " 3. pass an ad-hoc serializer through `serializers`, e.g. {'my_pkg.MyType': 'my_pkg:to_aiida_node'}." + ) + raise ValueError(msg) from exc + if store: + node.store() + return node + + +def serialize_to_aiida_nodes(inputs: dict[str, t.Any], serializers: dict[str, str] | None = None) -> dict[str, t.Any]: + """Serialize each value of a mapping to an AiiDA data node with :func:`general_serializer`.""" + return {key: general_serializer(value, serializers=serializers) for key, value in inputs.items()} diff --git a/tests/orm/nodes/data/test_serializer.py b/tests/orm/nodes/data/test_serializer.py new file mode 100644 index 0000000000..6252fc5a25 --- /dev/null +++ b/tests/orm/nodes/data/test_serializer.py @@ -0,0 +1,107 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Tests for :mod:`aiida.orm.nodes.data.serializer`.""" + +import numpy +import pytest + +from aiida import orm +from aiida.common.extendeddicts import AttributeDict +from aiida.orm.nodes.data.serializer import general_serializer, get_serializers, serialize_to_aiida_nodes + + +@pytest.mark.parametrize( + 'value, expected_type', + ( + (5, orm.Int), + (5.0, orm.Float), + ('x', orm.Str), + (True, orm.Bool), + ([1, 2], orm.List), + ({'a': 1}, orm.Dict), + (numpy.int64(3), orm.Int), + (numpy.array([1, 2]), orm.ArrayData), + (None, orm.NoneData), + ), +) +def test_core_owned_types_go_through_to_aiida_type(value, expected_type): + """Core-owned value types are dispatched by :func:`to_aiida_type` and stored.""" + node = general_serializer(value) + assert isinstance(node, expected_type) + assert node.is_stored is True + + +def test_existing_node_is_returned_unchanged(): + node = orm.Int(1).store() + assert general_serializer(node) is node + + +def test_attribute_dict_is_returned_unchanged(): + namespace = AttributeDict({'a': 1}) + assert general_serializer(namespace) is namespace + + +def test_store_false_does_not_store(): + node = general_serializer(5, store=False) + assert isinstance(node, orm.Int) + assert node.is_stored is False + + +class JsonableClass: + """Object exposing the ``as_dict`` / ``from_dict`` contract for the JSON fallback.""" + + def __init__(self, value): + self.value = value + + def as_dict(self): + return {'value': self.value} + + @classmethod + def from_dict(cls, dictionary): + return cls(dictionary['value']) + + +def test_json_able_object_falls_back_to_jsonable_data(): + node = general_serializer(JsonableClass(7)) + assert isinstance(node, orm.JsonableData) + assert node.obj.value == 7 + + +def test_unserializable_object_raises_with_guidance(): + class Opaque: + pass + + with pytest.raises(ValueError, match=r'cannot serialize the object of type `.*Opaque`'): + general_serializer(Opaque()) + + +def tuple_to_list(data, user=None): + """A serializer following the ``serializer(value, user=...)`` calling convention, used by the registry test.""" + return orm.List(list=list(data)) + + +def test_explicit_serializers_mapping_handles_foreign_type(): + """The registry branch: a type ``to_aiida_type`` does not own is resolved via the ``serializers`` mapping.""" + node = general_serializer((1, 2, 3), serializers={'builtins.tuple': f'{__name__}.tuple_to_list'}) + assert isinstance(node, orm.List) + assert node.get_list() == [1, 2, 3] + + +def test_serialize_to_aiida_nodes_maps_each_value(): + result = serialize_to_aiida_nodes({'i': 1, 's': 'x'}) + assert isinstance(result['i'], orm.Int) + assert isinstance(result['s'], orm.Str) + + +def test_get_serializers_is_cached_and_skips_non_type_keys(): + registry = get_serializers() + assert get_serializers() is registry + # ``core.dict`` yields the dot-less key ``dict`` and is skipped; only dotted (type-key-shaped) names are kept. + assert 'dict' not in registry + assert all('.' in key for key in registry) From 2304388c755cf9c225dce3777eb34ae52858535f Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Wed, 22 Jul 2026 10:39:35 +0200 Subject: [PATCH 09/19] =?UTF-8?q?=E2=9C=A8=20`workgraph`:=20node-graph=20d?= =?UTF-8?q?ep=20+=20serialize=5Fports?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Take `node-graph` as a hard aiida-core dependency and add the first subsystem module that uses it: `aiida/workgraph/serialization.py` with `serialize_ports`, which walks a `node_graph.SocketSpec` schema and serializes each leaf through `aiida.orm.general_serializer` (namespaces recurse, dynamic namespaces accept extra keys, metadata passes through, nodes are left unstored for the caller to store). It lives in the WorkGraph subsystem rather than in `aiida.orm` precisely because it imports node-graph: `aiida.orm` must stay node-graph-free so the base import path does not require it. A plain `import aiida` still pulls no node-graph; only importing `aiida.workgraph` does, which is acceptable now that node-graph is a hard dependency (the escape hatch of an optional `aiida-core[workgraph]` extra stays open for later). Adds `node-graph~=0.6.5` to the dependencies and the generated conda environment, registers `node_graph` as untyped for mypy, and refreshes the lock. cloudpickle enters the lock transitively through node-graph. Increment 4 of the serializer reconcile (step 5a): the node-graph-coupled half of the moved serializer, completing the core-side stack. Next the plugins repoint at it (aiida-pythonjob, then the WorkGraph engine). --- environment.yml | 1 + pyproject.toml | 2 + src/aiida/workgraph/__init__.py | 2 + src/aiida/workgraph/serialization.py | 103 +++++++++++++++++++++++ tests/workgraph/test_serialization.py | 52 ++++++++++++ uv.lock | 117 ++++++++++++++++++++++++++ 6 files changed, 277 insertions(+) create mode 100644 src/aiida/workgraph/serialization.py create mode 100644 tests/workgraph/test_serialization.py diff --git a/environment.yml b/environment.yml index 654095cb93..5d441e1e3a 100644 --- a/environment.yml +++ b/environment.yml @@ -21,6 +21,7 @@ dependencies: - jedi<0.19 - jinja2~=3.0 - kiwipy[rmq]~=0.9.0 +- node-graph~=0.6.5 - numpy<3,>=1.21 - paramiko~=3.0 - pgsu~=0.3.0 diff --git a/pyproject.toml b/pyproject.toml index 515895fcb7..8df0e8d060 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ 'jedi<0.19', 'jinja2~=3.0', 'kiwipy[rmq]~=0.9.0', + 'node-graph~=0.6.5', 'numpy>=1.21,<3', 'paramiko~=3.0', 'pgsu~=0.3.0', @@ -412,6 +413,7 @@ module = [ 'graphviz.*', 'kiwipy.*', 'mayavi.*', + 'node_graph.*', 'pgsu.*', 'pgtest.*', 'trogon.*', diff --git a/src/aiida/workgraph/__init__.py b/src/aiida/workgraph/__init__.py index cac16f23e8..e6e9290dbb 100644 --- a/src/aiida/workgraph/__init__.py +++ b/src/aiida/workgraph/__init__.py @@ -18,6 +18,7 @@ # fmt: off from .enums import * +from .serialization import * from .utils import * __all__ = ( @@ -28,6 +29,7 @@ 'TaskState', 'get_nested_dict', 'resolve_node_link_managers', + 'serialize_ports', 'update_nested_dict', 'update_nested_dict_with_special_keys', ) diff --git a/src/aiida/workgraph/serialization.py b/src/aiida/workgraph/serialization.py new file mode 100644 index 0000000000..1b2bd44a81 --- /dev/null +++ b/src/aiida/workgraph/serialization.py @@ -0,0 +1,103 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Schema-driven serialization of raw Python data into AiiDA data nodes. + +This is the node-graph-coupled layer of the serialization stack: it walks a :class:`node_graph.socket_spec.SocketSpec` +schema and serializes each leaf with :func:`~aiida.orm.general_serializer`, which is why it lives in the WorkGraph +subsystem (which may depend on node-graph) rather than in ``aiida.orm`` (which may not). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from node_graph.socket_meta import SocketMeta +from node_graph.socket_spec import SocketSpec +from node_graph.utils.struct_utils import is_structured_instance, structured_to_dict + +from aiida.orm import general_serializer + +if TYPE_CHECKING: + from aiida.orm import User + +__all__ = ('serialize_ports',) + + +def _ensure_spec(schema: SocketSpec | dict[str, Any]) -> SocketSpec: + """Return ``schema`` as a :class:`SocketSpec`, building one from a dict if needed.""" + if isinstance(schema, SocketSpec): + return schema + if isinstance(schema, dict): + return SocketSpec.from_dict(schema) + msg = f'unsupported schema type: {type(schema)}' + raise TypeError(msg) + + +def serialize_ports( + python_data: Any, + port_schema: SocketSpec | dict[str, Any], + serializers: dict[str, str] | None = None, + user: User | None = None, +) -> Any: + """Serialize raw Python data into AiiDA data nodes following a :class:`SocketSpec` schema. + + A namespace spec is walked recursively, serializing each declared field and (for a ``dynamic`` namespace) any extra + keys; metadata fields are passed through untouched. A leaf spec serializes the value directly. The produced nodes + are not stored (the caller stores them as part of the process inputs). + + :param python_data: the raw value or nested mapping to serialize. + :param port_schema: the socket spec describing the expected structure, or its dict form. + :param serializers: optional ``{type_key: import_path}`` override forwarded to :func:`general_serializer`. + :param user: the user to assign to newly created nodes. + :raises ValueError: if a namespace value is not a mapping, or holds a key the (non-dynamic) namespace does not + declare. + """ + spec = _ensure_spec(port_schema) + + # Leaf: serialize the value directly. + if not spec.is_namespace(): + return general_serializer(python_data, serializers=serializers, store=False, user=user) + + name = getattr(spec.meta, 'help', None) or '' + if is_structured_instance(python_data): + python_data = structured_to_dict(python_data) + if not isinstance(python_data, dict): + msg = f"expected a mapping for namespace '{name}', got {type(python_data)}" + raise ValueError(msg) + + out: dict[str, Any] = {} + fields = spec.fields or {} + + for key, value in python_data.items(): + if key in fields: + child_spec = fields[key] + if child_spec.meta.is_metadata: + # Metadata is not serialized; it is carried through verbatim. + out[key] = value + elif child_spec.is_namespace(): + out[key] = serialize_ports(value, child_spec, serializers=serializers, user=user) + else: + out[key] = general_serializer(value, serializers=serializers, store=False, user=user) + elif spec.meta.dynamic: + # Extra keys are allowed in a dynamic namespace; ``item`` (if given) types them, else ANY. + if spec.item is None: + if isinstance(value, dict): + item = SocketSpec(identifier='node_graph.namespace', meta=SocketMeta(dynamic=True)) + out[key] = serialize_ports(value, item, serializers=serializers, user=user) + else: + out[key] = general_serializer(value, serializers=serializers, store=False, user=user) + elif spec.item.is_namespace(): + out[key] = serialize_ports(value, spec.item, serializers=serializers, user=user) + else: + out[key] = general_serializer(value, serializers=serializers, store=False, user=user) + else: + msg = f"unexpected key '{key}' for namespace '{name}' (not dynamic)." + raise ValueError(msg) + + return out diff --git a/tests/workgraph/test_serialization.py b/tests/workgraph/test_serialization.py new file mode 100644 index 0000000000..84f70db0cf --- /dev/null +++ b/tests/workgraph/test_serialization.py @@ -0,0 +1,52 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Tests for :func:`aiida.workgraph.serialization.serialize_ports`.""" + +import pytest +from node_graph.socket_spec import SocketSpec, dynamic, namespace + +from aiida import orm +from aiida.workgraph import serialize_ports + + +def test_leaf_spec_serializes_the_value(): + """A non-namespace spec serializes the value directly and does not store it.""" + node = serialize_ports(5, SocketSpec(identifier='node_graph.int')) + assert isinstance(node, orm.Int) + assert node.value == 5 + assert node.is_stored is False + + +def test_namespace_serializes_each_declared_field(): + result = serialize_ports({'x': 1, 'y': 'a'}, namespace(x=int, y=str)) + assert isinstance(result['x'], orm.Int) + assert isinstance(result['y'], orm.Str) + assert (result['x'].value, result['y'].value) == (1, 'a') + + +def test_nested_namespace_recurses(): + result = serialize_ports({'inner': {'x': 1}}, namespace(inner=namespace(x=int))) + assert isinstance(result['inner']['x'], orm.Int) + assert result['inner']['x'].value == 1 + + +def test_dynamic_namespace_serializes_arbitrary_keys(): + result = serialize_ports({'a': 1, 'b': 2}, dynamic(int)) + assert set(result) == {'a', 'b'} + assert all(isinstance(node, orm.Int) for node in result.values()) + + +def test_non_mapping_for_namespace_raises(): + with pytest.raises(ValueError, match='expected a mapping'): + serialize_ports(5, namespace(x=int)) + + +def test_unexpected_key_in_static_namespace_raises(): + with pytest.raises(ValueError, match="unexpected key 'z'"): + serialize_ports({'z': 1}, namespace(x=int)) diff --git a/uv.lock b/uv.lock index ed9a4c3f7c..82fe8ba133 100644 --- a/uv.lock +++ b/uv.lock @@ -46,6 +46,7 @@ dependencies = [ { name = "jedi" }, { name = "jinja2" }, { name = "kiwipy", extra = ["rmq"] }, + { name = "node-graph" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "paramiko" }, @@ -228,6 +229,7 @@ requires-dist = [ { name = "myst-nb", marker = "extra == 'docs'", specifier = "~=1.0.0" }, { name = "nbclient", marker = "extra == 'tests'", specifier = "~=0.10" }, { name = "nbformat", marker = "extra == 'tests'", specifier = "~=5.10" }, + { name = "node-graph", specifier = "~=0.6.5" }, { name = "notebook", marker = "extra == 'notebook'", specifier = "~=6.1,>=6.1.5" }, { name = "numpy", specifier = ">=1.21,<3" }, { name = "packaging", marker = "extra == 'pre-commit'", specifier = "~=23.0" }, @@ -394,6 +396,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/b3/9b1a8074496371342ec1e796a96f99c82c945a339cd81a8e73de28b4cf9e/anyio-4.11.0-py3-none-any.whl", hash = "sha256:0287e96f4d26d4149305414d4e3bc32f0dcd0862365a4bddea19d7a1ec38c4fc", size = 109097, upload-time = "2025-09-23T09:19:10.601Z" }, ] +[[package]] +name = "anywidget" +version = "0.11.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ipywidgets" }, + { name = "psygnal" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/79/31/0491d707c674b34267f55d96d6a7148e55e7b6718a271686232cf295fbe2/anywidget-0.11.0.tar.gz", hash = "sha256:6695fbef9449cf8c27f421b96c5837aa37f909ec1f60cfa33add333e1b70b169", size = 426999, upload-time = "2026-04-27T23:42:09.576Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8e/c2/8fec8e8e2eb920cc2280f569144080cd58622a2eda83bfa4c0c354a63264/anywidget-0.11.0-py3-none-any.whl", hash = "sha256:c574d9acc6503ad27b37a9acea48f957a8ba7c9c9876cfcb37898931c098ce9d", size = 317341, upload-time = "2026-04-27T23:42:08.356Z" }, +] + [[package]] name = "appnope" version = "0.1.4" @@ -974,6 +990,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/93/2a/04893832bfeddc2d40a7de2e8153b3085f12d63507d91a9cf0157dc3a1c2/click_spinner-0.1.10-py2.py3-none-any.whl", hash = "sha256:d1ffcff1fdad9882396367f15fb957bcf7f5c64ab91927dee2127e0d2991ee84", size = 3986, upload-time = "2020-04-24T07:14:50.575Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -2040,6 +2065,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/6d/0d9848617b9f753b87f214f1c682592f7ca42de085f564352f10f0843026/ipywidgets-8.1.8-py3-none-any.whl", hash = "sha256:ecaca67aed704a338f88f67b1181b58f821ab5dc89c1f0f5ef99db43c1c2921e", size = 139808, upload-time = "2025-11-01T21:18:10.956Z" }, ] +[[package]] +name = "isodate" +version = "0.7.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/54/4d/e940025e2ce31a8ce1202635910747e5a87cc3a6a6bb2d00973375014749/isodate-0.7.2.tar.gz", hash = "sha256:4cd1aa0f43ca76f4a6c6c0292a85f40b35ec2e43e315b59f06e6d32171a953e6", size = 29705, upload-time = "2024-10-08T23:04:11.5Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/aa/0aca39a37d3c7eb941ba736ede56d689e7be91cab5d9ca846bde3999eba6/isodate-0.7.2-py3-none-any.whl", hash = "sha256:28009937d8031054830160fce6d409ed342816b543597cece116d966c6d99e15", size = 22320, upload-time = "2024-10-08T23:04:09.501Z" }, +] + [[package]] name = "isoduration" version = "20.11.0" @@ -3244,6 +3278,42 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/eb/8d/776adee7bbf76365fdd7f2552710282c79a4ead5d2a46408c9043a2b70ba/networkx-3.5-py3-none-any.whl", hash = "sha256:0030d386a9a06dee3565298b4a734b68589749a544acbb6c412dc9e2489ec6ec", size = 2034406, upload-time = "2025-05-29T11:35:04.961Z" }, ] +[[package]] +name = "node-graph" +version = "0.6.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "cloudpickle" }, + { name = "graphviz" }, + { name = "node-graph-widget" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "rdflib" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "typing-extensions" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/fa/f64cb06871fc56c43f34f799b68f84321451e0f821ccc5b9a76a8207d92a/node_graph-0.6.5.tar.gz", hash = "sha256:25114f96e3154aeccac8faf75cde16e8acaf772e0cc34857d446c1145ae30cde", size = 100369, upload-time = "2026-01-26T15:18:44.682Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/5a/fee76af3ad82de6ce8708846bc1c6217792cde04f8acc23991c8dc496731/node_graph-0.6.5-py3-none-any.whl", hash = "sha256:4efe5acd8900f9cf872b90f3bc0afccd63651c3bf7705bd6edfbd55e1dd2c471", size = 117763, upload-time = "2026-01-26T15:18:43.618Z" }, +] + +[[package]] +name = "node-graph-widget" +version = "0.0.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anywidget" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/e8/eec2320fcab1ac1861e1b031ba1d1545ef25c1607dd45d0c677535f0d098/node_graph_widget-0.0.5.tar.gz", hash = "sha256:6ce74d2e15faeb216cf8c79413b7140b64505f337c7f505d3f66a03927ed6841", size = 600443, upload-time = "2025-05-07T11:08:46.025Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/00/a754f8375b1470d0b058e0e53960843d2989d3ceb9bb8f62db12f4f8634f/node_graph_widget-0.0.5-py3-none-any.whl", hash = "sha256:90876cdf5f48f52c86c179929cefb83cc336af5cdae972df52ddd350e34b6606", size = 559939, upload-time = "2025-05-07T11:08:44.069Z" }, +] + [[package]] name = "nodeenv" version = "1.9.1" @@ -4131,6 +4201,40 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/53/cf/10c3e95827a3ca8af332dfc471befec86e15a14dc83cee893c49a4910dad/psycopg_binary-3.2.12-cp314-cp314-win_amd64.whl", hash = "sha256:48a8e29f3e38fcf8d393b8fe460d83e39c107ad7e5e61cd3858a7569e0554a39", size = 3005787, upload-time = "2025-10-26T00:36:06.783Z" }, ] +[[package]] +name = "psygnal" +version = "0.15.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/79/20c3e23e75272e9ddf018097cf872ab088bccba978888472656629efa4a3/psygnal-0.15.1.tar.gz", hash = "sha256:f64f62dee2306fc1c22050a59b6c6cdad126e04b0cf50e393ff858a1da719096", size = 123147, upload-time = "2026-01-04T16:38:41.959Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/44/ab13cb6147d010258826a43e574ad94599af0de29df13795fff9efee656c/psygnal-0.15.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8ee55e3997f796fd84d4fdbd829bb1b19d323e087c43d072744604a3016c8851", size = 587322, upload-time = "2026-01-04T16:38:04.827Z" }, + { url = "https://files.pythonhosted.org/packages/f2/a2/68c042a607ca613e9450dfee99cc5c2a4d10d95392fb1de2ba932dd0a605/psygnal-0.15.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:912bcf110bfe7b4aa121d24987b6a58afb967ff090a049dad136eaf3cbcc7bea", size = 576207, upload-time = "2026-01-04T16:38:06.183Z" }, + { url = "https://files.pythonhosted.org/packages/4b/86/123c7b169ad32994a0cd801cd1f11c1a2be84555807e9c8a8a4682c67a9f/psygnal-0.15.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b2e860c11fe075fd80c93a24081c577ef7ec5c9da41f0e75990aa4cccf3f79cf", size = 864261, upload-time = "2026-01-04T16:38:07.895Z" }, + { url = "https://files.pythonhosted.org/packages/20/f1/886cec7bec2f27fe453cfa32bfcaac08a83aab2a04895af68f93e1c493b8/psygnal-0.15.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:d5b8bebcf99699ef50b6ef572868a490f6d191dc4466e5bd9818ca27e17cd581", size = 872582, upload-time = "2026-01-04T16:38:09.745Z" }, + { url = "https://files.pythonhosted.org/packages/21/a3/da972a05568ee8a9dc6c6567bee2c0cc5af8c681baebcb9fdbbf3cceb4f7/psygnal-0.15.1-cp310-cp310-win_amd64.whl", hash = "sha256:06e0a90490e1205620d97ac52fbbe3282a22b126a26d02b3e1196bb46de16c7a", size = 411043, upload-time = "2026-01-04T16:38:11.588Z" }, + { url = "https://files.pythonhosted.org/packages/bb/a7/69495410025cc4298765545ce3b8c635cd4c8d3a362b7fbbc15b80e9fc8f/psygnal-0.15.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1adc41515f648696990964433f1e25d8dfd306813a3645366c85e01986ba57a0", size = 581002, upload-time = "2026-01-04T16:38:12.753Z" }, + { url = "https://files.pythonhosted.org/packages/75/1f/19a8126ccf3cd3974ba5d08a435a049b666961d90f5848ba83599d7a29de/psygnal-0.15.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:38ff18455b2ac73d4e8eea82ef298ce904b52e4dfdc603a24380c9c440e37519", size = 567775, upload-time = "2026-01-04T16:38:14.04Z" }, + { url = "https://files.pythonhosted.org/packages/54/c5/b1348880d603edb82128a721397a1ddcf3dfcf5384fe5689db6e471118ae/psygnal-0.15.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c923c322eeefb1140886927cfe7bda7c32341087e290e812b9c69a624ab72d54", size = 855961, upload-time = "2026-01-04T16:38:15.612Z" }, + { url = "https://files.pythonhosted.org/packages/e6/42/3da2d6f3583bd1a849f7faa2fd3492b14bfda05012519ceaea5992658af0/psygnal-0.15.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:2714ddaa41ea3134c0ee91cebd5fb11a88f254ea1d5948806ab0ad5f8be603d5", size = 862721, upload-time = "2026-01-04T16:38:17.059Z" }, + { url = "https://files.pythonhosted.org/packages/4d/14/6fc7e97fdecf7e8c5c105684bab784920312a3259800d8b53e3cf8783f42/psygnal-0.15.1-cp311-cp311-win_amd64.whl", hash = "sha256:877516056a5a383427a647fff2fad5179eaa3e12de2c083c273e748435414aef", size = 415696, upload-time = "2026-01-04T16:38:18.355Z" }, + { url = "https://files.pythonhosted.org/packages/76/65/b7bbca96bc477aa9ac2264e5907b2f4ccfcd1319f776dd1f35eec06cc2f4/psygnal-0.15.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8d56f0f35eaf4a21f660de76885222faf9e8c7112454528d3394d464f3d4d1a3", size = 598340, upload-time = "2026-01-04T16:38:19.752Z" }, + { url = "https://files.pythonhosted.org/packages/40/f2/56577465a1b42a5e6780bb5fab53fb68f8bfd72f0131ed397576529af724/psygnal-0.15.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0febcf757a1323d9b8bd75735ee3569213d8110012a7bf0f478e85c5ab459fc6", size = 575311, upload-time = "2026-01-04T16:38:21.137Z" }, + { url = "https://files.pythonhosted.org/packages/79/81/f642ac08104049383076f83480ed412c9626e068769a1c34873c595bec0e/psygnal-0.15.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b5e4837dfbfa4974dabe0795e32be9aadcd87603adf734738ce1114f72238a05", size = 889770, upload-time = "2026-01-04T16:38:22.629Z" }, + { url = "https://files.pythonhosted.org/packages/de/43/e571fa40b72780abed080ef829e5ad98017b6fe48d28c15a2404e006b676/psygnal-0.15.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:07b4c4e03bbf4e8cad7e25f4fbc1ba9575fb9c3d14991bc7edfeb8b09c8d6d54", size = 881105, upload-time = "2026-01-04T16:38:23.896Z" }, + { url = "https://files.pythonhosted.org/packages/e3/26/ef3ab825eb08eaecbbceeeb56383694fe64ce399dbfd1d0767bb85688785/psygnal-0.15.1-cp312-cp312-win_amd64.whl", hash = "sha256:4f0ce91b9c18e92281bf2c3fc4bb4e808d90f0b023d0a37b302d354188520338", size = 418969, upload-time = "2026-01-04T16:38:25.731Z" }, + { url = "https://files.pythonhosted.org/packages/46/21/5a142165d27063abf5921807d3c3d973f5d44ab414a13b210839a43ead4d/psygnal-0.15.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:2087aadc9404f007f79c2899e329932869e362c50de58b90631c5f49b4768cc5", size = 596768, upload-time = "2026-01-04T16:38:27.053Z" }, + { url = "https://files.pythonhosted.org/packages/e1/25/c1712931d61c118691e73daf29ef708c679ea9ba187c797dd5deee360411/psygnal-0.15.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0f3bf68ca42569dfdce20c6cf915d34b78b9e3ddddacb9f78728224fda6946b4", size = 574808, upload-time = "2026-01-04T16:38:28.779Z" }, + { url = "https://files.pythonhosted.org/packages/2d/4f/3593e5adb88a188c798604aed95fbc1479f30230e7f51e8f2c770e6a3832/psygnal-0.15.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e9fca977f5335deea39aed22e31d9795983e4f243e59a7d3c4105793adb7693d", size = 885616, upload-time = "2026-01-04T16:38:30.081Z" }, + { url = "https://files.pythonhosted.org/packages/58/4c/14779ed4c3a1d71fa1a9a87ecfb184ad3335dd64681067f77c1c47b14ae9/psygnal-0.15.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:0c85b7d05b92ccbec47c75ab8a5545eda462e81a492c82424aba5ab81a3ad89d", size = 876516, upload-time = "2026-01-04T16:38:31.422Z" }, + { url = "https://files.pythonhosted.org/packages/3e/bc/4f771e3cdcde4db4023dbf36d6f0aab44e02b9de719353c22954b655e2ff/psygnal-0.15.1-cp313-cp313-win_amd64.whl", hash = "sha256:ac0e693b29e0a429e97315a52313321855bef6140e9975b7ae78b4d93c8fbb42", size = 419172, upload-time = "2026-01-04T16:38:32.82Z" }, + { url = "https://files.pythonhosted.org/packages/f4/2e/975bd61727578d88df62797f78390965ca7905780cf01eb59cb095a13638/psygnal-0.15.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:803fc33c4280c822c6f4b22e6c3ea7c4483e190f3cc69e69350098b3799476f3", size = 595706, upload-time = "2026-01-04T16:38:34.139Z" }, + { url = "https://files.pythonhosted.org/packages/b8/55/e487f1d91497eb75e86c3fdfef69a21b1cab24d023383dd7648b08797d6a/psygnal-0.15.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4f53b4b83355b0a785b745987fd04e59bbf169a9028ed81a68ca7e05fb76d458", size = 575133, upload-time = "2026-01-04T16:38:35.448Z" }, + { url = "https://files.pythonhosted.org/packages/bf/2f/f286355accd0e68d3eef52e63c8b9ab6ba33ec3107177719a036b3319657/psygnal-0.15.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bcbca12190f5aa65c1f8fb04a81fa6f4463c5f5dde25cd74c3a56ceff6f37b02", size = 889565, upload-time = "2026-01-04T16:38:37.003Z" }, + { url = "https://files.pythonhosted.org/packages/fc/dc/40c6026c88d7f9220ecc913afe0501045a512c9b82f9b7e036bb089dc287/psygnal-0.15.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:1ac399566852fe4354ce26a1acbe12319232e8c2b615fe5ad1e114c547095cf6", size = 880863, upload-time = "2026-01-04T16:38:38.381Z" }, + { url = "https://files.pythonhosted.org/packages/b7/85/b4f45ec3057c473b5622fc002b3a636a698c34d3a0917a064ff5247f1984/psygnal-0.15.1-cp314-cp314-win_amd64.whl", hash = "sha256:d3a03055f331ce91d44581c71edb79938ccc133a94af2ce7ad3a18fa57ac7be5", size = 423654, upload-time = "2026-01-04T16:38:39.7Z" }, + { url = "https://files.pythonhosted.org/packages/46/49/7742544684bee728ec123515d2694cee859aa2a705951a461230b00f18cc/psygnal-0.15.1-py3-none-any.whl", hash = "sha256:4221140e633e45b076953c64bcb9b41a744833527f9a037c1ca98bc270798cbf", size = 90638, upload-time = "2026-01-04T16:38:40.841Z" }, +] + [[package]] name = "ptyprocess" version = "0.7.0" @@ -4839,6 +4943,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/01/1b/5dbe84eefc86f48473947e2f41711aded97eecef1231f4558f1f02713c12/pyzmq-27.1.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c9f7f6e13dff2e44a6afeaf2cf54cee5929ad64afaf4d40b50f93c58fc687355", size = 544862, upload-time = "2025-09-08T23:09:56.509Z" }, ] +[[package]] +name = "rdflib" +version = "7.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "isodate", marker = "python_full_version < '3.11'" }, + { name = "pyparsing" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/98/f5/18bb77b7af9526add0c727a3b2048959847dc5fb030913e2918bf384fec3/rdflib-7.6.0.tar.gz", hash = "sha256:6c831288d5e4a5a7ece85d0ccde9877d512a3d0f02d7c06455d00d6d0ea379df", size = 4943826, upload-time = "2026-02-13T07:15:55.938Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/c2/6604a71269e0c1bd75656d5a001432d16f2cc5b8c057140ec797155c295e/rdflib-7.6.0-py3-none-any.whl", hash = "sha256:30c0a3ebf4c0e09215f066be7246794b6492e054e782d7ac2a34c9f70a15e0dd", size = 615416, upload-time = "2026-02-13T07:15:46.487Z" }, +] + [[package]] name = "referencing" version = "0.37.0" From 363a9ac45bbbad6bc8e70d61f9c808e0fde2bdd3 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Wed, 22 Jul 2026 11:17:30 +0200 Subject: [PATCH 10/19] =?UTF-8?q?=E2=9C=A8=20`orm`:=20DateTime/Function=20?= =?UTF-8?q?data=20+=20deserializer?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move part of the serializer data layer out of aiida-pythonjob and into aiida-core so the serialize/deserialize stack lives in one place: - `DateTimeData` (`datetime.py`) and `FunctionData` (`function.py`), stdlib-only, each registered with `to_aiida_type` (for `datetime.datetime` and `types.FunctionType`) so they serialize through the same singledispatch path as the other core value types. - `deserialize_to_raw_python_data` (`deserializer.py`), the inverse of `general_serializer`: a node with a `value` is read directly, otherwise a registry maps its type to a deserializer, and mappings recurse. Registered as the `core.datetime` / `core.function` entry points, exported from `aiida.orm`, with focused tests and `test_fields` fixtures. `PickledData` is deliberately NOT moved here: aiida-shell already registers a `core.pickled` entry point for its own pickle data type, so core claiming `core.pickled` would collide until aiida-shell is updated. Since both aiida-shell and aiida-pythonjob carry their own pickle data type, consolidating it (and taking a direct `cloudpickle` dependency) belongs to the coordinated aiida-shell fold, not here. aiida-pythonjob keeps its own `PickledData` (which registers no entry point) for now. Increment 5 of the serializer reconcile (step 5a). With these, core owns the datetime/function data types and the deserializer; aiida-pythonjob repoints at this stack next, keeping only `AtomsData` (ASE) and `PickledData`. --- pyproject.toml | 2 + src/aiida/orm/__init__.py | 3 + src/aiida/orm/nodes/__init__.py | 3 + src/aiida/orm/nodes/data/__init__.py | 6 + src/aiida/orm/nodes/data/datetime.py | 45 ++++++++ src/aiida/orm/nodes/data/deserializer.py | 104 ++++++++++++++++++ src/aiida/orm/nodes/data/function.py | 79 +++++++++++++ tests/orm/nodes/data/test_datetime.py | 39 +++++++ tests/orm/nodes/data/test_deserializer.py | 41 +++++++ tests/orm/nodes/data/test_function.py | 47 ++++++++ ..._aiida.data.core.datetime.DateTimeData.yml | 21 ++++ ..._aiida.data.core.function.FunctionData.yml | 21 ++++ 12 files changed, 411 insertions(+) create mode 100644 src/aiida/orm/nodes/data/datetime.py create mode 100644 src/aiida/orm/nodes/data/deserializer.py create mode 100644 src/aiida/orm/nodes/data/function.py create mode 100644 tests/orm/nodes/data/test_datetime.py create mode 100644 tests/orm/nodes/data/test_deserializer.py create mode 100644 tests/orm/nodes/data/test_function.py create mode 100644 tests/orm/test_fields/fields_aiida.data.core.datetime.DateTimeData.yml create mode 100644 tests/orm/test_fields/fields_aiida.data.core.function.FunctionData.yml diff --git a/pyproject.toml b/pyproject.toml index 8df0e8d060..896a8da3c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,10 +119,12 @@ requires-python = '>=3.10' 'core.code.containerized' = 'aiida.orm.nodes.data.code.containerized:ContainerizedCode' 'core.code.installed' = 'aiida.orm.nodes.data.code.installed:InstalledCode' 'core.code.portable' = 'aiida.orm.nodes.data.code.portable:PortableCode' +'core.datetime' = 'aiida.orm.nodes.data.datetime:DateTimeData' 'core.dict' = 'aiida.orm.nodes.data.dict:Dict' 'core.enum' = 'aiida.orm.nodes.data.enum:EnumData' 'core.float' = 'aiida.orm.nodes.data.float:Float' 'core.folder' = 'aiida.orm.nodes.data.folder:FolderData' +'core.function' = 'aiida.orm.nodes.data.function:FunctionData' 'core.int' = 'aiida.orm.nodes.data.int:Int' 'core.jsonable' = 'aiida.orm.nodes.data.jsonable:JsonableData' 'core.list' = 'aiida.orm.nodes.data.list:List' diff --git a/src/aiida/orm/__init__.py b/src/aiida/orm/__init__.py index 32e97fd428..f690a8f50c 100644 --- a/src/aiida/orm/__init__.py +++ b/src/aiida/orm/__init__.py @@ -52,6 +52,7 @@ 'ComputerEntityLoader', 'ContainerizedCode', 'Data', + 'DateTimeData', 'Dict', 'Entity', 'EntityExtras', @@ -59,6 +60,7 @@ 'EnumData', 'Float', 'FolderData', + 'FunctionData', 'Group', 'GroupEntityLoader', 'ImportGroup', @@ -109,6 +111,7 @@ 'WorkflowNode', 'XyData', 'cif_from_ase', + 'deserialize_to_raw_python_data', 'find_bandgap', 'general_serializer', 'get_loader', diff --git a/src/aiida/orm/nodes/__init__.py b/src/aiida/orm/nodes/__init__.py index 8619c5fb69..70b91a72c9 100644 --- a/src/aiida/orm/nodes/__init__.py +++ b/src/aiida/orm/nodes/__init__.py @@ -31,10 +31,12 @@ 'Code', 'ContainerizedCode', 'Data', + 'DateTimeData', 'Dict', 'EnumData', 'Float', 'FolderData', + 'FunctionData', 'InstalledCode', 'Int', 'JsonableData', @@ -67,6 +69,7 @@ 'WorkflowNode', 'XyData', 'cif_from_ase', + 'deserialize_to_raw_python_data', 'find_bandgap', 'general_serializer', 'has_pycifrw', diff --git a/src/aiida/orm/nodes/data/__init__.py b/src/aiida/orm/nodes/data/__init__.py index 505682eb5d..ef69596bc4 100644 --- a/src/aiida/orm/nodes/data/__init__.py +++ b/src/aiida/orm/nodes/data/__init__.py @@ -18,10 +18,13 @@ from .cif import * from .code import * from .data import * +from .datetime import * +from .deserializer import * from .dict import * from .enum import * from .float import * from .folder import * +from .function import * from .int import * from .jsonable import * from .list import * @@ -45,10 +48,12 @@ 'Code', 'ContainerizedCode', 'Data', + 'DateTimeData', 'Dict', 'EnumData', 'Float', 'FolderData', + 'FunctionData', 'InstalledCode', 'Int', 'JsonableData', @@ -73,6 +78,7 @@ 'UpfData', 'XyData', 'cif_from_ase', + 'deserialize_to_raw_python_data', 'find_bandgap', 'general_serializer', 'has_pycifrw', diff --git a/src/aiida/orm/nodes/data/datetime.py b/src/aiida/orm/nodes/data/datetime.py new file mode 100644 index 0000000000..37094f0d7a --- /dev/null +++ b/src/aiida/orm/nodes/data/datetime.py @@ -0,0 +1,45 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""`Data` sub class to represent a :class:`datetime.datetime` value.""" + +from __future__ import annotations + +import datetime + +from .base import to_aiida_type +from .data import Data + +__all__ = ('DateTimeData',) + + +@to_aiida_type.register(datetime.datetime) +def _(value): + return DateTimeData(value) + + +class DateTimeData(Data): + """`Data` sub class to store a :class:`datetime.datetime` object. + + The value is stored as an ISO-8601 string for portability across backends and reconstructed on access. + """ + + def __init__(self, value, **kwargs): + if not isinstance(value, datetime.datetime): + msg = f'expected a datetime.datetime, got {type(value)}' + raise TypeError(msg) + super().__init__(**kwargs) + self.base.attributes.set('datetime', value.isoformat()) + + @property + def value(self) -> datetime.datetime: + """Return the stored value as a :class:`datetime.datetime`.""" + return datetime.datetime.fromisoformat(self.base.attributes.get('datetime')) + + def __str__(self) -> str: + return str(self.value) diff --git a/src/aiida/orm/nodes/data/deserializer.py b/src/aiida/orm/nodes/data/deserializer.py new file mode 100644 index 0000000000..c2f5425374 --- /dev/null +++ b/src/aiida/orm/nodes/data/deserializer.py @@ -0,0 +1,104 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Deserialize AiiDA data nodes back into raw Python values. + +This is the inverse of :func:`~aiida.orm.nodes.data.serializer.general_serializer`. A node that exposes a ``value`` +is returned directly; otherwise a registry keyed by the node's ``module.ClassName`` maps it to a deserializer. Nested +mappings are walked recursively. +""" + +from __future__ import annotations + +import functools +import typing as t + +#: Deserializers for node types that do not expose a ``value`` attribute, keyed by the node class import path and +#: mapping to the import path of a ``(node) -> value`` callable. +BUILTIN_DESERIALIZERS: dict[str, str] = { + 'aiida.orm.nodes.data.list.List': 'aiida.orm.nodes.data.deserializer.list_data_to_list', + 'aiida.orm.nodes.data.dict.Dict': 'aiida.orm.nodes.data.deserializer.dict_data_to_dict', + 'aiida.orm.nodes.data.array.array.ArrayData': 'aiida.orm.nodes.data.deserializer.array_data_to_array', + 'aiida.orm.nodes.data.structure.StructureData': 'aiida.orm.nodes.data.deserializer.structure_data_to_atoms', +} + +__all__ = ('deserialize_to_raw_python_data',) + + +def list_data_to_list(data): + return data.get_list() + + +def dict_data_to_dict(data): + return data.get_dict() + + +def array_data_to_array(data): + return data.get_array() + + +def structure_data_to_atoms(structure): + """Return the ASE ``Atoms`` of a ``StructureData`` (requires ``ase`` at call time).""" + return structure.get_ase() + + +def structure_data_to_pymatgen(structure): + """Return the pymatgen structure of a ``StructureData`` (requires ``pymatgen`` at call time).""" + return structure.get_pymatgen() + + +@functools.cache +def get_deserializers() -> dict[str, str]: + """Return the ``{node_type_path: deserializer_path}`` registry (cached).""" + return dict(BUILTIN_DESERIALIZERS) + + +def deserialize_to_raw_python_data( + data: t.Any, + deserializers: dict[str, str] | None = None, + dry_run: bool = False, +) -> t.Any: + """Deserialize an AiiDA data node (or nested mapping of nodes) back into raw Python values. + + A node exposing a ``value`` attribute is returned through it; otherwise the ``deserializers`` registry maps the + node type to a deserializer. Mappings are walked recursively. + + :param data: the node or nested mapping to deserialize. + :param deserializers: optional ``{node_type_path: deserializer_path}`` override; defaults to + :func:`get_deserializers`. + :param dry_run: if true, return ``None`` for ``value``-exposing leaves instead of reading them (used to validate + deserializability without materializing values). + :raises ValueError: if a node exposes no ``value`` and no matching deserializer is registered. + """ + from plumpy.utils import AttributesFrozendict + + from aiida import common, orm + + from .serializer import import_from_path + + if deserializers is None: + deserializers = get_deserializers() + + if isinstance(data, orm.Data): + if hasattr(data, 'value'): + return None if dry_run else data.value + type_key = f'{type(data).__module__}.{type(data).__name__}' + if type_key in deserializers: + deserializer = import_from_path(deserializers[type_key]) + return deserializer(data) + msg = ( + f'cannot deserialize an AiiDA data node of type `{type_key}`: it exposes no `value` attribute and no ' + f'matching deserializer is registered. Use a data type with a `value`, or provide one through ' + f"`deserializers`, e.g. {{'{type_key}': 'my_pkg:my_deserializer'}}." + ) + raise ValueError(msg) + + if isinstance(data, (common.extendeddicts.AttributeDict, AttributesFrozendict, dict)): + return {key: deserialize_to_raw_python_data(value, deserializers=deserializers) for key, value in data.items()} + + return None diff --git a/src/aiida/orm/nodes/data/function.py b/src/aiida/orm/nodes/data/function.py new file mode 100644 index 0000000000..eb7cd3a4fc --- /dev/null +++ b/src/aiida/orm/nodes/data/function.py @@ -0,0 +1,79 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""`Data` sub class to represent a reference to a Python function or class.""" + +from __future__ import annotations + +import types +from importlib import import_module + +from .base import to_aiida_type +from .data import Data + +__all__ = ('FunctionData',) + + +@to_aiida_type.register(types.FunctionType) +def _(value): + return FunctionData(value) + + +class FunctionData(Data): + """`Data` sub class that stores an importable reference to a Python function or class. + + Only the module path and qualified name are stored, so the referenced object must be importable in the environment + that later resolves it through :attr:`value`; the object itself is not serialized. + """ + + def __init__(self, value, **kwargs): + module = getattr(value, '__module__', None) + qualname = getattr(value, '__qualname__', None) or getattr(value, '__name__', None) + if not module or not qualname: + msg = f'expected a function-like object with a module and qualified name, got {type(value)}' + raise TypeError(msg) + super().__init__(**kwargs) + self.base.attributes.set('module_path', module) + self.base.attributes.set('qualname', qualname) + + @property + def module_path(self) -> str: + return self.base.attributes.get('module_path') + + @property + def qualname(self) -> str: + return self.base.attributes.get('qualname') + + @property + def path(self) -> str: + return f'{self.module_path}:{self.qualname}' + + @property + def value(self): + """Import and return the referenced function or class. + + :raises ImportError: if the module cannot be imported or the qualified name cannot be resolved within it. + """ + try: + module = import_module(self.module_path) + except Exception as exc: + msg = f"failed to import module '{self.module_path}' for FunctionData '{self.path}': {exc}" + raise ImportError(msg) from exc + + obj = module + try: + for part in self.qualname.split('.'): + obj = getattr(obj, part) + except AttributeError as exc: + msg = f"failed to resolve '{self.path}': attribute '{part}' not found." + raise ImportError(msg) from exc + + return obj + + def __str__(self) -> str: + return self.path diff --git a/tests/orm/nodes/data/test_datetime.py b/tests/orm/nodes/data/test_datetime.py new file mode 100644 index 0000000000..54a3db0436 --- /dev/null +++ b/tests/orm/nodes/data/test_datetime.py @@ -0,0 +1,39 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Tests for :class:`aiida.orm.nodes.data.datetime.DateTimeData`.""" + +import datetime + +import pytest + +from aiida import orm +from aiida.orm import DateTimeData, load_node + + +def test_value_roundtrip(): + value = datetime.datetime(2026, 7, 22, 13, 30, 15) + node = DateTimeData(value).store() + assert load_node(node.pk).value == value + + +def test_to_aiida_type_dispatches_datetime(): + value = datetime.datetime.now() + node = orm.to_aiida_type(value) + assert isinstance(node, DateTimeData) + assert node.value == value + + +def test_invalid_type_raises(): + with pytest.raises(TypeError, match='expected a datetime.datetime'): + DateTimeData('2026-07-22') + + +def test_str(): + value = datetime.datetime(2026, 7, 22) + assert str(DateTimeData(value)) == str(value) diff --git a/tests/orm/nodes/data/test_deserializer.py b/tests/orm/nodes/data/test_deserializer.py new file mode 100644 index 0000000000..f256907dba --- /dev/null +++ b/tests/orm/nodes/data/test_deserializer.py @@ -0,0 +1,41 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Tests for :mod:`aiida.orm.nodes.data.deserializer`.""" + +import pytest + +from aiida import orm +from aiida.orm.nodes.data.deserializer import deserialize_to_raw_python_data + + +def test_value_bearing_node_returns_its_value(): + assert deserialize_to_raw_python_data(orm.Int(5).store()) == 5 + + +def test_list_node_deserialized_through_registry(): + assert deserialize_to_raw_python_data(orm.List(list=[1, 2, 3]).store()) == [1, 2, 3] + + +def test_dict_node_deserialized_through_registry(): + assert deserialize_to_raw_python_data(orm.Dict(dict={'a': 1}).store()) == {'a': 1} + + +def test_nested_mapping_is_walked_recursively(): + data = {'x': orm.Int(1).store(), 'y': {'z': orm.Str('s').store()}} + assert deserialize_to_raw_python_data(data) == {'x': 1, 'y': {'z': 's'}} + + +def test_dry_run_returns_none_for_value_leaves(): + assert deserialize_to_raw_python_data(orm.Int(5).store(), dry_run=True) is None + + +def test_unknown_node_without_value_raises(): + node = orm.FolderData().store() + with pytest.raises(ValueError, match='cannot deserialize an AiiDA data node'): + deserialize_to_raw_python_data(node) diff --git a/tests/orm/nodes/data/test_function.py b/tests/orm/nodes/data/test_function.py new file mode 100644 index 0000000000..caf0c90b0e --- /dev/null +++ b/tests/orm/nodes/data/test_function.py @@ -0,0 +1,47 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Tests for :class:`aiida.orm.nodes.data.function.FunctionData`.""" + +import pytest + +from aiida import orm +from aiida.orm import FunctionData, load_node + + +def a_referenced_function(): + """A module-level function used as the target of a ``FunctionData`` reference.""" + return 42 + + +def test_stores_module_and_qualname_and_resolves_value(): + node = FunctionData(a_referenced_function).store() + loaded = load_node(node.pk) + assert loaded.module_path == __name__ + assert loaded.qualname == 'a_referenced_function' + assert loaded.path == f'{__name__}:a_referenced_function' + assert loaded.value is a_referenced_function + assert loaded.value() == 42 + + +def test_to_aiida_type_dispatches_function(): + node = orm.to_aiida_type(a_referenced_function) + assert isinstance(node, FunctionData) + assert node.value is a_referenced_function + + +def test_invalid_object_raises(): + with pytest.raises(TypeError, match='expected a function-like object'): + FunctionData(5) + + +def test_unresolvable_reference_raises_import_error(): + node = FunctionData(a_referenced_function) + node.base.attributes.set('qualname', 'does_not_exist') + with pytest.raises(ImportError, match="attribute 'does_not_exist' not found"): + _ = node.value diff --git a/tests/orm/test_fields/fields_aiida.data.core.datetime.DateTimeData.yml b/tests/orm/test_fields/fields_aiida.data.core.datetime.DateTimeData.yml new file mode 100644 index 0000000000..af6406519d --- /dev/null +++ b/tests/orm/test_fields/fields_aiida.data.core.datetime.DateTimeData.yml @@ -0,0 +1,21 @@ +attributes: QbAttributesField('attributes', dtype=, + doc='The node attributes') +computer: QbNumericField('computer', dtype=int | None, doc='The PK of the computer') +ctime: QbNumericField('ctime', dtype=, doc='The creation + time of the node') +description: QbStrField('description', dtype=, doc='The node description') +extras: QbDictField('extras', dtype=dict[str, typing.Any], doc='The node extras') +label: QbStrField('label', dtype=, doc='The node label') +mtime: QbNumericField('mtime', dtype=, doc='The modification + time of the node') +node_type: QbStrField('node_type', dtype=typing.Literal['data.core.datetime.DateTimeData.'], + doc='The type of the node.') +pk: QbNumericField('pk', dtype=, doc='The primary key of the entity') +process_type: QbStrField('process_type', dtype=str | None, doc='The process type of + the node') +repository_metadata: QbDictField('repository_metadata', dtype=dict[str, typing.Any], + doc='Virtual hierarchy of the file repository') +source: QbDictField('attributes.source', dtype=dict | None, doc='Source of the data') +user: QbNumericField('user', dtype=, doc='The PK of the user who owns + the node') +uuid: QbAnyField('uuid', dtype=, doc='The UUID of the node') diff --git a/tests/orm/test_fields/fields_aiida.data.core.function.FunctionData.yml b/tests/orm/test_fields/fields_aiida.data.core.function.FunctionData.yml new file mode 100644 index 0000000000..e374840b54 --- /dev/null +++ b/tests/orm/test_fields/fields_aiida.data.core.function.FunctionData.yml @@ -0,0 +1,21 @@ +attributes: QbAttributesField('attributes', dtype=, + doc='The node attributes') +computer: QbNumericField('computer', dtype=int | None, doc='The PK of the computer') +ctime: QbNumericField('ctime', dtype=, doc='The creation + time of the node') +description: QbStrField('description', dtype=, doc='The node description') +extras: QbDictField('extras', dtype=dict[str, typing.Any], doc='The node extras') +label: QbStrField('label', dtype=, doc='The node label') +mtime: QbNumericField('mtime', dtype=, doc='The modification + time of the node') +node_type: QbStrField('node_type', dtype=typing.Literal['data.core.function.FunctionData.'], + doc='The type of the node.') +pk: QbNumericField('pk', dtype=, doc='The primary key of the entity') +process_type: QbStrField('process_type', dtype=str | None, doc='The process type of + the node') +repository_metadata: QbDictField('repository_metadata', dtype=dict[str, typing.Any], + doc='Virtual hierarchy of the file repository') +source: QbDictField('attributes.source', dtype=dict | None, doc='Source of the data') +user: QbNumericField('user', dtype=, doc='The PK of the user who owns + the node') +uuid: QbAnyField('uuid', dtype=, doc='The UUID of the node') From f943b82a7d14abdf7df4a65564d3a745f98eb4a6 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Wed, 22 Jul 2026 14:41:57 +0200 Subject: [PATCH 11/19] =?UTF-8?q?=F0=9F=91=8C=20`JsonableData`:=20add=20`.?= =?UTF-8?q?value`=20alias=20for=20`.obj`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expose `.value` as an alias of `.obj`, mirroring the `value` accessor of the simple data types (`Int`, `Str`, ...). This is a backwards-compatible extension: it adds an accessor without changing any existing behaviour. It lets downstream code that follows the common ``node.value`` convention (aiida-pythonjob, repointing at core's `JsonableData` instead of its own copy) work unchanged, one more step toward a single `JsonableData` in core rather than duplicates in every plugin. --- src/aiida/orm/nodes/data/jsonable.py | 5 +++++ tests/orm/nodes/data/test_jsonable.py | 6 ++++++ 2 files changed, 11 insertions(+) diff --git a/src/aiida/orm/nodes/data/jsonable.py b/src/aiida/orm/nodes/data/jsonable.py index f559a547bd..20f335f787 100644 --- a/src/aiida/orm/nodes/data/jsonable.py +++ b/src/aiida/orm/nodes/data/jsonable.py @@ -241,6 +241,11 @@ def obj(self) -> JsonSerializableProtocol: """ return self._get_object() + @property + def value(self) -> JsonSerializableProtocol: + """Alias of :attr:`obj`, mirroring the ``value`` accessor of the simple data types (``Int``, ``Str``, ...).""" + return self._get_object() + @classmethod def _deserialize_float_constants(cls, data: typing.Any): """Deserialize the contents of a dictionary ``data`` deserializing infinity and NaN string constants. diff --git a/tests/orm/nodes/data/test_jsonable.py b/tests/orm/nodes/data/test_jsonable.py index e63adfd70e..b89c4f470a 100644 --- a/tests/orm/nodes/data/test_jsonable.py +++ b/tests/orm/nodes/data/test_jsonable.py @@ -181,6 +181,12 @@ class PydanticObj(BaseModel): b: str +def test_value_is_alias_of_obj(): + """``.value`` mirrors ``.obj`` for consistency with the simple data types.""" + node = JsonableData(JsonableClass({'a': 1})) + assert node.value is node.obj + + def test_wrap_object_with_to_dict(): """An object exposing ``to_dict`` (not ``as_dict``) round-trips.""" node = JsonableData(ToDictClass(7)).store() From f4860f40086e6f3f8f38a83adf0419edcdaf6184 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 6 Aug 2026 12:27:29 +0200 Subject: [PATCH 12/19] =?UTF-8?q?=F0=9F=94=A7=20`pre-commit`:=20exempt=20t?= =?UTF-8?q?he=20`aiida/workgraph`=20subtree?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The incoming WorkGraph subsystem keeps hand-maintained `__init__` files (a curated public API, lazy plugin imports, the node-graph optional-extra boundary) and is typed incrementally, so adjust core's tooling before it lands: - `autogenerate_all_imports.py`: skip the `aiida/workgraph/` subtree, so it is not star-imported and its `__all__` is not bubbled to parent packages. - mypy: add `src/aiida/workgraph/.*` to the exclude. - ruff: a scoped per-file-ignore for the CamelCase control-flow API (`If` / `While` / `Map` / `Zone` / `TaskPool`) and two pre-existing WorkGraph loop/argument patterns. --- .pre-commit-config.yaml | 2 ++ pyproject.toml | 7 +++++++ utils/autogenerate_all_imports.py | 5 +++++ 3 files changed, 14 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9492f95165..395e7c18d5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -136,6 +136,8 @@ repos: src/aiida/transports/cli.py| src/aiida/transports/plugins/local.py| src/aiida/transports/plugins/ssh.py| + + src/aiida/workgraph/.*| )$ - id: generate-conda-environment diff --git a/pyproject.toml b/pyproject.toml index 896a8da3c9..e06e42bbe6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -520,6 +520,13 @@ select = [ 'NPY201' # numpy compatibility ] +[tool.ruff.lint.per-file-ignores] +# The relocated WorkGraph subsystem uses CamelCase for its public control-flow / task-type API +# (``If`` / ``While`` / ``Map`` / ``Zone``, ``TaskPool``, the ``*_TaskSpec`` factory functions), +# which pep8-naming (N802/N803/N806) would otherwise flag. PLW2901/PLR1704 are pre-existing +# WorkGraph loop/argument patterns to clean up in a follow-up. +'src/aiida/workgraph/**' = ['N802', 'N803', 'N806', 'PLW2901', 'PLR1704'] + # Mark some classes as generic, per https://docs.astral.sh/ruff/settings/#lint_pyflakes_extend-generics # Needed due to https://github.com/astral-sh/ruff/issues/9298 [tool.ruff.lint.pyflakes] diff --git a/utils/autogenerate_all_imports.py b/utils/autogenerate_all_imports.py index 1f32538cfd..c10b955a8e 100755 --- a/utils/autogenerate_all_imports.py +++ b/utils/autogenerate_all_imports.py @@ -119,6 +119,11 @@ def write_inits(folder_path: Path, all_dict: dict, skip_children: dict[str, list rel_path = path.parent.relative_to(folder_path).as_posix() + if rel_path == 'workgraph' or rel_path.startswith('workgraph/'): + # The workgraph subsystem's ``__init__`` files are hand-maintained: a curated public API, + # lazy plugin imports, and the node-graph optional-extra boundary. Skip autogeneration. + continue + # get sub_dict for this folder path_all_dict = all_dict mod_path = path.parent.relative_to(folder_path).parts From 2461a3985c69abacae40a49ca18bfc1dc6b066cc Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 6 Aug 2026 12:27:49 +0200 Subject: [PATCH 13/19] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20`WorkChain`:=20extra?= =?UTF-8?q?ct=20`WorkflowProcess`=20base?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make WorkGraph a sibling of WorkChain under a shared base rather than having it inherit WorkChain. Per the review discussion (aiidateam#7479, aiidateam#7513): the two have different execution models, and bolting WorkGraph concepts (the awaitable-barrier policy) onto WorkChain both muddies WorkChain and couples the two, so any change to WorkChain would ripple into WorkGraph. Extract everything the two share into a new `WorkflowProcess(Process)`: the stepper seam, the awaitable-based waiting on child processes, the context (`ctx`), and the step lifecycle with its checkpointing. The seam `_create_stepper` / `_recreate_stepper` becomes abstract on the base; `WorkChain(WorkflowProcess)` keeps only the outline stepper and its spec/node binding. The `_awaitable_barrier` (a property of the stepping strategy, not of WorkChain) now lives on the base, so WorkChain no longer carries a WorkGraph concept. Also fix the `Protect` metaclass to scan each base's full MRO instead of only its direct bases, so a `@final` method stays protected when it is inherited through an intermediate class (e.g. `run`, now on `WorkflowProcess`, reached via `WorkChain`). Without this, overriding such a method in a `WorkChain` subclass would no longer raise. Pure refactor of the core side; WorkChain behaviour is unchanged (`test_work_chain.py` + `test_restart.py` green, mypy + ruff clean). The WorkGraph side re-parents onto `WorkflowProcess` when the engine moves. --- src/aiida/engine/__init__.py | 1 + src/aiida/engine/processes/__init__.py | 1 + .../engine/processes/workchains/__init__.py | 2 + .../engine/processes/workchains/workchain.py | 418 +-------------- .../processes/workchains/workflow_process.py | 483 ++++++++++++++++++ 5 files changed, 497 insertions(+), 408 deletions(-) create mode 100644 src/aiida/engine/processes/workchains/workflow_process.py diff --git a/src/aiida/engine/__init__.py b/src/aiida/engine/__init__.py index 0a855c6484..438c945cf0 100644 --- a/src/aiida/engine/__init__.py +++ b/src/aiida/engine/__init__.py @@ -55,6 +55,7 @@ 'WithNonDb', 'WithSerialize', 'WorkChain', + 'WorkflowProcess', 'append_', 'assign_', 'await_processes', diff --git a/src/aiida/engine/processes/__init__.py b/src/aiida/engine/processes/__init__.py index 59b8b3a987..8c28596c4d 100644 --- a/src/aiida/engine/processes/__init__.py +++ b/src/aiida/engine/processes/__init__.py @@ -51,6 +51,7 @@ 'WithNonDb', 'WithSerialize', 'WorkChain', + 'WorkflowProcess', 'append_', 'assign_', 'calcfunction', diff --git a/src/aiida/engine/processes/workchains/__init__.py b/src/aiida/engine/processes/workchains/__init__.py index f1def097d9..0424c6ac35 100644 --- a/src/aiida/engine/processes/workchains/__init__.py +++ b/src/aiida/engine/processes/workchains/__init__.py @@ -17,6 +17,7 @@ from .restart import * from .utils import * from .workchain import * +from .workflow_process import * __all__ = ( 'Awaitable', @@ -26,6 +27,7 @@ 'ProcessHandlerReport', 'ToContext', 'WorkChain', + 'WorkflowProcess', 'append_', 'assign_', 'construct_awaitable', diff --git a/src/aiida/engine/processes/workchains/workchain.py b/src/aiida/engine/processes/workchains/workchain.py index 6acaba102d..1e2a77abd3 100644 --- a/src/aiida/engine/processes/workchains/workchain.py +++ b/src/aiida/engine/processes/workchains/workchain.py @@ -10,28 +10,17 @@ from __future__ import annotations -import collections.abc -import functools import logging import typing as t -from plumpy import run_with_portal -from plumpy.persistence import auto_persist -from plumpy.process_states import Continue, Wait -from plumpy.processes import ProcessStateMachineMeta -from plumpy.workchains import Stepper, _PropagateReturn, if_, return_, while_ +from plumpy.workchains import Stepper, if_, return_, while_ from plumpy.workchains import WorkChainSpec as PlumpyWorkChainSpec from aiida.common import exceptions -from aiida.common.extendeddicts import AttributeDict -from aiida.common.lang import override -from aiida.orm import Node, ProcessNode, WorkChainNode -from aiida.orm.utils import load_node +from aiida.orm import WorkChainNode -from ..exit_code import ExitCode -from ..process import Process, ProcessState from ..process_spec import ProcessSpec -from .awaitable import Awaitable, AwaitableAction, AwaitableTarget, construct_awaitable +from .workflow_process import WorkflowProcess if t.TYPE_CHECKING: from aiida.engine.runners import Runner @@ -43,69 +32,16 @@ class WorkChainSpec(ProcessSpec, PlumpyWorkChainSpec): pass -MethodType = t.TypeVar('MethodType') +class WorkChain(WorkflowProcess): + """The `WorkChain` class is the principle component to implement workflows in AiiDA. - -class Protect(ProcessStateMachineMeta): - """Metaclass that allows protecting class methods from being overridden by subclasses. - - Usage as follows:: - - class SomeClass(metaclass=Protect): - - @Protect.final - def private_method(self): - "This method cannot be overridden by a subclass." - - If a subclass is imported that overrides the subclass, a ``RuntimeError`` is raised. + It is a :class:`~aiida.engine.processes.workchains.workflow_process.WorkflowProcess` whose stepper walks the + static outline declared on its spec. The shared workflow machinery (awaitables, context, the step lifecycle + and its checkpointing) lives on the base class; only the outline stepping is defined here. """ - __SENTINEL = object() - - def __new__(mcs, name, bases, namespace, **kwargs): - """Collect all methods that were marked as protected and raise if the subclass defines it. - - :raises RuntimeError: If the new class defines (i.e. overrides) a method that was decorated with ``final``. - """ - private = { - key for base in bases for key, value in vars(base).items() if callable(value) and mcs.__is_final(value) - } - for key in namespace: - if key in private: - raise RuntimeError(f'the method `{key}` is protected cannot be overridden.') - return super().__new__(mcs, name, bases, namespace, **kwargs) - - @classmethod - def __is_final(mcs, method) -> bool: # noqa: N804 - """Return whether the method has been decorated by the ``final`` classmethod. - - :return: Boolean, ``True`` if the method is marked as final, ``False`` otherwise. - """ - try: - return method.__final is mcs.__SENTINEL - except AttributeError: - return False - - @classmethod - def final(mcs, method: MethodType) -> MethodType: # noqa: N804 - """Decorate a method with this method to protect it from being overridden. - - Adds the ``__SENTINEL`` object as the ``__final`` private attribute to the given ``method`` and wraps it in - the ``typing.final`` decorator. The latter indicates to typing systems that it cannot be overridden in - subclasses. - """ - method.__final = mcs.__SENTINEL # type: ignore[attr-defined] - return t.final(method) - - -@auto_persist('_awaitables') -class WorkChain(Process, metaclass=Protect): - """The `WorkChain` class is the principle component to implement workflows in AiiDA.""" - _node_class = WorkChainNode _spec_class = WorkChainSpec - _STEPPER_STATE = 'stepper_state' - _CONTEXT = 'CONTEXT' def __init__( self, @@ -129,13 +65,6 @@ def __init__( super().__init__(inputs, logger, runner, enable_persistence=enable_persistence) - self._stepper: Stepper | None = None - self._awaitables: list[Awaitable] = [] - # The pks of awaitables whose completion callback is already registered. This is runtime state, callbacks - # do not survive a checkpoint, so it is not persisted and is reset in `load_instance_state`. - self._registered_awaitable_pks: set[int] = set() - self._context = AttributeDict() - @classmethod def spec(cls) -> WorkChainSpec: return super().spec() # type: ignore[return-value] @@ -144,340 +73,13 @@ def spec(cls) -> WorkChainSpec: def node(self) -> WorkChainNode: return super().node # type: ignore[return-value] - @property - def ctx(self) -> AttributeDict: - """Get the context.""" - return self._context - - @override - def save_instance_state(self, out_state, save_context): - """Save instance state. - - :param out_state: state to save in - - :param save_context: - :type save_context: :class:`!plumpy.persistence.LoadSaveContext` - - """ - super().save_instance_state(out_state, save_context) - # Save the context - out_state[self._CONTEXT] = self.ctx - - # Ask the stepper to save itself - if self._stepper is not None: - out_state[self._STEPPER_STATE] = self._stepper.save() - - @override - def load_instance_state(self, saved_state, load_context): - super().load_instance_state(saved_state, load_context) - # Load the context - self._context = saved_state[self._CONTEXT] - - # Recreate the stepper - self._stepper = None - stepper_state = saved_state.get(self._STEPPER_STATE, None) - if stepper_state is not None: - self._stepper = self._recreate_stepper(stepper_state) - - self.set_logger(self.node.logger) - - # Callbacks do not survive the checkpoint, so nothing is registered yet on the reloaded process. - self._registered_awaitable_pks = set() - if self._awaitables: - self._action_awaitables() - def _create_stepper(self) -> Stepper: - """Return the stepper that drives this work chain. - - This is the seam for supplying a different execution strategy. The default steps through the outline declared - on the spec, but a subclass may return any :class:`plumpy.workchains.Stepper`, for example one that derives the - order of execution from a graph of data dependencies instead of a static outline. - - A subclass that overrides this should also override :meth:`_recreate_stepper`, otherwise its processes cannot - be restored from a checkpoint. - """ + """Step through the outline declared on the spec.""" return self.spec().get_outline().create_stepper(self) # type: ignore[arg-type] def _recreate_stepper(self, saved_state: t.Any) -> Stepper: - """Restore the stepper from the state it wrote to the checkpoint. - - The counterpart of :meth:`_create_stepper`, called when a process is loaded from a checkpoint rather than - started fresh. + """Restore the outline stepper from the state it wrote to the checkpoint. :param saved_state: the state previously returned by ``Stepper.save()`` """ return self.spec().get_outline().recreate_stepper(saved_state, self) # type: ignore[arg-type] - - @property - def _awaitable_barrier(self) -> bool: - """Whether each step waits for everything it launched before the next one begins. - - This is the difference between the two execution models, and it is a property of the stepping strategy, so - the value is taken from the stepper. ``True``, the default, is the outline model: :meth:`_do_step` clears - the awaitables at the start of every step, so a step forms a barrier over the children it launched and the - process only resumes once all of them have finished. A stepper that schedules by data dependencies wants - ``False``: the awaitables persist across steps and the process resumes as each child finishes, so - independent branches stay in flight together. A stepper opts into the streaming model by defining - ``awaitable_barrier = False`` on itself. - """ - return getattr(self._stepper, 'awaitable_barrier', True) - - @Protect.final - def on_run(self): - super().on_run() - self.node.set_stepper_state_info(str(self._stepper)) - - def _resolve_nested_context(self, key: str) -> tuple[AttributeDict, str]: - """Returns a reference to a sub-dictionary of the context and the last key, - after resolving a potentially segmented key where required sub-dictionaries are created as needed. - - :param key: A key into the context, where words before a dot are interpreted as a key for a sub-dictionary - """ - ctx = self.ctx - ctx_path = key.split('.') - - for index, path in enumerate(ctx_path[:-1]): - try: - ctx = ctx[path] - except KeyError: # see below why this is the only exception we have to catch here - ctx[path] = AttributeDict() # create the sub-dict and update the context - ctx = ctx[path] - continue - - # Notes: - # * the first ctx (self.ctx) is guaranteed to be an AttributeDict, hence the post-"dereference" checking - # * the values can be many different things: on insertion they are either AtrributeDict, List or Awaitables - # (subclasses of AttributeDict) but after resolution of an Awaitable this will be the value itself - # * assumption: a resolved value is never a plain AttributeDict, on the other hand if a resolved Awaitable - # would be an AttributeDict we can append things to it since the order of tasks is maintained. - if type(ctx) is not AttributeDict: - raise ValueError( - f'Can not update the context for key `{key}`:' - f' found instance of `{type(ctx)}` at `{".".join(ctx_path[: index + 1])}`, expected AttributeDict' - ) - - return ctx, ctx_path[-1] - - def _insert_awaitable(self, awaitable: Awaitable) -> None: - """Insert an awaitable that should be terminated before before continuing to the next step. - - :param awaitable: the thing to await - """ - ctx, key = self._resolve_nested_context(awaitable.key) - - # Already assign the awaitable itself to the location in the context container where it is supposed to end up - # once it is resolved. This is especially important for the `APPEND` action, since it needs to maintain the - # order, but the awaitables will not necessarily be resolved in the order in which they are added. By using the - # awaitable as a placeholder, in the `_resolve_awaitable`, it can be found and replaced by the resolved value. - if awaitable.action == AwaitableAction.ASSIGN: - ctx[key] = awaitable - elif awaitable.action == AwaitableAction.APPEND: - ctx.setdefault(key, []).append(awaitable) - else: - raise AssertionError(f'Unsupported awaitable action: {awaitable.action}') - - self._awaitables.append( - awaitable - ) # add only if everything went ok, otherwise we end up in an inconsistent state - self._update_process_status() - - def _resolve_awaitable(self, awaitable: Awaitable, value: t.Any) -> None: - """Resolve an awaitable. - - Precondition: must be an awaitable that was previously inserted. - - :param awaitable: the awaitable to resolve - """ - ctx, key = self._resolve_nested_context(awaitable.key) - - if awaitable.action == AwaitableAction.ASSIGN: - ctx[key] = value - elif awaitable.action == AwaitableAction.APPEND: - # Find the same awaitable inserted in the context - container = ctx[key] - for index, placeholder in enumerate(container): - if isinstance(placeholder, Awaitable) and placeholder.pk == awaitable.pk: - container[index] = value - break - else: - raise AssertionError(f'Awaitable `{awaitable.pk} was not found in `ctx.{awaitable.key}`') - else: - raise AssertionError(f'Unsupported awaitable action: {awaitable.action}') - - awaitable.resolved = True - self._awaitables.remove(awaitable) # remove only if everything went ok, otherwise we may lose track - - if not self.has_terminated(): - # the process may be terminated, for example, if the process was killed or excepted - # then we should not try to update it - self._update_process_status() - - @Protect.final - def to_context(self, **kwargs: Awaitable | ProcessNode) -> None: - """Add a dictionary of awaitables to the context. - - This is a convenience method that provides syntactic sugar, for a user to add multiple intersteps that will - assign a certain value to the corresponding key in the context of the work chain. - """ - for key, value in kwargs.items(): - awaitable = construct_awaitable(value) - awaitable.key = key - self._insert_awaitable(awaitable) - - def _update_process_status(self) -> None: - """Set the process status with a message accounting the current sub processes that we are waiting for.""" - if self._awaitables: - status = f'Waiting for child processes: {", ".join([str(_.pk) for _ in self._awaitables])}' - else: - status = None - if self.paused: - # Update the pre-paused status so that when the process is played - # it will be set to the new status - self._pre_paused_status = status - else: - self.set_status(status) - - @override - @Protect.final - async def run(self) -> t.Any: - self._stepper = self._create_stepper() - return await run_with_portal(self._do_step) - - def _do_step(self) -> t.Any: - """Execute the next step in the outline and return the result. - - If the stepper returns a non-finished status and the return value is of type ToContext, the contents of the - ToContext container will be turned into awaitables if necessary. If any awaitables were created, the process - will enter in the Wait state, otherwise it will go to Continue. When the stepper returns that it is done, the - stepper result will be converted to None and returned, unless it is an integer or instance of ExitCode. - """ - from .context import ToContext - - # Under the barrier model the awaitables belong to a single step and are cleared before the next one, which - # is what forces every step to wait for all the children it launched. A streaming stepper keeps them, so - # children launched in earlier steps stay in flight while later steps run. - if self._awaitable_barrier: - self._awaitables = [] - result: t.Any = None - - try: - assert self._stepper is not None - finished, stepper_result = self._stepper.step() - except _PropagateReturn as exception: - finished, result = True, exception.exit_code - else: - # Set result to None unless stepper_result was non-zero positive integer or ExitCode with similar status - if isinstance(stepper_result, int) and stepper_result > 0: - result = ExitCode(stepper_result) - elif isinstance(stepper_result, ExitCode) and stepper_result.status > 0: - result = stepper_result - else: - result = None - - # If the stepper said we are finished or the result is an ExitCode, we exit by returning - if finished or isinstance(result, ExitCode): - return result - - if isinstance(stepper_result, ToContext): - self.to_context(**stepper_result) - - if self._awaitables: - return Wait(self._do_step, 'Waiting before next step') - - return Continue(self._do_step) - - def _store_nodes(self, data: t.Any) -> None: - """Recurse through a data structure and store any unstored nodes that are found along the way - - :param data: a data structure potentially containing unstored nodes - """ - if isinstance(data, Node) and not data.is_stored: - data.store() - elif isinstance(data, collections.abc.Mapping): - for _, value in data.items(): - self._store_nodes(value) - elif isinstance(data, collections.abc.Sequence) and not isinstance(data, str): - for value in data: - self._store_nodes(value) - - @override - @Protect.final - def on_exiting(self) -> None: - """Ensure that any unstored nodes in the context are stored, before the state is exited - - After the state is exited the next state will be entered and if persistence is enabled, a checkpoint will - be saved. If the context contains unstored nodes, the serialization necessary for checkpointing will fail. - """ - super().on_exiting() - try: - self._store_nodes(self.ctx) - except Exception: - # An uncaught exception here will have bizarre and disastrous consequences - self.logger.exception('exception in _store_nodes called in on_exiting') - - @Protect.final - def on_wait(self, awaitables: t.Sequence[t.Awaitable]): - """Entering the WAITING state.""" - super().on_wait(awaitables) - if self._awaitables: - self._action_awaitables() - else: - self.call_soon(self.resume) - - def _action_awaitables(self) -> None: - """Register the completion callback for each awaitable that does not already have one. - - Depending on the class type of the awaitable's target a different callback - function will be bound with the awaitable and the runner will be asked to - call it when the target is completed. - - The registration is guarded against duplicates: under the barrier model the awaitables are cleared each - step so the same one is never seen twice, but a streaming stepper keeps its awaitables across steps and - would otherwise register a further callback for the same awaitable on every pass through the waiting state. - """ - for awaitable in self._awaitables: - if awaitable.pk in self._registered_awaitable_pks: - continue - if awaitable.target == AwaitableTarget.PROCESS: - callback = functools.partial(self.call_soon, self._on_awaitable_finished, awaitable) - self.runner.call_on_process_finish(awaitable.pk, callback) - self._registered_awaitable_pks.add(awaitable.pk) - else: - raise AssertionError(f"invalid awaitable target '{awaitable.target}'") - - def _on_awaitable_resolved(self, awaitable: Awaitable) -> None: - """Hook called once a finished awaitable has been resolved onto the context, before the resume decision. - - Defaults to doing nothing. A subclass can use it to run bookkeeping that must see the resolved value and - must happen before the process is resumed, without having to reimplement :meth:`_on_awaitable_finished`. - - :param awaitable: the awaitable that has just been resolved - """ - - def _on_awaitable_finished(self, awaitable: Awaitable) -> None: - """Callback function, for when an awaitable process instance is completed. - - The awaitable will be effectuated on the context of the work chain and removed from the internal list. The - process is then resumed: under the barrier model only once every awaitable has finished, and under the - streaming model as soon as this one does, so a finished child can unblock its dependents while others run. - - :param awaitable: an Awaitable instance - """ - self.logger.info('received callback that awaitable %d has terminated', awaitable.pk) - - try: - node = load_node(awaitable.pk) - except (exceptions.MultipleObjectsError, exceptions.NotExistent): - raise ValueError(f'provided pk<{awaitable.pk}> could not be resolved to a valid Node instance') - - if awaitable.outputs: - value = {entry.link_label: entry.node for entry in node.base.links.get_outgoing()} - else: - value = node # type: ignore[assignment] - - self._resolve_awaitable(awaitable, value) - self._registered_awaitable_pks.discard(awaitable.pk) - self._on_awaitable_resolved(awaitable) - - if self.state == ProcessState.WAITING and (not self._awaitable_barrier or not self._awaitables): - self.resume() diff --git a/src/aiida/engine/processes/workchains/workflow_process.py b/src/aiida/engine/processes/workchains/workflow_process.py new file mode 100644 index 0000000000..7d9e040048 --- /dev/null +++ b/src/aiida/engine/processes/workchains/workflow_process.py @@ -0,0 +1,483 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""The :class:`WorkflowProcess`, the shared base for stepper-driven workflow processes.""" + +from __future__ import annotations + +import collections.abc +import functools +import logging +import typing as t + +from plumpy import run_with_portal +from plumpy.persistence import auto_persist +from plumpy.process_states import Continue, Wait +from plumpy.processes import ProcessStateMachineMeta +from plumpy.workchains import Stepper, _PropagateReturn + +from aiida.common import exceptions +from aiida.common.extendeddicts import AttributeDict +from aiida.common.lang import override +from aiida.orm import Node, ProcessNode +from aiida.orm.utils import load_node + +from ..exit_code import ExitCode +from ..process import Process, ProcessState +from .awaitable import Awaitable, AwaitableAction, AwaitableTarget, construct_awaitable + +if t.TYPE_CHECKING: + from aiida.engine.runners import Runner + +__all__ = ('WorkflowProcess',) + + +MethodType = t.TypeVar('MethodType') + + +class Protect(ProcessStateMachineMeta): + """Metaclass that allows protecting class methods from being overridden by subclasses. + + Usage as follows:: + + class SomeClass(metaclass=Protect): + + @Protect.final + def private_method(self): + "This method cannot be overridden by a subclass." + + If a subclass is imported that overrides the subclass, a ``RuntimeError`` is raised. + """ + + __SENTINEL = object() + + def __new__(mcs, name, bases, namespace, **kwargs): + """Collect all methods that were marked as protected and raise if the subclass defines it. + + The whole ancestry of each base is scanned (``base.__mro__``), not just the direct bases, so a ``final`` + method stays protected even when it is inherited through an intermediate class rather than defined on the + immediate parent (for example a ``final`` method on ``WorkflowProcess`` reached via ``WorkChain``). + + :raises RuntimeError: If the new class defines (i.e. overrides) a method that was decorated with ``final``. + """ + private = { + key + for base in bases + for klass in base.__mro__ + for key, value in vars(klass).items() + if callable(value) and mcs.__is_final(value) + } + for key in namespace: + if key in private: + raise RuntimeError(f'the method `{key}` is protected cannot be overridden.') + return super().__new__(mcs, name, bases, namespace, **kwargs) + + @classmethod + def __is_final(mcs, method) -> bool: # noqa: N804 + """Return whether the method has been decorated by the ``final`` classmethod. + + :return: Boolean, ``True`` if the method is marked as final, ``False`` otherwise. + """ + try: + return method.__final is mcs.__SENTINEL + except AttributeError: + return False + + @classmethod + def final(mcs, method: MethodType) -> MethodType: # noqa: N804 + """Decorate a method with this method to protect it from being overridden. + + Adds the ``__SENTINEL`` object as the ``__final`` private attribute to the given ``method`` and wraps it in + the ``typing.final`` decorator. The latter indicates to typing systems that it cannot be overridden in + subclasses. + """ + method.__final = mcs.__SENTINEL # type: ignore[attr-defined] + return t.final(method) + + +@auto_persist('_awaitables') +class WorkflowProcess(Process, metaclass=Protect): + """A :class:`~aiida.engine.processes.process.Process` whose execution is delegated to a pluggable stepper. + + This is the shared base for AiiDA's stepper-driven workflow processes: the :class:`~aiida.engine.WorkChain` + (which walks a static outline) and WorkGraph (which schedules by a graph of data dependencies). It owns + everything the two have in common, the stepper seam, the awaitable-based waiting on child processes, the + context (:attr:`ctx`), and the step lifecycle with its checkpointing, and leaves only the concrete stepping + strategy to the subclass via :meth:`_create_stepper` / :meth:`_recreate_stepper`. + + Making WorkChain and WorkGraph siblings under this base, rather than having WorkGraph inherit from WorkChain, + keeps the two execution models independent: neither carries the other's concepts, and a change to one + strategy (the outline or the DAG) cannot affect the other. + """ + + _STEPPER_STATE = 'stepper_state' + _CONTEXT = 'CONTEXT' + + def __init__( + self, + inputs: dict | None = None, + logger: logging.Logger | None = None, + runner: Runner | None = None, + enable_persistence: bool = True, + ) -> None: + """Construct a stepper process instance. + + :param inputs: process inputs + :param logger: aiida logger + :param runner: process runner + :param enable_persistence: whether to persist this process + """ + if self.__class__ == WorkflowProcess: + raise exceptions.InvalidOperation('cannot construct or launch a base `WorkflowProcess` class.') + + super().__init__(inputs, logger, runner, enable_persistence=enable_persistence) + + self._stepper: Stepper | None = None + self._awaitables: list[Awaitable] = [] + # The pks of awaitables whose completion callback is already registered. This is runtime state, callbacks + # do not survive a checkpoint, so it is not persisted and is reset in `load_instance_state`. + self._registered_awaitable_pks: set[int] = set() + self._context = AttributeDict() + + @property + def ctx(self) -> AttributeDict: + """Get the context.""" + return self._context + + @override + def save_instance_state(self, out_state, save_context): + """Save instance state. + + :param out_state: state to save in + + :param save_context: + :type save_context: :class:`!plumpy.persistence.LoadSaveContext` + + """ + super().save_instance_state(out_state, save_context) + # Save the context + out_state[self._CONTEXT] = self.ctx + + # Ask the stepper to save itself + if self._stepper is not None: + out_state[self._STEPPER_STATE] = self._stepper.save() + + @override + def load_instance_state(self, saved_state, load_context): + super().load_instance_state(saved_state, load_context) + # Load the context + self._context = saved_state[self._CONTEXT] + + # Recreate the stepper + self._stepper = None + stepper_state = saved_state.get(self._STEPPER_STATE, None) + if stepper_state is not None: + self._stepper = self._recreate_stepper(stepper_state) + + self.set_logger(self.node.logger) + + # Callbacks do not survive the checkpoint, so nothing is registered yet on the reloaded process. + self._registered_awaitable_pks = set() + if self._awaitables: + self._action_awaitables() + + def _create_stepper(self) -> Stepper: + """Return the stepper that drives this process. Subclasses must implement this. + + The stepper is the execution strategy: :class:`~aiida.engine.WorkChain` returns one that walks the outline + declared on its spec, WorkGraph one that derives the order from a graph of data dependencies. A subclass + that overrides this must also override :meth:`_recreate_stepper`, otherwise its processes cannot be restored + from a checkpoint. + """ + raise NotImplementedError + + def _recreate_stepper(self, saved_state: t.Any) -> Stepper: + """Restore the stepper from the state it wrote to the checkpoint. Subclasses must implement this. + + The counterpart of :meth:`_create_stepper`, called when a process is loaded from a checkpoint rather than + started fresh. + + :param saved_state: the state previously returned by ``Stepper.save()`` + """ + raise NotImplementedError + + @property + def _awaitable_barrier(self) -> bool: + """Whether each step waits for everything it launched before the next one begins. + + This is the difference between the two execution models, and it is a property of the stepping strategy, so + the value is taken from the stepper. ``True``, the default, is the outline model: :meth:`_do_step` clears + the awaitables at the start of every step, so a step forms a barrier over the children it launched and the + process only resumes once all of them have finished. A stepper that schedules by data dependencies wants + ``False``: the awaitables persist across steps and the process resumes as each child finishes, so + independent branches stay in flight together. A stepper opts into the streaming model by defining + ``awaitable_barrier = False`` on itself. + """ + return getattr(self._stepper, 'awaitable_barrier', True) + + @Protect.final + def on_run(self): + super().on_run() + # ``set_stepper_state_info`` is part of the node contract for a stepper-driven process; both WorkChainNode + # and WorkGraphNode provide it (the latter subclasses the former). + self.node.set_stepper_state_info(str(self._stepper)) + + def _resolve_nested_context(self, key: str) -> tuple[AttributeDict, str]: + """Returns a reference to a sub-dictionary of the context and the last key, + after resolving a potentially segmented key where required sub-dictionaries are created as needed. + + :param key: A key into the context, where words before a dot are interpreted as a key for a sub-dictionary + """ + ctx = self.ctx + ctx_path = key.split('.') + + for index, path in enumerate(ctx_path[:-1]): + try: + ctx = ctx[path] + except KeyError: # see below why this is the only exception we have to catch here + ctx[path] = AttributeDict() # create the sub-dict and update the context + ctx = ctx[path] + continue + + # Notes: + # * the first ctx (self.ctx) is guaranteed to be an AttributeDict, hence the post-"dereference" checking + # * the values can be many different things: on insertion they are either AtrributeDict, List or Awaitables + # (subclasses of AttributeDict) but after resolution of an Awaitable this will be the value itself + # * assumption: a resolved value is never a plain AttributeDict, on the other hand if a resolved Awaitable + # would be an AttributeDict we can append things to it since the order of tasks is maintained. + if type(ctx) is not AttributeDict: + raise ValueError( + f'Can not update the context for key `{key}`:' + f' found instance of `{type(ctx)}` at `{".".join(ctx_path[: index + 1])}`, expected AttributeDict' + ) + + return ctx, ctx_path[-1] + + def _insert_awaitable(self, awaitable: Awaitable) -> None: + """Insert an awaitable that should be terminated before before continuing to the next step. + + :param awaitable: the thing to await + """ + ctx, key = self._resolve_nested_context(awaitable.key) + + # Already assign the awaitable itself to the location in the context container where it is supposed to end up + # once it is resolved. This is especially important for the `APPEND` action, since it needs to maintain the + # order, but the awaitables will not necessarily be resolved in the order in which they are added. By using the + # awaitable as a placeholder, in the `_resolve_awaitable`, it can be found and replaced by the resolved value. + if awaitable.action == AwaitableAction.ASSIGN: + ctx[key] = awaitable + elif awaitable.action == AwaitableAction.APPEND: + ctx.setdefault(key, []).append(awaitable) + else: + raise AssertionError(f'Unsupported awaitable action: {awaitable.action}') + + self._awaitables.append( + awaitable + ) # add only if everything went ok, otherwise we end up in an inconsistent state + self._update_process_status() + + def _resolve_awaitable(self, awaitable: Awaitable, value: t.Any) -> None: + """Resolve an awaitable. + + Precondition: must be an awaitable that was previously inserted. + + :param awaitable: the awaitable to resolve + """ + ctx, key = self._resolve_nested_context(awaitable.key) + + if awaitable.action == AwaitableAction.ASSIGN: + ctx[key] = value + elif awaitable.action == AwaitableAction.APPEND: + # Find the same awaitable inserted in the context + container = ctx[key] + for index, placeholder in enumerate(container): + if isinstance(placeholder, Awaitable) and placeholder.pk == awaitable.pk: + container[index] = value + break + else: + raise AssertionError(f'Awaitable `{awaitable.pk} was not found in `ctx.{awaitable.key}`') + else: + raise AssertionError(f'Unsupported awaitable action: {awaitable.action}') + + awaitable.resolved = True + self._awaitables.remove(awaitable) # remove only if everything went ok, otherwise we may lose track + + if not self.has_terminated(): + # the process may be terminated, for example, if the process was killed or excepted + # then we should not try to update it + self._update_process_status() + + @Protect.final + def to_context(self, **kwargs: Awaitable | ProcessNode) -> None: + """Add a dictionary of awaitables to the context. + + This is a convenience method that provides syntactic sugar, for a user to add multiple intersteps that will + assign a certain value to the corresponding key in the context of the work chain. + """ + for key, value in kwargs.items(): + awaitable = construct_awaitable(value) + awaitable.key = key + self._insert_awaitable(awaitable) + + def _update_process_status(self) -> None: + """Set the process status with a message accounting the current sub processes that we are waiting for.""" + if self._awaitables: + status = f'Waiting for child processes: {", ".join([str(_.pk) for _ in self._awaitables])}' + else: + status = None + if self.paused: + # Update the pre-paused status so that when the process is played + # it will be set to the new status + self._pre_paused_status = status + else: + self.set_status(status) + + @override + @Protect.final + async def run(self) -> t.Any: + self._stepper = self._create_stepper() + return await run_with_portal(self._do_step) + + def _do_step(self) -> t.Any: + """Execute the next step in the outline and return the result. + + If the stepper returns a non-finished status and the return value is of type ToContext, the contents of the + ToContext container will be turned into awaitables if necessary. If any awaitables were created, the process + will enter in the Wait state, otherwise it will go to Continue. When the stepper returns that it is done, the + stepper result will be converted to None and returned, unless it is an integer or instance of ExitCode. + """ + from .context import ToContext + + # Under the barrier model the awaitables belong to a single step and are cleared before the next one, which + # is what forces every step to wait for all the children it launched. A streaming stepper keeps them, so + # children launched in earlier steps stay in flight while later steps run. + if self._awaitable_barrier: + self._awaitables = [] + result: t.Any = None + + try: + assert self._stepper is not None + finished, stepper_result = self._stepper.step() + except _PropagateReturn as exception: + finished, result = True, exception.exit_code + else: + # Set result to None unless stepper_result was non-zero positive integer or ExitCode with similar status + if isinstance(stepper_result, int) and stepper_result > 0: + result = ExitCode(stepper_result) + elif isinstance(stepper_result, ExitCode) and stepper_result.status > 0: + result = stepper_result + else: + result = None + + # If the stepper said we are finished or the result is an ExitCode, we exit by returning + if finished or isinstance(result, ExitCode): + return result + + if isinstance(stepper_result, ToContext): + self.to_context(**stepper_result) + + if self._awaitables: + return Wait(self._do_step, 'Waiting before next step') + + return Continue(self._do_step) + + def _store_nodes(self, data: t.Any) -> None: + """Recurse through a data structure and store any unstored nodes that are found along the way + + :param data: a data structure potentially containing unstored nodes + """ + if isinstance(data, Node) and not data.is_stored: + data.store() + elif isinstance(data, collections.abc.Mapping): + for _, value in data.items(): + self._store_nodes(value) + elif isinstance(data, collections.abc.Sequence) and not isinstance(data, str): + for value in data: + self._store_nodes(value) + + @override + @Protect.final + def on_exiting(self) -> None: + """Ensure that any unstored nodes in the context are stored, before the state is exited + + After the state is exited the next state will be entered and if persistence is enabled, a checkpoint will + be saved. If the context contains unstored nodes, the serialization necessary for checkpointing will fail. + """ + super().on_exiting() + try: + self._store_nodes(self.ctx) + except Exception: + # An uncaught exception here will have bizarre and disastrous consequences + self.logger.exception('exception in _store_nodes called in on_exiting') + + @Protect.final + def on_wait(self, awaitables: t.Sequence[t.Awaitable]): + """Entering the WAITING state.""" + super().on_wait(awaitables) + if self._awaitables: + self._action_awaitables() + else: + self.call_soon(self.resume) + + def _action_awaitables(self) -> None: + """Register the completion callback for each awaitable that does not already have one. + + Depending on the class type of the awaitable's target a different callback + function will be bound with the awaitable and the runner will be asked to + call it when the target is completed. + + The registration is guarded against duplicates: under the barrier model the awaitables are cleared each + step so the same one is never seen twice, but a streaming stepper keeps its awaitables across steps and + would otherwise register a further callback for the same awaitable on every pass through the waiting state. + """ + for awaitable in self._awaitables: + if awaitable.pk in self._registered_awaitable_pks: + continue + if awaitable.target == AwaitableTarget.PROCESS: + callback = functools.partial(self.call_soon, self._on_awaitable_finished, awaitable) + self.runner.call_on_process_finish(awaitable.pk, callback) + self._registered_awaitable_pks.add(awaitable.pk) + else: + raise AssertionError(f"invalid awaitable target '{awaitable.target}'") + + def _on_awaitable_resolved(self, awaitable: Awaitable) -> None: + """Hook called once a finished awaitable has been resolved onto the context, before the resume decision. + + Defaults to doing nothing. A subclass can use it to run bookkeeping that must see the resolved value and + must happen before the process is resumed, without having to reimplement :meth:`_on_awaitable_finished`. + + :param awaitable: the awaitable that has just been resolved + """ + + def _on_awaitable_finished(self, awaitable: Awaitable) -> None: + """Callback function, for when an awaitable process instance is completed. + + The awaitable will be effectuated on the context of the work chain and removed from the internal list. The + process is then resumed: under the barrier model only once every awaitable has finished, and under the + streaming model as soon as this one does, so a finished child can unblock its dependents while others run. + + :param awaitable: an Awaitable instance + """ + self.logger.info('received callback that awaitable %d has terminated', awaitable.pk) + + try: + node = load_node(awaitable.pk) + except (exceptions.MultipleObjectsError, exceptions.NotExistent): + raise ValueError(f'provided pk<{awaitable.pk}> could not be resolved to a valid Node instance') + + if awaitable.outputs: + value = {entry.link_label: entry.node for entry in node.base.links.get_outgoing()} + else: + value = node # type: ignore[assignment] + + self._resolve_awaitable(awaitable, value) + self._registered_awaitable_pks.discard(awaitable.pk) + self._on_awaitable_resolved(awaitable) + + if self.state == ProcessState.WAITING and (not self._awaitable_barrier or not self._awaitables): + self.resume() From 22e7fa35b1159184d69eb59b06f28d1870709cd7 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 6 Aug 2026 12:28:05 +0200 Subject: [PATCH 14/19] =?UTF-8?q?=E2=9C=A8=20`workgraph`:=20relocate=20the?= =?UTF-8?q?=20subsystem=20into=20core?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lift-and-shift the entire aiida-workgraph engine + authoring layer into `aiida/workgraph/`, so the WorkGraph framework lives in aiida-core. The authoring layer is a mutually-recursive knot (task / socket / registry / decorator / manager / workgraph / task_pool), so it moves as one unit and is refactored in place afterwards. - 49 modules relocated, imports repointed `aiida_workgraph` -> `aiida.workgraph`, and modernised to core's ruff; a scoped per-file-ignore keeps the CamelCase control-flow API (If / While / Map / Zone / TaskPool). - serializer: `AiidaSerializationAdapter` merged into `serialization.py`; the generic + authoring helpers merged into a `utils/` package. - `inspect_aiida_component_type` and `decorator` de-eagered (lazy plugin imports) so `import aiida.workgraph` pulls no plugin, and `import aiida` stays free of node-graph (the optional-extra boundary). - the two plugin task types (pythonjob / shelljob) fold in with lazy imports. - `aiida/workgraph/__init__.py` is hand-maintained (a curated public API), so it is excluded from the autogenerate-imports hook; the subsystem is added to the mypy exclude to be typed incrementally. Entry points still resolve through the aiida-workgraph shims; migrating them into core's `pyproject.toml` is a follow-up. Verified: the full WorkGraph suite passes through the shims at 181 passed / 11 failed, the exact pre-move baseline (the 11 are environmental `python3@localhost` failures), zero regressions. Requires node-graph, and the aiida-workgraph shim package during the transition. --- src/aiida/workgraph/__init__.py | 48 +- src/aiida/workgraph/cli/__init__.py | 13 + src/aiida/workgraph/cli/cmd_task.py | 81 +++ src/aiida/workgraph/cli/cmd_workgraph.py | 20 + src/aiida/workgraph/collection.py | 5 + src/aiida/workgraph/config.py | 29 + src/aiida/workgraph/decorator.py | 309 ++++++++ src/aiida/workgraph/engine/__init__.py | 0 .../workgraph/engine/error_handler_manager.py | 73 ++ src/aiida/workgraph/engine/process.py | 217 ++++++ src/aiida/workgraph/engine/stepper.py | 122 ++++ src/aiida/workgraph/engine/task_actions.py | 71 ++ src/aiida/workgraph/engine/task_manager.py | 578 +++++++++++++++ src/aiida/workgraph/engine/task_state.py | 376 ++++++++++ src/aiida/workgraph/executors/__init__.py | 0 src/aiida/workgraph/executors/builtins.py | 61 ++ src/aiida/workgraph/executors/test.py | 37 + src/aiida/workgraph/manager.py | 194 +++++ src/aiida/workgraph/orm/__init__.py | 0 src/aiida/workgraph/orm/mapping.py | 29 + src/aiida/workgraph/orm/utils.py | 36 + src/aiida/workgraph/properties/__init__.py | 3 + src/aiida/workgraph/properties/builtins.py | 120 +++ .../workgraph/properties/property_pool.py | 5 + src/aiida/workgraph/property.py | 28 + src/aiida/workgraph/registry.py | 12 + src/aiida/workgraph/schemas/__init__.py | 0 src/aiida/workgraph/serialization.py | 29 +- src/aiida/workgraph/socket.py | 57 ++ src/aiida/workgraph/socket_spec.py | 155 ++++ src/aiida/workgraph/sockets/__init__.py | 3 + src/aiida/workgraph/sockets/builtins.py | 78 ++ src/aiida/workgraph/sockets/socket_pool.py | 6 + src/aiida/workgraph/task.py | 254 +++++++ src/aiida/workgraph/tasks/__init__.py | 3 + src/aiida/workgraph/tasks/aiida.py | 138 ++++ src/aiida/workgraph/tasks/builtins.py | 287 ++++++++ src/aiida/workgraph/tasks/function_task.py | 110 +++ src/aiida/workgraph/tasks/graph_task.py | 123 ++++ src/aiida/workgraph/tasks/monitors.py | 58 ++ src/aiida/workgraph/tasks/pythonjob_tasks.py | 343 +++++++++ src/aiida/workgraph/tasks/shelljob_task.py | 205 ++++++ src/aiida/workgraph/tasks/subgraph_task.py | 92 +++ src/aiida/workgraph/tasks/task_pool.py | 6 + src/aiida/workgraph/tasks/tests.py | 59 ++ src/aiida/workgraph/utils.py | 147 ---- src/aiida/workgraph/utils/__init__.py | 682 ++++++++++++++++++ src/aiida/workgraph/utils/control.py | 157 ++++ src/aiida/workgraph/utils/logging.py | 14 + src/aiida/workgraph/utils/svg_to_html.py | 99 +++ src/aiida/workgraph/workgraph.py | 662 +++++++++++++++++ 51 files changed, 6077 insertions(+), 157 deletions(-) create mode 100644 src/aiida/workgraph/cli/__init__.py create mode 100644 src/aiida/workgraph/cli/cmd_task.py create mode 100644 src/aiida/workgraph/cli/cmd_workgraph.py create mode 100644 src/aiida/workgraph/collection.py create mode 100644 src/aiida/workgraph/config.py create mode 100644 src/aiida/workgraph/decorator.py create mode 100644 src/aiida/workgraph/engine/__init__.py create mode 100644 src/aiida/workgraph/engine/error_handler_manager.py create mode 100644 src/aiida/workgraph/engine/process.py create mode 100644 src/aiida/workgraph/engine/stepper.py create mode 100644 src/aiida/workgraph/engine/task_actions.py create mode 100644 src/aiida/workgraph/engine/task_manager.py create mode 100644 src/aiida/workgraph/engine/task_state.py create mode 100644 src/aiida/workgraph/executors/__init__.py create mode 100644 src/aiida/workgraph/executors/builtins.py create mode 100644 src/aiida/workgraph/executors/test.py create mode 100644 src/aiida/workgraph/manager.py create mode 100644 src/aiida/workgraph/orm/__init__.py create mode 100644 src/aiida/workgraph/orm/mapping.py create mode 100644 src/aiida/workgraph/orm/utils.py create mode 100644 src/aiida/workgraph/properties/__init__.py create mode 100644 src/aiida/workgraph/properties/builtins.py create mode 100644 src/aiida/workgraph/properties/property_pool.py create mode 100644 src/aiida/workgraph/property.py create mode 100644 src/aiida/workgraph/registry.py create mode 100644 src/aiida/workgraph/schemas/__init__.py create mode 100644 src/aiida/workgraph/socket.py create mode 100644 src/aiida/workgraph/socket_spec.py create mode 100644 src/aiida/workgraph/sockets/__init__.py create mode 100644 src/aiida/workgraph/sockets/builtins.py create mode 100644 src/aiida/workgraph/sockets/socket_pool.py create mode 100644 src/aiida/workgraph/task.py create mode 100644 src/aiida/workgraph/tasks/__init__.py create mode 100644 src/aiida/workgraph/tasks/aiida.py create mode 100644 src/aiida/workgraph/tasks/builtins.py create mode 100644 src/aiida/workgraph/tasks/function_task.py create mode 100644 src/aiida/workgraph/tasks/graph_task.py create mode 100644 src/aiida/workgraph/tasks/monitors.py create mode 100644 src/aiida/workgraph/tasks/pythonjob_tasks.py create mode 100644 src/aiida/workgraph/tasks/shelljob_task.py create mode 100644 src/aiida/workgraph/tasks/subgraph_task.py create mode 100644 src/aiida/workgraph/tasks/task_pool.py create mode 100644 src/aiida/workgraph/tasks/tests.py delete mode 100644 src/aiida/workgraph/utils.py create mode 100644 src/aiida/workgraph/utils/__init__.py create mode 100644 src/aiida/workgraph/utils/control.py create mode 100644 src/aiida/workgraph/utils/logging.py create mode 100644 src/aiida/workgraph/utils/svg_to_html.py create mode 100644 src/aiida/workgraph/workgraph.py diff --git a/src/aiida/workgraph/__init__.py b/src/aiida/workgraph/__init__.py index e6e9290dbb..f02748edd4 100644 --- a/src/aiida/workgraph/__init__.py +++ b/src/aiida/workgraph/__init__.py @@ -8,30 +8,60 @@ ########################################################################### """The AiiDA WorkGraph: a data-dependency workflow language and runtime. -This subpackage is the AiiDA-specific layer over the generic node-graph SDK. It is imported lazily and depends on -node-graph, so it must not be imported from aiida-core's own import path; a plain ``import aiida`` stays free of the -node-graph dependency, which the (optional) workgraph install provides. +This subpackage is the AiiDA-specific layer over the generic node-graph SDK. It depends on node-graph, so it is not +imported from aiida-core's own import path; a plain ``import aiida`` stays free of the node-graph dependency, which the +workgraph install provides. Hand-maintained (excluded from the autogenerate-imports hook) so that importing it exposes +the curated authoring API without pulling optional plugin task types. """ -# AUTO-GENERATED - -# fmt: off +from aiida import __version__ as __version__ +from . import decorator +from . import socket_spec as spec +from .collection import group from .enums import * +from .manager import If, Map, While, Zone, get_current_graph from .serialization import * -from .utils import * +from .socket_spec import dynamic, meta, namespace, select +from .task import Task +from .tasks import TaskPool +from .utils import ( + get_nested_dict, + resolve_node_link_managers, + update_nested_dict, + update_nested_dict_with_special_keys, +) +from .workgraph import WorkGraph + +# The ``task`` decorator must be the ``aiida.workgraph.task`` package attribute; the +# ``from .task import Task`` import above otherwise binds the ``task`` *submodule* there. +task = decorator.task __all__ = ( 'TERMINAL_TASK_STATES', + 'AiidaSerializationAdapter', + 'If', + 'Map', 'RuntimeInfoKey', + 'Task', 'TaskAction', 'TaskActionMessage', + 'TaskPool', 'TaskState', + 'While', + 'WorkGraph', + 'Zone', + 'dynamic', + 'get_current_graph', 'get_nested_dict', + 'group', + 'meta', + 'namespace', 'resolve_node_link_managers', + 'select', 'serialize_ports', + 'spec', + 'task', 'update_nested_dict', 'update_nested_dict_with_special_keys', ) - -# fmt: on diff --git a/src/aiida/workgraph/cli/__init__.py b/src/aiida/workgraph/cli/__init__.py new file mode 100644 index 0000000000..8c4b6b5044 --- /dev/null +++ b/src/aiida/workgraph/cli/__init__.py @@ -0,0 +1,13 @@ +"""Sub commands of the ``verdi`` command line interface. + +The commands need to be imported here for them to be registered with the top-level command group. +""" + +from aiida.plugins.entry_point import get_entry_points +from aiida.workgraph.cli import cmd_task + +eps = get_entry_points('workgraph.cmdline') +for ep in eps: + ep.load() + +__all__ = ['cmd_task'] diff --git a/src/aiida/workgraph/cli/cmd_task.py b/src/aiida/workgraph/cli/cmd_task.py new file mode 100644 index 0000000000..b002e75376 --- /dev/null +++ b/src/aiida/workgraph/cli/cmd_task.py @@ -0,0 +1,81 @@ +"""`verdi process` command.""" + +import click + +from aiida.cmdline.params import arguments, options +from aiida.cmdline.utils import decorators, echo +from aiida.workgraph.cli.cmd_workgraph import workgraph + +REPAIR_INSTRUCTIONS = """\ +If one ore more processes are unreachable, you can run the following commands to try and repair them: + + verdi daemon stop + verdi process repair + verdi daemon start +""" + + +@workgraph.group('task') +def workgraph_task(): + """Inspect and manage processes.""" + + +@workgraph_task.command('list') +@arguments.PROCESS() +@options.TIMEOUT() +@decorators.with_dbenv() +def task_show(process, timeout): + """List the tasks for one or multiple work graphs.""" + from aiida.workgraph import WorkGraph + + wg = WorkGraph.load(process.pk) + wg.show() + + +@workgraph_task.command('pause') +@arguments.PROCESS() +@click.argument('tasks', nargs=-1) +@options.TIMEOUT() +@decorators.with_dbenv() +def task_pause(process, tasks, timeout): + """Pause task.""" + from aiida.engine.processes import control + from aiida.workgraph.utils.control import pause_tasks + + try: + _, msg = pause_tasks(process.pk, tasks, timeout) + except control.ProcessTimeoutException as exception: + echo.echo_critical(f'{exception}\n{REPAIR_INSTRUCTIONS}') + + +@workgraph_task.command('play') +@arguments.PROCESS() +@click.argument('tasks', nargs=-1) +@options.TIMEOUT() +@decorators.with_dbenv() +def task_play(process, tasks, timeout): + """Play task.""" + from aiida.engine.processes import control + from aiida.workgraph.utils.control import play_tasks + + try: + _, msg = play_tasks(process.pk, tasks, timeout) + except control.ProcessTimeoutException as exception: + echo.echo_critical(f'{exception}\n{REPAIR_INSTRUCTIONS}') + + +@workgraph_task.command('kill') +@arguments.PROCESS() +@click.argument('tasks', nargs=-1) +@options.TIMEOUT() +@decorators.with_dbenv() +def task_kill(process, tasks, timeout): + """Kill task.""" + from aiida.engine.processes import control + from aiida.workgraph.utils.control import kill_tasks + + print('tasks', tasks) + try: + kill_tasks(process.pk, tasks, timeout) + except control.ProcessTimeoutException as exception: + echo.echo_critical(f'{exception}\n{REPAIR_INSTRUCTIONS}') diff --git a/src/aiida/workgraph/cli/cmd_workgraph.py b/src/aiida/workgraph/cli/cmd_workgraph.py new file mode 100644 index 0000000000..a6740619f6 --- /dev/null +++ b/src/aiida/workgraph/cli/cmd_workgraph.py @@ -0,0 +1,20 @@ +"""The main `workgraph` click group.""" + +import click + +from aiida.cmdline.groups import VerdiCommandGroup +from aiida.cmdline.params import options, types +from aiida.workgraph import __version__ + + +# Pass the version explicitly to ``version_option`` otherwise editable installs can show the wrong version number +@click.group(cls=VerdiCommandGroup, context_settings={'help_option_names': ['--help', '-h']}) +@options.PROFILE(type=types.ProfileParamType(load_profile=True), expose_value=False) +@options.VERBOSITY() +@click.version_option( + __version__, + package_name='aiida_core', + message='AiiDA-WorkGraph version %(version)s', +) +def workgraph(): + """The command line interface of AiiDA-WorkGraph.""" diff --git a/src/aiida/workgraph/collection.py b/src/aiida/workgraph/collection.py new file mode 100644 index 0000000000..f3c9925cc4 --- /dev/null +++ b/src/aiida/workgraph/collection.py @@ -0,0 +1,5 @@ +from node_graph.collection import group + +__all__ = [ + 'group', +] diff --git a/src/aiida/workgraph/config.py b/src/aiida/workgraph/config.py new file mode 100644 index 0000000000..f155f0c397 --- /dev/null +++ b/src/aiida/workgraph/config.py @@ -0,0 +1,29 @@ +import json +from pathlib import Path + +from aiida.engine import CalcJob, WorkChain +from aiida.manage import get_config +from aiida.orm.nodes.process.calculation.calcfunction import CalcFunctionNode +from aiida.orm.nodes.process.workflow.workfunction import WorkFunctionNode + +WORKGRAPH_EXTRA_KEY = '_workgraph' +WORKGRAPH_SHORT_EXTRA_KEY = '_workgraph_short' + +task_types = { + CalcFunctionNode: 'CALCFUNCTION', + WorkFunctionNode: 'WORKFUNCTION', + CalcJob: 'CALCJOB', + WorkChain: 'WORKCHAIN', +} + + +def load_config() -> dict: + """Load the configuration from the config file.""" + config = get_config() + config_file_path = Path(config.dirpath) / 'workgraph.json' + try: + with config_file_path.open('r') as f: + config = json.load(f) + except FileNotFoundError: + config = {} + return config diff --git a/src/aiida/workgraph/decorator.py b/src/aiida/workgraph/decorator.py new file mode 100644 index 0000000000..6cdcb3589d --- /dev/null +++ b/src/aiida/workgraph/decorator.py @@ -0,0 +1,309 @@ +from __future__ import annotations + +import inspect +from collections.abc import Callable + +from node_graph.error_handler import ErrorHandlerSpec, normalize_error_handlers +from node_graph.socket_spec import SocketSpec +from node_graph.task_spec import TaskSpec + +from aiida.engine import CalcJob, WorkChain, calcfunction, workfunction +from aiida.workgraph.task import Task +from aiida.workgraph.tasks.aiida import AiiDAProcessTask, _build_aiida_function_taskspec + +from .task import TaskHandle +from .workgraph import WorkGraph + + +def _spec_for( + obj, + *, + identifier: str | None, + inputs: SocketSpec | None = None, + outputs: SocketSpec | None = None, + catalog: str | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, +) -> TaskSpec: + # AiiDA process classes + if inspect.isclass(obj) and issubclass(obj, (CalcJob, WorkChain)): + return AiiDAProcessTask.build(obj, attached_error_handlers=error_handlers) + + # AiiDA process functions (calcfunction/workfunction) + if callable(obj) and getattr(obj, 'node_class', False): + return _build_aiida_function_taskspec( + obj, + identifier=identifier, + in_spec=inputs, + out_spec=outputs, + error_handlers=error_handlers, + catalog=catalog or 'Others', + ) + + # Plain Python function -> PyFunction + if callable(obj): + # Lazy: keeps ``import aiida.workgraph`` free of the aiida-pythonjob dependency. + from aiida.workgraph.tasks.pythonjob_tasks import build_pyfunction_taskspec + + spec = build_pyfunction_taskspec( + obj, + identifier=identifier, + in_spec=inputs, + out_spec=outputs, + error_handlers=error_handlers, + catalog=catalog or 'Others', + ) + return spec + + raise ValueError(f'Unsupported object for @task: {obj!r}') + + +def build_task_from_callable( + executor: Callable, + inputs: SocketSpec | list | None = None, + outputs: SocketSpec | list | None = None, +) -> Task: + """Build task from a callable object. + First, check if the executor is already a task. + If not, check if it is a function or a class. + If it is a function, build task from function. + If it is a class, it only supports CalcJob and WorkChain. + """ + from node_graph.task import Task + + # if it is already a task, return it + if ( + hasattr(executor, '_TaskCls') and inspect.isclass(executor._TaskCls) and issubclass(executor._TaskCls, Task) + ) or (inspect.isclass(executor) and issubclass(executor, Task)): + return executor + if inspect.isfunction(executor): + # calcfunction and workfunction + if getattr(executor, 'node_class', False): + return task(inputs=inputs, outputs=outputs)(executor) + else: + return task(inputs=inputs, outputs=outputs)(executor) + elif issubclass(executor, CalcJob) or issubclass(executor, WorkChain): + if inputs is not None or outputs is not None: + raise ValueError('Can not override inputs or outputs of an AiiDA process classes.') + return task()(executor) + raise ValueError(f'The executor {executor} is not supported.') + + +def nonfunctional_usage(callable: Callable): + """ + This is a decorator for a decorator factory (a function that returns a decorator). + It allows the usage of the decorator factory in a nonfunctional way. So a decorator + factory that has been decorated by this decorator that could only be used befor like + this + + .. code-block:: python + + @decorator_factory() + def foo(): + pass + + can now be also used like this + + .. code-block:: python + + @decorator_factory + def foo(): + pass + + """ + + def decorator_task_wrapper(*args, **kwargs): + if len(args) == 1 and isinstance(args[0], Callable) and len(kwargs) == 0: + return callable()(args[0]) + else: + return callable(*args, **kwargs) + + return decorator_task_wrapper + + +class TaskDecoratorCollection: + """Collection of task decorators.""" + + @staticmethod + @nonfunctional_usage + def decorator_task( + identifier: str | None = None, + inputs: SocketSpec | list | None = None, + outputs: SocketSpec | list | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, + catalog: str = 'Others', + ) -> Callable: + """Generate a decorator that register a function as a task. + + Attributes: + indentifier (str): task identifier + catalog (str): task catalog + inputs (list): task inputs + outputs (list): task outputs + """ + + def decorator(obj: WorkGraph | type | callable) -> TaskHandle: + normalized_handlers = normalize_error_handlers(error_handlers) + spec = _spec_for( + obj, + identifier=identifier, + catalog=catalog, + inputs=inputs, + outputs=outputs, + error_handlers=normalized_handlers, + ) + + handle = TaskHandle(spec) + handle._callable = obj + return handle + + return decorator + + @staticmethod + @nonfunctional_usage + def decorator_graph( + identifier: str | None = None, + catalog: str | None = None, + inputs: SocketSpec | list | None = None, + outputs: SocketSpec | list | None = None, + max_depth: int = 100, + max_number_jobs: int | None = None, + ) -> Callable: + """Generate a decorator that register a function as a graph task. + Attributes: + indentifier (str): task identifier + catalog (str): task catalog + inputs (list): task inputs + outputs (list): task outputs + """ + + def decorator(func) -> TaskHandle: + from aiida.workgraph.tasks.graph_task import _build_graph_task_taskspec + + handle = TaskHandle( + _build_graph_task_taskspec( + func, + identifier=identifier, + catalog=catalog, + in_spec=inputs, + out_spec=outputs, + max_depth=max_depth, + max_number_jobs=max_number_jobs, + ) + ) + handle._callable = func + return handle + + return decorator + + @staticmethod + @nonfunctional_usage + def calcfunction( + inputs: SocketSpec | list | None = None, + outputs: SocketSpec | list | None = None, + catalog: str | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, + ) -> Callable: + def decorator(func) -> TaskHandle: + func_decorated = calcfunction(func) + handle = TaskHandle( + _build_aiida_function_taskspec( + func_decorated, + in_spec=inputs, + out_spec=outputs, + catalog=catalog, + error_handlers=error_handlers, + ) + ) + handle._callable = func_decorated + return handle + + return decorator + + @staticmethod + @nonfunctional_usage + def workfunction( + inputs: SocketSpec | list | None = None, + outputs: SocketSpec | list | None = None, + catalog: str | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, + ) -> Callable: + def decorator(func) -> TaskHandle: + func_decorated = workfunction(func) + handle = TaskHandle( + _build_aiida_function_taskspec( + func_decorated, + in_spec=inputs, + out_spec=outputs, + catalog=catalog, + error_handlers=error_handlers, + ) + ) + handle._callable = func_decorated + return handle + + return decorator + + @staticmethod + @nonfunctional_usage + def pythonjob( + inputs: SocketSpec | list | None = None, + outputs: SocketSpec | list | None = None, + catalog: str | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, + ) -> Callable: + def decorator(func) -> TaskHandle: + from aiida.workgraph.tasks.pythonjob_tasks import build_pythonjob_taskspec + + spec = build_pythonjob_taskspec( + func, + in_spec=inputs, + out_spec=outputs, + catalog=catalog, + error_handlers=error_handlers, + ) + handle = TaskHandle(spec) + handle._callable = func + return handle + + return decorator + + @staticmethod + @nonfunctional_usage + def monitor( + inputs: SocketSpec | list | None = None, + outputs: SocketSpec | list | None = None, + catalog: str | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, + ) -> Callable: + def decorator(func) -> TaskHandle: + from aiida.workgraph.tasks.pythonjob_tasks import build_monitor_function_taskspec + + handle = TaskHandle( + build_monitor_function_taskspec( + func, + in_spec=inputs, + out_spec=outputs, + catalog=catalog, + error_handlers=error_handlers, + ) + ) + handle._callable = func + return handle + + return decorator + + # Making decorator_task accessible as 'task' + task = decorator_task + + # Making decorator_graph accessible as 'graph' + graph = decorator_graph + + def __call__(self, *args, **kwargs): + # This allows using '@task' to directly apply the decorator_task functionality + if len(args) == 1 and isinstance(args[0], Callable) and len(kwargs) == 0: + return self.decorator_task()(args[0]) + else: + return self.decorator_task(*args, **kwargs) + + +task = TaskDecoratorCollection() diff --git a/src/aiida/workgraph/engine/__init__.py b/src/aiida/workgraph/engine/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/aiida/workgraph/engine/error_handler_manager.py b/src/aiida/workgraph/engine/error_handler_manager.py new file mode 100644 index 0000000000..c43dd2ae3e --- /dev/null +++ b/src/aiida/workgraph/engine/error_handler_manager.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import logging +import traceback +from typing import Any + +from node_graph.error_handler import ErrorHandlerSpec + + +class ErrorHandlerManager: + def __init__(self, process: Any, logger: logging.Logger) -> None: + self.process = process + self.logger = logger + + @property + def ctx(self) -> Any: + """Read the context off the process, which replaces it wholesale when loading from a checkpoint.""" + return self.process.ctx + + def run_error_handlers(self, task_name: str) -> None: + """Run error handlers for a task.""" + + self.process.report(f'Run error handlers for {task_name}') + + node = self.process.task_manager.state_manager.get_task_runtime_info(task_name, 'process') + if not node or not node.exit_status: + return + # error_handlers from the task + error_handlers = self.process.wg.tasks[task_name].get_error_handlers() + for data in error_handlers.values(): + if node.exit_status in data.exit_codes: + self.run_error_handler(data, task_name) + return + # error_handlers from the workgraph + for data in self.process.wg._error_handlers.values(): + if node.exit_code.status in data['tasks'].get(task_name, {}).get('exit_codes', []): + self.run_error_handler(data, task_name) + return + + def run_error_handler(self, handler: ErrorHandlerSpec, task_name: str) -> None: + """Run the error handler for a task.""" + from inspect import signature + + from node_graph.executor import RuntimeExecutor + + executor = RuntimeExecutor(**handler.executor.to_dict()).callable + executor_sig = signature(executor) + self.process.report(f'Run error handler: {executor.__name__}') + if handler.retry < handler.max_retries: + task = self.process.task_manager.get_task(task_name) + prev_allow_overrides = getattr(task, '_allow_input_overrides', False) + try: + task.set_input_resolver(self.process.task_manager.get_socket_value) + task._allow_input_overrides = True + # Run the error handler to update the inputs of the task + if 'engine' in executor_sig.parameters: + msg = executor(task, engine=self, **(handler.kwargs or {})) + else: + msg = executor(task, **(handler.kwargs or {})) + # Reset the task to rerun it + self.process.task_manager.state_manager.reset_task(task.name) + # Save the updated task into self.ctx._wgdata + tdata = task.to_dict() + self.ctx._wgdata['tasks'][task.name] = tdata + if msg: + self.process.report(msg) + handler.retry += 1 + except Exception as e: + error_traceback = traceback.format_exc() # Capture the full traceback + self.logger.error(f'Error in running error handler for {task_name}: {e}\n{error_traceback}') + self.process.report(f'Error in running error handler for {task_name}: {e}\n{error_traceback}') + finally: + task._allow_input_overrides = prev_allow_overrides diff --git a/src/aiida/workgraph/engine/process.py b/src/aiida/workgraph/engine/process.py new file mode 100644 index 0000000000..dda68621ec --- /dev/null +++ b/src/aiida/workgraph/engine/process.py @@ -0,0 +1,217 @@ +"""The AiiDA process that executes a work graph.""" + +from __future__ import annotations + +import logging +import typing as t + +import kiwipy +from plumpy import process_comms +from plumpy.workchains import Stepper + +from aiida.common.lang import override +from aiida.engine.processes.workchains.awaitable import Awaitable +from aiida.engine.processes.workchains.workchain import WorkChain, WorkChainSpec +from aiida.orm import WorkGraphNode +from aiida.workgraph.engine.error_handler_manager import ErrorHandlerManager +from aiida.workgraph.engine.stepper import DagStepper +from aiida.workgraph.engine.task_manager import TaskManager +from aiida.workgraph.enums import TaskActionMessage + +if t.TYPE_CHECKING: + from aiida.engine.runners import Runner + from aiida.workgraph import WorkGraph + +__all__ = ('WorkGraphProcess', 'WorkGraphSpec') + + +class WorkGraphSpec(WorkChainSpec): + WORKGRAPH_DATA_KEY = 'workgraph_data' + + +class WorkGraphProcess(WorkChain): + """Execute a work graph, scheduling its tasks by their data dependencies. + + A work chain declares its execution order up front as an outline; a work graph derives it from the links + between tasks, which are only known once the graph is built. Everything else a work chain provides (context, + awaitables, checkpointing, node lifecycle) applies unchanged, so this supplies a :class:`DagStepper` through the + stepper hooks and inherits the rest. + + The one way it departs from :class:`~aiida.engine.processes.workchains.workchain.WorkChain` is concurrency: a + work chain waits for everything a step launched before starting the next, whereas here independent branches + stay in flight together. That is not overridden here; it follows from :class:`DagStepper` declaring + ``awaitable_barrier = False``, which the work chain honours. Only two small work-graph-specific hooks remain, + :meth:`_action_awaitables` (surface the waiting status in the report) and :meth:`_on_awaitable_resolved` + (record the finished child's outcome on its task). + """ + + # Narrowing the node and spec classes is how every AiiDA process specialises its base; mypy sees plain mutable + # class attributes and flags the covariance. + _node_class = WorkGraphNode # type: ignore[mutable-override] + _spec_class = WorkGraphSpec # type: ignore[mutable-override] + + def __init__( + self, + inputs: dict[str, t.Any] | None = None, + logger: logging.Logger | None = None, + runner: Runner | None = None, + enable_persistence: bool = True, + ) -> None: + super().__init__(inputs, logger, runner, enable_persistence=enable_persistence) + self._init_runtime_state() + self._init_managers() + + def _init_runtime_state(self) -> None: + """Initialise the state that is rebuilt on every load rather than restored from the checkpoint.""" + self._wg: WorkGraph | None = None + + def _init_managers(self) -> None: + self.task_manager = TaskManager(self.logger, self.runner, self) + self.error_handler_manager = ErrorHandlerManager(self, self.logger) + + @classmethod + def define(cls, spec: WorkGraphSpec) -> None: # type: ignore[override] + super().define(spec) + spec.input_namespace( + 'graph_inputs', + dynamic=True, + required=False, + help='Graph level inputs', + ) + spec.input_namespace( + 'tasks', + dynamic=True, + required=False, + help='Tasks inputs', + ) + spec.input_namespace( + spec.WORKGRAPH_DATA_KEY, + dynamic=True, + required=False, + help='WorkGraph data', + ) + spec.exit_code(2, 'ERROR_SUBPROCESS', message='A subprocess has failed.') + spec.outputs.dynamic = True + spec.exit_code(201, 'UNKNOWN_MESSAGE_TYPE', message='The message type is unknown.') + spec.exit_code(202, 'UNKNOWN_TASK_TYPE', message='The task type is unknown.') + spec.exit_code( + 301, + 'OUTPUS_NOT_MATCH_RESULTS', + message='The outputs of the process do not match the results.', + ) + spec.exit_code( + 302, + 'TASK_FAILED', + message='Some of the tasks failed.', + ) + spec.exit_code( + 303, + 'TASK_NON_ZERO_EXIT_STATUS', + message='Some of the tasks exited with non-zero status.', + ) + + @property + def wg(self) -> WorkGraph: + """The work graph being executed, rebuilt from the context the first time it is needed after a reload.""" + if self._wg is None: + from aiida.workgraph import WorkGraph + + self._wg = WorkGraph.from_dict(self.ctx._wgdata) + return self._wg + + def set_workgraph_data(self, wgdata: dict[str, t.Any]) -> None: + """Install the graph data, keeping the live graph and the checkpointed copy in step.""" + from aiida.workgraph import WorkGraph + + self.ctx._wgdata = wgdata + self._wg = WorkGraph.from_dict(wgdata) + + def _create_stepper(self) -> Stepper: + stepper = DagStepper(self) # type: ignore[arg-type] + stepper.setup() + return stepper + + def _recreate_stepper(self, saved_state: t.Any) -> Stepper: + """Restore the stepper after a checkpoint. + + ``saved_state`` is unused: :class:`DagStepper` holds no state of its own, everything it needs lives in the + context, which the base class has already restored. Notably :meth:`DagStepper.setup` must not run again, + as it would reset the execution bookkeeping and rerun the finished tasks. + """ + return DagStepper(self) # type: ignore[arg-type] + + @override + def load_instance_state(self, saved_state: t.MutableMapping[str, t.Any], load_context: t.Any) -> None: + from aiida.orm.utils.log import create_logger_adapter + + # `WorkChain.load_instance_state` re-registers the awaitable callbacks before returning, so the runtime + # state it consults has to be in place first. + self._init_runtime_state() + + # `no-untyped-call` here and below: both are unannotated aiida-core internals. + super().load_instance_state(saved_state, load_context) # type: ignore[no-untyped-call] + + # TODO: avoid hardcoding the logger + self.node._logger = logging.getLogger('aiida.orm.nodes.process.workflow.workchain.WorkChainNode') # type: ignore[assignment] + # First time the property is called after the node is stored, create the logger adapter + self.node._logger_adapter = create_logger_adapter(self.node._logger, self.node) # type: ignore[no-untyped-call] + self.set_logger(self.node._logger_adapter) + + self._init_managers() + + def _action_awaitables(self) -> None: + """Register the awaitable callbacks (via `WorkChain`), then surface the waiting status in the report log. + + `WorkChain` records "Waiting for child processes: ..." only as the process status; echoing it to the + report makes it visible in `verdi process report` when a graph pauses for its children. + """ + super()._action_awaitables() + if self._awaitables: + self.report(f'Process status: {self.status}') + + def _on_awaitable_resolved(self, awaitable: Awaitable) -> None: + """Record a finished child's outcome on its task before the process decides whether to resume. + + This is the only work-graph-specific step in the awaitable lifecycle. The rest, including resuming as soon + as any child finishes rather than only once all do, comes from `WorkChain`, because :class:`DagStepper` + declares ``awaitable_barrier = False``. + + :param awaitable: the awaitable whose target process has terminated + """ + self.task_manager.state_manager.update_task_state(awaitable.key) + + def _build_process_label(self) -> str: + """Use the workgraph name as the process label.""" + return f'WorkGraph<{self.inputs[WorkGraphSpec.WORKGRAPH_DATA_KEY]["name"]}>' + + def on_create(self) -> None: + """Called when a Process is created.""" + from aiida.workgraph.utils import save_workgraph_data + + super().on_create() + raw_inputs = dict(self.inputs) + self.node.label = raw_inputs[WorkGraphSpec.WORKGRAPH_DATA_KEY]['name'] + save_workgraph_data(self.node, raw_inputs) + + def apply_action(self, msg: TaskActionMessage) -> None: + if msg['catalog'] == 'task': + self.task_manager.action_manager.apply_task_actions(msg) + else: + self.report(f'Unknow message type {msg}') + + def message_receive(self, _comm: kiwipy.Communicator, msg: dict[str, t.Any]) -> t.Any: + """Handle the work-graph-specific ``custom`` intent and defer every standard intent to the base class. + + Only the ``custom`` intent, which carries task actions such as pausing or skipping an individual task, is + particular to a work graph. The rest belong to plumpy's message protocol, which decides such things as + which key holds the pause text and whether a kill is forced; that protocol changes, so reimplementing it + here means silently drifting out of step with it. + + :param _comm: the communicator that sent the message + :param msg: the message + :return: the outcome of processing the message, sent back as the response to the sender + """ + if msg[process_comms.INTENT_KEY] == 'custom': + return self._schedule_rpc(self.apply_action, msg=msg) + + return super().message_receive(_comm, msg) diff --git a/src/aiida/workgraph/engine/stepper.py b/src/aiida/workgraph/engine/stepper.py new file mode 100644 index 0000000000..4bd054e305 --- /dev/null +++ b/src/aiida/workgraph/engine/stepper.py @@ -0,0 +1,122 @@ +"""Data-dependency stepping for a work graph.""" + +from __future__ import annotations + +import typing as t + +from node_graph.config import BUILTIN_TASKS +from plumpy.workchains import Stepper + +from aiida.engine.processes.exit_code import ExitCode +from aiida.workgraph.enums import TaskState + +if t.TYPE_CHECKING: + from aiida.workgraph.engine.process import WorkGraphProcess + + +class DagStepper(Stepper): + """Drive a process by the data dependencies between its tasks instead of a declared outline. + + A :class:`~aiida.engine.processes.workchains.workchain.WorkChain` outline fixes at class-definition time both + which steps exist and in what order they run. Here neither is known until the graph is built: a step means + "launch whatever became ready since the last one", and readiness is derived from the links between tasks. Two + consequences follow, and they are why this cannot be expressed as an outline: + + * A step may launch several tasks at once, and independent branches stay in flight together. The outline + stepper's implicit barrier, waiting for everything launched in a step before starting the next, would + serialise them. + * The number of steps is a property of the graph, not of the class. + + The stepper deliberately keeps no state of its own. Everything it needs (the graph data, task results and + execution bookkeeping) lives in the process context, which is checkpointed with the process, so restoring from + a checkpoint is a matter of constructing a fresh instance against the restored context. + """ + + # Opt out of the outline barrier: awaitables persist across steps and the process resumes as each child + # finishes, so a task whose inputs are ready starts while an unrelated task is still running. + awaitable_barrier = False + + @property + def process(self) -> WorkGraphProcess: + """The process being stepped (:class:`~plumpy.workchains.Stepper` names it ``_workchain``).""" + return t.cast('WorkGraphProcess', self._workchain) + + def setup(self) -> None: + """Materialise the graph from the process inputs and seed the execution bookkeeping in the context. + + Called once, when the process starts. Restoring from a checkpoint must *not* call this: the context it + writes is part of the checkpoint, so re-running it would discard the progress made so far. + """ + from aiida.workgraph.utils import restore_workgraph_data_from_raw_inputs + + process = self.process + # Track which awaitables already have a completion callback registered, see + # `WorkGraphProcess._action_awaitables`. + process.ctx._awaitable_actions = [] + process.ctx._new_data = {} + process.ctx._executed_tasks = [] + + wgdata = restore_workgraph_data_from_raw_inputs(dict(process.inputs or {})) + process.set_workgraph_data(wgdata) + + # Expose the graph-level containers as if they were the results of tasks, so that links pointing at them + # resolve through the same lookup as links between real tasks. + process.ctx._task_results = { + 'graph_ctx': process.wg.ctx._value, + 'graph_inputs': process.wg.inputs._value, + 'graph_outputs': process.wg.outputs._value, + } + process.task_manager.set_task_results() + process.task_manager.state_manager.update_meta_tasks('graph_inputs') + for task_name in BUILTIN_TASKS: + process.task_manager.state_manager.set_task_runtime_info(task_name, 'state', TaskState.FINISHED) + + def step(self) -> tuple[bool, t.Any]: + """Launch every task whose inputs are now available, and report whether the graph is complete. + + :return: ``(finished, result)``. While tasks remain, ``(False, None)``. Once the graph is complete, + ``result`` is the exit code reporting which tasks failed if any did, and otherwise whatever + :meth:`finalize` returns after mapping the outputs. + """ + process = self.process + process.task_manager.continue_workgraph() + finished, failure = process.task_manager.is_workgraph_finished() + + if not finished: + return False, None + + # A graph that finished with failed tasks reports them through ``failure``, naming the tasks; the + # outputs are not mapped in that case, since the graph did not produce them. + if isinstance(failure, ExitCode): + return True, failure + + return True, self.finalize() + + def finalize(self) -> ExitCode | None: + """Map the graph outputs onto the process outputs once every task has reached a terminal state. + + :return: the ``TASK_FAILED`` exit code if any task failed, otherwise ``None``. + """ + from aiida.workgraph.utils import resolve_node_link_managers + + process = self.process + # Refresh the meta-tasks so that a graph exposing its context or inputs directly as outputs sees the + # final values rather than those captured at setup. + process.task_manager.state_manager.update_meta_tasks('graph_ctx') + process.task_manager.state_manager.update_meta_tasks('graph_inputs') + process.out_many(resolve_node_link_managers(process.ctx._task_results['graph_outputs'])) + + if process.ctx._new_data: + process.out('new_data', process.ctx._new_data) + + process.report('Finalize workgraph.') + + for task in process.wg.tasks: + if process.task_manager.state_manager.get_task_runtime_info(task.name, 'state') == TaskState.FAILED: + return t.cast(ExitCode, process.exit_codes.TASK_FAILED) + + return None + + def __str__(self) -> str: + """Report progress, which ends up in ``verdi process status`` via ``WorkChain.on_run``.""" + return f'{len(self.process.ctx._executed_tasks)}/{len(self.process.wg.tasks)} tasks launched' diff --git a/src/aiida/workgraph/engine/task_actions.py b/src/aiida/workgraph/engine/task_actions.py new file mode 100644 index 0000000000..1ecf4abf53 --- /dev/null +++ b/src/aiida/workgraph/engine/task_actions.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from collections.abc import Callable + +from typing_extensions import assert_never + +from aiida.workgraph.enums import TaskAction, TaskActionMessage, TaskState + + +class TaskActionManager: + """ + Handles externally-triggered task actions (RESET, PAUSE, PLAY, SKIP, KILL) + to change task states or runtime behavior. + """ + + def __init__(self, state_manager, logger, process): + """ + :param state_manager: A reference to TaskStateManager for updating states. + :param logger: A logger instance. + """ + self.state_manager = state_manager + self.logger = logger + self.process = process + + def apply_task_actions(self, msg: TaskActionMessage) -> None: + """ + Apply task actions to the workgraph based on user or external messages. + + :raises ValueError: if the message carries an unknown task action. + """ + action = TaskAction(msg['action'].upper()) + tasks = msg['tasks'] + self.process.report(f'Action: {action}. Tasks: {tasks}') + + # Each action maps to a callable that takes a single task name. + handler: Callable[[str], None] + match action: + case TaskAction.RESET: + handler = self.state_manager.reset_task + case TaskAction.PAUSE: + handler = self.pause_task + case TaskAction.PLAY: + handler = self.play_task + case TaskAction.SKIP: + handler = self.skip_task + case TaskAction.KILL: + handler = self.kill_task + case _: + assert_never(action) + for name in tasks: + handler(name) + + def pause_task(self, name: str) -> None: + """Mark the task to be paused.""" + self.state_manager.set_task_runtime_info(name, 'action', TaskAction.PAUSE) + self.process.report(f'Task {name} action: PAUSE.') + + def play_task(self, name: str) -> None: + """Remove the PAUSE action on a task.""" + self.state_manager.set_task_runtime_info(name, 'action', '') + self.process.report(f'Task {name} action: PLAY.') + + def skip_task(self, name: str) -> None: + """Force a task to be SKIPPED.""" + self.state_manager.set_task_runtime_info(name, 'state', TaskState.SKIPPED) + self.process.report(f'Task {name} action: SKIP.') + + def kill_task(self, name: str) -> None: + """KILL a running task. + This is not needed for task with AiiDA process, because one can kill the AiiDA process directly. + """ diff --git a/src/aiida/workgraph/engine/task_manager.py b/src/aiida/workgraph/engine/task_manager.py new file mode 100644 index 0000000000..7fde2f8192 --- /dev/null +++ b/src/aiida/workgraph/engine/task_manager.py @@ -0,0 +1,578 @@ +from __future__ import annotations + +import traceback +from typing import Any + +from node_graph.link import TaskLink + +from aiida.engine.processes import Process +from aiida.engine.processes.exit_code import ExitCode +from aiida.workgraph.enums import TaskAction, TaskState +from aiida.workgraph.socket import TaskSocketNamespace +from aiida.workgraph.task import Task +from aiida.workgraph.utils import get_nested_dict + +from .task_actions import TaskActionManager +from .task_state import TaskStateManager + +MAX_NUMBER_AWAITABLES_MSG = 'The maximum number of subprocesses has been reached: {}. Cannot launch the job: {}.' + +process_task_types = [ + 'CALCJOB', + 'WORKCHAIN', + 'GRAPH', + 'SUBGRAPH', + 'PYTHONJOB', + 'SHELLJOB', +] + + +class TaskManager: + """Manages task execution, state updates, and error handling.""" + + def __init__(self, logger, runner, process: Process): + """ + :param logger: A logger instance. + :param runner: An AiiDA runner. + :param process: The AiiDA process object that orchestrates the entire WorkGraph. + """ + self.logger = logger + self.runner = runner + self.process = process + # Sub-managers + self.state_manager = TaskStateManager(logger, process) + self.action_manager = TaskActionManager(self.state_manager, logger, process) + + @property + def ctx(self): + """Read the context off the process, which replaces it wholesale when loading from a checkpoint.""" + return self.process.ctx + + def get_task(self, name: str): + """Get task from the context.""" + task = self.process.wg.tasks[name] + task.set_input_resolver(self.get_socket_value) + task.action = self.state_manager.get_task_runtime_info(name, 'action') + # update task results + # namespace socket does not have a value, but _value + for socket in task.outputs: + if socket._identifier == 'workgraph.namespace': + socket._value = get_nested_dict(self.ctx._task_results[name], socket._name, default=None) + else: + socket.value = get_nested_dict(self.ctx._task_results[name], socket._name, default=None) + return task + + def set_task_results(self) -> None: + from node_graph.config import BUILTIN_TASKS + + for task in self.process.wg.tasks: + if task.name in BUILTIN_TASKS: + # skip built-in nodes, they are not executed + continue + if self.state_manager.get_task_runtime_info(task.name, 'action') == TaskAction.RESET: + self.state_manager.reset_task(task.name) + self.state_manager.update_task_state(task.name) + + def is_workgraph_finished(self) -> tuple[bool, ExitCode | None]: + """Check if the workgraph is finished. + For `while` workgraph, we need check its conditions""" + is_finished = True + failed_tasks = [] + not_finished_tasks = [] + for task in self.process.wg.tasks: + # if the task is in mapped state, we need to check its children (mapped tasks) + if self.state_manager.get_task_runtime_info(task.name, 'state') == TaskState.MAPPED: + self.state_manager.update_template_task_state(task.name) + elif self.state_manager.get_task_runtime_info(task.name, 'state') in { + TaskState.RUNNING, + TaskState.CREATED, + TaskState.PLANNED, + TaskState.READY, + }: + not_finished_tasks.append(task.name) + is_finished = False + elif self.state_manager.get_task_runtime_info(task.name, 'state') == TaskState.FAILED: + failed_tasks.append(task.name) + if is_finished and len(failed_tasks) > 0: + message = f'WorkGraph finished, but tasks: {failed_tasks} failed. Thus all their child tasks are skipped.' + self.process.report(message) + result = ExitCode(302, message) + else: + result = None + # print("not_finished_tasks: ", not_finished_tasks) + return is_finished, result + + def continue_workgraph(self) -> None: + """ + Resume the WorkGraph by looking for tasks that are ready to run. + """ + # self.process.report("Continue workgraph.") + task_to_run = [] + for task in self.process.wg.tasks: + # update task state + if ( + self.state_manager.get_task_runtime_info(task.name, 'state') + in { + TaskState.CREATED, + TaskState.RUNNING, + TaskState.FINISHED, + TaskState.FAILED, + TaskState.SKIPPED, + TaskState.MAPPED, + } + or task.name in self.ctx._executed_tasks + ): + continue + ready, _ = self.state_manager.is_task_ready_to_run(task.name) + if ready: + task_to_run.append(task.name) + self.process.report('tasks ready to run: {}'.format(','.join(task_to_run))) + self.run_tasks(task_to_run) + + def should_run_task(self, task: Task) -> bool: + """Check if the task should run.""" + name = task.name + # skip if the max number of awaitables is reached + if task.task_type.upper() in process_task_types: + if len(self.process._awaitables) >= self.process.wg.max_number_jobs: + self.process.report(MAX_NUMBER_AWAITABLES_MSG.format(self.process.wg.max_number_jobs, name)) + return False + # skip if the task is already executed or if the task is in a skippped state + if ( + name in self.ctx._executed_tasks + or self.state_manager.get_task_runtime_info(name, 'state') == TaskState.SKIPPED + ): + return False + return True + + def run_tasks(self, names: list[str], continue_workgraph: bool = True) -> None: + """Run tasks. + Task type includes: Node, Data, CalcFunction, WorkFunction, CalcJob, WorkChain, GraphBuilder, + WorkGraph, PythonJob, ShellJob, While, If, Zone, GetContext, SetContext, Normal. + + """ + for name in names: + # skip if the max number of awaitables is reached + task = self.process.wg.tasks[name] + task.action = self.state_manager.get_task_runtime_info(name, 'action') + if not self.should_run_task(task): + continue + + self.ctx._executed_tasks.append(name) + # print("-" * 60) + + self.logger.info(f'Run task: {name}, type: {task.task_type}') + inputs = self.get_inputs(name) + # print("kwargs: ", inputs["kwargs"]) + self.ctx._task_results[task.name] = {} + task_type = task.task_type.upper() + if task_type == 'PYFUNCTION': + if task.spec.metadata.get('is_coroutine', False): + self.execute_process_task(task, **inputs) + else: + self.execute_function_task(task, continue_workgraph, **inputs) + elif task_type in ['CALCFUNCTION', 'WORKFUNCTION']: + self.execute_function_task(task, continue_workgraph, **inputs) + elif task_type in [ + 'CALCJOB', + 'WORKCHAIN', + 'SHELLJOB', + 'PYTHONJOB', + 'SUBGRAPH', + 'GRAPH', + 'MONITOR', + ]: + self.execute_process_task(task, **inputs) + elif task_type == 'WHILE': + self.execute_while_task(task) + elif task_type == 'IF': + self.execute_if_task(task) + elif task_type == 'ZONE': + self.execute_zone_task(task) + elif task_type == 'MAP': + self.execute_map_task(task, inputs['kwargs']) + elif task_type == 'NORMAL': + self.execute_normal_task( + task, + continue_workgraph, + **inputs, + ) + else: + self.process.report(f'Unknown task type {task_type}') + self.state_manager.set_task_runtime_info(name, 'state', TaskState.FAILED) + + def execute_function_task(self, task, continue_workgraph=None, args=None, kwargs=None, var_kwargs=None): + """Execute a CalcFunction or WorkFunction task.""" + + try: + process, _ = task.execute(args, kwargs, var_kwargs) + self.state_manager.set_task_runtime_info(task.name, 'process', process) + self.state_manager.update_task_state(task.name) + except Exception as e: + error_traceback = traceback.format_exc() # Capture the full traceback + self.logger.error(f'Error in task {task.name}: {e}\n{error_traceback}') + self.state_manager.update_task_state(task.name, success=False) + # exclude the current tasks from the next run + if continue_workgraph: + self.continue_workgraph() + + def execute_process_task(self, task, args=None, kwargs=None, var_kwargs=None): + """Execute a CalcJob or WorkChain task.""" + try: + process, state = task.execute( + engine_process=self.process, + args=args, + kwargs=kwargs, + var_kwargs=var_kwargs, + ) + self.state_manager.set_task_runtime_info(task.name, 'state', state) + self.state_manager.set_task_runtime_info(task.name, 'action', '') + self.state_manager.set_task_runtime_info(task.name, 'process', process) + # update the parent task state of mappped tasks + if self.process.wg.tasks[task.name].map_data: + parent_task_name = self.process.wg.tasks[task.name].map_data['parent'] + if self.process.node.get_task_state(parent_task_name) == TaskState.PLANNED: + self.process.node.set_task_state(parent_task_name, state) + self.process.to_context(**{task.name: process}) + except Exception as e: + error_traceback = traceback.format_exc() # Capture the full traceback + self.logger.error(f'Error in task {task.name}: {e}\n{error_traceback}') # Log the error with traceback + self.state_manager.update_task_state(task.name, success=False) + + def execute_while_task(self, task): + """Execute a WHILE task.""" + # TODO refactor this for while, if and zone + # in case of an empty zone, it will finish immediately + name = task.name + if self.state_manager.are_childen_finished(name)[0]: + self.state_manager.update_while_task_state(name) + else: + # check the conditions of the while task + should_run = self.should_run_while_task(name) + if not should_run: + self.state_manager.set_task_runtime_info(name, 'state', TaskState.FINISHED) + self.state_manager.set_tasks_state( + [child.name for child in self.process.wg.tasks[name].children], + TaskState.SKIPPED, + ) + self.state_manager.update_parent_task_state(name) + self.process.report( + f'While Task {name}: Condition not fullilled, task finished. Skip all its children.' + ) + else: + execution_count = self.state_manager.get_task_runtime_info(name, 'execution_count') + self.state_manager.set_task_runtime_info(name, 'state', TaskState.RUNNING) + self.state_manager.set_task_runtime_info(name, 'execution_count', execution_count + 1) + self.continue_workgraph() + + def execute_if_task(self, task): + # in case of an empty zone, it will finish immediately + name = task.name + if self.state_manager.are_childen_finished(name)[0]: + self.state_manager.update_zone_task_state(name) + else: + should_run = self.should_run_if_task(name) + if should_run: + self.state_manager.set_task_runtime_info(name, 'state', TaskState.RUNNING) + else: + self.state_manager.set_tasks_state([child.name for child in task.children], TaskState.SKIPPED) + self.state_manager.update_zone_task_state(name) + self.continue_workgraph() + + def execute_zone_task(self, task): + # in case of an empty zone, it will finish immediately + name = task.name + if self.state_manager.are_childen_finished(name)[0]: + self.state_manager.update_zone_task_state(name) + else: + self.state_manager.set_task_runtime_info(name, 'state', TaskState.RUNNING) + self.continue_workgraph() + + def execute_map_task(self, task, kwargs): + """ + 1. Clone the subgraph tasks for each item in `source`. + 2. Mark this MAP node as running and schedule a continuation. + """ + name = task.name + # we also store the links, so that we can load it in the GUI + map_info = {'prefix': [], 'children': [], 'links': []} + if self.state_manager.are_childen_finished(name)[0]: + self.state_manager.update_zone_task_state(name) + else: + self.state_manager.set_task_runtime_info(name, 'state', TaskState.RUNNING) + item_task = next(child for child in task.children if child.identifier == 'workgraph.map_item') + source = kwargs['source'] + map_info['prefix'] = list(source.keys()) + for prefix, value in source.items(): + new_tasks, new_links = self.generate_mapped_tasks(task, prefix=prefix) + self.update_map_item_task_state(item_task, prefix, value) + map_info['children'] = list(new_tasks.keys()) + map_info['links'] = new_links + self.state_manager.set_task_runtime_info(name, 'map_info', map_info) + # gather task finishes immediately + gather_task = task.gather_item_task + self.state_manager.set_task_runtime_info(gather_task.name, 'state', TaskState.FINISHED) + + self.continue_workgraph() + + def execute_normal_task(self, task, continue_workgraph=None, args=None, kwargs=None, var_kwargs=None): + """Execute a Normal task.""" + name = task.name + + # A "context" key is special and should be passed to the context manager + # TODO this is hard coded for now, need to be refactored + if 'context' in task.args_data['kwargs']: + self.ctx.task_name = name + kwargs.update({'context': self.ctx}) + for key in task.args_data['args']: + kwargs.pop(key, None) + try: + results, _ = task.execute(args, kwargs, var_kwargs) + self.state_manager.update_normal_task_state(name, results) + except Exception as e: + error_traceback = traceback.format_exc() + self.logger.error(f'Error in task {task.name}: {e}\n{error_traceback}') + self.state_manager.update_normal_task_state(name, results=None, success=False) + if continue_workgraph: + self.continue_workgraph() + + def get_socket_value(self, socket) -> Any: + """Get the value of the socket recursively.""" + socket_value = None + if isinstance(socket, TaskSocketNamespace): + socket_value = {} + for name, sub_socket in socket._sockets.items(): + value = self.get_socket_value(sub_socket) + if value is None or (isinstance(value, dict) and value == {}): + continue + socket_value[name] = value + else: + socket_value = socket.property.value + if ( + socket._task is not None + and socket._full_name.split('.')[0] == 'inputs' + and socket._metadata.extras.get('value_source') == 'property' + ): + return socket_value + links = socket._links + if len(links) == 1: + link = links[0] + if self.ctx._task_results.get(link.from_task.name): + # handle the special socket _wait, _outputs + if link.from_socket._scoped_name == '_wait': + return socket_value + elif link.from_socket._scoped_name == '_outputs': + socket_value = self.ctx._task_results[link.from_task.name] + else: + socket_value = get_nested_dict( + self.ctx._task_results[link.from_task.name], + link.from_socket._scoped_name, + default=None, + ) + # handle the case of multiple outputs + elif len(links) > 1: + socket_value = {} + for link in links: + item_name = f'{link.from_task.name}_{link.from_socket._scoped_name}' + # handle the special socket _wait, _outputs + if link.from_socket._scoped_name in ['_wait', '_outputs']: + continue + if self.ctx._task_results[link.from_task.name] is None: + socket_value[item_name] = None + else: + socket_value[item_name] = self.ctx._task_results[link.from_task.name][link.from_socket._scoped_name] + return socket_value + + def get_inputs( + self, name: str + ) -> tuple[ + list[Any], + dict[str, Any], + list[Any] | None, + dict[str, Any] | None, + dict[str, Any], + ]: + """Get input based on the links.""" + from aiida.workgraph.utils import update_nested_dict_with_special_keys + + args = [] + kwargs = {} + var_kwargs = None + task = self.process.wg.tasks[name] + inputs = {} + for prop in task.properties: + inputs[prop.name] = prop.value + + inputs.update(self.get_socket_value(task.inputs)) + + for name, input in inputs.items(): + # only need to check the top level key + key = name.split('.')[0] + if key in task.args_data['args']: + args.append(input) + elif key in task.args_data['kwargs']: + kwargs[name] = input + elif key == task.args_data['var_kwargs']: + var_kwargs = input + for i, key in enumerate(task.args_data['args']): + kwargs[key] = args[i] + # update the port namespace + kwargs = update_nested_dict_with_special_keys(kwargs) + return { + 'args': args, + 'kwargs': kwargs, + 'var_kwargs': var_kwargs, + } + + def should_run_while_task(self, name: str) -> tuple[bool, Any]: + """Check if the while task should run.""" + # check the conditions of the while task + execution_count = self.state_manager.get_task_runtime_info(name, 'execution_count') + not_excess_max_iterations = execution_count < self.process.wg.tasks[name].inputs.max_iterations.property.value + conditions = [not_excess_max_iterations] + inputs = self.get_inputs(name) + kwargs = inputs['kwargs'] + if isinstance(kwargs['conditions'], list): + for condition in kwargs['conditions']: + value = get_nested_dict(self.ctx, condition) + conditions.append(value) + elif isinstance(kwargs['conditions'], dict): + for _, value in kwargs['conditions'].items(): + conditions.append(value) + else: + conditions.append(kwargs['conditions']) + return False not in conditions + + def should_run_if_task(self, name: str) -> tuple[bool, Any]: + """Check if the IF task should run.""" + inputs = self.get_inputs(name) + kwargs = inputs['kwargs'] + flag = kwargs['conditions'] + if kwargs['invert_condition']: + return not flag + return flag + + def get_all_children(self, name: str) -> list[str]: + """Find all children of the zone_task, and their children recursively""" + child_tasks = [] + task = self.process.wg.tasks[name] + if not hasattr(task, 'children'): + return child_tasks + for child_task in task.children: + child_tasks.append(child_task.name) + child_tasks.extend(self.get_all_children(child_task.name)) + return child_tasks + + def generate_mapped_tasks(self, zone_task: Task, prefix: str) -> None: + """ + Recursively clone the subgraph starting from zone_children, + rewriting references to old tasks with new task names. + """ + # keep track of the mapped tasks + new_tasks = {} + all_links = [] + child_tasks = self.get_all_children(zone_task.name) + for child_task in child_tasks: + # since the child task is mapped, it should be skipped + self.state_manager.set_task_runtime_info(child_task, 'state', TaskState.MAPPED) + task = self.copy_task(child_task, prefix) + new_tasks[child_task] = task + links = self.process.wg.tasks[child_task].inputs._all_links + all_links.extend(links) + # fix references in the newly mapped tasks (children, input_links, etc.) + new_links = self._patch_cloned_tasks(new_tasks, all_links) + # update process.wg.connectivity so the new tasks are recognized in child_node, zone references, etc. + self._patch_connectivity(new_tasks) + return new_tasks, new_links + + def update_map_item_task_state(self, item_task, prefix, value: Any): + new_name = f'{prefix}_{item_task.name}' + self.ctx._task_results[new_name]['key'] = prefix + self.ctx._task_results[new_name]['value'] = value + self.state_manager.set_task_runtime_info(new_name, 'state', TaskState.FINISHED) + + def copy_task(self, name: str, prefix: str) -> Task: + import uuid + + # keep track of the mapped tasks + if not self.process.wg.tasks[name].mapped_tasks: + self.process.wg.tasks[name].mapped_tasks = {} + task_data = self.process.wg.tasks[name].to_dict() + new_name = f'{prefix}_{name}' + task_data['name'] = new_name + task_data['map_data'] = {'parent': name, 'prefix': prefix} + task_data['uuid'] = str(uuid.uuid4()) + # Reset runtime states + self.ctx._task_results[new_name] = {} + self.state_manager.set_task_runtime_info(new_name, 'state', TaskState.PLANNED) + self.state_manager.set_task_runtime_info(new_name, 'action', '') + # Insert new_data in ctx._tasks + task = self.process.wg.add_task_from_dict(task_data) + self.process.wg.tasks[name].mapped_tasks[prefix] = task + return task + + def _patch_cloned_tasks( + self, + new_tasks: dict[str, Task], + all_links: list[TaskLink], + ): + """ + For each newly mapped task, fix references (children, input_links, etc.) + from old_name -> new_name. + """ + for orginal_name, task in new_tasks.items(): + orginal_task = self.process.wg.tasks[orginal_name] + # fix children references + if hasattr(orginal_task, 'children'): + for child_task in orginal_task.children: + task.children.add(new_tasks[child_task.name]) + # since this is a newly created task, it should not have any mapped tasks + task.mapped_tasks = None + # fix parent reference + if orginal_task.parent is not None: + if orginal_task.parent.name in new_tasks: + task.parent = new_tasks[orginal_task.parent.name] + else: + task.parent = orginal_task.parent + # fix links references + new_links = [] + for link in all_links: + if link.to_task.name in new_tasks: + to_node = new_tasks[link.to_task.name] + to_socket = to_node.inputs[link.to_socket._scoped_name] + else: + # if the to_node is not in the new_tasks, skip + continue + new_links.append(link.to_dict()) + if link.from_task.name in new_tasks: + from_node = new_tasks[link.from_task.name] + from_socket = from_node.outputs[link.from_socket._scoped_name] + else: + from_socket = link.from_socket + self.process.wg.add_link( + from_socket, + to_socket, + ) + return new_links + + def _patch_connectivity(self, new_tasks: dict[str, Task]) -> None: + """ + Update the global connectivity for newly created tasks. + """ + for name, task in new_tasks.items(): + # child_node + new_child_node = [] + for child_task in self.process.wg.connectivity['child_node'][name]: + if child_task in new_tasks: + new_child_node.append(new_tasks[child_task].name) + else: + new_child_node.append(child_task) + self.process.wg.connectivity['child_node'][task.name] = new_child_node + # input_tasks + new_input_tasks = [] + for input_task in self.process.wg.connectivity['zone'][name]['input_tasks']: + if input_task in new_tasks: + new_input_tasks.append(new_tasks[input_task].name) + else: + new_input_tasks.append(input_task) + self.process.wg.connectivity['zone'][task.name] = {'input_tasks': new_input_tasks} diff --git a/src/aiida/workgraph/engine/task_state.py b/src/aiida/workgraph/engine/task_state.py new file mode 100644 index 0000000000..52740ea9b6 --- /dev/null +++ b/src/aiida/workgraph/engine/task_state.py @@ -0,0 +1,376 @@ +from __future__ import annotations + +from typing import Any + +from node_graph.socket import BaseSocket, TaskSocketNamespace +from typing_extensions import assert_never + +from aiida.orm import Data, ProcessNode +from aiida.orm.utils.serialize import serialize +from aiida.workgraph.enums import TERMINAL_TASK_STATES, RuntimeInfoKey, TaskState +from aiida.workgraph.orm.utils import deserialize_safe + + +class TaskStateManager: + """ + Handles all low-level operations on tasks' states, runtime info, + and relationships (parent/child). + """ + + def __init__(self, logger, process): + """ + :param logger: Logger instance. + :param process: The current AiiDA process. + """ + self.logger = logger + self.process = process + + @property + def ctx(self): + """Read the context off the process, which replaces it wholesale when loading from a checkpoint.""" + return self.process.ctx + + def get_task_runtime_info(self, name: str, key: RuntimeInfoKey) -> Any: + """Fetch a task runtime property (e.g. process, state, action).""" + match key: + case 'process': + value = self.process.node.get_task_process(name) + return deserialize_safe(value) if value else None + case 'state': + return self.process.node.get_task_state(name) + case 'action': + return self.process.node.get_task_action(name) + case 'execution_count': + return self.process.node.get_task_execution_count(name) + case _: + raise ValueError(f'Invalid key: {key}') + + def set_task_runtime_info(self, name: str, key: RuntimeInfoKey, value: Any) -> None: + """Set a task runtime property (e.g. process, state, action). + All the runtime info are store into the process node, which allow us + access this info outside the engine + """ + match key: + case 'process': + serialized = serialize(value) + self.process.node.set_task_process(name, serialized) + case 'state': + self.process.node.set_task_state(name, value) + case 'action': + self.process.node.set_task_action(name, value) + case 'execution_count': + self.process.node.set_task_execution_count(name, value) + case 'map_info': + self.process.node.set_task_map_info(name, value) + case _: + assert_never(key) + + def set_tasks_state(self, tasks: list[str], value: str) -> None: + """ + Set the state for a list of tasks (and their children) to `value`. + Typically used for skip or reset tasks. + """ + for name in tasks: + self.set_task_runtime_info(name, 'state', value) + if hasattr(self.process.wg.tasks[name], 'children'): + self.set_tasks_state([task.name for task in self.process.wg.tasks[name].children], value) + # TODO should we also reset the mapped tasks? + + def update_task_state(self, name: str, success=True) -> None: + """Update task state when the task is finished.""" + from aiida.workgraph.utils import resolve_node_link_managers + + task = self.process.wg.tasks[name] + self.ctx._task_results.setdefault(name, {}) + if success: + node = self.get_task_runtime_info(name, 'process') + if isinstance(node, ProcessNode): + state = node.process_state.value.upper() + if node.is_finished_ok: + self.set_task_runtime_info(task.name, 'state', state) + + self.ctx._task_results[name] = resolve_node_link_managers(node.outputs) + self.set_task_runtime_info(task.name, 'state', TaskState.FINISHED) + self.update_meta_tasks(name) + self.process.report(f'Task: {name}, type: {task.task_type}, finished.') + self.apply_socket_spec_extras_to_aiida_node(name, node) + # all other states are considered as failed + else: + self.ctx._task_results[name] = resolve_node_link_managers(node.outputs) + self.on_task_failed(name) + elif isinstance(node, Data): + output_name = next( + output_name for output_name in task.outputs._get_keys() if output_name not in ['_wait', '_outputs'] + ) + self.ctx._task_results[name] = {output_name: node} + self.set_task_runtime_info(task.name, 'state', TaskState.FINISHED) + self.update_meta_tasks(name) + self.process.report(f'Task: {name} finished.') + else: + self.on_task_failed(name) + # After finishing, inform the parent + self.update_parent_task_state(name) + + def update_normal_task_state(self, name, results, success=True): + """Set the results of a normal task. + A normal task is created by decorating a function with @task(). + """ + + if success: + task = self.process.wg.tasks[name] + if isinstance(results, tuple): + # there are two built-in outputs: _wait and _outputs + if len(task.outputs) - 2 != len(results): + self.on_task_failed(name) + return self.process.exit_codes.OUTPUS_NOT_MATCH_RESULTS + output_names = [ + output._name for output in task.outputs if output._metadata.extra.get('builtin_socket') is not True + ] + for i, output_name in enumerate(output_names): + self.ctx._task_results[name][output_name] = results[i] + elif isinstance(results, dict): + self.ctx._task_results[name] = results + else: + output_names = [ + output_name for output_name in task.outputs._get_keys() if output_name not in ['_wait', '_outputs'] + ] + # some task does not have any output + if len(output_names) == 1: + self.ctx._task_results[name][output_names[0]] = results + if isinstance(results, Data): + results.store() + self.set_task_runtime_info(task.name, 'process', results) + elif len(output_names) > 1: + self.process.exit_codes.OUTPUS_NOT_MATCH_RESULTS + self.update_meta_tasks(name) + self.set_task_runtime_info(name, 'state', TaskState.FINISHED) + self.process.report(f'Task: {name} finished.') + else: + self.on_task_failed(name) + self.update_parent_task_state(name) + + def update_meta_tasks(self, name: str) -> None: + """Export task results to the context based on context mapping.""" + from aiida.workgraph.utils import get_nested_dict, resolve_node_link_managers, update_nested_dict + + for link in self.process.wg.links: + if link.from_task.name == name and link.to_task.name in [ + 'graph_ctx', + 'graph_outputs', + ]: + key = link.to_socket._scoped_name + result_key = link.from_socket._scoped_name + # built-in "_outputs" means the whole task result + if result_key == '_outputs': + result = self.ctx._task_results[name] + else: + result = get_nested_dict(self.ctx._task_results[name], result_key, default=None) + result = resolve_node_link_managers(result) + update_nested_dict(self.ctx._task_results[link.to_task.name], key, result) + + def reset_task( + self, + name: str, + reset_process: bool = True, + recursive: bool = True, + reset_execution_count: bool = True, + ) -> None: + """ + Reset the task's state to PLANNED, optionally clearing the process reference + and recursing to children. If the task is a WHILE, reset its execution_count. + """ + self.logger.debug(f'Resetting task {name}.') + self.set_task_runtime_info(name, 'state', TaskState.PLANNED) + if reset_process: + self.set_task_runtime_info(name, 'process', None) + self.remove_executed_task(name) + + task_type = self.process.wg.tasks[name].task_type.upper() + if task_type == 'WHILE': + if reset_execution_count: + self.set_task_runtime_info(name, 'execution_count', 0) + for child_task in self.process.wg.tasks[name].children: + self.reset_task(child_task.name, reset_process=False, recursive=False) + elif task_type in ['IF', 'ZONE']: + for child_task in self.process.wg.tasks[name].children: + self.reset_task(child_task.name, reset_process=False, recursive=False) + + if recursive: + # reset its child tasks + child_names = self.process.wg.connectivity['child_node'][name] + for child_name in child_names: + self.reset_task(child_name, recursive=False) + + self.logger.debug(f'Task {name} was reset.') + + def remove_executed_task(self, name: str) -> None: + """ + Remove tasks from `ctx._executed_tasks` if they match this name (or name.*). + """ + self.ctx._executed_tasks = [label for label in self.ctx._executed_tasks if label.split('.')[0] != name] + + def is_task_ready_to_run(self, name: str) -> tuple[bool, str | None]: + """ + Check if the task is ready to run. We consider parent states, input tasks, etc. + For tasks inside a ZONE or with a parent task, we require the parent + to be in a running state, and the zone's input tasks finished or failed. + """ + parent_task = self.process.wg.tasks[name].parent + parent_states = [True, True] + + # If the task has a parent zone + if parent_task: + state = self.get_task_runtime_info(parent_task.name, 'state') + if state != TaskState.RUNNING: + parent_states[1] = False + + # Check input tasks from the zone connectivity + for child_task_name in self.process.wg.connectivity['zone'][name]['input_tasks']: + child_state = self.get_task_runtime_info(child_task_name, 'state') + if child_state not in TERMINAL_TASK_STATES: + parent_states[0] = False + break + return all(parent_states), parent_states + + def on_task_failed(self, name: str) -> None: + """ + Mark a task as FAILED, skip its children, and run any error handlers. + """ + task_type = self.process.wg.tasks[name].task_type + self.set_task_runtime_info(name, 'state', TaskState.FAILED) + self.set_tasks_state(self.process.wg.connectivity['child_node'][name], TaskState.SKIPPED) + msg = f'Task, {name}, type: {task_type}, failed.' + process = self.get_task_runtime_info(name, 'process') + if isinstance(process, ProcessNode): + msg += f' Error message: {process.exit_message}' + self.process.report(msg) + self.process.error_handler_manager.run_error_handlers(name) + + def update_parent_task_state(self, name: str) -> None: + """ + If a task has a parent (WHILE, IF, ZONE, MAP), notify the parent to update + its own state. Also handle mapped tasks referencing a 'map_data.parent' node. + """ + parent_task = self.process.wg.tasks[name].parent + if parent_task: + task_type = parent_task.task_type.upper() + if task_type == 'WHILE': + self.update_while_task_state(parent_task.name) + elif task_type in ['IF', 'ZONE']: + self.update_zone_task_state(parent_task.name) + elif task_type == 'MAP': + self.update_map_task_state(parent_task.name) + + # If the task is a mapped child, update its parent's "template" (the original map node) + if self.process.wg.tasks[name].map_data: + map_parent = self.process.wg.tasks[name].map_data['parent'] + self.update_template_task_state(map_parent) + + def update_while_task_state(self, name: str) -> None: + """ + Called when a child of a WHILE task finishes. If all children are done, we decide + whether to reset for the next iteration or finalize the WHILE. + """ + finished, _ = self.are_childen_finished(name) + + if finished: + self.process.report(f'While Task {name}: this iteration finished. Try to reset for the next iteration.') + # reset the condition tasks + for link in self.process.wg.tasks[name].inputs.conditions._links: + self.reset_task(link.from_task.name, recursive=False) + # reset the task and all its children, so that the task can run again + # do not reset the execution count + self.reset_task(name, reset_execution_count=False) + + def update_zone_task_state(self, name: str) -> None: + """ + Update the state of an IF or ZONE block. Mark it FINISHED if children are done. + """ + finished, _ = self.are_childen_finished(name) + if finished: + self.set_task_runtime_info(name, 'state', TaskState.FINISHED) + self.process.report(f'Task: {name} finished.') + self.update_parent_task_state(name) + + def update_map_task_state(self, name: str) -> None: + """Update the map task state. + 1) check if all child tasks are finished. + 2) gather the results of all the mapped tasks. + 3) update the parent task state. + """ + finished, _ = self.are_childen_finished(name) + if finished: + map_zone = self.process.wg.tasks[name] + # gather the results of all the mapped tasks + gather_task = map_zone.gather_item_task + for input in gather_task.inputs: + if input._name.startswith('_'): + continue + results = {} + link = input._links[0] + for prefix, mapped_task in self.process.wg.tasks[gather_task.name].mapped_tasks.items(): + results[prefix] = self.ctx._task_results[mapped_task.name][link.to_socket._name] + self.ctx._task_results[name][link.to_socket._name] = results + self.set_task_runtime_info(name, 'state', TaskState.FINISHED) + # self.update_meta_tasks(name) + self.process.report(f'Task: {name} finished.') + self.update_meta_tasks(name) + self.update_parent_task_state(name) + + def update_template_task_state(self, name: str) -> None: + """Update the template task state. + 1) check if all child tasks are finished. + 2) gather the results of all the mapped tasks. + 3) update the parent task state. + """ + finished, _ = self.are_childen_finished(name) + if finished: + # # gather the results of all the mapped tasks + # results = {} + # for prefix, mapped_task in self.process.wg.tasks[name].mapped_tasks.items(): + # for output in mapped_task.outputs: + # if output._name in self.ctx._task_results[mapped_task.name]: + # results.setdefault(output._name, {}) + # results[output._name][prefix] = self.ctx._task_results[mapped_task.name][output._name] + # self.ctx._task_results[name] = results + self.set_task_runtime_info(name, 'state', TaskState.FINISHED) + # self.update_meta_tasks(name) + self.process.report(f'Task: {name} finished.') + self.update_parent_task_state(name) + + def are_childen_finished(self, name: str) -> tuple[bool, Any]: + """Check if the child tasks are finished.""" + task = self.process.wg.tasks[name] + finished = True + if hasattr(task, 'children'): + for child in task.children: + if self.get_task_runtime_info(child.name, 'state') not in TERMINAL_TASK_STATES: + finished = False + break + # check the mapped tasks + mapped_tasks = task.mapped_tasks or {} + for mapped_task in mapped_tasks.values(): + if self.get_task_runtime_info(mapped_task.name, 'state') not in TERMINAL_TASK_STATES: + finished = False + break + return finished, None + + def apply_socket_spec_extras_to_aiida_node(self, name: str, node: ProcessNode) -> None: + """Apply the socket spec extras to the AiiDA process node for a task.""" + task = self.process.wg.tasks[name] + task.set_outputs_from_process_node(node) + self.set_socket_spec_extra(task.outputs) + + @classmethod + def set_socket_spec_extra(cls, socket: BaseSocket) -> None: + """Set the socket spec extra to the AiiDA process node for a task.""" + if isinstance(socket, TaskSocketNamespace): + for sub_socket in socket._sockets.values(): + cls.set_socket_spec_extra(sub_socket) + elif isinstance(socket.value, Data): + extras = { + key: value + for key, value in socket._metadata.extras.items() + if key not in ['identifier', 'builtin_socket', 'function_socket'] + } + socket.value.base.extras.set_many(extras) diff --git a/src/aiida/workgraph/executors/__init__.py b/src/aiida/workgraph/executors/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/aiida/workgraph/executors/builtins.py b/src/aiida/workgraph/executors/builtins.py new file mode 100644 index 0000000000..dc89c18d73 --- /dev/null +++ b/src/aiida/workgraph/executors/builtins.py @@ -0,0 +1,61 @@ +import logging +from typing import Any + +from aiida import orm + +LOGGER = logging.getLogger(__name__) + + +def UnavailableExecutor(*args, **kwargs): + raise RuntimeError('This executor was defined dynamically and is not available from the database snapshot.') + + +def get_context(context: dict, key: str) -> Any: + """Get the context value.""" + key = key.value if isinstance(key, orm.Str) else key + results = {'result': context._task_results['graph_ctx'].get(key)} + return results + + +def update_ctx(context: dict, key: str, value: Any) -> None: + """Set the context value.""" + key = key.value if isinstance(key, orm.Str) else key + context._task_results['graph_ctx'][key] = value + + +def select(condition, true=None, false=None): + """Select the data based on the condition.""" + if condition: + return true + return false + + +def get_item(data: dict, key: str) -> Any: + """Get an item from a dictionary.""" + return data.get(key, None) + + +def return_input(**kwargs: Any) -> dict: + """Return the input""" + return kwargs + + +def load_node(pk: int | None = None, uuid: str | None = None) -> orm.Node: + """Load an AiiDA node by its primary key or UUID.""" + if uuid is not None: + pk = uuid.value if isinstance(uuid, orm.Str) else uuid + else: + pk = pk.value if isinstance(pk, orm.Int) else pk + return orm.load_node(pk) + + +def load_code(pk: int | None = None, uuid: str | None = None, label: str | None = None) -> orm.Code: + """Load an AiiDA code by its primary key or UUID.""" + if uuid is not None: + pk = uuid.value if isinstance(uuid, orm.Str) else uuid + elif label is not None: + pk = label.value if isinstance(label, orm.Str) else label + else: + pk = pk.value if isinstance(pk, orm.Int) else pk + LOGGER.info('Loading code with pk: %s', pk) + return orm.load_code(pk) diff --git a/src/aiida/workgraph/executors/test.py b/src/aiida/workgraph/executors/test.py new file mode 100644 index 0000000000..15ff4907c6 --- /dev/null +++ b/src/aiida/workgraph/executors/test.py @@ -0,0 +1,37 @@ +import time + +from aiida.calculations.arithmetic.add import ArithmeticAddCalculation +from aiida.orm import Int +from aiida.workgraph import task +from aiida.workgraph.socket_spec import namespace + +ArithmeticAddTask = task(ArithmeticAddCalculation) + + +@task +def add(x: Int = 0, y: Int = 0, t: Int = 1) -> Int: + """Add node.""" + time.sleep(t) + return x + y + + +@task +def sum_diff(x: Int = 0, y: Int = 0, t: Int = 1) -> namespace(sum=Int, diff=Int): + """Add node.""" + time.sleep(t) + return {'sum': x + y, 'diff': x - y} + + +@task.pythonjob() +def add_pythonjob(x: int, y: int) -> int: + return x + y + + +@task.graph +def Fibonacci(n, a=0, b=1): + """Fibonacci sequence.""" + if n == 0: + return a + if n == 1: + return b + return Fibonacci(n=n - 1, a=b, b=add(x=a, y=b).result) diff --git a/src/aiida/workgraph/manager.py b/src/aiida/workgraph/manager.py new file mode 100644 index 0000000000..4ec3c08b66 --- /dev/null +++ b/src/aiida/workgraph/manager.py @@ -0,0 +1,194 @@ +""" +Simple global variable approach. +Note pitfalls: + - lack of concurrency control + - collisions in library code + - difficulty testing in isolation + - etc. +""" + +from contextlib import contextmanager + +from aiida.workgraph.socket import TaskSocket +from aiida.workgraph.tasks.task_pool import TaskPool + +DEFAULT_MAP_PLACEHOLDER = 'map_input' + + +class CurrentGraphManager: + _instance = None + + def __new__(cls, *args, **kwargs): + """ + Enforce the singleton pattern. Only one instance of + CurrentGraphManager is created for the entire process. + """ + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._graph = None # Storage for the active graph + return cls._instance + + def get_current_graph(self): + """ + Retrieve the current graph, or create a new one if none is set. + """ + from aiida.workgraph import WorkGraph + + if self._graph is None: + self._graph = WorkGraph() + return self._graph + + def set_current_graph(self, graph): + """ + Set the active graph to the given instance. + """ + self._graph = graph + + @contextmanager + def active_graph(self, graph): + """ + Context manager that temporarily overrides the current graph + with `graph`, restoring the old graph when exiting the context. + """ + old_graph = self._graph + self._graph = graph + try: + yield graph + finally: + self._graph = old_graph + + +# Create a global manager instance +_manager = CurrentGraphManager() + + +def get_current_graph(): + """ + Helper function to retrieve the graph + through the global manager instance. + """ + return _manager.get_current_graph() + + +def set_current_graph(graph): + """ + Helper function to set the graph through the + global manager instance. + """ + _manager.set_current_graph(graph) + + +@contextmanager +def active_graph(graph): + """ + Top-level context manager that defers to + the manager's `active_graph` method. + """ + with _manager.active_graph(graph) as g: + yield g + + +@contextmanager +def Zone(): + """ + Context manager to create a "zone" in the current graph. + """ + + wg = get_current_graph() + + zone_task = wg.add_task( + TaskPool.workgraph.zone, + ) + + old_zone = getattr(wg, '_active_zone', None) + if old_zone: + old_zone.children.add(zone_task) + wg._active_zone = zone_task + + try: + yield zone_task + finally: + wg._active_zone = old_zone + + +@contextmanager +def If(condition_socket: TaskSocket, invert_condition: bool = False): + """ + Context manager to create a "conditional zone" in the current graph. + + :param condition_socket: A TaskSocket or boolean-like object (e.g. sum_ > 0) + :param invert_condition: Whether to invert the condition (useful for else-zones) + """ + + wg = get_current_graph() + + zone_task = wg.add_task( + TaskPool.workgraph.if_zone, + conditions=condition_socket, + invert_condition=invert_condition, + ) + + old_zone = getattr(wg, '_active_zone', None) + if old_zone: + old_zone.children.add(zone_task) + wg._active_zone = zone_task + + try: + yield zone_task + finally: + wg._active_zone = old_zone + + +@contextmanager +def While(condition_socket: TaskSocket, max_iterations: int = 10000): + """ + Context manager to create a "while zone" in the current graph. + + :param condition_socket: A TaskSocket or boolean-like object (e.g. sum_ > 0) + :param max_iterations: Maximum number of iterations before breaking the loop + """ + + wg = get_current_graph() + + zone_task = wg.add_task( + TaskPool.workgraph.while_zone, + conditions=condition_socket, + max_iterations=max_iterations, + ) + + old_zone = getattr(wg, '_active_zone', None) + if old_zone: + old_zone.children.add(zone_task) + wg._active_zone = zone_task + + try: + yield zone_task + finally: + wg._active_zone = old_zone + + +@contextmanager +def Map(source_socket: TaskSocket, placeholder: str = DEFAULT_MAP_PLACEHOLDER): + """ + Context manager to create a "map zone" in the current graph. + + :param source_socket: A TaskSocket or boolean-like object (e.g. sum_ > 0) + :param placeholder: The placeholder string to use as the input for the mapped tasks + """ + + wg = get_current_graph() + + zone_task = wg.add_task( + TaskPool.workgraph.map_zone, + source=source_socket, + ) + + old_zone = getattr(wg, '_active_zone', None) + if old_zone: + old_zone.children.add(zone_task) + wg._active_zone = zone_task + + try: + yield zone_task + finally: + wg._active_zone = old_zone diff --git a/src/aiida/workgraph/orm/__init__.py b/src/aiida/workgraph/orm/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/aiida/workgraph/orm/mapping.py b/src/aiida/workgraph/orm/mapping.py new file mode 100644 index 0000000000..21ecd143de --- /dev/null +++ b/src/aiida/workgraph/orm/mapping.py @@ -0,0 +1,29 @@ +from typing import Any + +from aiida import orm + +builtins_type_mapping = { + 'any': 'workgraph.any', + 'default': 'workgraph.any', + 'namespace': 'workgraph.namespace', + int: 'workgraph.int', + float: 'workgraph.float', + str: 'workgraph.string', + bool: 'workgraph.bool', + list: 'workgraph.list', + dict: 'workgraph.dict', + orm.Int: 'workgraph.int', + orm.Float: 'workgraph.float', + orm.Str: 'workgraph.string', + orm.Bool: 'workgraph.bool', + orm.List: 'workgraph.list', + orm.Dict: 'workgraph.dict', + orm.StructureData: 'workgraph.aiida_structuredata', + Any: 'workgraph.any', +} + +TYPE_PROMOTIONS: set[tuple[str, str]] = { + ('workgraph.bool', 'workgraph.int'), + ('workgraph.bool', 'workgraph.float'), + ('workgraph.int', 'workgraph.float'), +} diff --git a/src/aiida/workgraph/orm/utils.py b/src/aiida/workgraph/orm/utils.py new file mode 100644 index 0000000000..a9540c845b --- /dev/null +++ b/src/aiida/workgraph/orm/utils.py @@ -0,0 +1,36 @@ +from typing import Any + +import yaml + +from aiida.orm.utils.serialize import ( + _COMPUTER_TAG, + _GROUP_TAG, + _NODE_LINKS_MANAGER_TAG, + _NODE_TAG, + computer_constructor, + group_constructor, + node_constructor, + node_links_manager_constructor, +) + + +class AiiDASafeLoader(yaml.SafeLoader): + """ + A “safe” AiiDA-specific YAML loader. + + Since we are extending SafeLoader, we need to carefully add only those + constructors we consider safe. Anything we omit will raise an error if + encountered in the YAML. + """ + + pass + + +AiiDASafeLoader.add_constructor(_NODE_TAG, node_constructor) +AiiDASafeLoader.add_constructor(_COMPUTER_TAG, computer_constructor) +AiiDASafeLoader.add_constructor(_NODE_LINKS_MANAGER_TAG, node_links_manager_constructor) +AiiDASafeLoader.add_constructor(_GROUP_TAG, group_constructor) + + +def deserialize_safe(serialized: str) -> Any: + return yaml.load(serialized, Loader=AiiDASafeLoader) diff --git a/src/aiida/workgraph/properties/__init__.py b/src/aiida/workgraph/properties/__init__.py new file mode 100644 index 0000000000..d160bd982a --- /dev/null +++ b/src/aiida/workgraph/properties/__init__.py @@ -0,0 +1,3 @@ +from .property_pool import PropertyPool + +__all__ = ['PropertyPool'] diff --git a/src/aiida/workgraph/properties/builtins.py b/src/aiida/workgraph/properties/builtins.py new file mode 100644 index 0000000000..02d1e431bd --- /dev/null +++ b/src/aiida/workgraph/properties/builtins.py @@ -0,0 +1,120 @@ +from typing import Any + +from aiida import orm +from aiida.orm import NoneData +from aiida.workgraph.property import TaskProperty + + +class PropertyAny(TaskProperty): + """A new class for Any type.""" + + identifier: str = 'workgraph.any' + + def validate(self, _: Any) -> None: + """No validation needed.""" + + +class PropertyInt(TaskProperty): + """A new class for integer type.""" + + identifier: str = 'workgraph.int' + allowed_types = (int, orm.Int, type(None), NoneData) + + +class PropertyFloat(TaskProperty): + """A new class for float type.""" + + identifier: str = 'workgraph.float' + allowed_types = (int, float, orm.Int, orm.Float, type(None), NoneData) + + +class PropertyBool(TaskProperty): + """A new class for bool type.""" + + identifier: str = 'workgraph.bool' + allowed_types = (bool, int, orm.Bool, orm.Int, type(None), NoneData) + + +class PropertyString(TaskProperty): + """A new class for string type.""" + + identifier: str = 'workgraph.string' + allowed_types = (str, orm.Str, type(None), NoneData) + + +class PropertyList(TaskProperty): + """A new class for List type.""" + + identifier: str = 'workgraph.list' + allowed_types = (list, tuple, orm.List, type(None), NoneData) + + +class PropertyDict(TaskProperty): + """A new class for Dict type.""" + + identifier: str = 'workgraph.dict' + allowed_types = (dict, orm.Dict, type(None), NoneData) + + +# ==================================== +class PropertyVector(TaskProperty): + """Vector property""" + + identifier: str = 'workgraph.vector' + allowed_item_types = (object, type(None), NoneData) + + def __init__(self, name, description='', size=3, default=None, update=None) -> None: + self.size = size + default = [] if default is None else default + super().__init__(name, description, default, update) + + def validate(self, value: Any) -> None: + """Validate the given value based on allowed types.""" + if value is not None: + if len(value) != self.size: + raise ValueError(f'Invalid size: Expected {self.size}, got {len(value)} instead.') + for item in value: + if not isinstance(item, self.allowed_item_types): + raise ValueError( + f'Invalid item type: Expected {self.allowed_item_types}, got {type(item)} instead.' + ) + + super().validate(value) + + def set_value(self, value: list) -> None: + self.validate(value) + self._value = value + if self.update: + self.update() + + def copy(self): + p = self.__class__(self.name, self.description, self.size, self.value, self.update) + p.value = self.value + return p + + def get_metadata(self): + metadata = {'default': self.default, 'size': self.size} + return metadata + + +class PropertyAiiDAIntVector(PropertyVector): + """A new class for integer vector type.""" + + identifier: str = 'workgraph.aiida_int_vector' + allowed_types = (list, orm.List, type(None), NoneData) + allowed_item_types = (int, type(None), NoneData) + + +class PropertyAiiDAFloatVector(PropertyVector): + """A new class for float vector type.""" + + identifier: str = 'workgraph.aiida_float_vector' + allowed_types = (list, orm.List, type(None), NoneData) + allowed_item_types = (int, float, type(None), NoneData) + + +class PropertyStructureData(TaskProperty): + """A new class for Any type.""" + + identifier: str = 'workgraph.aiida_structuredata' + allowed_types = (orm.StructureData, type(None), NoneData) diff --git a/src/aiida/workgraph/properties/property_pool.py b/src/aiida/workgraph/properties/property_pool.py new file mode 100644 index 0000000000..e401e46d4c --- /dev/null +++ b/src/aiida/workgraph/properties/property_pool.py @@ -0,0 +1,5 @@ +from node_graph.registry import EntryPointPool + +# global instance +PropertyPool = EntryPointPool(entry_point_group='aiida_workgraph.property') +PropertyPool['any'] = PropertyPool.workgraph.any diff --git a/src/aiida/workgraph/property.py b/src/aiida/workgraph/property.py new file mode 100644 index 0000000000..ba379c3b89 --- /dev/null +++ b/src/aiida/workgraph/property.py @@ -0,0 +1,28 @@ +from collections.abc import Callable + +from node_graph.property import TaskProperty as BaseTaskProperty + + +class TaskProperty(BaseTaskProperty): + """Represent a property of a Task in the AiiDA WorkGraph.""" + + def validate(self, value: any) -> None: + super().validate(value) + + @classmethod + def new(cls, identifier: Callable | str, name: str | None = None, **kwargs) -> 'TaskProperty': + """Create a property from a identifier.""" + # use PropertyPool from aiida.workgraph.properties + # to override the default PropertyPool from node_graph + from aiida.workgraph.properties import PropertyPool + + return super().new(identifier, name=name, PropertyPool=PropertyPool, **kwargs) + + +def unwrap_aiida_node(value): + if hasattr(value, 'value'): + return value.value + return TaskProperty.NOT_ADAPTED + + +TaskProperty.register_validation_adapter(unwrap_aiida_node) diff --git a/src/aiida/workgraph/registry.py b/src/aiida/workgraph/registry.py new file mode 100644 index 0000000000..25ec75e2cd --- /dev/null +++ b/src/aiida/workgraph/registry.py @@ -0,0 +1,12 @@ +from node_graph.registry import RegistryHub + +registry_hub = RegistryHub.from_prefix( + task_group='aiida_workgraph.task', + socket_group='aiida_workgraph.socket', + property_group='aiida_workgraph.property', + type_mapping_group='aiida_workgraph.type_mapping', + type_promotion_group='aiida_workgraph.type_promotion', + identifier_prefix='workgraph', +) + +type_mapping = registry_hub.type_mapping diff --git a/src/aiida/workgraph/schemas/__init__.py b/src/aiida/workgraph/schemas/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/aiida/workgraph/serialization.py b/src/aiida/workgraph/serialization.py index 1b2bd44a81..0e31de9142 100644 --- a/src/aiida/workgraph/serialization.py +++ b/src/aiida/workgraph/serialization.py @@ -17,16 +17,19 @@ from typing import TYPE_CHECKING, Any +from node_graph.serializer import SerializationAdapter from node_graph.socket_meta import SocketMeta from node_graph.socket_spec import SocketSpec +from node_graph.utils import resolve_tagged_values from node_graph.utils.struct_utils import is_structured_instance, structured_to_dict from aiida.orm import general_serializer +from aiida.orm.nodes.data.serializer import get_serializers if TYPE_CHECKING: from aiida.orm import User -__all__ = ('serialize_ports',) +__all__ = ('AiidaSerializationAdapter', 'serialize_ports') def _ensure_spec(schema: SocketSpec | dict[str, Any]) -> SocketSpec: @@ -101,3 +104,27 @@ def serialize_ports( raise ValueError(msg) return out + + +class AiidaSerializationAdapter(SerializationAdapter): + """node-graph serialization adapter that serializes socket values into AiiDA data nodes.""" + + id: str = 'aiida' + name: str = 'AiiDA' + + def __init__(self, serializers: dict[str, str] | None = None, user: Any = None) -> None: + self.serializers = serializers or get_serializers() + self.user = user + + def serialize(self, value: Any, socket: Any, *, store: bool) -> Any: + if socket is None: + return value + spec = socket._to_spec() + resolve_tagged_values(value) + return serialize_ports(python_data=value, port_schema=spec, serializers=self.serializers, user=self.user) + + def serialize_ports(self, python_data: Any, port_schema: Any, *, store: bool) -> Any: + resolve_tagged_values(python_data) + return serialize_ports( + python_data=python_data, port_schema=port_schema, serializers=self.serializers, user=self.user + ) diff --git a/src/aiida/workgraph/socket.py b/src/aiida/workgraph/socket.py new file mode 100644 index 0000000000..fd9dd8df2d --- /dev/null +++ b/src/aiida/workgraph/socket.py @@ -0,0 +1,57 @@ +from node_graph.socket import ( + TaskSocket as BaseTaskSocket, +) +from node_graph.socket import ( + TaskSocketNamespace as BaseTaskSocketNamespace, +) + +from aiida import orm +from aiida.workgraph.property import TaskProperty +from aiida.workgraph.registry import type_mapping + + +class TaskSocket(BaseTaskSocket): + """Represent a socket of a Task in the AiiDA WorkGraph.""" + + # use TaskProperty from aiida.workgraph.property + # to override the default TaskProperty from node_graph + _socket_property_class = TaskProperty + + @property + def _decorator(self): + from aiida.workgraph.decorator import task + + return task + + @property + def node_value(self): + return self.get_node_value() + + def get_node_value(self): + """Obtain the actual Python `value` of the object attached to the Socket.""" + if isinstance(self.value, orm.Data): + if hasattr(self.value, 'value'): + return self.value.value + else: + raise ValueError( + 'Data node does not have a value attribute. We do not know how to extract the raw Python value.' + ) + else: + return self.value + + +class TaskSocketNamespace(BaseTaskSocketNamespace): + """Represent a namespace of a Task in the AiiDA WorkGraph.""" + + _identifier = 'workgraph.namespace' + _socket_property_class = TaskProperty + _type_mapping: dict = type_mapping + + @property + def _decorator(self): + from aiida.workgraph.decorator import task + + return task + + def __init__(self, *args, **kwargs): + super().__init__(*args, entry_point='aiida_workgraph.socket', **kwargs) diff --git a/src/aiida/workgraph/socket_spec.py b/src/aiida/workgraph/socket_spec.py new file mode 100644 index 0000000000..9808238ac5 --- /dev/null +++ b/src/aiida/workgraph/socket_spec.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +from dataclasses import replace +from typing import Any + +from node_graph.socket_spec import ( + Leaf, + SocketMeta, + SocketSpec, + SocketSpecSelect, + meta, + select, +) +from node_graph.socket_spec import ( + SocketSpecAPI as _SocketSpecAPI, +) +from plumpy.ports import Port, PortNamespace + +from aiida.engine import Process +from aiida.engine.processes.process_spec import ProcessSpec +from aiida.workgraph.registry import type_mapping + +from .socket import TaskSocketNamespace + +__all__ = [ + 'Leaf', + 'SocketSpecAPI', + 'SocketSpecSelect', + 'dynamic', + 'from_aiida_process', + 'infer_specs_from_callable', + 'meta', + 'namespace', + 'select', + 'socket', + 'validate_socket_data', +] + + +class SocketSpecAPI(_SocketSpecAPI): + MAP: dict[Any, str] = type_mapping + NAMESPACE: str = 'workgraph.namespace' + DEFAULT: str = 'workgraph.any' + ANNOTATED: str = 'workgraph.annotated' + + SocketNamespace = TaskSocketNamespace + + @classmethod + def _identifier_from_valid_type(cls, valid_type: Any) -> str: + """Map AiiDA Port.valid_type -> identifier with our mapping. + - tuple of types: if len==1 map that; else -> any/default + - None/empty: any/default + - single type: mapped identifier + """ + if isinstance(valid_type, tuple): + if len(valid_type) == 1: + return cls._map_identifier(valid_type[0]) + return cls.DEFAULT + if valid_type in (None, Ellipsis): + return cls.DEFAULT + return cls._map_identifier(valid_type) + + @classmethod + def _from_port(cls, port: Port | PortNamespace, *, parent_required: bool, role: str) -> SocketSpec: + """Recursively convert an AiiDA Port/PortNamespace to a SocketSpec. + `role` is "input" or "output" (affects call_role metadata). + """ + if isinstance(port, PortNamespace): + required_here = bool(getattr(port, 'required', True)) and bool(parent_required) + + # Build child fields by iterating explicit .ports mapping + fields: dict[str, SocketSpec] = {} + for name, child in port.ports.items(): + fields[name] = cls._from_port(child, parent_required=required_here, role=role) + + ns = SocketSpec( + identifier=cls.NAMESPACE, + fields=fields, + meta=SocketMeta( + required=required_here, + is_metadata=getattr(port, 'is_metadata', False), + call_role=('kwargs' if role == 'input' else None), + ), + ) + + # Dynamic namespace? (DynamicPortNamespace derives from PortNamespace) + is_dyn = bool(getattr(port, 'dynamic', False)) + if is_dyn: + valid_type = getattr(port, 'valid_type', None) + if valid_type: + item_ident = cls._identifier_from_valid_type(valid_type) + ns = replace(ns, meta=replace(ns.meta, dynamic=True), item=SocketSpec(identifier=item_ident)) + else: + ns = replace(ns, meta=replace(ns.meta, dynamic=True), item=None) + return ns + + # Leaf Port (InputPort/OutputPort) + required_here = bool(getattr(port, 'required', True)) and bool(parent_required) + valid_type = getattr(port, 'valid_type', None) + ident = cls._identifier_from_valid_type(valid_type) + return SocketSpec( + identifier=ident, + meta=SocketMeta( + required=required_here, + is_metadata=getattr(port, 'is_metadata', False), + call_role=('kwargs' if role == 'input' else None), + ), + ) + + @classmethod + def from_aiida_process( + cls, process_or_spec: type[Process] | Process | ProcessSpec + ) -> tuple[SocketSpec, SocketSpec]: + """Return (inputs_spec, outputs_spec) for an AiiDA Process or its ProcessSpec. + + Accepts: + - AiiDA Process subclass (e.g. CalcJob, WorkChain) + - AiiDA Process instance + - ProcessSpec object (as returned by `.spec()`) + """ + # Normalize to a ProcessSpec + if isinstance(process_or_spec, ProcessSpec): + spec = process_or_spec + elif isinstance(process_or_spec, type) and issubclass(process_or_spec, Process): + spec = process_or_spec.spec() + elif isinstance(process_or_spec, Process): + spec = process_or_spec.spec() + else: + raise TypeError( + 'from_aiida_process expects an AiiDA Process class/instance or a ProcessSpec; ' + f'got {type(process_or_spec)!r}' + ) + + # Validate spec structure + if not isinstance(spec, ProcessSpec): + raise TypeError(f'.spec() did not return a ProcessSpec; got {type(spec)!r}') + + inputs_ns = getattr(spec, 'inputs', None) + outputs_ns = getattr(spec, 'outputs', None) + if not isinstance(inputs_ns, PortNamespace) or not isinstance(outputs_ns, PortNamespace): + raise TypeError('Spec does not expose PortNamespace for inputs/outputs') + + in_spec = cls._from_port(inputs_ns, parent_required=True, role='input') + out_spec = cls._from_port(outputs_ns, parent_required=True, role='output') + # tag top-level outputs with 'return' + out_spec = replace(out_spec, meta=replace(out_spec.meta, call_role='return')) + return in_spec, out_spec + + +socket = SocketSpecAPI.socket +namespace = SocketSpecAPI.namespace +dynamic = SocketSpecAPI.dynamic +validate_socket_data = SocketSpecAPI.validate_socket_data +infer_specs_from_callable = SocketSpecAPI.infer_specs_from_callable +from_aiida_process = SocketSpecAPI.from_aiida_process diff --git a/src/aiida/workgraph/sockets/__init__.py b/src/aiida/workgraph/sockets/__init__.py new file mode 100644 index 0000000000..f9f0a09e8e --- /dev/null +++ b/src/aiida/workgraph/sockets/__init__.py @@ -0,0 +1,3 @@ +from .socket_pool import SocketPool + +__all__ = ['SocketPool'] diff --git a/src/aiida/workgraph/sockets/builtins.py b/src/aiida/workgraph/sockets/builtins.py new file mode 100644 index 0000000000..c22571f512 --- /dev/null +++ b/src/aiida/workgraph/sockets/builtins.py @@ -0,0 +1,78 @@ +from aiida.workgraph.socket import TaskSocket + + +class SocketAny(TaskSocket): + """Any socket.""" + + _identifier: str = 'workgraph.any' + _socket_property_identifier: str = 'workgraph.any' + + +class SocketAnnotated(TaskSocket): + """Socket for annotated Python types stored in metadata.""" + + _identifier: str = 'workgraph.annotated' + _socket_property_identifier: str = 'workgraph.any' + + +class SocketFloat(TaskSocket): + """Float socket.""" + + _identifier: str = 'workgraph.float' + _socket_property_identifier: str = 'workgraph.float' + + +class SocketInt(TaskSocket): + """Int socket.""" + + _identifier: str = 'workgraph.int' + _socket_property_identifier: str = 'workgraph.int' + + +class SocketString(TaskSocket): + """String socket.""" + + _identifier: str = 'workgraph.string' + _socket_property_identifier: str = 'workgraph.string' + + +class SocketBool(TaskSocket): + """Bool socket.""" + + _identifier: str = 'workgraph.bool' + _socket_property_identifier: str = 'workgraph.bool' + + +class SocketList(TaskSocket): + """List socket.""" + + _identifier: str = 'workgraph.list' + _socket_property_identifier: str = 'workgraph.list' + + +class SocketDict(TaskSocket): + """Dict socket.""" + + _identifier: str = 'workgraph.dict' + _socket_property_identifier: str = 'workgraph.dict' + + +class SocketAiiDAIntVector(TaskSocket): + """Socket with a AiiDAIntVector property.""" + + _identifier: str = 'workgraph.aiida_int_vector' + _socket_property_identifier: str = 'workgraph.aiida_int_vector' + + +class SocketAiiDAFloatVector(TaskSocket): + """Socket with a FloatVector property.""" + + _identifier: str = 'workgraph.aiida_float_vector' + _socket_property_identifier: str = 'workgraph.aiida_float_vector' + + +class SocketStructureData(TaskSocket): + """Any socket.""" + + _identifier: str = 'workgraph.aiida_structuredata' + _socket_property_identifier: str = 'workgraph.aiida_structuredata' diff --git a/src/aiida/workgraph/sockets/socket_pool.py b/src/aiida/workgraph/sockets/socket_pool.py new file mode 100644 index 0000000000..b042506753 --- /dev/null +++ b/src/aiida/workgraph/sockets/socket_pool.py @@ -0,0 +1,6 @@ +from node_graph.registry import EntryPointPool + +# global instance +SocketPool = EntryPointPool(entry_point_group='aiida_workgraph.socket') +SocketPool['any'] = SocketPool.workgraph.any +SocketPool['namespace'] = SocketPool.workgraph.namespace diff --git a/src/aiida/workgraph/task.py b/src/aiida/workgraph/task.py new file mode 100644 index 0000000000..9da22d70e8 --- /dev/null +++ b/src/aiida/workgraph/task.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from node_graph.task import Task as GraphTask +from node_graph.task import TaskSet +from node_graph.task_spec import BaseHandle, TaskSpec + +import aiida +from aiida.workgraph.enums import TaskState +from aiida.workgraph.socket_spec import SocketSpecAPI + +from .registry import RegistryHub, registry_hub + +if TYPE_CHECKING: + pass + + +class Task(GraphTask): + """Represent a Task in the AiiDA WorkGraph. + + The class extends from node_graph.task.Task and add new + attributes to it. + """ + + _REGISTRY: RegistryHub | None = registry_hub + _SOCKET_SPEC_API = SocketSpecAPI + + _default_spec = TaskSpec( + identifier='workgraph.task', + task_type='Normal', + inputs=_SOCKET_SPEC_API.namespace(), + outputs=_SOCKET_SPEC_API.namespace(), + catalog='Base', + base_class_path='aiida.workgraph.task.Task', + ) + + def __init__( + self, + process: aiida.orm.ProcessNode | None = None, + pk: int | None = None, + **kwargs: Any, + ) -> None: + """ + Initialize a Task instance. + """ + super().__init__( + **kwargs, + ) + self.waiting_on = WaitingTaskSet(parent=self) + self.process = process + self.pk = pk + self.state = TaskState.PLANNED + self.action = '' + self.show_socket_depth = 0 + self.parent = None + self.map_data = None + self.mapped_tasks = None + self.execution_count = 0 + + def to_dict(self, include_sockets: bool = False, should_serialize: bool = False) -> dict[str, Any]: + from aiida.orm.utils.serialize import serialize + + tdata = super().to_dict(include_sockets=include_sockets, should_serialize=should_serialize) + tdata['wait'] = [task.name for task in self.waiting_on] + tdata['children'] = [] + tdata['execution_count'] = self.execution_count + tdata['parent_task'] = [self.parent.name] if self.parent else [None] + tdata['process'] = serialize(self.process) if self.process else serialize(None) + tdata['metadata']['pk'] = self.process.pk if self.process else None + + return tdata + + def set_from_builder(self, builder: Any) -> None: + """Set the task inputs from a AiiDA ProcessBuilder.""" + from aiida.workgraph.utils import get_dict_from_builder + + data = get_dict_from_builder(builder) + self.set_inputs(data) + + @classmethod + def new(cls, identifier: str | Callable, name: str | None = None) -> Task: + """Create a task from a identifier.""" + from aiida.workgraph.tasks import TaskPool + + return super().new(identifier, name=name, TaskPool=TaskPool) + + @classmethod + def from_dict(cls, data: dict[str, Any], TaskPool: Any | None = None) -> Task: + """Create a task from a dictionary. This method initializes a Task instance with properties and settings + defined within the provided data dictionary. If TaskPool is not specified, the default TaskPool from + aiida.workgraph.tasks is used. + + Args: + data (Dict[str, Any]): A dictionary containing the task's configuration. + TaskPool (Optional[Any]): A pool of task configurations, defaults to None + which will use the global TaskPool. + + Returns: + Task: An instance of Task initialized with the provided data.""" + from aiida.workgraph.tasks import TaskPool as workgraph_TaskPool + + if TaskPool is None: + TaskPool = workgraph_TaskPool + task = GraphTask.from_dict(data, TaskPool=TaskPool) + + return task + + def update_from_dict(self, data: dict[str, Any]) -> None: + from aiida.workgraph.orm.utils import deserialize_safe + + super().update_from_dict(data) + process = data.get('process', None) + if process and isinstance(process, str): + process = deserialize_safe(process) + self.process = process + self.waiting_on.add(data.get('wait', [])) + self.map_data = data.get('map_data', None) + + def reset(self) -> None: + self.process = None + self.state = TaskState.PLANNED + + def update_state(self, data: dict[str, Any]) -> None: + """Set the outputs of the task from a dictionary.""" + self.state = data['state'] + self.ctime = data['ctime'] + self.mtime = data['mtime'] + self.pk = data['pk'] + if data['pk'] is not None: + node = aiida.orm.load_node(data['pk']) + self.process = self.node = node + if isinstance(node, aiida.orm.ProcessNode): + self.set_outputs_from_process_node(node) + elif isinstance(node, aiida.orm.Data): + self.set_outputs_from_data_node(node) + + def set_outputs_from_process_node(self, node: aiida.orm.ProcessNode) -> None: + from aiida.workgraph.utils import resolve_node_link_managers + + # if the process is finished ok, update the output sockets + # note the task.state may not be the same as the node.process_state + # for example, task.state can be `SKIPPED` if it is inside a conditional block, + # even if the node.is_finished_ok is True + self.process = node + if node.is_finished_ok: + self.outputs._set_socket_value(resolve_node_link_managers(node.outputs)) + + def set_outputs_from_data_node(self, node: aiida.orm.Data) -> None: + self.outputs[0].value = node + + def execute(self, args=None, kwargs=None, var_kwargs=None): + """Execute the task.""" + from node_graph.task_spec import BaseHandle + + executor = self.get_executor().callable + # the imported executor could be a wrapped function + if isinstance(executor, BaseHandle) and hasattr(executor, '_callable'): + executor = getattr(executor, '_callable') + if var_kwargs is None: + result = executor(*args, **kwargs) + else: + result = executor(*args, **kwargs, **var_kwargs) + return result, TaskState.FINISHED + + def to_widget_value(self): + from aiida.workgraph.utils import workgraph_to_short_json + + tdata = self.to_dict(include_sockets=True) + wgdata = {'name': self.name, 'tasks': {self.name: tdata}, 'links': []} + wgdata = workgraph_to_short_json(wgdata) + return wgdata + + +class WaitingTaskSet(TaskSet): + def add(self, tasks: list[str | Task] | str | Task) -> None: + """Add tasks to the collection. Tasks can be a list or a single Task or task name.""" + normalize_tasks = super().add(tasks) + for task in normalize_tasks: + source = task.outputs._wait + target = self.parent.inputs._wait + self.graph.add_link(source, target) + + +class TaskHandle(BaseHandle): + def __init__(self, spec): + from aiida.workgraph import WorkGraph + from aiida.workgraph.manager import get_current_graph + + super().__init__(spec, get_current_graph, graph_class=WorkGraph) + + def __call__(self, *args, **kwargs): + """Build a task into the current graph; forbid calling a task from inside + another running @task/@task.calcfunction/@task.workfunction body (i.e., during process execution).""" + + from aiida.engine import FunctionProcess, Process + + try: + current = Process.current() + except Exception: + current = None + + # Forbid the nested call only when the current process runs a user Python + # function in this interpreter: core's @calcfunction/@workfunction + # (FunctionProcess), or any plugin process that opts in via the + # `_runs_python_function_locally` marker (e.g. aiida-pythonjob's PyFunction). + # Keying on the marker rather than importing the plugin keeps this module + # free of any downstream dependency, a prerequisite for it living in core. + if current is not None and ( + isinstance(current, FunctionProcess) or getattr(current, '_runs_python_function_locally', False) + ): + running = getattr(current, 'process_label', current.__class__.__name__) + raise RuntimeError( + 'Invalid nested task call.\n\n' + f"• You invoked task '{self.identifier}' from inside the running process " + f"'{running}' ({current.__class__.__name__}).\n" + '• Tasks must not call other tasks directly inside a ' + '@task/@task.calcfunction/@task.workfunction body.\n\n' + 'Do one of the following instead:\n' + ' 1) Compose tasks in a @task.graph function (build a graph and connect tasks), or\n' + ' 2) Move shared logic into a plain Python helper function and call that.' + ) + + outputs = super().__call__(*args, **kwargs) + # if "metadata.call_link_label" is set, use it as the name of the task + if outputs._task.inputs.metadata.call_link_label.value is not None: + from aiida.workgraph.utils import _validate_task_name + + graph = outputs._graph + new_name = outputs._task.inputs.metadata.call_link_label.value + # `call_link_label` overrides the (already validated) function-derived name, so it + # must be validated too; otherwise an invalid override slips past `add_task` and + # only fails silently inside the engine at run time (see issue #784). + _validate_task_name(new_name, source='call_link_label') + outputs._task.name = new_name + # update the names of tasks and links collections in the graph + graph.tasks._items = {task.name: task for task in graph.tasks._items.values()} + graph.links._items = {link.name: link for link in graph.links._items.values()} + + return outputs + + def run(self, /, *args, **kwargs): + graph = self.build(*args, **kwargs) + return graph.run() + + def run_get_graph(self, /, *args, **kwargs): + graph = self.build(*args, **kwargs) + return graph.run(), graph + + def submit(self, /, *args, **kwargs): + graph = self.build(*args, **kwargs) + graph.submit() + return graph diff --git a/src/aiida/workgraph/tasks/__init__.py b/src/aiida/workgraph/tasks/__init__.py new file mode 100644 index 0000000000..57b10b1986 --- /dev/null +++ b/src/aiida/workgraph/tasks/__init__.py @@ -0,0 +1,3 @@ +from .task_pool import TaskPool + +__all__ = ['TaskPool'] diff --git a/src/aiida/workgraph/tasks/aiida.py b/src/aiida/workgraph/tasks/aiida.py new file mode 100644 index 0000000000..62124a1550 --- /dev/null +++ b/src/aiida/workgraph/tasks/aiida.py @@ -0,0 +1,138 @@ +from collections.abc import Callable + +from node_graph.error_handler import ErrorHandlerSpec, normalize_error_handlers +from node_graph.executor import RuntimeExecutor +from node_graph.socket_spec import SocketSpec +from node_graph.task_spec import SchemaSource, TaskSpec + +from aiida.engine import Process +from aiida.workgraph.enums import TaskAction, TaskState +from aiida.workgraph.socket_spec import from_aiida_process +from aiida.workgraph.task import Task +from aiida.workgraph.utils import inspect_aiida_component_type + +from .function_task import build_callable_TaskSpec + + +class AiiDAFunctionTask(Task): + """Task with AiiDA calcfunction/workfunction as executor.""" + + identifier = 'workgraph.aiida_functions' + name = 'aiida_function' + task_type = 'function' + catalog = 'AIIDA' + + def execute(self, args=None, kwargs=None, var_kwargs=None): + from node_graph.task_spec import BaseHandle + + from aiida.engine import run_get_node + + executor = RuntimeExecutor(**self.get_executor().to_dict()).callable + # the imported executor could be a wrapped function + if isinstance(executor, BaseHandle) and hasattr(executor, '_callable'): + executor = getattr(executor, '_callable') + kwargs.setdefault('metadata', {}) + kwargs['metadata'].update({'call_link_label': self.name}) + # since aiida 2.5.0, we need to use args_dict to pass the args to the run_get_node + if var_kwargs is None: + _, process = run_get_node(executor, **kwargs) + else: + _, process = run_get_node(executor, **kwargs, **var_kwargs) + + return process, TaskState.FINISHED + + +class AiiDAProcessTask(Task): + """Task with AiiDA calcfunction/workfunction as executor.""" + + identifier = 'workgraph.aiida_process' + name = 'aiida_process' + task_type = 'Process' + catalog = 'AIIDA' + + @classmethod + def build( + cls, + callable, + attached_error_handlers: dict[str, ErrorHandlerSpec] | None = None, + ): + attached_error_handlers = normalize_error_handlers(attached_error_handlers) + in_spec, out_spec = from_aiida_process(callable) + return TaskSpec( + identifier=callable.__name__, + schema_source=SchemaSource.CALLABLE, + catalog='AIIDA', + inputs=in_spec, + outputs=out_spec, + executor=RuntimeExecutor.from_callable(callable), + attached_error_handlers=attached_error_handlers, + base_class=cls, + task_type=inspect_aiida_component_type(callable), + ) + + def execute(self, engine_process, args=None, kwargs=None, var_kwargs=None): + from aiida.workgraph.utils import create_and_pause_process + + executor = RuntimeExecutor(**self.get_executor().to_dict()).callable + + kwargs.setdefault('metadata', {}) + kwargs['metadata'].update({'call_link_label': self.name}) + if self.action == TaskAction.PAUSE: + engine_process.report(f'Task {self.name} is created and paused.') + process = create_and_pause_process( + engine_process.runner, + executor, + kwargs, + state_msg='Paused through WorkGraph', + ) + state = TaskState.CREATED + process = process.node + else: + process = engine_process.submit(executor, **kwargs) + state = TaskState.RUNNING + + return process, state + + +class CalcJobTask(AiiDAProcessTask): + identifier = 'workgraph.calcjob' + name = 'calcjob' + task_type = 'CalcJob' + catalog = 'AIIDA' + + +class WorkChainTask(AiiDAProcessTask): + identifier = 'workgraph.workchain' + name = 'workchain' + task_type = 'WorkChain' + catalog = 'AIIDA' + + +def _build_aiida_function_taskspec( + obj: Callable, + identifier: str | None = None, + catalog: str = 'AIIDA', + in_spec: SocketSpec | None = None, + out_spec: SocketSpec | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, +) -> TaskSpec: + from dataclasses import replace + + from aiida.workgraph.utils import inspect_aiida_component_type + + spec = build_callable_TaskSpec( + obj=obj, + task_type=inspect_aiida_component_type(obj), + catalog=catalog, + base_class=AiiDAFunctionTask, + identifier=identifier, + process_cls=Process, + in_spec=in_spec, + out_spec=out_spec, + error_handlers=error_handlers, + ) + # the outputs of calcfunctions/workfunctions are always dynamic + spec = replace(spec, outputs=replace(spec.outputs, meta=replace(spec.outputs.meta, dynamic=True))) + if obj.spec().inputs.dynamic: + spec = replace(spec, inputs=replace(spec.inputs, meta=replace(spec.inputs.meta, dynamic=True))) + return spec diff --git a/src/aiida/workgraph/tasks/builtins.py b/src/aiida/workgraph/tasks/builtins.py new file mode 100644 index 0000000000..96baed9b26 --- /dev/null +++ b/src/aiida/workgraph/tasks/builtins.py @@ -0,0 +1,287 @@ +from __future__ import annotations + +from typing import Annotated, Any + +from node_graph import RuntimeExecutor +from node_graph.socket import BaseSocket +from node_graph.socket_spec import SocketMeta, SocketSpec +from node_graph.task import BuiltinPolicy, ChildTaskSet +from node_graph.task_spec import TaskSpec +from node_graph.tasks.builtins import _GraphIOSharedMixin + +from aiida import orm +from aiida.workgraph import dynamic, meta, namespace, task +from aiida.workgraph.executors.builtins import get_context, load_code, load_node, return_input, select, update_ctx +from aiida.workgraph.task import Task + + +class GraphLevelTask(_GraphIOSharedMixin, Task): + """Graph level task variant with shared IO.""" + + _default_spec = TaskSpec( + identifier='workgraph.graph_level_task', + catalog='Builtins', + base_class_path='aiida.workgraph.tasks.builtins.GraphLevelTask', + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._unify_io() + + +class Zone(Task): + """ + Extend the Task class to include a 'children' attribute. + """ + + _default_spec = TaskSpec( + identifier='workgraph.zone', + task_type='ZONE', + catalog='Control', + inputs=namespace(), + outputs=namespace(), + base_class_path='aiida.workgraph.tasks.builtins.Zone', + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.children = ChildTaskSet(parent=self) + + def add_task(self, *args, **kwargs) -> Task: + """Syntactic sugar to add a task to the zone.""" + task = self.graph.add_task(*args, **kwargs) + self.children.add(task) + task.parent = self + return task + + def to_dict(self, **kwargs) -> dict[str, Any]: + tdata = super().to_dict(**kwargs) + tdata['children'] = [task.name for task in self.children] + return tdata + + +class While(Zone): + """While""" + + _default_spec = TaskSpec( + identifier='workgraph.while_zone', + task_type='WHILE', + catalog='Control', + inputs=namespace( + max_iterations=Annotated[int, SocketSpec('workgraph.any', default=10000)], + conditions=Annotated[Any, SocketSpec('workgraph.any', link_limit=100000)], + ), + base_class_path='aiida.workgraph.tasks.builtins.While', + ) + + +class If(Zone): + """If task""" + + _default_spec = TaskSpec( + identifier='workgraph.if_zone', + task_type='IF', + catalog='Control', + inputs=namespace( + invert_condition=Annotated[bool, SocketSpec('workgraph.bool', default=False)], + conditions=Annotated[Any, SocketSpec('workgraph.any', link_limit=100000)], + ), + base_class_path='aiida.workgraph.tasks.builtins.If', + ) + + +class Map(Zone): + """Map""" + + _default_spec = TaskSpec( + identifier='workgraph.map_zone', + task_type='MAP', + catalog='Control', + inputs=namespace( + source=dynamic(Any), + ), + outputs=namespace(), + base_class_path='aiida.workgraph.tasks.builtins.Map', + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @property + def item(self): + for child in self.children: + if child.identifier == 'workgraph.map_item': + return child.outputs + # create a child map_item_task if it does not exist + map_item_task = self.add_task('workgraph.map_item') + return map_item_task.outputs + + @property + def gather_item_task(self) -> Task | None: + for child in self.children: + if child.identifier == 'workgraph.gather_item': + return child + gather_item = self.add_task('workgraph.gather_item') + return gather_item + + def gather(self, sockets: dict[str, BaseSocket]) -> None: + gather_item = self.gather_item_task + for name in sockets: + gather_item.add_input_spec('workgraph.any', name=name) + self.add_output_spec('workgraph.namespace', name=name) + gather_item.set_inputs(sockets) + return gather_item.outputs + + +class MapItem(Task): + """MapItem""" + + # turn off framework builtins for these graph-level nodes + _BUILTINS_POLICY = BuiltinPolicy(input_wait=False, output_wait=False, default_output=False) + + _default_spec = TaskSpec( + identifier='workgraph.map_item', + task_type='Normal', + catalog='Control', + inputs=namespace( + source=SocketSpec('workgraph.any', link_limit=100000, meta=SocketMeta(required=False)), + key=SocketSpec('workgraph.string', meta=SocketMeta(required=False)), + ), + outputs=namespace(key=str, value=Any), + base_class_path='aiida.workgraph.tasks.builtins.MapItem', + ) + + +class GatherItem(Task): + """GatherItem""" + + # turn off framework builtins for these graph-level nodes + _BUILTINS_POLICY = BuiltinPolicy(input_wait=True, output_wait=False, default_output=False) + + _default_spec = TaskSpec( + identifier='workgraph.gather_item', + task_type='Normal', + catalog='Control', + inputs=namespace(), + outputs=namespace(), + executor=RuntimeExecutor.from_callable(return_input), + base_class_path='aiida.workgraph.tasks.builtins.GatherItem', + ) + + +class SetContext(Task): + """SetContext""" + + _default_spec = TaskSpec( + identifier='workgraph.set_context', + task_type='Normal', + catalog='Control', + inputs=namespace( + context=SocketSpec('workgraph.any', meta=SocketMeta(required=False)), + key=Any, + value=Any, + ), + executor=RuntimeExecutor.from_callable(update_ctx), + base_class_path='aiida.workgraph.tasks.builtins.SetContext', + ) + + +class GetContext(Task): + """GetContext""" + + _default_spec = TaskSpec( + identifier='workgraph.get_context', + task_type='Normal', + catalog='Control', + inputs=namespace( + context=SocketSpec('workgraph.any', meta=SocketMeta(required=False)), + key=Any, + ), + outputs=namespace(result=Any), + executor=RuntimeExecutor.from_callable(get_context), + base_class_path='aiida.workgraph.tasks.builtins.GetContext', + ) + + +class Select(Task): + """Select""" + + _default_spec = TaskSpec( + identifier='workgraph.select', + task_type='Normal', + catalog='Control', + inputs=namespace( + condition=Any, + true=Any, + false=Any, + ), + outputs=namespace(result=Any), + executor=RuntimeExecutor.from_callable(select), + base_class_path='aiida.workgraph.tasks.builtins.Select', + ) + + +@task(identifier='workgraph.aiida_int') +def aiida_int(value: int) -> orm.Int: + return orm.Int(value) + + +@task(identifier='workgraph.aiida_float') +def aiida_float(value: float) -> orm.Float: + return orm.Float(value) + + +@task(identifier='workgraph.aiida_string') +def aiida_string(value: str) -> orm.Str: + return orm.Str(value) + + +@task(identifier='workgraph.aiida_list') +def aiida_list(value: list) -> orm.List: + return orm.List(value) + + +@task(identifier='workgraph.aiida_dict') +def aiida_dict(value: dict) -> orm.Dict: + return orm.Dict(value) + + +class AiiDANode(Task): + """AiiDANode""" + + identifier = 'workgraph.load_node' + name = 'AiiDANode' + catalog = 'Test' + + _default_spec = TaskSpec( + identifier=identifier, + task_type='Normal', + inputs=namespace( + pk=Annotated[int, meta(required=False)], + uuid=Annotated[str, meta(required=False)], + ), + outputs=namespace(node=orm.Node), + executor=RuntimeExecutor.from_callable(load_node), + base_class_path='aiida.workgraph.task.Task', + ) + + +class AiiDACode(Task): + """AiiDACode""" + + identifier = 'workgraph.load_code' + name = 'AiiDACode' + catalog = 'Test' + + _default_spec = TaskSpec( + identifier=identifier, + task_type='Normal', + inputs=namespace( + pk=Annotated[int, meta(required=False)], + uuid=Annotated[str, meta(required=False)], + label=Annotated[str, meta(required=False)], + ), + outputs=namespace(code=orm.Code), + executor=RuntimeExecutor.from_callable(load_code), + base_class_path='aiida.workgraph.task.Task', + ) diff --git a/src/aiida/workgraph/tasks/function_task.py b/src/aiida/workgraph/tasks/function_task.py new file mode 100644 index 0000000000..c63ae57578 --- /dev/null +++ b/src/aiida/workgraph/tasks/function_task.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from node_graph.error_handler import ErrorHandlerSpec, normalize_error_handlers +from node_graph.executor import RuntimeExecutor +from node_graph.socket_spec import SocketSpec, merge_specs +from node_graph.task_spec import SchemaSource, TaskSpec + +from aiida.workgraph.socket_spec import ( + from_aiida_process, + infer_specs_from_callable, +) + +if TYPE_CHECKING: + from node_graph import Node + + +def build_callable_TaskSpec( + *, + obj: Callable, + task_type: str, + base_class: type[Node], + identifier: str | None = None, + catalog: str = 'Others', + in_spec: SocketSpec | list[str] | None = None, + out_spec: SocketSpec | list[str] | None = None, + process_cls: type | None = None, # e.g. PythonJob, PyFunction, or aiida.engine.Process + add_inputs: SocketSpec | list[str] | None = None, + add_outputs: SocketSpec | list[str] | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, + metadata: dict | None = None, +) -> TaskSpec: + """ + - infers function I/O + - optionally merges process-contributed I/O + - optionally merges additional I/O + - records *each* contribution in metadata + """ + from aiida.workgraph.socket_spec import validate_socket_data + + error_handlers = normalize_error_handlers(error_handlers) + + in_spec = validate_socket_data(in_spec) + out_spec = validate_socket_data(out_spec) + + # 1) infer from the callable (keep a snapshot before augmentation) + func_in, func_out = infer_specs_from_callable(obj, in_spec, out_spec) + # "metadata" is reserved for AiiDA process, so raise error if user tries to use it + if 'metadata' in func_in.fields: + fn = getattr(obj, '__name__', 'the task function') + raise ValueError( + "Invalid input name: 'metadata'\n" + "Reason: In AiiDA, 'metadata' is reserved for process-level settings " + '(e.g., call_link_label, description) and cannot be used as a task input.\n\n' + f'How to fix: Rename the argument in {fn} to something else, e.g.: task_metadata.\n\n' + 'Example:\n' + ' # before\n' + f' def {fn}(metadata: dict, x: int):\n' + ' ...\n\n' + ' # after\n' + f' def {fn}(task_metadata: dict, x: int):\n' + ' ...\n' + ) + + # 2) process-contributed I/O (if any) + proc_in = proc_out = None + if process_cls is not None: + proc_in, proc_out = from_aiida_process(process_cls) + func_in = merge_specs(func_in, proc_in) + func_out = merge_specs(func_out, proc_out) + + # 3) additional fields (if any) + if add_inputs is not None: + func_in = merge_specs(func_in, add_inputs) + if add_outputs is not None: + func_out = merge_specs(func_out, add_outputs) + + # 4) metadata: keep a record of each contribution + metadata = metadata or {} + metadata.update( + { + 'non_function_inputs': list( + set((proc_in and proc_in.fields.keys()) or []) | set((add_inputs and add_inputs.fields.keys()) or []) + ), + 'non_function_outputs': list( + set((proc_out and proc_out.fields.keys()) or []) + | set((add_outputs and add_outputs.fields.keys()) or []) + ), + } + ) + # We always use EMBEDDED schema for function tasks + # but when store the spec in the DB, we will check if the + # callable is a BaseHandler, and switch the schema_source to HANDLER accordingly. + # This avoid cyclic import. + schema_source = SchemaSource.EMBEDDED + + return TaskSpec( + identifier=identifier or obj.__name__, + schema_source=schema_source, + task_type=task_type, + catalog=catalog, + inputs=func_in, + outputs=func_out, + executor=RuntimeExecutor.from_callable(obj), + error_handlers=error_handlers, + base_class=base_class, + metadata=metadata, + ) diff --git a/src/aiida/workgraph/tasks/graph_task.py b/src/aiida/workgraph/tasks/graph_task.py new file mode 100644 index 0000000000..0c164c775b --- /dev/null +++ b/src/aiida/workgraph/tasks/graph_task.py @@ -0,0 +1,123 @@ +from collections.abc import Callable + +from node_graph.executor import RuntimeExecutor +from node_graph.socket_spec import SocketSpec +from node_graph.task_spec import TaskSpec + +from aiida.engine import Process +from aiida.workgraph.enums import TaskAction, TaskState +from aiida.workgraph.task import Task + +from .function_task import build_callable_TaskSpec + + +class GraphTask(Task): + """Graph builder task""" + + identifier = 'workgraph.graph_task' + name = 'graph_task' + task_type = 'graph_task' + catalog = 'builtins' + + def execute(self, engine_process, args=None, kwargs=None, var_kwargs=None): + from node_graph.utils.graph import materialize_graph + + from aiida.workgraph import WorkGraph + from aiida.workgraph.engine.process import WorkGraphProcess + from aiida.workgraph.task import TaskHandle + from aiida.workgraph.utils import call_depth_from_node, create_and_pause_process + + executor = RuntimeExecutor(**self.get_executor().to_dict()).callable + max_depth = self.spec.metadata.get('max_depth', 100) + metadata = kwargs.pop('metadata', {}) if kwargs else {} + metadata.setdefault('call_link_label', self.name) + # Cloudpickle doesn't restore the function's own name in its globals after unpickling, + # so any recursive calls would raise NameError. We re-insert a task handle into its + # globals under its original name. We reuse the spec built at decoration time rather + # than re-decorating the function: re-decoration would re-infer the signature, which + # fails under PEP 563 once cloudpickle has dropped the names used only in stringized + # annotations (issue #783). + # Downside: this mutates the module globals at runtime, if another symbol with the same name exists, + # we may introduce hard-to-trace bugs or collisions. + if isinstance(executor, TaskHandle) and hasattr(executor, '_callable'): + executor = executor._callable + recursion_handle = TaskHandle(self.spec) + recursion_handle._callable = executor + executor.__globals__[executor.__name__] = recursion_handle + depth = call_depth_from_node(engine_process.node) + if depth >= max_depth: + if depth >= max_depth: + msg = ( + f"Graph task '{self.name}' exceeded the recursion safeguard.\n" + f'- Current AiiDA process call depth (approx.): {depth}\n' + f'- Allowed maximum : {max_depth}\n' + f'- Process UUID: {engine_process.node.uuid}\n\n' + f'Deeply nested process calls (>100) are generally discouraged. ' + f'Prefer wrapping iterative logic inside a single task instead of ' + f'recursively spawning new graph tasks.\n\n' + f'However, if you are confident that recursion is the right design, ' + f'you can explicitly set a higher limit in your decorator, e.g.:\n' + f' @task.graph(max_depth=200)\n' + ) + engine_process.report(msg) + raise RecursionError(msg) + wg = materialize_graph( + executor, + self.spec.inputs, + self.spec.outputs, + self.name, + WorkGraph, + args=args, + kwargs=kwargs, + var_kwargs=var_kwargs, + ) + # Set the maximum number of concurrent jobs + max_number_jobs = self.spec.metadata.get('max_number_jobs') + if max_number_jobs is not None: + wg.max_number_jobs = max_number_jobs + wg.parent_uuid = engine_process.node.uuid + inputs = wg.to_engine_inputs(metadata=metadata) + if self.action == TaskAction.PAUSE: + engine_process.report(f'Task {self.name} is created and paused.') + process = create_and_pause_process( + engine_process.runner, + WorkGraphProcess, + inputs, + state_msg='Paused through WorkGraph', + ) + state = TaskState.CREATED + process = process.node + else: + process = engine_process.submit(WorkGraphProcess, **inputs) + state = TaskState.RUNNING + + return process, state + + +def _build_graph_task_taskspec( + obj: Callable, + identifier: str | None = None, + in_spec: SocketSpec | None = None, + out_spec: SocketSpec | None = None, + max_depth: int = 100, + max_number_jobs: int = 1000000, + catalog: str = 'Others', +) -> TaskSpec: + # defaults for max depth + metadata = {'max_depth': max_depth, 'max_number_jobs': max_number_jobs} + # We use Process as the process class here, so that the task inherits the metadata + # inputs from the base Process class, such as 'call_link_label'. + # While the actual process class will be the WorkGraphProcess, + # which is set at runtime in the execute() method + + return build_callable_TaskSpec( + obj=obj, + task_type='GRAPH', + base_class=GraphTask, + identifier=identifier, + catalog=catalog, + process_cls=Process, + in_spec=in_spec, + out_spec=out_spec, + metadata=metadata, + ) diff --git a/src/aiida/workgraph/tasks/monitors.py b/src/aiida/workgraph/tasks/monitors.py new file mode 100644 index 0000000000..858d305d2d --- /dev/null +++ b/src/aiida/workgraph/tasks/monitors.py @@ -0,0 +1,58 @@ +import datetime +import logging + +from aiida.workgraph import task +from aiida.workgraph.enums import TERMINAL_TASK_STATES + +LOGGER = logging.getLogger(__name__) + + +@task.monitor +def monitor_file(filepath: str): + """Return `True` when the file is detected.""" + import os + + return os.path.exists(filepath) + + +@task.monitor +def monitor_time(time: str | datetime.datetime): + """Return `True` when the given moment in time has passed. + + If given as a string, `time` should be in ISO format (e.g., '2025-07-30T12:00:00'). + """ + + if isinstance(time, str): + try: + time = datetime.datetime.fromisoformat(time) + except ValueError as err: + raise ValueError(f'Invalid time format: {time}. Expected ISO format.') from err + + return datetime.datetime.now() > time + + +@task.monitor +def monitor_task(task_name: str, workgraph_pk: int | None = None, workgraph_name: str | None = None): + """Return `True` if the task in the WorkGraph is completed.""" + from aiida import orm + from aiida.workgraph.engine.process import WorkGraphProcess + + if workgraph_pk: + try: + node = orm.load_node(workgraph_pk) + except Exception: + return False + else: + builder = orm.QueryBuilder() + builder.append( + WorkGraphProcess, + filters={'attributes.process_label': {'==': f'WorkGraph<{workgraph_name}>'}}, + tag='process', + ) + if builder.count() == 0: + return False + LOGGER.debug('Found workgraph') + node = builder.first()[0] + state = node.task_states.get(task_name, '') + LOGGER.debug('Task state: %s', state) + return state in TERMINAL_TASK_STATES diff --git a/src/aiida/workgraph/tasks/pythonjob_tasks.py b/src/aiida/workgraph/tasks/pythonjob_tasks.py new file mode 100644 index 0000000000..b80bd8d49c --- /dev/null +++ b/src/aiida/workgraph/tasks/pythonjob_tasks.py @@ -0,0 +1,343 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Annotated, Any + +from aiida_pythonjob import MonitorPyFunction, PyFunction, PythonJob, pyfunction +from node_graph.error_handler import ErrorHandlerSpec +from node_graph.executor import RuntimeExecutor +from node_graph.socket_spec import SocketMeta, SocketSpec, SocketSpecSelect +from node_graph.task_spec import BaseHandle, TaskSpec + +from aiida import orm +from aiida.common.extendeddicts import AttributeDict +from aiida.engine import run_get_node +from aiida.workgraph.enums import TaskAction, TaskState +from aiida.workgraph.socket_spec import namespace +from aiida.workgraph.task import Task +from aiida.workgraph.utils import create_and_pause_process + +from .function_task import build_callable_TaskSpec + + +class BaseSerializablePythonTask(Task): + """ + A base Task that handles serialization and deserialization + of Python data into AiiDA Data nodes, so that raw Python data + can be stored/passed around by the WorkGraph engine. + Subclasses must implement their own `execute` method. + """ + + @property + def non_function_inputs(self): + return self.spec.metadata.get('non_function_inputs', []) + + @property + def non_function_outputs(self): + return self.spec.metadata.get('non_function_outputs', []) + + @property + def function_inputs_spec(self): + inputs_spec = namespace( + _=Annotated[ + Any, + self.spec.inputs, + SocketSpecSelect(exclude=self.non_function_inputs), + ] + ).fields['_'] + return inputs_spec + + @property + def function_outputs_spec(self): + outputs_spec = namespace( + _=Annotated[ + Any, + self.spec.outputs, + SocketSpecSelect(exclude=self.non_function_outputs), + ] + ).fields['_'] + return outputs_spec + + def get_function_inputs(self, kwargs, var_kwargs): + function_inputs = kwargs.pop('function_inputs', {}) or {} + for key in list(kwargs.keys()): + if key not in self.non_function_inputs: + function_inputs[key] = kwargs.pop(key) + # Handle var_kwargs + if self.get_args_data()['var_kwargs'] is not None: + var_key = self.get_args_data()['var_kwargs'] + function_inputs.pop(var_key, None) + if var_kwargs: + if isinstance(var_kwargs, (dict, AttributeDict)): + function_inputs.update(var_kwargs) + elif isinstance(var_kwargs, orm.Data): + function_inputs.update(var_kwargs.value) + else: + raise ValueError(f'Invalid var_kwargs type: {type(var_kwargs)}') + return function_inputs + + def get_process_metadata(self, kwargs): + metadata = kwargs.pop('metadata', {}) + metadata.update({'call_link_label': self.name}) + return metadata + + +class PythonJobTask(BaseSerializablePythonTask): + """PythonJob Task.""" + + identifier = 'workgraph.pythonjob' + + def execute(self, engine_process, args=None, kwargs=None, var_kwargs=None): + """ + Here is the specialized 'execute' method for PythonJobTask, + including the 'prepare_for_python_task' logic. + """ + from aiida_pythonjob import prepare_pythonjob_inputs + + # Pull out code, computer, etc + computer = kwargs.pop('computer', 'localhost') + if isinstance(computer, orm.Str): + computer = computer.value + command_info = kwargs.pop('command_info', {}) + register_pickle_by_value = kwargs.pop('register_pickle_by_value', False) + upload_files = kwargs.pop('upload_files', {}) + metadata = self.get_process_metadata(kwargs) + function_inputs = self.get_function_inputs(kwargs, var_kwargs) + func = RuntimeExecutor(**self.get_executor().to_dict()).callable + # If it's a wrapped function, unwrap + if isinstance(func, BaseHandle) and hasattr(func, '_callable'): + func = func._callable + + if hasattr(func, 'is_process_function'): + func = func.func + + # Prepare the final inputs for PythonJob + inputs = prepare_pythonjob_inputs( + function=func, + function_inputs=function_inputs, + inputs_spec=self.function_inputs_spec, + outputs_spec=self.function_outputs_spec, + code=kwargs.pop('code', None), + command_info=command_info, + computer=computer, + metadata=metadata, + upload_files=upload_files, + process_label=f'PythonJob<{self.name}>', + register_pickle_by_value=register_pickle_by_value, + **kwargs, + ) + + # If we want to pause + if self.action == TaskAction.PAUSE: + engine_process.report(f'Task {self.name} is created and paused.') + process = create_and_pause_process( + engine_process.runner, + PythonJob, + inputs, + state_msg='Paused through WorkGraph', + ) + state = TaskState.CREATED + process = process.node + else: + process = engine_process.submit(PythonJob, **inputs) + state = TaskState.RUNNING + + return process, state + + +class PyFunctionTask(BaseSerializablePythonTask): + """PyFunction Task.""" + + identifier = 'workgraph.pyfunction' + + def execute(self, args=None, kwargs=None, var_kwargs=None, engine_process=None): + from aiida_pythonjob import prepare_pyfunction_inputs + + kwargs = kwargs or {} + metadata = self.get_process_metadata(kwargs) + func = RuntimeExecutor(**self.get_executor().to_dict()).callable + # If it's a wrapped function, unwrap + if isinstance(func, BaseHandle) and hasattr(func, '_callable'): + func = func._callable + + if self.spec.metadata.get('is_coroutine', False): + function_inputs = self.get_function_inputs(kwargs, var_kwargs) + inputs = prepare_pyfunction_inputs( + function=func, + function_inputs=function_inputs, + inputs_spec=self.function_inputs_spec, + outputs_spec=self.function_outputs_spec, + metadata=metadata, + process_label=kwargs.pop('process_label', None), + deserializers=kwargs.pop('deserializers', None), + serializers=kwargs.pop('serializers', None), + register_pickle_by_value=kwargs.pop('register_pickle_by_value', False), + ) + if self.action == TaskAction.PAUSE: + engine_process.report(f'Task {self.name} is created and paused.') + process = create_and_pause_process( + engine_process.runner, + PyFunction, + inputs, + state_msg='Paused through WorkGraph', + ) + state = TaskState.CREATED + process = process.node + else: + process = engine_process.submit(PyFunction, **inputs) + state = TaskState.RUNNING + + return process, state + else: + # Make sure it's process_function-decorated + if not hasattr(func, 'is_process_function'): + func = pyfunction()(func) + + # If we have var_kwargs, pass them in + if var_kwargs is None: + _, process = run_get_node( + func, + inputs_spec=self.function_inputs_spec, + outputs_spec=self.function_outputs_spec, + metadata=metadata, + **kwargs, + ) + else: + _, process = run_get_node(func, **kwargs, **var_kwargs) + + return process, TaskState.FINISHED + + +class MonitorFunctionTask(BaseSerializablePythonTask): + """Monitor Function Task.""" + + identifier = 'workgraph.monitor_function' + + def execute(self, args=None, kwargs=None, var_kwargs=None, engine_process=None): + from aiida_pythonjob import prepare_monitor_function_inputs + + kwargs = kwargs or {} + metadata = self.get_process_metadata(kwargs) + func = RuntimeExecutor(**self.get_executor().to_dict()).callable + # If it's a wrapped function, unwrap + if isinstance(func, BaseHandle) and hasattr(func, '_callable'): + func = func._callable + function_inputs = self.get_function_inputs(kwargs, var_kwargs) + inputs = prepare_monitor_function_inputs( + function=func, + function_inputs=function_inputs, + inputs_spec=self.function_inputs_spec, + outputs_spec=self.function_outputs_spec, + metadata=metadata, + process_label=kwargs.pop('process_label', None), + deserializers=kwargs.pop('deserializers', None), + serializers=kwargs.pop('serializers', None), + register_pickle_by_value=kwargs.pop('register_pickle_by_value', False), + **kwargs, + ) + if self.action == TaskAction.PAUSE: + engine_process.report(f'Task {self.name} is created and paused.') + process = create_and_pause_process( + engine_process.runner, + MonitorPyFunction, + inputs, + state_msg='Paused through WorkGraph', + ) + state = TaskState.CREATED + process = process.node + else: + process = engine_process.submit(MonitorPyFunction, **inputs) + state = TaskState.RUNNING + return process, state + + +def build_pythonjob_taskspec( + obj: Callable, + identifier: str | None = None, + catalog: str = 'Others', + in_spec: SocketSpec | None | list = None, + out_spec: SocketSpec | None | list = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, +) -> TaskSpec: + # allow list specs just for PythonJob (keep existing behavior) + from aiida.workgraph.socket_spec import validate_socket_data + + in_spec = validate_socket_data(in_spec) + out_spec = validate_socket_data(out_spec) + + # additions specific to PythonJob + add_in = namespace( + computer=Annotated[str, SocketMeta(required=False)], + command_info=Annotated[dict, SocketMeta(required=False)], + register_pickle_by_value=Annotated[bool, SocketMeta(required=False)], + ) + + return build_callable_TaskSpec( + obj=obj, + task_type='PYTHONJOB', + catalog=catalog, + base_class=PythonJobTask, + identifier=identifier, + process_cls=PythonJob, + in_spec=in_spec, + out_spec=out_spec, + add_inputs=add_in, + error_handlers=error_handlers, + ) + + +def build_pyfunction_taskspec( + obj: Callable, + identifier: str | None = None, + catalog: str = 'Others', + in_spec: SocketSpec | None = None, + out_spec: SocketSpec | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, +) -> TaskSpec: + import asyncio + + if asyncio.iscoroutinefunction(obj): + metadata = {'is_coroutine': True} + else: + metadata = {} + return build_callable_TaskSpec( + obj=obj, + task_type='PYFUNCTION', + catalog=catalog, + base_class=PyFunctionTask, + identifier=identifier, + process_cls=PyFunction, + in_spec=in_spec, + out_spec=out_spec, + error_handlers=error_handlers, + metadata=metadata, + ) + + +def build_monitor_function_taskspec( + obj: Callable, + identifier: str | None = None, + catalog: str = 'Others', + in_spec: SocketSpec | None = None, + out_spec: SocketSpec | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, +) -> TaskSpec: + add_in = namespace( + interval=(int, 5), + timeout=(int, 3600), + ) + add_out = namespace(exit_code=Any) + + return build_callable_TaskSpec( + obj=obj, + task_type='MONITOR', + catalog=catalog, + base_class=MonitorFunctionTask, + identifier=identifier, + process_cls=MonitorPyFunction, + in_spec=in_spec, + out_spec=out_spec, + add_inputs=add_in, + add_outputs=add_out, + error_handlers=error_handlers, + ) diff --git a/src/aiida/workgraph/tasks/shelljob_task.py b/src/aiida/workgraph/tasks/shelljob_task.py new file mode 100644 index 0000000000..98ba8309d5 --- /dev/null +++ b/src/aiida/workgraph/tasks/shelljob_task.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import inspect +from collections.abc import Callable +from dataclasses import replace +from typing import Annotated, Any + +from aiida_shell import ShellJob +from aiida_shell.launch import prepare_shell_job_inputs +from node_graph.executor import RuntimeExecutor +from node_graph.socket_spec import SocketMeta, SocketSpec, merge_specs +from node_graph.task_spec import TaskSpec + +from aiida import orm +from aiida.workgraph.enums import TaskAction, TaskState +from aiida.workgraph.socket_spec import from_aiida_process, namespace +from aiida.workgraph.task import Task, TaskHandle + + +def _serialize_value(self, store: bool = False) -> Any: + from node_graph.utils import resolve_tagged_values + + value = resolve_tagged_values(self._value) + if value is None: + return None + return RuntimeExecutor.from_callable(value).to_dict() + + +class ShellJobTask(Task): + """Runtime for ShellJob nodes. + + This class is referenced by TaskSpec.base_class_path so the engine can import + it and call `execute`. + """ + + identifier = 'workgraph.shelljob' + name = 'shelljob' + task_type = 'SHELLJOB' + catalog = 'AIIDA' + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # override the _serialize_value + self.inputs['parser'].set_serializer(_serialize_value) + + def execute(self, engine_process, args=None, kwargs=None, var_kwargs=None): + """Submit/launch the AiiDA ShellJob. + + - Translates friendly inputs (command, resolve_command, parser, ...) + using `prepare_shell_job_inputs`. + - Submits or runs under the engine's runner. + """ + from aiida.workgraph.utils import create_and_pause_process + + kwargs = dict(kwargs or {}) + + # Detect and translate aiida-shell convenience arguments + signature = inspect.signature(prepare_shell_job_inputs) + aiida_shell_keys = signature.parameters.keys() + + subset = {k: kwargs[k] for k in list(kwargs) if k in aiida_shell_keys} + + parser = subset.get('parser', None) + if isinstance(parser, dict) and {'module_path', 'callable_name'} <= set(parser): + # already a Executor dict -> build executor instance + subset['parser'] = RuntimeExecutor(**parser).callable + elif inspect.isfunction(parser): + subset['parser'] = parser + + if subset: + if 'command' in subset: + subset['command'] = ( + subset['command'].value if isinstance(subset['command'], orm.Str) else subset['command'] + ) + if 'resolve_command' in subset: + subset['resolve_command'] = ( + subset['resolve_command'].value + if isinstance(subset['resolve_command'], orm.Bool) + else subset['resolve_command'] + ) + if 'arguments' in subset: + subset['arguments'] = ( + subset['arguments'].get_list() if isinstance(subset['arguments'], orm.List) else subset['arguments'] + ) + prepared = prepare_shell_job_inputs(**subset) + # drop original keys so they won't clash with launch kwargs + for k in subset.keys(): + kwargs.pop(k, None) + # merge translated inputs + kwargs.update(prepared) + + # metadata + md = kwargs.setdefault('metadata', {}) + md.setdefault('call_link_label', self.name) + + if getattr(self, 'action', None) == TaskAction.PAUSE: + engine_process.report(f'Task {self.name} is created and paused.') + process = create_and_pause_process( + engine_process.runner, + ShellJob, + kwargs, + state_msg='Paused through WorkGraph', + ) + state = TaskState.CREATED + process = process.node + else: + process = engine_process.submit(ShellJob, **kwargs) + state = TaskState.RUNNING + return process, state + + +def _build_shelljob_TaskSpec( + *, + identifier: str | None = None, + outputs: SocketSpec | list[str] | None = None, + parser_outputs: SocketSpec | list[str] | None = None, +) -> TaskSpec: + """Create a `TaskSpec` for a ShellJob, augmenting inputs/outputs as needed. + + - Start from AiiDA Process spec inference + - Add inputs: command, resolve_command + - Ensure stdout/stderr outputs exist + - Optionally add user-declared outputs and parser_outputs (as leaf-any) + """ + from aiida_shell.parsers.shell import ShellParser + + from aiida.workgraph.socket_spec import validate_socket_data + + outputs = validate_socket_data(outputs) + parser_outputs = validate_socket_data(parser_outputs) + + in_spec, out_spec = from_aiida_process(ShellJob) + # the code socket is not required in the task + # as we can build it from the command input + code_spec = in_spec.fields['code'] + patched_code = replace(code_spec, meta=replace(code_spec.meta, required=False)) + in_spec = replace(in_spec, fields={**in_spec.fields, 'code': patched_code}) + + # Add additional inputs + additions_in = namespace(command=Any, resolve_command=Annotated[bool, SocketMeta(required=False)]) + in_spec = merge_specs(in_spec, additions_in) + + # Ensure stdout/stderr outputs + additions_out = namespace(stdout=Any, stderr=Any) + out_spec = merge_specs(out_spec, additions_out) + + # add extra outputs requested by user + if outputs: + # make sure the key are AiiDA compatible + fields = {ShellParser.format_link_label(key): value for key, value in outputs.fields.items()} + outputs = replace(outputs, fields=fields) + out_spec = merge_specs(out_spec, outputs) + + if parser_outputs: + out_spec = merge_specs(out_spec, parser_outputs) + + exec_payload = RuntimeExecutor.from_callable(ShellJob) + + return TaskSpec( + identifier=identifier or 'ShellJob', + catalog='AIIDA', + task_type='SHELLJOB', + inputs=in_spec, + outputs=out_spec, + executor=exec_payload, + base_class=ShellJobTask, + metadata={'task_type': 'SHELLJOB'}, + ) + + +# Public factory used by users inside a WorkGraph + + +def shelljob( + *, + command: str, + arguments: list[str] | None = None, + nodes: dict[str, Any] | None = None, + filenames: dict[str, str] | None = None, + outputs: list[str | dict[str, Any]] | None = None, + parser: Callable | None = None, + parser_outputs: SocketSpec | list[str] | None = None, + metadata: dict[str, Any] | None = None, + resolve_command: bool = True, +): + """Create a ShellJob node in the active WorkGraph and return its outputs handle. + + Usage: + with WorkGraph(name="test_shell_date_with_arguments") as wg: + outs = shelljob(command="date", arguments=["--iso-8601"]) # returns handle + wg.run() + """ + spec = _build_shelljob_TaskSpec(outputs=outputs, parser_outputs=parser_outputs) + + handle = TaskHandle(spec) + return handle( + command=command, + arguments=arguments, + nodes=nodes, + filenames=filenames, + outputs=outputs, + parser=parser, + metadata=metadata, + resolve_command=resolve_command, + ) diff --git a/src/aiida/workgraph/tasks/subgraph_task.py b/src/aiida/workgraph/tasks/subgraph_task.py new file mode 100644 index 0000000000..49d590c7c9 --- /dev/null +++ b/src/aiida/workgraph/tasks/subgraph_task.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from node_graph.task_spec import TaskSpec + +from aiida.workgraph.enums import TaskAction, TaskState +from aiida.workgraph.task import Task + +if TYPE_CHECKING: + from aiida.workgraph import WorkGraph + + +class SubGraphTask(Task): + """Task created from WorkGraph.""" + + identifier = 'workgraph.workgraph_task' + name = 'SubGraphTask' + task_type = 'Normal' + catalog = 'Builtins' + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._subgraph = None + + @property + def subgraph(self): + from copy import deepcopy + + from aiida.workgraph import WorkGraph + + if not self._subgraph: + graph_data = deepcopy(self.get_executor().graph_data) + self._subgraph = WorkGraph.from_dict(graph_data) + return self._subgraph + + @property + def tasks(self): + return self.subgraph.tasks + + @property + def links(self): + return self.subgraph.links + + def prepare_for_subgraph_task(self, kwargs: dict) -> tuple: + """Prepare the inputs for SubGraph task""" + # update the subgraph inputs by the kwargs + for name, data in kwargs.items(): + input_socket = self.subgraph.inputs[name] + input_socket._set_socket_value(data) + # merge the properties + metadata = {'call_link_label': self.name} + inputs = self.subgraph.to_engine_inputs(metadata=metadata) + return inputs + + def execute(self, engine_process, args=None, kwargs=None, var_kwargs=None): + from aiida.workgraph.engine.process import WorkGraphProcess + from aiida.workgraph.utils import create_and_pause_process + + inputs = self.prepare_for_subgraph_task(kwargs) + + if self.action == TaskAction.PAUSE: + engine_process.report(f'Task {self.name} is created and paused.') + process = create_and_pause_process( + engine_process.runner, + WorkGraphProcess, + inputs, + state_msg='Paused through WorkGraph', + ) + state = TaskState.CREATED + process = process.node + else: + process = engine_process.submit(WorkGraphProcess, **inputs) + state = TaskState.RUNNING + + return process, state + + +def _build_subgraph_task_TaskSpec( + graph: WorkGraph, + name: str | None = None, +) -> TaskSpec: + from node_graph.executor import SafeExecutor + + return TaskSpec( + identifier=name or graph.name, + task_type='SubGraph', + inputs=graph.spec.inputs, + outputs=graph.spec.outputs, + executor=SafeExecutor.from_graph(graph), + base_class=SubGraphTask, + ) diff --git a/src/aiida/workgraph/tasks/task_pool.py b/src/aiida/workgraph/tasks/task_pool.py new file mode 100644 index 0000000000..b8cb90a90f --- /dev/null +++ b/src/aiida/workgraph/tasks/task_pool.py @@ -0,0 +1,6 @@ +from node_graph.registry import EntryPointPool + +# global instance +TaskPool = EntryPointPool(entry_point_group='aiida_workgraph.task') +TaskPool['any'] = TaskPool.workgraph.any +TaskPool['graph_level'] = TaskPool.workgraph.graph_level diff --git a/src/aiida/workgraph/tasks/tests.py b/src/aiida/workgraph/tasks/tests.py new file mode 100644 index 0000000000..d3b1801921 --- /dev/null +++ b/src/aiida/workgraph/tasks/tests.py @@ -0,0 +1,59 @@ +from pydantic import BaseModel + +from aiida.calculations.arithmetic.add import ArithmeticAddCalculation +from aiida.workgraph import Task, task + + +class BlobModel(BaseModel): + model_config = {'leaf': True} # always a leaf blob + + a: int + b: int + + +class AnotherModel(BaseModel): + a: int + b: int + + +@task +def add(x, y): + """Add two numbers.""" + return x + y + + +@task +def multiply(x, y): + """Multiply two numbers.""" + return x * y + + +def handle_negative_sum(task: Task): + """Handle negative sum by resetting the task and changing the inputs. + self is the WorkGraph instance, thus we can access the tasks and the context. + """ + from aiida import orm + + # modify task inputs + task.set_inputs( + { + 'x': orm.Int(abs(task.inputs.x.value)), + 'y': orm.Int(abs(task.inputs.y.value)), + } + ) + msg = 'Run error handler: handle_negative_sum.' + return msg + + +BaseAddTask = task( + error_handlers={ + 'handle_negative_sum': { + 'executor': handle_negative_sum, + 'exit_codes': [ + ArithmeticAddCalculation.exit_codes.ERROR_NEGATIVE_NUMBER.status, + ], + 'max_retries': 5, + 'kwargs': {}, + } + } +)(ArithmeticAddCalculation) diff --git a/src/aiida/workgraph/utils.py b/src/aiida/workgraph/utils.py deleted file mode 100644 index e166ad8e28..0000000000 --- a/src/aiida/workgraph/utils.py +++ /dev/null @@ -1,147 +0,0 @@ -########################################################################### -# Copyright (c), The AiiDA team. All rights reserved. # -# This file is part of the AiiDA code. # -# # -# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # -# For further information on the license, see the LICENSE.txt file # -# For further information please visit http://www.aiida.net # -########################################################################### -"""Generic helpers used by the WorkGraph runtime. - -These are the node-graph-free, plugin-free utilities the engine relies on: dotted-key access into nested -dictionaries, and resolving AiiDA ``NodeLinksManager`` structures into plain dictionaries. -""" - -from __future__ import annotations - -from typing import Any - -from aiida.orm.utils.managers import NodeLinksManager - -__all__ = ( - 'get_nested_dict', - 'resolve_node_link_managers', - 'update_nested_dict', - 'update_nested_dict_with_special_keys', -) - - -def get_nested_dict(d: Any, name: str, **kwargs: Any) -> Any: - """Get the value from a nested dictionary. - - ``d`` is deliberately ``Any``: the traversal descends through both plain dicts and AiiDA - ``NodeLinksManager`` containers, whose values are heterogeneous. - - If default is provided, return the default value if the key is not found. - Otherwise, raise ValueError. - For example: - d = {"base": {"pw": {"parameters": 2}}} - name = "base.pw.parameters" - """ - keys = name.split('.') - current = d - for key in keys: - if key not in current: - if 'default' in kwargs: - return kwargs.get('default') - if isinstance(current, dict): - avaiable_keys = list(current.keys()) - elif isinstance(current, NodeLinksManager): - avaiable_keys = list(current._get_keys()) - else: - avaiable_keys = [] - raise ValueError(f'{name} not exist. Available keys: {avaiable_keys}') - current = current[key] - return current - - -def merge_dicts(dict1: Any, dict2: Any) -> Any: - """Recursively merges two dictionaries.""" - for key, value in dict2.items(): - if key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict): - # Recursively merge dictionaries - dict1[key] = merge_dicts(dict1[key], value) - else: - # Overwrite or add the key - dict1[key] = value - return dict1 - - -def update_nested_dict(base: dict[str, Any] | None, key_path: str, value: Any) -> dict[str, Any]: - """ - Update or create a nested dictionary structure based on a dotted key path. - - This function allows updating a nested dictionary or creating one if `d` is `None`. - Given a dictionary and a key path (e.g., "base.pw.parameters"), it will traverse - or create the necessary nested structure to set the provided value at the specified - key location. If intermediate dictionaries do not exist, they will be created. - If the resulting dictionary is empty, it is set to `None`. - - Args: - base (Dict[str, Any] | None): The dictionary to update, which can be `None`. - If `None`, an empty dictionary will be created. - key (str): A dotted key path string representing the nested structure. - value (Any): The value to set at the specified key. - - Example: - base = None - key = "scf.pw.parameters" - value = 2 - After running: - update_nested_dict(d, key, value) - The result will be: - base = {"scf": {"pw": {"parameters": 2}}} - - Edge Case: - If the resulting dictionary is empty after the update, it will be set to `None`. - - """ - if base is None: - base = {} - keys = key_path.split('.') - current_key = keys[0] - if len(keys) == 1: - # Base case: Merge dictionaries or set the value directly. - if isinstance(base.get(current_key), dict) and isinstance(value, dict): - base[current_key] = merge_dicts(base[current_key], value) - else: - base[current_key] = value - else: - # Recursive case: Ensure the key exists and is a dictionary, then recurse. - if current_key not in base or not isinstance(base[current_key], dict): - base[current_key] = {} - base[current_key] = update_nested_dict(base[current_key], '.'.join(keys[1:]), value) - - return base - - -def update_nested_dict_with_special_keys(data: dict[str, Any]) -> dict[str, Any]: - """Update the nested dictionary with special keys like "base.pw.parameters".""" - # Remove None - data = {k: v for k, v in data.items() if v is not None} - special_keys = [k for k in data.keys() if '.' in k] - for key in special_keys: - value = data.pop(key) - update_nested_dict(data, key, value) - return data - - -def resolve_node_link_managers(data: Any) -> Any: - """Recursively resolve all NodeLinksManagers either in a dictionary or a NodeLinksManager.""" - if isinstance(data, dict): - return {key: resolve_node_link_managers(value) for key, value in data.items()} - if isinstance(data, NodeLinksManager): - return convert_node_link_manager_to_dict(data) - return data - - -def convert_node_link_manager_to_dict(node_link_manager: NodeLinksManager) -> dict[str, Any]: - """Recursively convert a NodeLinksManager to a dictionary representation.""" - data = {} - for name in node_link_manager._get_keys(): - item = node_link_manager._get_node_by_link_label(name) - if isinstance(item, NodeLinksManager): - data[name] = convert_node_link_manager_to_dict(item) - else: - data[name] = item - return data diff --git a/src/aiida/workgraph/utils/__init__.py b/src/aiida/workgraph/utils/__init__.py new file mode 100644 index 0000000000..09b468d4d1 --- /dev/null +++ b/src/aiida/workgraph/utils/__init__.py @@ -0,0 +1,682 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""WorkGraph runtime and authoring helpers (relocated from aiida-workgraph).""" + +from __future__ import annotations + +import inspect +import logging +from collections.abc import Callable +from copy import deepcopy +from typing import Any, Literal, TypeAlias + +import yaml +from node_graph.socket import TaggedValue +from node_graph.socket_spec import SocketSpec + +from aiida import orm +from aiida.common.exceptions import NotExistent +from aiida.common.links import validate_link_label +from aiida.engine import CalcJob, WorkChain +from aiida.engine.processes import Process +from aiida.engine.runners import Runner +from aiida.orm.utils.serialize import serialize +from aiida.workgraph.config import task_types +from aiida.workgraph.orm.utils import deserialize_safe + +LOGGER = logging.getLogger(__name__) + + +# A task name is set through three mechanisms; each gets the fix hint for *that* mechanism, +# looked up by ``source`` rather than branched on (see ``_validate_task_name``): +# +# * ``'explicit_name'``: an explicit ``name=`` passed to :meth:`WorkGraph.add_task`. Only the +# low-level API can set it, so the hint may safely point at ``WorkGraph.add_task``. +# * ``'derived_name'``: a name derived from the callable when no explicit name is given. This +# is the usual high-level case (a task called inside ``@task.graph``, which never touches +# ``add_task`` itself) and also low-level ``add_task(callable)`` without a ``name=``. The +# only fix common to both is to rename the callable, so the hint must not mention ``name=`` +# (wrong for the high-level user) nor ``call_link_label`` (wrong for the low-level one). +# * ``'call_link_label'``: the ``metadata={'call_link_label': ...}`` override applied in +# :meth:`TaskHandle.__call__`, which renames the task *after* ``add_task`` ran (high-level). +TaskNameSource: TypeAlias = Literal['explicit_name', 'derived_name', 'call_link_label'] + +_TASK_NAME_FIX_HINTS: dict[TaskNameSource, str] = { + 'explicit_name': ( + 'How to fix: pass a valid `name=...` to `WorkGraph.add_task`, or omit it and rename ' + 'the function/callable the task is built from.' + ), + 'derived_name': ( + 'How to fix: rename the function/callable the task is built from so its name is a valid link label.' + ), + 'call_link_label': ( + "How to fix: choose a valid `call_link_label` (e.g. `metadata={'call_link_label': ''}`), or " + 'omit it and rename the function/callable the task is built from.' + ), +} + + +def _validate_task_name(name: str, *, source: TaskNameSource) -> None: + """Validate that a task name can be used as an AiiDA provenance link label. + + A task's name becomes the ``call_link_label`` of the process it launches, so it has to + satisfy AiiDA's link-label rules: a valid Python identifier, containing only letters, + digits and underscores, that does not start or end with an underscore. This is stricter + than ``node_graph``'s own name check, and it is enforced at build time so that an invalid + name fails fast instead of silently producing no output when the task runs. + + :param name: the task name to validate. + :param source: which mechanism set the name; selects the API-appropriate fix hint in the + error message (see :data:`_TASK_NAME_FIX_HINTS`). + :raises ValueError: if ``name`` cannot be used as a link label, with guidance on how to + fix it. + """ + try: + validate_link_label(name) + except ValueError as exc: + msg = ( + f"Invalid task name '{name}': {exc}.\n" + 'A task name is used as the AiiDA provenance link label of the process it ' + 'launches, so it must be a valid Python identifier containing only letters, ' + f'digits and underscores, and may not start or end with an underscore.\n' + f'{_TASK_NAME_FIX_HINTS[source]}' + ) + raise ValueError(msg) from exc + + +def inspect_aiida_component_type(executor: Callable) -> str: + task_type = None + if isinstance(executor, type): + # Lazy plugin imports so that ``import aiida.workgraph`` never pulls a downstream plugin. + # TODO: invert onto a ``_workgraph_task_type`` marker declared by each plugin process. + try: + from aiida_pythonjob import PythonJob + from aiida_pythonjob.calculations.pyfunction import PyFunction + except ImportError: + PythonJob = PyFunction = None + try: + from aiida_shell.calculations.shell import ShellJob + except ImportError: + ShellJob = None + if PythonJob is not None and executor == PythonJob: + task_type = 'PYTHONJOB' + elif PyFunction is not None and executor == PyFunction: + task_type = 'PYFUNCTION' + elif ShellJob is not None and executor == ShellJob: + task_type = 'SHELLJOB' + elif issubclass(executor, CalcJob): + task_type = task_types[CalcJob] + elif issubclass(executor, WorkChain): + task_type = task_types[WorkChain] + elif inspect.isfunction(executor): + if getattr(executor, 'node_class', False): + task_type = task_types[executor.node_class] + return task_type + + +def generate_provenance_graph(pk: int, output: str | None = None, width: str = '100%', height: str = '600px') -> Any: + """Generate the node graph for the given node pk. + If in Jupyter, return the graphviz object. + Otherwise, save the graph to an HTML file. + """ + + import pathlib + + from IPython.display import IFrame + + from aiida import orm + from aiida.tools.visualization import Graph + + from .svg_to_html import svg_to_html + + in_jupyter = False + try: + from IPython import get_ipython + + if get_ipython() is not None: + in_jupyter = True + except NameError: + pass + + graph = Graph() + calc_node = orm.load_node(pk) + graph.recurse_ancestors(calc_node, annotate_links='both') + graph.recurse_descendants(calc_node, annotate_links='both') + g = graph.graphviz + if not in_jupyter: + html_content = svg_to_html(g._repr_image_svg_xml(), width, height) + if output is None: + pathlib.Path('html').mkdir(exist_ok=True) + output = f'html/node_graph_{pk}.html' + with open(output, 'w') as f: + f.write(html_content) + return IFrame(output, width=width, height=height) + return g + + +def get_dict_from_builder(builder: Any) -> dict: + """Transform builder to pure dict.""" + from aiida.engine.processes.builder import ProcessBuilderNamespace + + if isinstance(builder, ProcessBuilderNamespace): + return {k: get_dict_from_builder(v) for k, v in builder.items()} + else: + return builder + + +def clean_pickled_task_executor(tdata: dict[str, Any]) -> None: + """Clean the pickled executor in the task data.""" + from node_graph.executor import RuntimeExecutor + + from aiida.workgraph.executors.builtins import UnavailableExecutor + + # spec + if 'spec' in tdata: + executor = tdata['spec'].get('executor', {}) + if executor.get('mode', '') == 'pickled_callable': + tdata['spec']['executor'] = RuntimeExecutor.from_callable(UnavailableExecutor).to_dict() + if executor.get('mode', '') == 'graph': + wgdata = executor['graph_data'] + for task in wgdata['tasks'].values(): + clean_pickled_task_executor(task) + # error handler + for name, handler in tdata.get('error_handlers', {}).items(): + if handler.get('mode', '') == 'pickled_callable': + tdata['error_handlers'][name] = RuntimeExecutor.from_callable(UnavailableExecutor).to_dict() + + +def save_workgraph_data(node: int | orm.Node, inputs: dict[str, Any]) -> None: + from aiida.workgraph.engine.process import WorkGraphSpec + + inputs = shallow_copy_nested_dict(inputs) + wgdata = inputs.pop(WorkGraphSpec.WORKGRAPH_DATA_KEY, {}) + task_states = {} + task_processes = {} + task_actions = {} + short_wgdata = workgraph_to_short_json(wgdata) + for name, task in wgdata['tasks'].items(): + task_states[name] = task['state'] + task_processes[name] = task['process'] + task_actions[name] = task['action'] + # clean pickled executor before save to database + clean_pickled_task_executor(task) + node.task_states = task_states + node.task_processes = task_processes + node.task_actions = task_actions + node.workgraph_data = wgdata + node.workgraph_data_short = short_wgdata + node.workgraph_error_handlers = wgdata.pop('error_handlers', {}) + graph_inputs = dict(inputs.pop('graph_inputs', {})) + tasks = dict(inputs.pop('tasks', {})) + tasks['graph_inputs'] = graph_inputs + node.task_inputs = serialize(tasks) + + +def restore_workgraph_data_from_raw_inputs(raw_inputs: dict[str, Any]) -> dict[str, Any]: + """Restore the workgraph data from the raw inputs.""" + from aiida.workgraph.engine.process import WorkGraphSpec + + raw_inputs = dict(raw_inputs) + wgdata = dict(raw_inputs.pop(WorkGraphSpec.WORKGRAPH_DATA_KEY, {})) + task_inputs = dict(raw_inputs.pop('tasks', {})) + graph_inputs = dict(raw_inputs.pop('graph_inputs', {})) + task_inputs['graph_inputs'] = graph_inputs + for name, data in task_inputs.items(): + wgdata['tasks'][name]['inputs'] = data + return wgdata + + +def load_workgraph_data(node: int | orm.Node) -> dict[str, Any] | None: + """ + Get the workgraph data from the given process node. + """ + from aiida.orm import load_node + from aiida.workgraph.engine.process import WorkGraphSpec + + if isinstance(node, int): + node = load_node(node) + wgdata = node.base.attributes.get(WorkGraphSpec.WORKGRAPH_DATA_KEY) + try: + task_inputs = deserialize_safe(node.task_inputs or '') + except (yaml.constructor.ConstructorError, yaml.YAMLError): + LOGGER.info('Could not deserialize inputs. The workgraph is still loaded and tasks/outputs remain inspectable.') + task_inputs = {} + + for name, data in task_inputs.items(): + wgdata['tasks'][name]['inputs'] = data + wgdata['error_handlers'] = node.workgraph_error_handlers + return wgdata + + +def get_parent_workgraphs(pk: int) -> list[list[str, int]]: + """Get the list of parent workgraphs. + Use aiida incoming links to find the parent workgraphs. + the parent workgraph is the workgraph that has a link (type CALL_WORK) to the current workgraph. + """ + from aiida import orm + from aiida.common.links import LinkType + + node = orm.load_node(pk) + parent_workgraphs = [[node.process_label, node.pk]] + links = node.base.links.get_incoming(link_type=LinkType.CALL_WORK).all() + if len(links) > 0: + parent_workgraphs.extend(get_parent_workgraphs(links[0].node.pk)) + return parent_workgraphs + + +def get_processes_latest( + pk: int, task_name: str | None = None, item_type: str = 'task' +) -> dict[str, dict[str, int | str]]: + """Get the latest info of all tasks from the process.""" + import aiida + from aiida.orm import WorkGraphNode + + tasks = {} + if pk is None: + return tasks + node = aiida.orm.load_node(pk) + if item_type == 'called_process': + # fetch the process that called by the workgraph + for link in node.base.links.get_outgoing().all(): + if isinstance(link.node, aiida.orm.ProcessNode): + tasks[f'{link.link_label}-{link.node.pk}'] = { + 'pk': link.node.pk, + 'process_type': link.node.process_type, + 'state': link.node.process_state.value, + 'ctime': link.node.ctime, + 'mtime': link.node.mtime, + } + elif item_type == 'task': + if not isinstance(node, WorkGraphNode): + return tasks + task_states = node.task_states + task_processes = node.task_processes + task_names = [task_name] if task_name else task_states.keys() + for name in task_names: + state = task_states[name] + task_process = deserialize_safe(task_processes.get(name, '')) + tasks[name] = { + 'pk': task_process.pk if task_process else None, + 'process_type': task_process.process_type if task_process else '', + 'state': state, + 'ctime': task_process.ctime if task_process else None, + 'mtime': task_process.mtime if task_process else None, + } + + return tasks + + +def get_or_create_code( + computer: str = 'localhost', + code_label: str = 'python3', + code_path: str | None = None, + prepend_text: str = '', +): + """Try to load code, create if not exit.""" + from aiida.orm.nodes.data.code.installed import InstalledCode + + try: + return orm.load_code(f'{code_label}@{computer}') + except NotExistent: + description = f'Code on computer: {computer}' + computer = orm.load_computer(computer) + code_path = code_path or code_label + code = InstalledCode( + computer=computer, + label=code_label, + description=description, + filepath_executable=code_path, + default_calc_job_plugin='workgraph.python', + prepend_text=prepend_text, + ) + + code.store() + return code + + +def create_and_pause_process( + runner: Runner = None, + process_class: Callable | None = None, + inputs: dict | None = None, + state_msg: str = '', +) -> Process: + from aiida.engine.utils import instantiate_process + + process_inited = instantiate_process(runner, process_class, **inputs) + process_inited.pause(state_msg) + process_inited.runner.persister.save_checkpoint(process_inited) + process_inited.close() + runner.controller.continue_process(process_inited.pid, nowait=True, no_reply=True) + return process_inited + + +def get_raw_value(identifier, value: Any) -> Any: + """Get the raw value from a Data node.""" + if identifier in [ + 'workgraph.int', + 'workgraph.float', + 'workgraph.string', + 'workgraph.bool', + 'workgraph.aiida_int', + 'workgraph.aiida_float', + 'workgraph.aiida_string', + 'workgraph.aiida_bool', + ]: + if isinstance(value, TaggedValue): + value = value.__wrapped__ + if value is not None and isinstance(value, orm.Data): + return value.value + else: + return value + elif isinstance(value, orm.Data): + # avoid modifying the original attributes + content = deepcopy(value.backend_entity.attributes) + content['node_type'] = value.node_type + return content + + +def process_properties(task: dict) -> dict: + """Extract raw values.""" + result = {} + for name, prop in task.get('properties', {}).items(): + identifier = prop['identifier'] + value = prop.get('value') + result[name] = { + 'identifier': identifier, + 'value': get_raw_value(identifier, value), + } + for name, input in task.get('input_sockets', {}).get('sockets', {}).items(): + if input.get('property'): + prop = input['property'] + identifier = prop['identifier'] + value = prop.get('value') + result[name] = { + 'identifier': identifier, + 'value': get_raw_value(identifier, value), + } + + return result + + +def workgraph_to_short_json(wgdata: dict[str, str | list | dict]) -> dict[str, str | dict]: + """Export a workgraph to a rete js editor data.""" + + wgdata_short = { + 'name': wgdata['name'], + 'uuid': wgdata.get('uuid', ''), + 'state': wgdata.get('state', ''), + 'nodes': {}, + 'links': deepcopy(wgdata.get('links', [])), + } + for name, task in wgdata['tasks'].items(): + # Add required inputs to tasks + inputs = [] + for input in task.get('input_sockets', {}).get('sockets', {}).values(): + metadata = input.get('metadata', {}) or {} + if metadata.get('required', False): + inputs.append({'name': input['name'], 'identifier': input['identifier']}) + + properties = process_properties(task) + wgdata_short['nodes'][name] = { + 'identifier': task['identifier'], + 'label': task['name'], + 'node_type': task['spec']['task_type'].upper(), + 'inputs': inputs, + 'properties': properties, + 'outputs': [], + 'position': task.get('position', [0, 0]), + 'children': task.get('children', []), + } + + # Add links to tasks + for link in wgdata_short.get('links', []): + wgdata_short['nodes'][link['to_task']]['inputs'].append( + { + 'name': link['to_socket'], + } + ) + wgdata_short['nodes'][link['from_task']]['outputs'].append( + { + 'name': link['from_socket'], + } + ) + + # remove the inputs socket of "graph_inputs" + if 'graph_inputs' in wgdata_short['nodes']: + wgdata_short['nodes']['graph_inputs']['inputs'] = [] + # remove the empty graph-level tasks + for name in ['graph_inputs', 'graph_outputs', 'graph_ctx']: + if name in wgdata_short['nodes']: + node = wgdata_short['nodes'][name] + if len(node['inputs']) == 0 and len(node['outputs']) == 0: + del wgdata_short['nodes'][name] + for link in wgdata_short['links']: + link['from_node'] = link.pop('from_task') + link['to_node'] = link.pop('to_task') + link['from_socket'] = link.pop('from_socket') + link['to_socket'] = link.pop('to_socket') + return wgdata_short + + +def wait_to_link(wgdata: dict[str, Any]) -> None: + """Convert wait attribute to link.""" + for name, task in wgdata['tasks'].items(): + for wait_task in task['wait']: + if wait_task in wgdata['tasks']: + wgdata['links'].append( + { + 'from_task': wait_task, + 'from_socket': '_wait', + 'to_task': name, + 'to_socket': '_wait', + } + ) + + +def shallow_copy_nested_dict(d): + """Recursively copies only the dictionary structure but keeps value references.""" + from plumpy.utils import AttributesFrozendict + + if isinstance(d, (dict, AttributesFrozendict)): + return {key: shallow_copy_nested_dict(value) for key, value in d.items()} + return d + + +def make_json_serializable(data): + """Recursively convert AiiDA objects to JSON-serializable structures.""" + from collections.abc import Mapping, Sequence + + if isinstance(data, orm.Data): + # Return an int if it's an orm.Int, or a more general dict for other data + return { + '__aiida_class__': data.__class__.__name__, + 'uuid': str(data.uuid), + } + elif isinstance(data, Mapping): + return {k: make_json_serializable(v) for k, v in data.items()} + elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)): + return [make_json_serializable(item) for item in data] + else: + return data + + +def resolve_tagged_values(inputs: dict[str, Any]) -> None: + """Recursively resolve all TaggedValue either in a dictionary or a TaggedValue.""" + from node_graph.utils import resolve_tagged_values as _resolve_tagged_values + + _resolve_tagged_values(inputs) + + +def serialize_graph_level_data( + input_socket: dict[str, Any], + port_schema: SocketSpec | dict[str, Any], + serializers: dict[str, str] | None = None, +) -> dict[str, Any]: + """Recursively walk over the sockets and convert raw Python + values to AiiDA Data nodes, if needed. + """ + from aiida.workgraph import serialize_ports + + resolve_tagged_values(input_socket) + return serialize_ports( + python_data=input_socket, + port_schema=port_schema, + serializers=serializers or {}, + ) + + +def get_process_summary(node: orm.ProcessNode | int, data: str = ['outputs']) -> None: + """Get the outputs of a process node.""" + from aiida.cmdline.utils.common import format_nested_links + from aiida.common.links import LinkType + + node = orm.load_node(node) if isinstance(node, int) else node + result = '' + if 'inputs' in data: + nodes_input = node.base.links.get_incoming(link_type=(LinkType.INPUT_CALC, LinkType.INPUT_WORK)) + result += f'\n{format_nested_links(nodes_input.nested(), headers=["Inputs", "PK", "Type"])}' + + if 'outputs' in data: + nodes_output = node.base.links.get_outgoing(link_type=(LinkType.CREATE, LinkType.RETURN)) + result += f'\n{format_nested_links(nodes_output.nested(), headers=["Outputs", "PK", "Type"])}' + return result + + +def call_depth_from_node(node: str | int | orm.Node) -> int: + node = orm.load_node(node) if not isinstance(node, orm.Node) else node + depth = 0 + while getattr(node, 'caller', None) is not None: + depth += 1 + node = node.caller + return depth + + +# --- generic dict / link-manager helpers (moved to core in the serializer step) --- +from aiida.orm.utils.managers import NodeLinksManager as _NodeLinksManager # noqa: E402 + + +def get_nested_dict(d: Any, name: str, **kwargs: Any) -> Any: + """Get the value from a nested dictionary. + + ``d`` is deliberately ``Any``: the traversal descends through both plain dicts and AiiDA + ``_NodeLinksManager`` containers, whose values are heterogeneous. + + If default is provided, return the default value if the key is not found. + Otherwise, raise ValueError. + For example: + d = {"base": {"pw": {"parameters": 2}}} + name = "base.pw.parameters" + """ + keys = name.split('.') + current = d + for key in keys: + if key not in current: + if 'default' in kwargs: + return kwargs.get('default') + if isinstance(current, dict): + avaiable_keys = list(current.keys()) + elif isinstance(current, _NodeLinksManager): + avaiable_keys = list(current._get_keys()) + else: + avaiable_keys = [] + raise ValueError(f'{name} not exist. Available keys: {avaiable_keys}') + current = current[key] + return current + + +def merge_dicts(dict1: Any, dict2: Any) -> Any: + """Recursively merges two dictionaries.""" + for key, value in dict2.items(): + if key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict): + # Recursively merge dictionaries + dict1[key] = merge_dicts(dict1[key], value) + else: + # Overwrite or add the key + dict1[key] = value + return dict1 + + +def update_nested_dict(base: dict[str, Any] | None, key_path: str, value: Any) -> dict[str, Any]: + """ + Update or create a nested dictionary structure based on a dotted key path. + + This function allows updating a nested dictionary or creating one if `d` is `None`. + Given a dictionary and a key path (e.g., "base.pw.parameters"), it will traverse + or create the necessary nested structure to set the provided value at the specified + key location. If intermediate dictionaries do not exist, they will be created. + If the resulting dictionary is empty, it is set to `None`. + + Args: + base (Dict[str, Any] | None): The dictionary to update, which can be `None`. + If `None`, an empty dictionary will be created. + key (str): A dotted key path string representing the nested structure. + value (Any): The value to set at the specified key. + + Example: + base = None + key = "scf.pw.parameters" + value = 2 + After running: + update_nested_dict(d, key, value) + The result will be: + base = {"scf": {"pw": {"parameters": 2}}} + + Edge Case: + If the resulting dictionary is empty after the update, it will be set to `None`. + + """ + if base is None: + base = {} + keys = key_path.split('.') + current_key = keys[0] + if len(keys) == 1: + # Base case: Merge dictionaries or set the value directly. + if isinstance(base.get(current_key), dict) and isinstance(value, dict): + base[current_key] = merge_dicts(base[current_key], value) + else: + base[current_key] = value + else: + # Recursive case: Ensure the key exists and is a dictionary, then recurse. + if current_key not in base or not isinstance(base[current_key], dict): + base[current_key] = {} + base[current_key] = update_nested_dict(base[current_key], '.'.join(keys[1:]), value) + + return base + + +def update_nested_dict_with_special_keys(data: dict[str, Any]) -> dict[str, Any]: + """Update the nested dictionary with special keys like "base.pw.parameters".""" + # Remove None + data = {k: v for k, v in data.items() if v is not None} + special_keys = [k for k in data.keys() if '.' in k] + for key in special_keys: + value = data.pop(key) + update_nested_dict(data, key, value) + return data + + +def resolve_node_link_managers(data: Any) -> Any: + """Recursively resolve all NodeLinksManagers either in a dictionary or a _NodeLinksManager.""" + if isinstance(data, dict): + return {key: resolve_node_link_managers(value) for key, value in data.items()} + if isinstance(data, _NodeLinksManager): + return convert_node_link_manager_to_dict(data) + return data + + +def convert_node_link_manager_to_dict(node_link_manager: _NodeLinksManager) -> dict[str, Any]: + """Recursively convert a _NodeLinksManager to a dictionary representation.""" + data = {} + for name in node_link_manager._get_keys(): + item = node_link_manager._get_node_by_link_label(name) + if isinstance(item, _NodeLinksManager): + data[name] = convert_node_link_manager_to_dict(item) + else: + data[name] = item + return data diff --git a/src/aiida/workgraph/utils/control.py b/src/aiida/workgraph/utils/control.py new file mode 100644 index 0000000000..c8d55db92f --- /dev/null +++ b/src/aiida/workgraph/utils/control.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +import logging + +from aiida import orm +from aiida.engine.processes import control +from aiida.manage import get_manager +from aiida.workgraph.enums import RuntimeInfoKey, TaskAction, TaskActionMessage, TaskState + +LOGGER = logging.getLogger(__name__) + + +def create_task_action( + pk: int, + tasks: list, + action: TaskAction = TaskAction.PAUSE, +): + """Send task action to Process.""" + + controller = get_manager().get_process_controller() + # Send the canonical action value as a plain string; the engine re-validates it + # into a TaskAction on receipt. + message: TaskActionMessage = {'intent': 'custom', 'catalog': 'task', 'action': str(action), 'tasks': tasks} + controller._communicator.rpc_send(pk, message) + + +def get_task_runtime_info(node, name: str, key: RuntimeInfoKey) -> str: + """Get task state info from attributes.""" + from aiida.workgraph.orm.utils import deserialize_safe + + match key: + case 'process': + return deserialize_safe(node.task_processes.get(name, '')) + case 'state': + return node.task_states.get(name, '') + case 'action': + return node.task_actions.get(name, '') + case _: + raise ValueError(f'Invalid key: {key}') + + +def pause_tasks(pk: int, tasks: list[str], timeout: int = 5): + """Pause task.""" + node = orm.load_node(pk) + if node.is_finished: + message = 'WorkGraph is finished. Cannot pause tasks.' + LOGGER.warning(message) + return False, message + elif node.process_state.value.upper() in [ + 'CREATED', + 'RUNNING', + 'WAITING', + 'PAUSED', + ]: + for name in tasks: + if get_task_runtime_info(node, name, 'state') == TaskState.PLANNED: + create_task_action(pk, tasks, action=TaskAction.PAUSE) + elif get_task_runtime_info(node, name, 'state') == TaskState.RUNNING: + try: + control.pause_processes( + [get_task_runtime_info(node, name, 'process')], + all_entries=None, + timeout=timeout, + ) + except Exception as e: + LOGGER.exception('Pause task %s failed: %s', name, e) + return True, '' + + +def play_tasks(pk: int, tasks: list, timeout: int = 5): + node = orm.load_node(pk) + if node.is_finished: + message = 'WorkGraph is finished. Cannot kill tasks.' + LOGGER.warning(message) + return False, message + elif node.process_state.value.upper() in [ + 'CREATED', + 'RUNNING', + 'WAITING', + 'PAUSED', + ]: + for name in tasks: + state = get_task_runtime_info(node, name, 'state') + if state == TaskState.PLANNED: + create_task_action(pk, tasks, action=TaskAction.PLAY) + break + process = get_task_runtime_info(node, name, 'process') + if process.is_finished: + raise ValueError(f'Task {name} is already finished.') + elif process.process_state.value.upper() in ['CREATED', 'WAITING']: + try: + control.play_processes( + [process], + all_entries=None, + timeout=timeout, + ) + except Exception as e: + LOGGER.exception('Play task %s failed: %s', name, e) + return True, '' + + +def kill_tasks(pk: int, tasks: list, timeout: int = 5): + node = orm.load_node(pk) + if node.is_finished: + message = 'WorkGraph is finished. Cannot kill tasks.' + LOGGER.warning(message) + return False, message + elif node.process_state.value.upper() in [ + 'CREATED', + 'RUNNING', + 'WAITING', + 'PAUSED', + ]: + for name in tasks: + state = get_task_runtime_info(node, name, 'state') + process = get_task_runtime_info(node, name, 'process') + if state == TaskState.PLANNED: + create_task_action(pk, tasks, action=TaskAction.SKIP) + # A live task to kill is either CREATED or RUNNING; WAITING/PAUSED are + # AiiDA process states a *task* state never takes, so they were dead + # entries in this list. + elif state in {TaskState.CREATED, TaskState.RUNNING}: + if process is None: + LOGGER.warning('Task %s is not an AiiDA process.', name) + create_task_action(pk, tasks, action=TaskAction.KILL) + else: + try: + control.kill_processes( + [process], + all_entries=None, + timeout=timeout, + ) + except Exception as e: + LOGGER.exception('Kill task %s failed: %s', name, e) + return True, '' + + +def reset_tasks(pk: int, tasks: list) -> None: + """Reset tasks + Args: + tasks (list): a list of task names. + """ + node = orm.load_node(pk) + if node.is_finished: + message = 'WorkGraph is finished. Cannot kill tasks.' + LOGGER.warning(message) + return False, message + elif node.process_state.value.upper() in [ + 'CREATED', + 'RUNNING', + 'WAITING', + 'PAUSED', + ]: + for name in tasks: + create_task_action(pk, tasks, action=TaskAction.RESET) + + return True, '' diff --git a/src/aiida/workgraph/utils/logging.py b/src/aiida/workgraph/utils/logging.py new file mode 100644 index 0000000000..b5501143db --- /dev/null +++ b/src/aiida/workgraph/utils/logging.py @@ -0,0 +1,14 @@ +import subprocess + +from aiida.manage import get_manager +from aiida.manage.configuration import reset_config + + +def set_aiida_loglevel(level: str): + """Set the AiiDA log level.""" + subprocess.run( + ['verdi', 'config', 'set', 'logging.aiida_loglevel', level], + check=True, + ) + get_manager().unload_profile() + reset_config() diff --git a/src/aiida/workgraph/utils/svg_to_html.py b/src/aiida/workgraph/utils/svg_to_html.py new file mode 100644 index 0000000000..d16d6e28a9 --- /dev/null +++ b/src/aiida/workgraph/utils/svg_to_html.py @@ -0,0 +1,99 @@ +def svg_to_html(svg_xml: str, width: str = '100%', height: str = '100%') -> str: + """ + Converts an SVG XML string into an HTML string with embedded SVG, + scaled to the specified width and height using CSS and includes functionality + for panning and zooming the SVG based on the mouse point. + """ + html_template = f""" + + + + + + Interactive SVG Viewer + + + + + {svg_xml} + + + + """ + return html_template diff --git a/src/aiida/workgraph/workgraph.py b/src/aiida/workgraph/workgraph.py new file mode 100644 index 0000000000..208a645bd8 --- /dev/null +++ b/src/aiida/workgraph/workgraph.py @@ -0,0 +1,662 @@ +from __future__ import annotations + +import logging +import time +from typing import Any + +import node_graph +from node_graph.analysis import GraphAnalysis +from node_graph.config import BUILTIN_TASKS +from node_graph.error_handler import ErrorHandlerSpec +from node_graph.socket import BaseSocket, TaskSocketNamespace + +import aiida +from aiida.workgraph.enums import TaskAction, TaskState +from aiida.workgraph.socket_spec import SocketSpecAPI +from aiida.workgraph.task import Task + +from .registry import RegistryHub, registry_hub + +LOGGER = logging.getLogger(__name__) + + +class WorkGraph(node_graph.Graph): + """Build flexible workflows with AiiDA. + + The class extends from NodeGraph and provides methods to run, + submit tasks, wait for tasks to finish, and update the process status. + It is used to handle various states of a workgraph process and provides + convenient operations to interact with it. + + Attributes: + process (aiida.orm.ProcessNode): The process node that represents the process status and other details. + state (str): The current state of the workgraph process. + pk (int): The primary key of the process node. + """ + + _REGISTRY: RegistryHub | None = registry_hub + _SOCKET_SPEC_API = SocketSpecAPI + + platform: str = 'aiida.workgraph' + + def __init__( + self, + name: str = 'WorkGraph', + inputs: type | list[str] | None = None, + outputs: type | list[str] | None = None, + error_handlers: dict[str, ErrorHandlerSpec] | None = None, + serialization: object | None = None, + serialization_policy: str = 'off', + **kwargs, + ) -> None: + """ + Initialize a WorkGraph instance. + + Args: + name (str, optional): The name of the WorkGraph. Defaults to 'WorkGraph'. + **kwargs: Additional keyword arguments to be passed to the WorkGraph class. + """ + from aiida.workgraph.serialization import AiidaSerializationAdapter + + if serialization is None: + serialization = AiidaSerializationAdapter() + super().__init__( + name, + inputs=inputs, + outputs=outputs, + serialization=serialization, + serialization_policy=serialization_policy, + **kwargs, + ) + self.process = None + self.restart_process = None + self.max_number_jobs = 1000000 + self.max_iteration = 1000000 + self._error_handlers = error_handlers or {} + self.analyzer = GraphAnalysis(self) + + def to_engine_inputs(self, metadata: dict[str, Any] | None = None) -> dict[str, Any]: + wgdata = self.to_dict(should_serialize=True) + metadata = metadata or {} + task_inputs = self.gather_task_inputs(wgdata['tasks']) + graph_inputs = task_inputs.pop('graph_inputs', {}) + inputs = { + 'metadata': metadata, + 'workgraph_data': wgdata, + 'tasks': task_inputs, + 'graph_inputs': graph_inputs, + } + return inputs + + def gather_task_inputs(self, data: dict[str, Any] | None = None) -> dict[str, Any]: + """Gather the inputs of all tasks.""" + inputs = {} + for name, task in data.items(): + inputs[name] = task.pop('inputs', {}) + return inputs + + def check_before_run(self) -> bool: + self.check_required_inputs() + self.check_modified_tasks() + + def check_required_inputs(self) -> None: + """Check if all required inputs are provided.""" + missing_inputs = self.find_missing_inputs(self.inputs) + for task in self.tasks: + if task.name in BUILTIN_TASKS: + continue + missing_inputs.extend(self.find_missing_inputs(task.inputs)) + if missing_inputs: + bullets = '\n'.join(f' • {p}' for p in sorted(missing_inputs)) + raise ValueError( + 'Missing required inputs:\n' + f'{bullets}\n\n' + 'How to fix:\n' + ' 1) Provide these values (at build time or by linking from upstream task outputs).\n' + ' 2) If some are intentionally unused, exclude them from the namespace at the call site, e.g.:\n' + ' Annotated[dict, some_task.inputs, SocketSpecSelect(exclude=["pw.structure", ...])]\n\n' + "Note: exclude paths are relative to the task's input namespace (e.g. 'pw.structure')." + ) + + def find_missing_inputs(self, socket: BaseSocket) -> list[str]: + """Check if all required inputs are provided.""" + missing_inputs = [] + for sub_socket in socket: + if isinstance(sub_socket, TaskSocketNamespace): + missing_inputs.extend(self.find_missing_inputs(sub_socket)) + elif sub_socket._metadata.required and sub_socket.value is None and len(sub_socket._links) == 0: + missing_inputs.append(f'{sub_socket._task.name}.{sub_socket._scoped_name}') + return missing_inputs + + def check_modified_tasks(self) -> None: + """Check if there are modified tasks compared to the existing process. + If there are modified tasks, reset them and their descendants to be re-run. + """ + existing_process = self._load_existing_process() + if existing_process: + diffs = self.analyzer.compare_graphs(existing_process, self) + self.reset_tasks(diffs['modified_tasks']) + + def run( + self, + inputs: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, + ) -> Any: + """ + Run the AiiDA workgraph process and update the process status. The method uses AiiDA's engine to run + the process, when the process is finished, update the status of the tasks + """ + from aiida.workgraph.engine.process import WorkGraphProcess + + # set task inputs + if inputs is not None: + self.set_inputs(inputs) + + # One can not run again if the process is alreay created. otherwise, a new process node will + # be created again. + if self.process is not None: + raise ValueError(f'Process {self.process.pk} has already been created. Please use the submit() method.') + self.check_before_run() + inputs = self.to_engine_inputs(metadata=metadata) + _, node = aiida.engine.run_get_node(WorkGraphProcess, inputs=inputs) + self.process = node + self.update() + return self.outputs._value + + def submit( + self, + inputs: dict[str, Any] | None = None, + wait: bool = False, + timeout: int = 600, + interval: int = 5, + metadata: dict[str, Any] | None = None, + ) -> aiida.orm.ProcessNode: + """Submit the AiiDA workgraph process and optionally wait for it to finish. + Args: + wait (bool): Wait for the process to finish. + timeout (int): The maximum time in seconds to wait for the process to finish. Defaults to 600. + restart (bool): Restart the process, and reset the modified tasks, then only re-run the modified tasks. + new (bool): Submit a new process. + """ + + # set task inputs + if inputs is not None: + self.set_inputs(inputs) + + # save the workgraph to the process node + self.save(metadata=metadata) + if self.process.process_state.value.upper() not in ['CREATED']: + raise ValueError(f'Process {self.process.pk} has already been submitted.') + self.continue_process() + # as long as we submit the process, it is a new submission, we should set restart_process to None + self.restart_process = None + if wait: + self.wait(timeout=timeout, interval=interval) + return self.process + + def save(self, metadata: dict[str, Any] | None = None) -> None: + """Save the udpated workgraph to the process + This is only used for a running workgraph. + Save the AiiDA workgraph process and update the process status. + """ + from aiida.engine.utils import instantiate_process + from aiida.manage import manager + from aiida.workgraph.engine.process import WorkGraphProcess + + self.check_before_run() + inputs = self.to_engine_inputs(metadata) + if self.process is None: + runner = manager.get_manager().get_runner() + # init a process node + process_inited = instantiate_process(runner, WorkGraphProcess, **inputs) + process_inited.runner.persister.save_checkpoint(process_inited) + self.process = process_inited.node + self.process_inited = process_inited + process_inited.close() + LOGGER.info('WorkGraph process created, PK: %s', self.process.pk) + else: + self.save_to_base(inputs) + self.update() + + def save_to_base(self, wgdata: dict[str, Any]) -> None: + """Save new wgdata to attribute. + It will first check the difference, and reset tasks if needed. + """ + from aiida.workgraph.utils import save_workgraph_data + + save_workgraph_data(self.process, wgdata) + + def _load_existing_process(self): + """Load an existing workgraph process if available.""" + if self.process: + return WorkGraph.load(self.process) + if self.restart_process: + return WorkGraph.load(self.restart_process) + return None + + def build_connectivity(self) -> None: + """Analyze the connectivity of workgraph and save it into dict.""" + connectivity = self.analyzer.build_connectivity() + return connectivity + + def to_dict(self, include_sockets: bool = False, should_serialize: bool = False) -> dict[str, Any]: + """Convert the workgraph to a dictionary.""" + from aiida.orm.utils.serialize import serialize + + wgdata = super().to_dict(include_sockets=include_sockets, should_serialize=should_serialize) + wgdata.update( + { + 'restart_process': self.restart_process.pk if self.restart_process else None, + 'max_iteration': self.max_iteration, + 'max_number_jobs': self.max_number_jobs, + } + ) + # save error handlers + wgdata['error_handlers'] = {name: eh.to_dict() for name, eh in self.get_error_handlers().items()} + wgdata['connectivity'] = self.build_connectivity() + wgdata['process'] = serialize(self.process) if self.process else serialize(None) + wgdata['metadata']['pk'] = self.process.pk if self.process else None + + return wgdata + + def wait(self, timeout: int = 600, tasks: dict | None = None, interval: int = 5) -> None: + """ + Periodically checks and waits for the AiiDA workgraph process to finish until a given timeout. + + Args: + timeout (int): The maximum time in seconds to wait for the process to finish. Defaults to 600. + tasks (dict): Optional; specifies task states to wait for in the format {task_name: [acceptable_states]}. + interval (int): The time interval in seconds between checks. Defaults to 5. + + Raises: + TimeoutError: If the process does not finish within the given timeout. + """ + terminating_states = ( + 'KILLED', + 'PAUSED', + 'FINISHED', + 'FAILED', + 'CANCELLED', + 'EXCEPTED', + ) + start = time.time() + self.update() + finished = False + + while not finished: + self.update() + + if tasks is not None: + states = [] + for name, value in tasks.items(): + flag = self.tasks[name].state in value + states.append(flag) + finished = all(states) + else: + finished = self.state in terminating_states + + if finished: + LOGGER.info('Process %s finished with state: %s', self.process.pk, self.state) + return + + time.sleep(interval) + + if time.time() - start > timeout: + raise TimeoutError( + f'Timeout reached after {timeout} seconds while waiting for the WorkGraph: {self.process.pk}. ' + ) + + def update(self) -> None: + """ + Update the current state and primary key of the process node as well as the state, node and primary key + of the tasks that are outgoing from the process node. This includes updating the state of process nodes + linked to the current process, and data nodes linked to the current process. + """ + from aiida.workgraph.utils import get_processes_latest, resolve_node_link_managers + + if self.process is None: + return + + self.state = self.process.process_state.value.upper() + processes_data = get_processes_latest(self.pk) + for name, data in processes_data.items(): + # the mapped tasks are not in the workgraph + if name not in self.tasks: + continue + self.tasks[name].update_state(data) + + if self.widget is not None: + states = {name: data['state'] for name, data in processes_data.items()} + self.widget.states = states + + if self.process.is_finished_ok: + self.outputs._set_socket_value(resolve_node_link_managers(self.process.outputs)) + + @property + def pk(self) -> int | None: + return self.process.pk if self.process else None + + @classmethod + def from_dict(cls, wgdata: dict[str, Any]) -> WorkGraph: + wg = super().from_dict(wgdata) + for key in [ + 'max_iteration', + 'max_number_jobs', + 'connectivity', + ]: + if key in wgdata: + setattr(wg, key, wgdata[key]) + if 'error_handlers' in wgdata: + wg._error_handlers = { + name: ErrorHandlerSpec.from_dict(eh) for name, eh in wgdata.get('error_handlers', {}).items() + } + # for zone tasks, add their children + for task in wg.tasks: + if hasattr(task, 'children'): + task.children.add(wgdata['tasks'][task.name].get('children', [])) + return wg + + @classmethod + def from_yaml(cls, filename: str | None = None, string: str | None = None) -> WorkGraph: + """Build WrokGraph from yaml file.""" + import yaml + + # import json + # from aiida.workgraph.utils import make_json_serializable + from node_graph.utils import yaml_to_dict + + # import importlib.resources + # import jsonschema + + if filename: + with open(filename) as f: + wgdata = yaml.safe_load(f) + elif string: + wgdata = yaml.safe_load(string) + else: + raise Exception('Please specific a filename or yaml string.') + wgdata = yaml_to_dict(wgdata) + # serialized_data = make_json_serializable(wgdata) + # with importlib.resources.open_text( + # "aiida.workgraph.schemas", "aiida.workgraph.schema.json" + # ) as f: + # schema = json.load(f) + # jsonschema.validate(instance=serialized_data, schema=schema) + + nt = cls.from_dict(wgdata) + return nt + + @classmethod + def load(cls, pk: int | str | aiida.orm.ProcessNode) -> WorkGraph | None: + """ + Load WorkGraph from the process node with the given primary key. + + Args: + pk (int, str, orm.ProcessNode): The primary key or uuid of the process node, + or the process node itself. + """ + from aiida.orm import WorkGraphNode + from aiida.workgraph.utils import load_workgraph_data + + if isinstance(pk, (int, str)): + process = aiida.orm.load_node(pk) + elif isinstance(pk, aiida.orm.ProcessNode): + process = pk + else: + raise ValueError(f'Invalid pk type: {type(pk)}, requires int, str or ProcessNode.') + if not isinstance(process, WorkGraphNode): + raise ValueError(f'Process {pk} is not a WorkGraph') + wgdata = load_workgraph_data(process) + wg = cls.from_dict(wgdata) + wg.process = process + wg.update() + return wg + + def show(self) -> None: + """ + Print the current state of the workgraph process. + """ + from tabulate import tabulate + + table = [] + self.update() + for task in self.tasks: + table.append([task.name, task.pk, task.state]) + print('-' * 80) + print(f'WorkGraph: {self.name}, PK: {self.pk}, State: {self.state}') + print('-' * 80) + print('Tasks:') + print(tabulate(table, headers=['Name', 'PK', 'State'])) + print('-' * 80) + + # def pause(self) -> None: + # """Pause the workgraph.""" + # from aiida.engine.processes import control + # try: + # control.pause_processes([self.process]) + # except Exception as e: + # print(f"Pause process failed: {e}") + + def pause_tasks(self, tasks: list[str]) -> None: + """Pause the given tasks.""" + from aiida.workgraph.utils.control import pause_tasks + + if self.process is None: + for name in tasks: + self.tasks[name].action = TaskAction.PAUSE + else: + _, msg = pause_tasks(self.process.pk, tasks) + + return 'Send message to pause tasks.' + + def play_tasks(self, tasks: list[str]) -> None: + """Play the given tasks""" + + from aiida.workgraph.utils.control import play_tasks + + if self.process is None: + for name in tasks: + self.tasks[name].action = '' + else: + _, msg = play_tasks(self.process.pk, tasks) + return 'Send message to play tasks.' + + def kill_tasks(self, tasks: list[str]) -> None: + """Kill the given tasks""" + + from aiida.workgraph.utils.control import kill_tasks + + if self.process is None: + for name in tasks: + self.tasks[name].action = TaskAction.KILL + else: + _, msg = kill_tasks(self.process.pk, tasks) + return 'Send message to kill tasks.' + + def reset_tasks(self, tasks: list[str]) -> None: + from aiida.workgraph.utils.control import reset_tasks + + LOGGER.info('Reset tasks: %s', tasks) + + if self.process is None: + for name in tasks: + self.tasks[name].state = TaskState.PLANNED + self.tasks[name].process = None + child_tasks = self.analyzer.get_all_descendants(self.tasks[name]) + for name in child_tasks: + self.tasks[name].state = TaskState.PLANNED + self.tasks[name].process = None + else: + _, msg = reset_tasks(self.process.pk, tasks) + return 'Send message to reset tasks.' + + def continue_process(self): + """Continue a saved process by sending the task to RabbitMA. + Use with caution, this may launch duplicate processes.""" + from aiida.manage import get_manager + + process_controller = get_manager().get_process_controller() + process_controller.continue_process(self.pk) + + def play(self): + import os + + os.system(f'verdi process play {self.process.pk}') + + def restart(self): + """Create a restart submission.""" + if self.process is None: + raise ValueError('No process found. One can not restart from a non-existing process.') + # save the current process node as restart_process + # so that the WorkGraphSaver can compare the difference, and reset the modified tasks + self.restart_process = self.process + self.process = None + self.state = 'PLANNED' + + def reset(self) -> None: + """Reset the workgraph to create a new submission.""" + + self.process = None + for task in self.tasks: + task.reset() + self.state = 'PLANNED' + + def extend(self, wg: WorkGraph, prefix: str = '') -> None: + """Append a workgraph to the current workgraph. + prefix is used to add a prefix to the task names. + """ + for task in wg.tasks: + # skip the built-in tasks + # need to fix this in the future + if task.name in BUILTIN_TASKS: + continue + task.name = prefix + task.name + task.graph = self + self.tasks._append(task) + self.update_ctx(wg.ctx._value) + # links + for link in wg.links: + # skip the links that are from or to built-in tasks + if link.from_task.name in BUILTIN_TASKS or link.to_task.name in BUILTIN_TASKS: + link.unmount() + continue + self.links._append(link) + + def get_error_handlers(self) -> dict[str, ErrorHandlerSpec]: + """Get the error handlers.""" + return self._error_handlers + + def add_task( + self, + identifier: str | callable, + name: str | None = None, + include_builtins: bool = False, + **kwargs, + ) -> Task: + """Add a task to the workgraph.""" + from node_graph.task_spec import TaskSpec + + from aiida.engine import ProcessBuilder + from aiida.workgraph.decorator import build_task_from_callable + from aiida.workgraph.task import Task, TaskHandle + from aiida.workgraph.tasks.shelljob_task import ( + _build_shelljob_TaskSpec, + shelljob, + ) + from aiida.workgraph.tasks.subgraph_task import _build_subgraph_task_TaskSpec + from aiida.workgraph.utils import get_dict_from_builder + + if name in BUILTIN_TASKS and not include_builtins: + raise ValueError(f'Task name {name} can not be used, it is reserved.') + + if isinstance(identifier, str): + identifier = self._REGISTRY.task_pool[identifier.lower()].load() + if isinstance(identifier, WorkGraph): + identifier = _build_subgraph_task_TaskSpec(identifier, name=name) + elif isinstance(identifier, ProcessBuilder): + kwargs = {**kwargs, **get_dict_from_builder(identifier)} + identifier = build_task_from_callable(identifier.process_class) + # todo + elif identifier is shelljob: + spec = _build_shelljob_TaskSpec( + outputs=kwargs.get('outputs'), + parser_outputs=kwargs.pop('parser_outputs', None), + ) + identifier = TaskHandle(spec) + # build the task on the fly if the identifier is a callable + elif callable(identifier) and not isinstance(identifier, (TaskSpec, TaskHandle, Task)): + identifier = build_task_from_callable(identifier) + node = self.tasks._new(identifier, name, **kwargs) + # A task name becomes the AiiDA `call_link_label` of the process it launches, so it + # must be a valid link label. Validate the resolved name (which may have been derived + # from the function name) here at build time; otherwise an invalid name only fails + # inside the engine at run time, where the error is swallowed (see issue #784). The + # fix hint differs depending on whether the user passed `name=` or it was derived. + from aiida.workgraph.utils import _validate_task_name + + try: + _validate_task_name(node.name, source='explicit_name' if name is not None else 'derived_name') + except ValueError: + # Roll back the partially-added task so the workgraph is not left in an + # inconsistent state, then re-raise the actionable error. + if node.name in self.tasks: + del self.tasks[node.name] + raise + self._version += 1 + return node + + def to_widget_value(self) -> dict[str, Any]: + """Convert the workgraph to a dictionary that can be used by the widget.""" + from aiida.workgraph.utils import wait_to_link, workgraph_to_short_json + + wgdata = self.to_dict(include_sockets=True) + wait_to_link(wgdata) + wgdata = workgraph_to_short_json(wgdata) + + return wgdata + + def generate_provenance_graph(self): + """Generate the provenance graph of the workgraph process.""" + from aiida.workgraph.utils import generate_provenance_graph + + if self.process is None: + raise ValueError('No process found. Please run or submit the workgraph first.') + return generate_provenance_graph(self.process.pk) + + def __repr__(self) -> str: + return f'WorkGraph(name="{self.name}", uuid="{self.uuid}")' + + def __str__(self) -> str: + return f'WorkGraph(name="{self.name}", uuid="{self.uuid}")' + + def __enter__(self): + """Called when entering the `with NodeGraph() as ng:` block.""" + from aiida.workgraph.manager import get_current_graph, set_current_graph + + self._previous_graph = get_current_graph() + set_current_graph(self) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Called upon leaving the `with NodeGraph() as ng:` block.""" + from aiida.workgraph.manager import set_current_graph + + set_current_graph(self._previous_graph) + self._previous_graph = None + + def __call__(self, inputs: dict[str, Any] | None = None) -> Any: + """Call the graph with inputs and return as a task. + + Used in context managers as a simple assignment. + + >>> wg1 = WorkGraph() + >>> with WorkGraph() as wg2: + >>> task_outputs = wg1({'input1': 42, 'input2': 'hello'}) + """ + from aiida.workgraph.manager import get_current_graph + + graph = get_current_graph() + task = graph.add_task(self) + inputs = inputs or {} + task.set_inputs(inputs) + return task.outputs From 45eaa7b8af90bbecb2e77bb2a3101b54ec0b2df2 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 6 Aug 2026 13:04:44 +0200 Subject: [PATCH 15/19] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20`WorkGraphProcess`:?= =?UTF-8?q?=20subclass=20`WorkflowProcess`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Put the shared base to use: `WorkGraphProcess` now inherits `WorkflowProcess` directly, as a sibling of `WorkChain` rather than a subclass of it. Each supplies its own stepper (outline vs `DagStepper`) and they share only the awaitable / context / checkpointing machinery on `WorkflowProcess`, so a change to `WorkChain`'s outline model can no longer ripple into WorkGraph (the review concern in aiidateam#7479 / aiidateam#7513). `WorkGraphSpec` likewise drops `WorkChainSpec` for `ProcessSpec`, WorkGraph declares no outline. No behaviour change: the full WorkGraph suite passes at 181 passed / 11 failed, the pre-move baseline. --- src/aiida/workgraph/engine/process.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/aiida/workgraph/engine/process.py b/src/aiida/workgraph/engine/process.py index dda68621ec..f81be5664f 100644 --- a/src/aiida/workgraph/engine/process.py +++ b/src/aiida/workgraph/engine/process.py @@ -10,8 +10,9 @@ from plumpy.workchains import Stepper from aiida.common.lang import override +from aiida.engine.processes.process_spec import ProcessSpec from aiida.engine.processes.workchains.awaitable import Awaitable -from aiida.engine.processes.workchains.workchain import WorkChain, WorkChainSpec +from aiida.engine.processes.workchains.workflow_process import WorkflowProcess from aiida.orm import WorkGraphNode from aiida.workgraph.engine.error_handler_manager import ErrorHandlerManager from aiida.workgraph.engine.stepper import DagStepper @@ -25,11 +26,11 @@ __all__ = ('WorkGraphProcess', 'WorkGraphSpec') -class WorkGraphSpec(WorkChainSpec): +class WorkGraphSpec(ProcessSpec): WORKGRAPH_DATA_KEY = 'workgraph_data' -class WorkGraphProcess(WorkChain): +class WorkGraphProcess(WorkflowProcess): """Execute a work graph, scheduling its tasks by their data dependencies. A work chain declares its execution order up front as an outline; a work graph derives it from the links @@ -144,7 +145,7 @@ def _recreate_stepper(self, saved_state: t.Any) -> Stepper: def load_instance_state(self, saved_state: t.MutableMapping[str, t.Any], load_context: t.Any) -> None: from aiida.orm.utils.log import create_logger_adapter - # `WorkChain.load_instance_state` re-registers the awaitable callbacks before returning, so the runtime + # `WorkflowProcess.load_instance_state` re-registers the awaitable callbacks before returning, so the runtime # state it consults has to be in place first. self._init_runtime_state() @@ -160,9 +161,9 @@ def load_instance_state(self, saved_state: t.MutableMapping[str, t.Any], load_co self._init_managers() def _action_awaitables(self) -> None: - """Register the awaitable callbacks (via `WorkChain`), then surface the waiting status in the report log. + """Register the awaitable callbacks (via `WorkflowProcess`), then surface the waiting status in the report log. - `WorkChain` records "Waiting for child processes: ..." only as the process status; echoing it to the + `WorkflowProcess` records "Waiting for child processes: ..." only as the process status; echoing it to the report makes it visible in `verdi process report` when a graph pauses for its children. """ super()._action_awaitables() @@ -173,7 +174,7 @@ def _on_awaitable_resolved(self, awaitable: Awaitable) -> None: """Record a finished child's outcome on its task before the process decides whether to resume. This is the only work-graph-specific step in the awaitable lifecycle. The rest, including resuming as soon - as any child finishes rather than only once all do, comes from `WorkChain`, because :class:`DagStepper` + as any child finishes rather than only once all do, comes from `WorkflowProcess`, because :class:`DagStepper` declares ``awaitable_barrier = False``. :param awaitable: the awaitable whose target process has terminated From 0c7a3cc75aa8c43a81cb10e4620ee97a725abe57 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 6 Aug 2026 15:56:42 +0200 Subject: [PATCH 16/19] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20engine:=20name=20the?= =?UTF-8?q?=20shared=20base=20`Workflow`,=20move=20it=20up?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Relocate the shared stepper-driven base out of `workchains/` (where it was just introduced) to a neutral `aiida/engine/processes/workflow.py`, and rename it `WorkflowProcess` -> `Workflow`. It is the common base of `WorkChain` and WorkGraph, not a workchain-specific thing, so it does not belong under `workchains/`, and `Workflow` mirrors the node side (`WorkChain`/`WorkChainNode`, `Workflow`/`WorkflowNode`). Unreleased, so no compat break. No import cycle (awaitable loads before workchain in the package `__init__`). `test_work_chain.py` green, mypy + ruff clean. --- src/aiida/engine/__init__.py | 2 +- src/aiida/engine/processes/__init__.py | 3 ++- .../engine/processes/workchains/__init__.py | 2 -- .../engine/processes/workchains/workchain.py | 6 +++--- .../workflow_process.py => workflow.py} | 20 +++++++++---------- src/aiida/workgraph/engine/process.py | 12 +++++------ 6 files changed, 22 insertions(+), 23 deletions(-) rename src/aiida/engine/processes/{workchains/workflow_process.py => workflow.py} (97%) diff --git a/src/aiida/engine/__init__.py b/src/aiida/engine/__init__.py index 438c945cf0..6d76e4a381 100644 --- a/src/aiida/engine/__init__.py +++ b/src/aiida/engine/__init__.py @@ -55,7 +55,7 @@ 'WithNonDb', 'WithSerialize', 'WorkChain', - 'WorkflowProcess', + 'Workflow', 'append_', 'assign_', 'await_processes', diff --git a/src/aiida/engine/processes/__init__.py b/src/aiida/engine/processes/__init__.py index 8c28596c4d..4f25afa397 100644 --- a/src/aiida/engine/processes/__init__.py +++ b/src/aiida/engine/processes/__init__.py @@ -21,6 +21,7 @@ from .process import * from .process_spec import * from .workchains import * +from .workflow import * __all__ = ( 'PORT_NAMESPACE_SEPARATOR', @@ -51,7 +52,7 @@ 'WithNonDb', 'WithSerialize', 'WorkChain', - 'WorkflowProcess', + 'Workflow', 'append_', 'assign_', 'calcfunction', diff --git a/src/aiida/engine/processes/workchains/__init__.py b/src/aiida/engine/processes/workchains/__init__.py index 0424c6ac35..f1def097d9 100644 --- a/src/aiida/engine/processes/workchains/__init__.py +++ b/src/aiida/engine/processes/workchains/__init__.py @@ -17,7 +17,6 @@ from .restart import * from .utils import * from .workchain import * -from .workflow_process import * __all__ = ( 'Awaitable', @@ -27,7 +26,6 @@ 'ProcessHandlerReport', 'ToContext', 'WorkChain', - 'WorkflowProcess', 'append_', 'assign_', 'construct_awaitable', diff --git a/src/aiida/engine/processes/workchains/workchain.py b/src/aiida/engine/processes/workchains/workchain.py index 1e2a77abd3..405532b99d 100644 --- a/src/aiida/engine/processes/workchains/workchain.py +++ b/src/aiida/engine/processes/workchains/workchain.py @@ -20,7 +20,7 @@ from aiida.orm import WorkChainNode from ..process_spec import ProcessSpec -from .workflow_process import WorkflowProcess +from ..workflow import Workflow if t.TYPE_CHECKING: from aiida.engine.runners import Runner @@ -32,10 +32,10 @@ class WorkChainSpec(ProcessSpec, PlumpyWorkChainSpec): pass -class WorkChain(WorkflowProcess): +class WorkChain(Workflow): """The `WorkChain` class is the principle component to implement workflows in AiiDA. - It is a :class:`~aiida.engine.processes.workchains.workflow_process.WorkflowProcess` whose stepper walks the + It is a :class:`~aiida.engine.processes.workflow.Workflow` whose stepper walks the static outline declared on its spec. The shared workflow machinery (awaitables, context, the step lifecycle and its checkpointing) lives on the base class; only the outline stepping is defined here. """ diff --git a/src/aiida/engine/processes/workchains/workflow_process.py b/src/aiida/engine/processes/workflow.py similarity index 97% rename from src/aiida/engine/processes/workchains/workflow_process.py rename to src/aiida/engine/processes/workflow.py index 7d9e040048..64d7a34a6b 100644 --- a/src/aiida/engine/processes/workchains/workflow_process.py +++ b/src/aiida/engine/processes/workflow.py @@ -6,7 +6,7 @@ # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### -"""The :class:`WorkflowProcess`, the shared base for stepper-driven workflow processes.""" +"""The :class:`Workflow`, the shared base for stepper-driven workflow processes.""" from __future__ import annotations @@ -27,14 +27,14 @@ from aiida.orm import Node, ProcessNode from aiida.orm.utils import load_node -from ..exit_code import ExitCode -from ..process import Process, ProcessState -from .awaitable import Awaitable, AwaitableAction, AwaitableTarget, construct_awaitable +from .exit_code import ExitCode +from .process import Process, ProcessState +from .workchains.awaitable import Awaitable, AwaitableAction, AwaitableTarget, construct_awaitable if t.TYPE_CHECKING: from aiida.engine.runners import Runner -__all__ = ('WorkflowProcess',) +__all__ = ('Workflow',) MethodType = t.TypeVar('MethodType') @@ -61,7 +61,7 @@ def __new__(mcs, name, bases, namespace, **kwargs): The whole ancestry of each base is scanned (``base.__mro__``), not just the direct bases, so a ``final`` method stays protected even when it is inherited through an intermediate class rather than defined on the - immediate parent (for example a ``final`` method on ``WorkflowProcess`` reached via ``WorkChain``). + immediate parent (for example a ``final`` method on ``Workflow`` reached via ``WorkChain``). :raises RuntimeError: If the new class defines (i.e. overrides) a method that was decorated with ``final``. """ @@ -101,7 +101,7 @@ def final(mcs, method: MethodType) -> MethodType: # noqa: N804 @auto_persist('_awaitables') -class WorkflowProcess(Process, metaclass=Protect): +class Workflow(Process, metaclass=Protect): """A :class:`~aiida.engine.processes.process.Process` whose execution is delegated to a pluggable stepper. This is the shared base for AiiDA's stepper-driven workflow processes: the :class:`~aiida.engine.WorkChain` @@ -132,8 +132,8 @@ def __init__( :param runner: process runner :param enable_persistence: whether to persist this process """ - if self.__class__ == WorkflowProcess: - raise exceptions.InvalidOperation('cannot construct or launch a base `WorkflowProcess` class.') + if self.__class__ == Workflow: + raise exceptions.InvalidOperation('cannot construct or launch a base `Workflow` class.') super().__init__(inputs, logger, runner, enable_persistence=enable_persistence) @@ -351,7 +351,7 @@ def _do_step(self) -> t.Any: will enter in the Wait state, otherwise it will go to Continue. When the stepper returns that it is done, the stepper result will be converted to None and returned, unless it is an integer or instance of ExitCode. """ - from .context import ToContext + from .workchains.context import ToContext # Under the barrier model the awaitables belong to a single step and are cleared before the next one, which # is what forces every step to wait for all the children it launched. A streaming stepper keeps them, so diff --git a/src/aiida/workgraph/engine/process.py b/src/aiida/workgraph/engine/process.py index f81be5664f..0b35b7fd50 100644 --- a/src/aiida/workgraph/engine/process.py +++ b/src/aiida/workgraph/engine/process.py @@ -12,7 +12,7 @@ from aiida.common.lang import override from aiida.engine.processes.process_spec import ProcessSpec from aiida.engine.processes.workchains.awaitable import Awaitable -from aiida.engine.processes.workchains.workflow_process import WorkflowProcess +from aiida.engine.processes.workflow import Workflow from aiida.orm import WorkGraphNode from aiida.workgraph.engine.error_handler_manager import ErrorHandlerManager from aiida.workgraph.engine.stepper import DagStepper @@ -30,7 +30,7 @@ class WorkGraphSpec(ProcessSpec): WORKGRAPH_DATA_KEY = 'workgraph_data' -class WorkGraphProcess(WorkflowProcess): +class WorkGraphProcess(Workflow): """Execute a work graph, scheduling its tasks by their data dependencies. A work chain declares its execution order up front as an outline; a work graph derives it from the links @@ -145,7 +145,7 @@ def _recreate_stepper(self, saved_state: t.Any) -> Stepper: def load_instance_state(self, saved_state: t.MutableMapping[str, t.Any], load_context: t.Any) -> None: from aiida.orm.utils.log import create_logger_adapter - # `WorkflowProcess.load_instance_state` re-registers the awaitable callbacks before returning, so the runtime + # `Workflow.load_instance_state` re-registers the awaitable callbacks before returning, so the runtime # state it consults has to be in place first. self._init_runtime_state() @@ -161,9 +161,9 @@ def load_instance_state(self, saved_state: t.MutableMapping[str, t.Any], load_co self._init_managers() def _action_awaitables(self) -> None: - """Register the awaitable callbacks (via `WorkflowProcess`), then surface the waiting status in the report log. + """Register the awaitable callbacks (via `Workflow`), then surface the waiting status in the report log. - `WorkflowProcess` records "Waiting for child processes: ..." only as the process status; echoing it to the + `Workflow` records "Waiting for child processes: ..." only as the process status; echoing it to the report makes it visible in `verdi process report` when a graph pauses for its children. """ super()._action_awaitables() @@ -174,7 +174,7 @@ def _on_awaitable_resolved(self, awaitable: Awaitable) -> None: """Record a finished child's outcome on its task before the process decides whether to resume. This is the only work-graph-specific step in the awaitable lifecycle. The rest, including resuming as soon - as any child finishes rather than only once all do, comes from `WorkflowProcess`, because :class:`DagStepper` + as any child finishes rather than only once all do, comes from `Workflow`, because :class:`DagStepper` declares ``awaitable_barrier = False``. :param awaitable: the awaitable whose target process has terminated From 78ac84fbb5c1decdf99a8250e4bd97933a9bcef4 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 6 Aug 2026 16:00:41 +0200 Subject: [PATCH 17/19] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20`workgraph`:=20detec?= =?UTF-8?q?t=20task=20types=20via=20a=20marker,=20not=20imports?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `inspect_aiida_component_type` mapped an executor class to its task-type string by importing `PythonJob`/`PyFunction`/`ShellJob` and comparing. Invert it (GRASP information expert): each plugin process declares `_workgraph_task_type`, and the host reads the marker with `getattr`. Core now imports no downstream plugin to recognise them; the lazy-import interim is gone. Suite parity (181/11). --- src/aiida/workgraph/utils/__init__.py | 22 +++++----------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/src/aiida/workgraph/utils/__init__.py b/src/aiida/workgraph/utils/__init__.py index 09b468d4d1..6752dbb00a 100644 --- a/src/aiida/workgraph/utils/__init__.py +++ b/src/aiida/workgraph/utils/__init__.py @@ -93,23 +93,11 @@ def _validate_task_name(name: str, *, source: TaskNameSource) -> None: def inspect_aiida_component_type(executor: Callable) -> str: task_type = None if isinstance(executor, type): - # Lazy plugin imports so that ``import aiida.workgraph`` never pulls a downstream plugin. - # TODO: invert onto a ``_workgraph_task_type`` marker declared by each plugin process. - try: - from aiida_pythonjob import PythonJob - from aiida_pythonjob.calculations.pyfunction import PyFunction - except ImportError: - PythonJob = PyFunction = None - try: - from aiida_shell.calculations.shell import ShellJob - except ImportError: - ShellJob = None - if PythonJob is not None and executor == PythonJob: - task_type = 'PYTHONJOB' - elif PyFunction is not None and executor == PyFunction: - task_type = 'PYFUNCTION' - elif ShellJob is not None and executor == ShellJob: - task_type = 'SHELLJOB' + # A plugin process declares its WorkGraph task type via the ``_workgraph_task_type`` marker, so the + # host recognises it without importing the plugin class (GRASP: the process is the information expert). + declared = getattr(executor, '_workgraph_task_type', None) + if declared is not None: + task_type = declared elif issubclass(executor, CalcJob): task_type = task_types[CalcJob] elif issubclass(executor, WorkChain): From c1418c4af461d2310de824f2ce99961870ba75f0 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 6 Aug 2026 16:00:43 +0200 Subject: [PATCH 18/19] =?UTF-8?q?=E2=9C=A8=20`workgraph`:=20register=20the?= =?UTF-8?q?=20entry=20points=20in=20core?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move the WorkGraph entry points (task / property / socket / type-mapping groups, the `aiida.workflows` process, the `verdi workgraph` command) into core's `pyproject.toml`, pointing at `aiida.workgraph.*`. With these in core, aiida-workgraph no longer needs to register anything, so it can be archived rather than kept as a shim. Group names are kept as `aiida_workgraph.*` for now (the registry reads them); renaming the groups is a separate cleanup. --- pyproject.toml | 64 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index e06e42bbe6..ea805cbbe4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,6 +87,9 @@ requires-python = '>=3.10' [project.entry-points.'aiida.calculations.monitors'] 'core.always_kill' = 'aiida.calculations.monitors.base:always_kill' +[project.entry-points."aiida.cmdline"] +"workgraph" = "aiida.workgraph.cli.cmd_workgraph:workgraph" + [project.entry-points.'aiida.cmdline.computer.configure'] 'core.local' = 'aiida.transports.plugins.local:CONFIGURE_LOCAL_CMD' 'core.ssh' = 'aiida.transports.plugins.ssh:CONFIGURE_SSH_CMD' @@ -216,6 +219,67 @@ requires-python = '>=3.10' [project.entry-points.'aiida.workflows'] 'core.arithmetic.add_multiply' = 'aiida.workflows.arithmetic.add_multiply:add_multiply' 'core.arithmetic.multiply_add' = 'aiida.workflows.arithmetic.multiply_add:MultiplyAddWorkChain' +"workgraph.process" = "aiida.workgraph.engine.process:WorkGraphProcess" + +[project.entry-points."aiida_workgraph.property"] +"workgraph.aiida_float_vector" = "aiida.workgraph.properties.builtins:PropertyAiiDAFloatVector" +"workgraph.aiida_int_vector" = "aiida.workgraph.properties.builtins:PropertyAiiDAIntVector" +"workgraph.aiida_structuredata" = "aiida.workgraph.properties.builtins:PropertyStructureData" +"workgraph.any" = "aiida.workgraph.properties.builtins:PropertyAny" +"workgraph.bool" = "aiida.workgraph.properties.builtins:PropertyBool" +"workgraph.dict" = "aiida.workgraph.properties.builtins:PropertyDict" +"workgraph.float" = "aiida.workgraph.properties.builtins:PropertyFloat" +"workgraph.int" = "aiida.workgraph.properties.builtins:PropertyInt" +"workgraph.list" = "aiida.workgraph.properties.builtins:PropertyList" +"workgraph.string" = "aiida.workgraph.properties.builtins:PropertyString" + +[project.entry-points."aiida_workgraph.socket"] +"workgraph.aiida_float_vector" = "aiida.workgraph.sockets.builtins:SocketAiiDAFloatVector" +"workgraph.aiida_int_vector" = "aiida.workgraph.sockets.builtins:SocketAiiDAIntVector" +"workgraph.aiida_structuredata" = "aiida.workgraph.sockets.builtins:SocketStructureData" +"workgraph.annotated" = "aiida.workgraph.sockets.builtins:SocketAnnotated" +"workgraph.any" = "aiida.workgraph.sockets.builtins:SocketAny" +"workgraph.bool" = "aiida.workgraph.sockets.builtins:SocketBool" +"workgraph.dict" = "aiida.workgraph.sockets.builtins:SocketDict" +"workgraph.float" = "aiida.workgraph.sockets.builtins:SocketFloat" +"workgraph.int" = "aiida.workgraph.sockets.builtins:SocketInt" +"workgraph.list" = "aiida.workgraph.sockets.builtins:SocketList" +"workgraph.namespace" = "aiida.workgraph.socket:TaskSocketNamespace" +"workgraph.string" = "aiida.workgraph.sockets.builtins:SocketString" + +[project.entry-points."aiida_workgraph.task"] +"workgraph.aiida_dict" = "aiida.workgraph.tasks.builtins:aiida_dict" +"workgraph.aiida_float" = "aiida.workgraph.tasks.builtins:aiida_float" +"workgraph.aiida_int" = "aiida.workgraph.tasks.builtins:aiida_int" +"workgraph.aiida_list" = "aiida.workgraph.tasks.builtins:aiida_list" +"workgraph.aiida_process" = "aiida.workgraph.tasks.aiida:AiiDAProcessTask" +"workgraph.aiida_string" = "aiida.workgraph.tasks.builtins:aiida_string" +"workgraph.any" = "aiida.workgraph.task:Task" +"workgraph.gather_item" = "aiida.workgraph.tasks.builtins:GatherItem" +"workgraph.get_context" = "aiida.workgraph.tasks.builtins:GetContext" +"workgraph.graph_level" = "aiida.workgraph.tasks.builtins:GraphLevelTask" +"workgraph.if_zone" = "aiida.workgraph.tasks.builtins:If" +"workgraph.load_code" = "aiida.workgraph.tasks.builtins:AiiDACode" +"workgraph.load_node" = "aiida.workgraph.tasks.builtins:AiiDANode" +"workgraph.map_item" = "aiida.workgraph.tasks.builtins:MapItem" +"workgraph.map_zone" = "aiida.workgraph.tasks.builtins:Map" +"workgraph.monitor_file" = "aiida.workgraph.tasks.monitors:monitor_file" +"workgraph.monitor_task" = "aiida.workgraph.tasks.monitors:monitor_task" +"workgraph.monitor_time" = "aiida.workgraph.tasks.monitors:monitor_time" +"workgraph.select" = "aiida.workgraph.tasks.builtins:Select" +"workgraph.set_context" = "aiida.workgraph.tasks.builtins:SetContext" +"workgraph.subgraph_task" = "aiida.workgraph.tasks.subgraph_task:SubGraphTask" +"workgraph.test_add" = "aiida.workgraph.executors.test:add" +"workgraph.test_arithmetic_multiply_add" = "aiida.workgraph.tasks.test:TestArithmeticMultiplyAdd" +"workgraph.test_sum_diff" = "aiida.workgraph.executors.test:sum_diff" +"workgraph.while_zone" = "aiida.workgraph.tasks.builtins:While" +"workgraph.zone" = "aiida.workgraph.tasks.builtins:Zone" + +[project.entry-points."aiida_workgraph.type_mapping"] +"workgraph.builtins_mapping" = "aiida.workgraph.orm.mapping:builtins_type_mapping" + +[project.entry-points."aiida_workgraph.type_promotion"] +"workgraph.builtins_mapping" = "aiida.workgraph.orm.mapping:TYPE_PROMOTIONS" [project.optional-dependencies] atomic_tools = [ From 174a9fccc0e6000f8a278b6e7ccda8f8b5505fe4 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Fri, 7 Aug 2026 12:38:58 +0200 Subject: [PATCH 19/19] =?UTF-8?q?=F0=9F=93=9A=20docs:=20add=20the=20WorkGr?= =?UTF-8?q?aph-into-core=20AEP?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bring the AEP for moving WorkGraph's engine + framework into aiida-core into the docs (drafted as issue aiidateam#7479), now that AEPs live in `docs/`. Adds a `docs/source/internals/aep/` section (index + this proposal) wired into the internals toctree. Scoped to the WorkGraph move (the `Workflow` shared base, the subsystem relocation, the serializer reconcile, entry points); a Scope section frames the overall migration as three tracks (WorkGraph engine, then shell + pythonjob, then node-graph) and marks tracks 2 and 3 as separate follow-up AEPs. Snippet-driven: MWEs for the outline barrier vs WorkGraph streaming, the `Workflow` base, `general_serializer` dispatch, the `_workgraph_task_type` marker, the `TaskState` enum, and the archived shim; a component-to-home table; one Mermaid class diagram, via `sphinxcontrib-mermaid` (conf.py + docs extra). The implementation status lists the eleven self-contained commits and why the work splits that way. AEP number left to assign. --- docs/source/conf.py | 1 + docs/source/internals/aep/index.rst | 13 + .../internals/aep/workgraph_into_core.md | 234 ++++++++++++++++++ docs/source/internals/index.rst | 1 + pyproject.toml | 1 + uv.lock | 15 ++ 6 files changed, 265 insertions(+) create mode 100644 docs/source/internals/aep/index.rst create mode 100644 docs/source/internals/aep/workgraph_into_core.md diff --git a/docs/source/conf.py b/docs/source/conf.py index 3ea486ff76..bd8414d7df 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -75,6 +75,7 @@ 'IPython.sphinxext.ipython_directive', 'aiida.sphinxext', 'sphinx_design', + 'sphinxcontrib.mermaid', 'sphinx_copybutton', 'sphinxext.rediraffe', 'notfound.extension', diff --git a/docs/source/internals/aep/index.rst b/docs/source/internals/aep/index.rst new file mode 100644 index 0000000000..f8f0da7b6f --- /dev/null +++ b/docs/source/internals/aep/index.rst @@ -0,0 +1,13 @@ +.. _internals:aep: + +=========================== +AiiDA Enhancement Proposals +=========================== + +Design proposals for substantial changes to AiiDA. Historically tracked in a separate repository and, more recently, +as GitHub issues; going forward they live here alongside the code they describe. + +.. toctree:: + :maxdepth: 1 + + workgraph_into_core diff --git a/docs/source/internals/aep/workgraph_into_core.md b/docs/source/internals/aep/workgraph_into_core.md new file mode 100644 index 0000000000..86e1ef5297 --- /dev/null +++ b/docs/source/internals/aep/workgraph_into_core.md @@ -0,0 +1,234 @@ +# Move the WorkGraph engine and framework into aiida-core + +| | | +|---|---| +| **AEP number** | to be assigned | +| **Authors** | Julian Geiger ([@GeigerJ2](https://github.com/GeigerJ2)) | +| **Status** | draft | +| **Type** | Standard | +| **Created** | 2026-08-06 | +| **Targets** | AiiDA v3 ([#7406](https://github.com/aiidateam/aiida-core/issues/7406)) | +| **Discussion** | [#7479](https://github.com/aiidateam/aiida-core/issues/7479), [#7533](https://github.com/aiidateam/aiida-core/pull/7533) | + +## Scope + +Bringing aiida-core's workflow ecosystem in-house splits into **three independent tracks**, each its own AEP: + +1. **WorkGraph engine + framework → core** (this AEP): the native DAG workflow model on a shared `Workflow` process base, the shared serializer foundation it uses, and the entry points, so aiida-workgraph becomes an archivable shim. +2. **aiida-shell + aiida-pythonjob → core**: `ShellJob` and `PythonJob`/`PyFunction` as first-class CalcJobs in `aiida/calculations/`, consuming track 1's serializer. +3. **node-graph → core**: vendor the generic graph / task / socket spec as ABCs, then archive the package. + +This AEP covers **track 1**. The dependencies between tracks: track 1 keeps node-graph as an external hard dependency (the way it depends on plumpy) until track 3 vendors it in, and it lands the serializer foundation (`general_serializer`, the datetime / function data types, the deserializer) that track 2 builds on. Landing the whole serializer in track 1 keeps core's serialize/deserialize stack in one place and lets aiida-pythonjob drop its copy immediately, even though only `general_serializer` + `serialize_ports` are exercised by the WorkGraph engine itself. + +Further, later work (each its own design): the optional **semantic / knowledge-graph layer** (`[semantics]` extra, node-graph's `knowledge` module, part of track 3) and a pluggable **execution-backend ABC**. + +## Motivation + +- **One source of truth.** aiida-workgraph and its satellites reimplement things core already has: their own `JsonableData`, a `builtin_serializers` table duplicating `to_aiida_type`, `PickledData`, `NoneData`, their own graph engine. node-graph was in fact *extracted out of* aiida-workgraph (same author, one month apart in 2023) and never fully separated, so much of the cross-package duplication is unfinished-extraction residue rather than a designed boundary. +- **A package that cannot be built on the public API is already core in practice.** WorkGraph is the only package (of ~40 audited in [#7410](https://github.com/aiidateam/aiida-core/issues/7410)) that reaches into engine / daemon / config internals no third-party plugin uses. +- **Keep the recommended authoring API in core**, not in a separate external package (not even under `aiidateam`), and drop the dependency-inversion gymnastics the split forces (registry / entry-point indirection so core never imports downstream). +- **Raise code quality:** fold loosely-gated plugin code under core's strict mypy + test bar. + +Guiding principle throughout: *generalize / extend aiida-core and drop the plugins' hand-written duplicates.* Backwards-compatible extension (additive accessors, more-permissive validation) is preferred; since this targets v3, incompatible changes are acceptable where genuinely required. + +## Why a graph engine belongs in core + +Core's only execution model is the plumpy **outline stepper**: a lexically-ordered sequence, static and sealed at first instantiation (grepping `topological` / `networkx` / DAG across `aiida/` turns up only SQLAlchemy's `declarative_base`). It expresses neither data-dependency scheduling nor sub-step concurrency, which is what a research workflow usually wants. + +```python +# WorkChain: the outline is a chain of hard barriers. +spec.outline(cls.submit_a, cls.submit_b, cls.combine) +# submit_b runs only after EVERY child launched in submit_a has finished. + +# WorkGraph: each task wakes when its own inputs are ready; siblings stream. +wg = WorkGraph() +a = wg.add_task(calc, x=1) +b = wg.add_task(calc, x=a.outputs.result) # starts the moment `a` is done +wg.run() +``` + +Concretely, the outline stepper cannot: + +- **Order by data dependency** (outline order is lexical; nothing declares "B consumes A's output"). +- **Run sub-steps concurrently** (the core gap): `_do_step` clears awaitables at every step boundary and resumes only when *all* are done. WorkGraph wakes each task on its own deps. +- **Map / fan-out at runtime** (`Map`/`GatherItem`); outline vocabulary is only `while_`/`if_`/`return_`. +- **Mutate the graph mid-run**, or run **parameterised steps that return data** (an outline step takes only `self`; data flows through untyped `self.ctx`). +- **Treat the workflow as data**: WorkGraph stores the whole graph on the node and rebuilds it on restore (inspect, diff, restart-with-modification, GUI). A WorkChain checkpoint only restores a position in an outline rebuilt from the class. + +## The enabling refactor: a shared `Workflow` base + +WorkGraph was an external fork because `WorkChain` hard-binds execution to the outline stepper and seals `run`/`on_run`/`to_context`/`on_exiting`/`on_wait` with `@Protect.final`. Swapping the scheduler meant going around `WorkChain` and re-copying its awaitable/context/checkpoint machinery, and that copy is the root of the duplication. + +The fix extracts the shared machinery into a new `Workflow(Process)` base; `WorkChain` and `WorkGraphProcess` become siblings, each supplying only its own stepper. + +```{mermaid} +classDiagram + Process <|-- CalcJob + Process <|-- FunctionProcess + Process <|-- Workflow + Workflow <|-- WorkChain + Workflow <|-- WorkGraphProcess + note for Workflow "stepper seam + awaitables + ctx + step lifecycle" + note for WorkChain "outline stepper" + note for WorkGraphProcess "DagStepper (streams; awaitable_barrier=False)" +``` + +```python +class Workflow(Process, metaclass=Protect): + """Shared base: stepper seam, awaitable-based waiting, ctx, step lifecycle + checkpointing.""" + def _create_stepper(self) -> Stepper: # abstract: the subclass picks the strategy + raise NotImplementedError + +class WorkChain(Workflow): # walks a static outline + def _create_stepper(self): + return self.spec().get_outline().create_stepper(self) + +class WorkGraphProcess(Workflow): # schedules by data dependencies + def _create_stepper(self): + return DagStepper(self) # sets awaitable_barrier = False to stream +``` + +- The `_create_stepper`/`_recreate_stepper` seam is abstract and left overridable, unlike the `@Protect.final` lifecycle methods around it. +- The awaitable-clearing barrier belongs to the stepping strategy, so it lives on the base and reads from the stepper: `awaitable_barrier = False` streams (launch when a task's inputs are terminal, resume on the *first* child to finish); the outline default barriers. WorkChain carries no WorkGraph concept. A falsification-checked test guards it: flip the flag and a ready task waits for its slow sibling, so the test fails. +- WorkGraph's copied `AwaitableManager` + `ContextManager` are deleted; both siblings use the base's. +- The extraction also fixes `Protect`: it now scans each base's full MRO, so a `@final` method reached through an intermediate class (e.g. `run`, now on `Workflow`) stays protected. + +The first prototype had `WorkGraphProcess` subclass `WorkChain`; review ([#7479](https://github.com/aiidateam/aiida-core/issues/7479)) flagged that it bolts WorkGraph's barrier policy onto `WorkChain` and couples the two execution models, hence the sibling split. + +Two boundaries the refactor keeps deliberate: + +- **`TaskState` stays its own enum.** It tracks a DAG slot (`PLANNED`/`SKIPPED`/`MAPPED`, states a task may hold without ever becoming a process); `ProcessState` is the lifecycle of one live process. The overlapping names (`RUNNING`, ...) roll up a child's real state. +- **No `CalculationProcess` base** for symmetry with `CalculationNode`. A base is introduced only where there is shared implementation. `CalcJob` and a calcfunction share almost none, and one `FunctionProcess` already backs both `@calcfunction` (a `CalculationNode`) and `@workfunction` (a `WorkflowNode`), differing only in `_node_class`. The calculation/workflow split stays enforced on the nodes and in the engine's link rules. + +`TaskState`'s members make the first boundary concrete: + +```python +class TaskState(str, Enum): # a DAG slot, separate from plumpy's ProcessState + PLANNED = 'PLANNED' # a task may sit here and never become a live + READY = 'READY' # process; PLANNED / SKIPPED / MAPPED have no + CREATED = 'CREATED' # ProcessState analogue + RUNNING = 'RUNNING' # RUNNING / FINISHED / FAILED share the name but + FINISHED = 'FINISHED' # only roll up the child process's real state + FAILED = 'FAILED' + SKIPPED = 'SKIPPED' + MAPPED = 'MAPPED' +``` + +## Serialization onto core's machinery + +The moved serializer is rebuilt on core, so no `JsonableData`/`to_aiida_type` copy is carried: + +- `general_serializer` (`aiida/orm/nodes/data/serializer.py`) dispatches value types through core's `to_aiida_type`, foreign types through an `aiida.data` entry-point registry, and JSON-able fallbacks to `JsonableData`; `serialize_to_aiida_nodes` maps it over a dict. Core's `to_aiida_type` already subsumes the plugin's `builtin_serializers` table (the scalars, list, dict, numpy, enum, `None`), so no value-type mapping is copied. The registry is lazy/cached; custom serializers come from a `serializers=` argument (the `pythonjob.json` config is dropped). +- New core data nodes `DateTimeData` and `FunctionData`, each registered with `to_aiida_type`; `deserialize_to_raw_python_data` is the inverse. Core's `JsonableData` gains a `.value` alias rather than a duplicate. +- `serialize_ports` (`aiida/workgraph/serialization.py`) walks a node-graph `SocketSpec` and serializes each leaf through `general_serializer`. It sits in the subsystem because it imports node-graph; `aiida.orm` and a plain `import aiida` stay node-graph-free. + +```python +from aiida.orm import general_serializer +general_serializer(3.14) # Float via to_aiida_type +general_serializer({'a': [1, 2]}) # Dict via to_aiida_type +general_serializer(datetime.now()) # DateTimeData via to_aiida_type +general_serializer(MyDataclass(x=1)) # JsonableData JSON-able fallback +general_serializer(open('f')) # ValueError with guidance (no serializer, not JSON-able) +``` + +`PickledData` stays in the plugins: aiida-shell already registers `core.pickled`, so core claiming it collides until the track-2 fold consolidates both onto one type (with a `cloudpickle` dependency). + +## Core imports no plugin + +WorkGraph has to recognise plugin processes (`PythonJob`/`ShellJob`/...) as task types. Each plugin process declares a marker and core reads it with `getattr`, so `aiida.workgraph` recognises them while importing no downstream package (GRASP: the process is the information expert). The earlier version imported each plugin class to compare against, which coupled core to its own plugins. + +```python +# in the plugin (aiida-pythonjob), one class attribute: +class PythonJob(CalcJob): + _workgraph_task_type = 'PYTHONJOB' + +# in core, importing nothing from the plugins: +def inspect_aiida_component_type(executor): + declared = getattr(executor, '_workgraph_task_type', None) + return declared or _core_fallback(executor) # CalcJob / WorkChain / process function +``` + +## Where things land + +Layered, batteries-included-but-extensible, the pattern core already uses for `Data` / `Transport` / `Scheduler` / `StorageBackend`: + +| Component | Home | Note | +|---|---|---| +| Graph spec (graph/task/socket/link/registry ABCs) | vendored from node-graph | track 3; external dep for now | +| AiiDA dialect (`Task`/socket subclasses, `WorkGraph` authoring, decorator, registry) | `aiida/workgraph/` | subclass core's spec | +| Execution (`WorkGraphProcess` + `DagStepper`) | `aiida/workgraph/engine/` | sibling of `WorkChain` | +| Serialization (`general_serializer`, deserializer, `serialize_ports`) | `aiida/orm/nodes/data/` + `aiida/workgraph/` | on `to_aiida_type` | +| Data types (`DateTimeData`/`FunctionData`/extended `JsonableData`) | `aiida/orm/nodes/data/` | | +| Calc jobs (`ShellJob`, `PythonJob`/`PyFunction`) | `aiida/calculations/` | track 2 | +| Duplicated scaffolding (zones, `*_pool.py`, forked helpers, dead `validate()` overrides) | delete | replace with imports/thin subclasses | +| Optional extras (semantics/rdflib, viz widget, ASE, cloudpickle) | `[extras]` | keep core lean | + +Control-flow constructs (`Map`/`Select`/`SetContext`/...) are each two-part: a declaration and a runtime half. Route both halves together to their homes (declaration to the spec/dialect, semantics to the engine) so a construct is never split across the move. + +Two of the scaffolding forks already drifted into real bugs (a lost `ContextVar` isolation in the copied context manager, a bypassed validation adapter), the usual failure mode of copied code and the reason to delete rather than re-fork. The one seam worth preserving is the plumpy `Port` → node-graph `SocketSpec` bridge, the real boundary between the two type systems. + +## Alternatives considered + +- **Keep WorkGraph a plugin.** Viable once the stepper seam exists, but it leaves the duplication in place and the "cannot be built on the public API" problem unsolved; [#7410](https://github.com/aiidateam/aiida-core/issues/7410) already assumes absorption. +- **Make WorkGraph subclass `WorkChain`** (the first prototype). Rejected: it bolts WorkGraph's barrier policy onto `WorkChain` and couples the two execution models. The shared `Workflow` base keeps them independent. +- **Keep node-graph external, or move it to `aiidateam` without folding** (track 3's question). Rejected: either leaves core's recommended authoring API in a separate package and keeps the dependency-inversion indirection; the genericity that would justify a standalone package never materialised (its one real second consumer, the multi-engine POC, is itself AiiDA-fused). + +## Roadmap + +Track 1, leaves-first; each a reviewable slice, no boil-the-ocean branch. + +1. Extract `Workflow(Process)`; re-parent `WorkChain`. +2. `WorkGraphProcess(Workflow)` + `DagStepper`, sibling of `WorkChain`; drop the copied managers. +3. The serializer foundation (`general_serializer`, datetime / function data, deserializer, `serialize_ports`) onto `to_aiida_type`/`JsonableData`; repoint aiida-pythonjob and the engine. +4. Relocate the WorkGraph subsystem into `aiida/workgraph/`; marker-based task detection; entry points in core. +5. AiiDA `Task` base + sockets subclass core's own spec; de-pluginise the `Task.__call__` guard. +6. Reduce aiida-workgraph to a shim, then archive. + +Steps 1 to 4 are landing now (see Implementation status). The engine may instead go to `aiida/engine/processes/workgraphs/` mirroring `workchains/`, a relocation either way. + +The other two tracks are separate AEPs. **Track 2:** fold `ShellJob`, then `PythonJob`/`PyFunction`, into `aiida/calculations/` (resolve the `core.pickled`/`cloudpickle` collision, `ase` behind an extra); archive both plugins. **Track 3:** vendor node-graph's spec into core as ABCs, move the `Map`/`Select`/ctx control-flow constructs into it, add the `[semantics]` extra, archive node-graph. + +## Governance and record + +- **Author on board.** node-graph's / aiida-workgraph's author approved consuming the scinode repos into core. He has since left the team, but the team is in direct contact with him, so there is no handover concern. Confirm no external dependents on node-graph before archiving. +- **Prior record:** the v3 public/private-API AEP [#7410](https://github.com/aiidateam/aiida-core/issues/7410) already states WorkGraph is "slated to move into core"; the v3 parent [#7406](https://github.com/aiidateam/aiida-core/issues/7406) invites AEP sub-issues and has none on workflows; the shell capability has a standing request [#5287](https://github.com/aiidateam/aiida-core/issues/5287). +- **`process_type` data break** (cumulative, from ORM entry-point + class-path moves) is kept soft by [#7386](https://github.com/aiidateam/aiida-core/issues/7386) (unknown types fall back to `ProcessNode`); note once in v3 release notes. +- **On a standalone aiida-workgraph 1.0:** it mostly buys a backwards-compat obligation for a package meant to dissolve; prefer targeting core v3 directly. + +## Open questions + +- node-graph external users: confirm none depend on it as a standalone SDK before archiving. +- Release cadence: folding couples shell/pythonjob/node-graph to core's slower cadence. +- Extras policy: `ase`, `cloudpickle`/`PickledData`, `rdflib`/semantics, the widget. +- Control-flow ownership: which layer defines `If`/`While`/`Map` and their meaning (interacts with aiida-workgraph [#601](https://github.com/aiidateam/aiida-workgraph/issues/601)). +- Core API gaps the fork exposed, each a small core PR: no public `deserialize_safe`; no supported process-control RPC / extensible `Intent`; no per-plugin mutable-state slot on `ProcessNode`; private `instantiate_process`; hardcoded `WorkChainNode` logger name. + +## Implementation status + +Landing in PR [#7533](https://github.com/aiidateam/aiida-core/pull/7533) (`refactor/workgraph-into-core`), targeting v3, as a chain of self-contained, individually-verified commits (leaves-first, so each reviews on its own): + +1. `general_serializer` / `serialize_to_aiida_nodes` on `to_aiida_type` + `JsonableData` (the generic serializer, no copied value-type table). +2. node-graph as a hard dependency + `serialize_ports` over `SocketSpec` (its node-graph-coupled half, kept out of `aiida.orm`). +3. `DateTimeData` / `FunctionData` + `deserialize_to_raw_python_data` (the data nodes and the inverse). +4. `JsonableData.value` alias (a backwards-compatible accessor, so aiida-pythonjob can repoint at core's `JsonableData`). +5. pre-commit / mypy / ruff exemptions for `aiida/workgraph/` (pure tooling, so the relocation commit stays code-only). +6. extract `Workflow(Process)`, re-parent `WorkChain`, fix `Protect`'s MRO scan (core-only; WorkChain behaviour unchanged). +7. lift-and-shift the subsystem into `aiida/workgraph/` (49 modules as one unit, imports repointed, no logic change). +8. `WorkGraphProcess(Workflow)`, sibling of `WorkChain` (drops the inherit-WorkChain coupling review flagged). +9. rename the base to `Workflow`, relocate it to `aiida/engine/processes/workflow.py`. +10. marker-based task detection (core imports no plugin). +11. entry points registered in core. + +Why the split: the serializer reconcile (1 to 4) lands and tests bottom-up before anything consumes it; the tooling prep (5) is isolated so the big move is code-only; the enabling `Workflow` refactor (6) carries zero WorkGraph code, so it reads as a pure aiida-core change against WorkChain's own suite; the relocation (7) is mechanical with no behaviour change; re-parenting (8), naming (9), dependency inversion (10) and registration (11) each move exactly one thing. Every non-docs commit is verified against the WorkGraph suite at 181 passed / 11 pre-existing environmental failures, zero regressions. + +After (11), aiida-workgraph is a forwarding shim, archivable: + +```python +# aiida_workgraph/__init__.py, the whole package after the move +import aiida.workgraph as _core +from aiida.workgraph import * # forward every symbol +task = _core.task # rebind the decorator past the `task` submodule +# entry points now register `aiida.workgraph.*`, so nothing here needs to exist +``` + +node-graph is still a hard dependency (not yet vendored). diff --git a/docs/source/internals/index.rst b/docs/source/internals/index.rst index 6e9742241e..33e378af7b 100644 --- a/docs/source/internals/index.rst +++ b/docs/source/internals/index.rst @@ -9,6 +9,7 @@ Internal architecture engine broker rest_api + aep/index .. todo:: diff --git a/pyproject.toml b/pyproject.toml index ea805cbbe4..7f3c66dbfb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -295,6 +295,7 @@ bpython = [ 'bpython~=0.20' ] docs = [ + 'sphinxcontrib-mermaid~=1.0', 'pydata-sphinx-theme~=0.15.1', 'sphinx~=7.2.0', 'sphinx-copybutton~=0.5.0', diff --git a/uv.lock b/uv.lock index 82fe8ba133..b51922c3e7 100644 --- a/uv.lock +++ b/uv.lock @@ -88,6 +88,7 @@ docs = [ { name = "sphinx-intl" }, { name = "sphinx-notfound-page" }, { name = "sphinx-sqlalchemy" }, + { name = "sphinxcontrib-mermaid" }, { name = "sphinxext-rediraffe" }, ] notebook = [ @@ -272,6 +273,7 @@ requires-dist = [ { name = "sphinx-intl", marker = "extra == 'docs'", specifier = "~=2.1.0" }, { name = "sphinx-notfound-page", marker = "extra == 'docs'", specifier = "~=1.0" }, { name = "sphinx-sqlalchemy", marker = "extra == 'docs'", specifier = "~=0.2.0" }, + { name = "sphinxcontrib-mermaid", marker = "extra == 'docs'", specifier = "~=1.0" }, { name = "sphinxext-rediraffe", marker = "extra == 'docs'", specifier = "~=0.2.4" }, { name = "sqlalchemy", specifier = ">=2.0.20,<3" }, { name = "sqlalchemy", extras = ["mypy"], marker = "extra == 'pre-commit'", specifier = "~=2.0" }, @@ -5610,6 +5612,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c2/42/4c8646762ee83602e3fb3fbe774c2fac12f317deb0b5dbeeedd2d3ba4b77/sphinxcontrib_jsmath-1.0.1-py2.py3-none-any.whl", hash = "sha256:2ec2eaebfb78f3f2078e73666b1415417a116cc848b72e5172e596c871103178", size = 5071, upload-time = "2019-01-21T16:10:14.333Z" }, ] +[[package]] +name = "sphinxcontrib-mermaid" +version = "1.2.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, + { name = "sphinx" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f5/49/c6ddfe709a4ab76ac6e5a00e696f73626b2c189dc1e1965a361ec102e6cc/sphinxcontrib_mermaid-1.2.3.tar.gz", hash = "sha256:358699d0ec924ef679b41873d9edd97d0773446daf9760c75e18dc0adfd91371", size = 18885, upload-time = "2025-11-26T04:18:32.43Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/39/8b54299ffa00e597d3b0b4d042241a0a0b22cb429ad007ccfb9c1745b4d1/sphinxcontrib_mermaid-1.2.3-py3-none-any.whl", hash = "sha256:5be782b27026bef97bfb15ccb2f7868b674a1afc0982b54cb149702cfc25aa02", size = 13413, upload-time = "2025-11-26T04:18:31.269Z" }, +] + [[package]] name = "sphinxcontrib-qthelp" version = "2.0.0"