✨ Store a modelled socket through its own model

[elinscott]

Stacked on #800 (enum flattening). Depends on scinode/node-graph#182, which introduces input_model= / output_model=; against the pinned node-graph release this branch does not import.

The idea

scinode/node-graph#182 lets a task declare its inputs and outputs as Pydantic models: @task(input_model=PhInputs, output_model=PhOutputs). The model declares the sockets, and node-graph runs the model's validators. That PR is the baseline for this PR; read it first.

What that PR does not implement is serialization. node-graph never stores a socket value: it stays a Python object from the moment it is written to the moment the body runs.

aiida-workgraph cares about serialization. Every socket value becomes an AiiDA node on the way into a process and is read back from a node on the way out. That path has its own type table — general_serializer, _flatten_enums, builtins_type_mapping — and knows nothing about Pydantic models.

If we just take the new input model implementation as-is, we run into three errors

1. a Decimal input cannot be stored at all
2. an Enum input works only through a package-specific workaround
3. a body receives whatever node the engine stored, not the type its model declares.

This PR makes the Pydantic model the authority on serialization, too.

• On the way in, a modelled value is rendered by its model's own serializer
• On the way out, the AiiDA node is taken off for fields the model declares as plain Python

The checks node-graph defines are hooked into the two places this engine runs them — graph expansion inside a submitted process, and a task's own process for output_model.

What it looks like

class MoneyInputs(BaseModel):
    amount: Decimal

    @field_serializer("amount")
    def _dump_amount(self, value: Decimal) -> str:
        return str(value)

@task(input_model=MoneyInputs)
def double_money(amount):
    return {"kind": type(amount).__name__, "doubled": str(amount * 2)}

amount is stored as '0.10', the string the model's serializer rendered, and the body is handed a Decimal back — so str(amount * 2) is exactly '0.20', never a float's '0.2000000000000000041...'.

Changes

On the way in, storage goes through the model. AiidaSerializationAdapter.serialize asks node-graph which input model owns a socket and, if one does, renders the value with model_dump(mode='json') instead of _flatten_enums. The lookup follows the socket's full path through nested models and typed mappings, so a Decimal at items['a'].inner.value is rendered by the model that declares it. An Enum is stored as its bare value and rebuilt by model_validate; no enum-specific code remains.

On the way out, the read edge unwraps what the model declares as Python. deserialize takes the AiiDA node off a socket whose model says the body receives Python, and leaves it on for an AiiDA type or Any. The model must decide this, not the socket identifier: int and orm.Int both map to workgraph.int. The value keeps its tag and the node's uuid either way, so a graph body draws a link from it rather than holding a copy. A leaf body therefore gets the model's validated values (Enum as the member, Decimal as a Decimal). A graph body gets them too: the expansion check hands back the values the model made, re-tagged leaf by leaf with the original tag and uuid, so spin == Spin.COLLINEAR holds inside a graph body and the link still draws.

A model that declares an AiiDA type makes the task a calcfunction. A PyFunction body is handed the value a node carries, never the node, so @task(input_model=M) with a field declaring orm.StructureData (or Any, at any depth, mapping items included) is refused at decoration:

Structure declares 'structure' as an AiiDA type, and a task declared with @task runs its body as a PyFunction, which is handed the value a node carries, never the node.
How to fix: declare the task with @task.calcfunction, whose body is handed the node; or declare the field as the Python type the body reads.

@task.calcfunction(input_model=M) takes the model, AiiDA's calcfunction wraps the validated callable, and the body receives the nodes. A namespace-shaped model — a nested model or a dict[str, T] — is refused on a calcfunction, because aiida-core gives a process function one port carrying one node per named parameter; a rule broken inside a calcfunction body excepts the node, with the model's report on node.exception.

The checks run at this engine's two process boundaries. node-graph's graph-expansion check sits in materialize_graph, which GraphTask.execute calls, so a @task.graph(input_model=M) inside a submitted process is checked when its inputs are known — including a subgraph whose bound an upstream task produced. A leaf task runs as an AiiDA process rather than through node-graph's executor wrapper, so the run-edge check is hooked into the process runner: a return value output_model rejects fails that process with exit status 323 and the model's report as the exit message (verdi process show).

Nothing changes for a task without a model. A socket no model owns takes the base adapter, which is also what lets #799's identifier-keyed unwrap coexist with this one.

Testing

tests/test_input_model.py, 64 tests against a live profile; the contract's own tests are node-graph's. All reproduced on this tip.

• Storage. Each modelled leaf's stored string and the type the body receives, over a nested model, a mapping of models, and a mapping of models with a model inside. Control: the same task without a model — the enum control's body reports str, the Decimal control cannot be stored.
• What a body receives, per field kind in one run: str/int/float/bool/dict/list as Python, Any and orm.Int as nodes, with a spec assertion that int and orm.Int share an identifier. Control: the base deserialize makes the same graph fail at expansion.
• Provenance. The node a subtask reads is the node the graph was given, by uuid. Control: an unwrap that drops the tag leaves the subtask reading an orphan node with no incoming links.
• Where refusals land. A validator broken inside a graph body leaves the workgraph with no called descendants (exit 302; the model's text reaches the console only). A validator broken at the run edge gives exit 323 with the model's report as the exit message. Control: the same graph at collinear expands, runs, and hands the body the enum member.
• A graph body sees the member. A modelled graph body reports Spin.COLLINEAR under its tag, the same as an annotated graph body; before this branch it reported 'collinear'. (reproduced)
• Calcfunction. A StructureData field and a field under Any are each refused on @task with the message above, and the same two models on @task.calcfunction run with the bodies reading back StructureData and Dict; a namespace-shaped model is refused on a calcfunction at decoration; a rule broken in a calcfunction body leaves the node excepted with the report on it. Control: the unmodelled @task still loses the node in silence. (reproduced)
• Written members. A leaf handed system={'nbnd': 20} out of a five-member namespace reads back ['nbnd'] from storage, and a mapping's items likewise. (reproduced)
• Full suite: 19 failed, 271 passed, 6 skipped paired with this branch's node-graph; the same 19 by name, all AttributeError: 'TaskHandle' object has no attribute 'build' — the known TaskHandle/GraphTaskHandle pairing noise, untouched here. Paired with both patched branches, which carry the fix for that noise: 373 passed, 6 skipped, 0 failed.

Open decision — how a graph body gets the form its field declares

The read edge is one commit and can be dropped on its own. Measured on one harness against #799's identifier-keyed unwrap (what the body receives per field kind; whether a graph input's node is still the subtask's node by uuid):

	graph expands	body receives	graph-input link
#799's deserialize alone	no — Decimal, Enum, dict, list, Any reach the model as nodes and it refuses them	—	lost (subtask node has 0 incoming links)
the model-keyed unwrap here	yes	str/int/float/bool/dict/list and the rendered Decimal/Enum as Python; Any and AiiDA-typed as nodes	preserved

The link loss is #799's return value.value, which drops the TaggedValue a graph body turns into a link; keeping the tag is what this row does differently. Measured against #799 at the revision the prototype carries; patched also unwraps workgraph.dict/workgraph.list, which changes nothing above.

Not yet covered

A leaf taking a namespace of nodes — pseudos: dict[str, UpfData], one node per element — has no model that can declare it: aiida-core delivers that shape through **kwargs and nothing else (a mapping written into a named parameter meets to_aiida_type and collapses into one orm.Dict), and no model field can name **kwargs. Under a model it lives on a @task.graph, whose own socket may be dict[str, Node], or each node takes a field of its own. Giving the contract a spelling for it is separate work.

Note for the maintainers

The mypy pre-commit hook is language: system, so it runs whatever mypy is first on PATH — an interpreter without aiida or node_graph — and reports 16 errors in files this branch does not touch, identically on the unmodified tip. Against the project's interpreter (mypy --config-file=pyproject.toml --python-executable .venv/bin/python src) the tree is clean, so these commits skipped the hook. Fix is yours to pick: args: [--python-executable, .venv/bin/python], or language: python with additional_dependencies: [aiida-core, node-graph].

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}