v3: workgraph into core - #7533
Draft
GeigerJ2 wants to merge 20 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.
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.
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.
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.
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.
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.
`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).
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.
|
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 #7533 +/- ##
===========================================
- Coverage 80.68% 28.56% -52.11%
===========================================
Files 581 635 +54
Lines 47068 50416 +3348
===========================================
- Hits 37972 14397 -23575
- Misses 9096 36019 +26923 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Comment on lines
+245
to
+247
| def value(self) -> JsonSerializableProtocol: | ||
| """Alias of :attr:`obj`, mirroring the ``value`` accessor of the simple data types (``Int``, ``Str``, ...).""" | ||
| return self._get_object() |
Member
There was a problem hiding this comment.
@GeigerJ2 does WorkGraph expect all Data nodes to have .value?
| return NoneData() | ||
|
|
||
|
|
||
| class NoneData(Data): |
Member
There was a problem hiding this comment.
Very Python. Also, aiida-pythonjob already has a NoneData. Do we need another in core? Bigger question - should AiiDA hold on to anything other than the absolute bases (e.g., ProcessNode, Data, etc.), letting instead aiida-pythonjob (should be aiida-python really, though one exists 🥲) extend with BaseType and all Python primitives?
As I write this, I recall that we also marked aiida-pythonjob for core integration 😅 TBD
Collaborator
Author
There was a problem hiding this comment.
As I write this, I recall that we also marked aiida-pythonjob for core integration 😅 TBD
Yeah, I think that change was probably necessary, as it's used in WG, so, it's already moved here to core. It's a bit hard to disentagle everything, tbh, as WG also interfaces with aiida-shell and aiida-pythonjob.
Member
There was a problem hiding this comment.
Okay, no worries. As discussed, those are all marked for integration. If we need to already pull in things to get things working, so be it.
Member
|
@GeigerJ2 I started looking over your commits 🙂 Will take a closer look over the weekend. Tag me for review when ready 🙏 |
GeigerJ2
force-pushed
the
refactor/workgraph-into-core
branch
2 times, most recently
from
August 7, 2026 11:15
82cf24a to
aa92fc7
Compare
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.
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.
So far, pure AI output; mostly a mechanical move from the
aiida-workgraphrepo.