Refactor workgraph into core - #7513
Draft
agoscinski wants to merge 11 commits into
Draft
Conversation
`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 aiidateam#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.
`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 aiidateam#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.
`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.
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.
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.
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.
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.
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.
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).
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`.
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.
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #7513 +/- ##
==========================================
+ Coverage 80.63% 80.75% +0.13%
==========================================
Files 580 591 +11
Lines 46756 47471 +715
==========================================
+ Hits 37695 38331 +636
- Misses 9061 9140 +79 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
agoscinski
commented
Aug 2, 2026
| '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', | ||
| } |
Collaborator
Author
There was a problem hiding this comment.
This needs a proper integration with the orm serialization mechanisms. There have been several discussions. I am not sure what the current state of the proposed solution is.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.