✨ Take a task's input and output contracts from Pydantic models

[elinscott]

Stacked on #178 (enum and literal sockets). The commits to review here are the ones after it; #178's own commits are in the diff because this branch builds on them.

The idea

Today a task declares its sockets in one of two ways: from its signature — def ph(spin: str, structure: str) gives two input sockets typed str — or explicitly, @task(inputs=namespace(spin=str, structure=str)), which can also nest, mark a socket required, and give it a default. Either way a socket carries a type and a default and nothing else, which is not expressive enough for some things we want to do (see below).

Pydantic has much richer validation vocabulary. node-graph already reads models — SocketSpec.from_model expands pydantic models into sockets — but it keeps only the field names, types and defaults and throws the rules away.

This branch keeps the rules. @task(input_model=PhInputs) makes the model the declaration of ph's inputs, and node-graph enforces the model's validators wherever a value arrives at ph. @task(output_model=PhOutputs) does the same for what ph returns.

What it looks like

from enum import Enum
from pydantic import BaseModel, Field, field_validator
from node_graph import task

class SpinType(str, Enum):
    NONE = "none"
    COLLINEAR = "collinear"
    NON_COLLINEAR = "non_collinear"
    SPIN_ORBIT = "spin_orbit"

class PhInputs(BaseModel):
    spin: SpinType = SpinType.NONE
    structure: str
    ecutwfc: float = Field(default=60.0, le=200)

    @field_validator("spin")
    @classmethod
    def _supported(cls, value):
        if value in (SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT):
            raise ValueError("ph.x has no electric-field perturbation for noncollinear magnetism")
        return value

class PhOutputs(BaseModel):
    dielectric: float

@task(input_model=PhInputs, output_model=PhOutputs)
def ph(spin, structure, ecutwfc):
    ...
    return {"dielectric": dielectric}

PhInputs does two jobs. It declares the sockets — nothing new; other annotations do that. It also carries rules: ecutwfc has a bound, and two of the four SpinType members are ones ph.x cannot run. Those rules travel with the sockets. Put ph inside a workflow that accepts every spin treatment:

class WorkflowInputs(BaseModel):
    spin: SpinType = SpinType.NONE
    structure: str

@task.graph(input_model=WorkflowInputs)
def eps(spin, structure):
    return ph(spin=spin, structure=structure).dielectric

eps.build(spin=SpinType.NON_COLLINEAR, structure="si")

node_graph.input_model.TaskInputValidationError: Task 'ph' got inputs PhInputs rejects:
1 validation error for PhInputs__Rules
spin
  Value error, ph.x has no electric-field perturbation for noncollinear magnetism [type=value_error, input_value=<SpinType.NON_COLLINEAR: 'non_collinear'>, input_type=SpinType]

eps does not have to know which spins ph supports, and ph's body never ran: the value was refused while the graph was being built, at the line inside eps that handed it to ph, before anything was stored or submitted. A bad value written directly fails at that line too — graph.add_task(ph, "ph", structure="si", ecutwfc=500) raises Input should be less than or equal to 200 — and a return value the output model rejects fails the task that produced it:

node_graph.input_model.TaskOutputValidationError: Task 'ph_forgets' returned outputs PhOutputs rejects:
1 validation error for PhOutputs
dielectric
  Field required [type=missing, input_value={}, input_type=dict]

What we tried first

This ph example comes from aiida-koopmans, whose dielectric task runs ph.x, which aborts on an electric-field perturbation under noncollinear magnetism, while the workflow around it takes all four of aiida-quantumespresso's SpinTypes. I tried to impose the spin restriction three ways before this pydantic-based rewrite:

• A raise in the body. Inside DielectricTask, add if spin in (NON_COLLINEAR, SPIN_ORBIT): raise NotImplementedError(...). This only fires when the body runs, so inside a deferred @task.graph, that raises mid-run in the daemon, after the value was stored and the process submitted. It also means that we are describing ph-specific restrictions outside of the body of ph.
• Narrowing the annotation to Literal[SpinType.NONE, SpinType.COLLINEAR] fell victim to the missing Literal support (Literal[...] not supported #175). 🐛 Decide enum/Literal socket membership once, at assignment #178 fixes that at _set_socket_value, and what it gives is membership: a value outside the Literal is refused at the socket, with a message that names the socket. Nevertheless, every other kind of rule a pydantic model can state stays out of reach.
• Restating the rule at the caller. I already use pydantic to parse input yaml files that parametrise the workflow to run. This started to grow route-conditional validators (e.g. "if requesting eps to be calculated, then you can't have non-collinear spin"). This is a re-implementation of knowledge that should be local to the ph graph, not the input file parser, and is prone to drift.

When a rule is checked

A value can arrive at a task at three moments, and the model is consulted at each. The rest of this description calls them checkpoints A, B and C.

• A — when an input is written: calling the task, graph.add_task(...), task.set_inputs(...). Field types, constraints and mode='after' field validators run on every field whose value is present. A field holding a socket or a task gets its type checked only; model validators wait, since the other fields may not exist yet. A rule that reads a sibling nobody has written yet waits alone; the other rules in the same write still run (info.data['x'] waits, info.data.get('x') returns None as dict.get always has). One route bypasses A: assigning to a socket directly (task.inputs.x.value = ...), which the socket layer never checked either.
• B — when a @task.graph is expanded: its inputs are values by then, so the full model runs, model validators included, on an untagged copy that is then discarded; the body keeps its tagged originals, which it turns into links.
• C — when a leaf task runs: the assembled inputs are validated once more and the body receives the result. The return value is validated against output_model here and only here; a missing or mistyped field fails the task that produced it, an undeclared key is refused by name.

output_model is refused on a @task.graph: a graph returns socket references, values that do not exist yet. Put it on the function tasks whose outputs the graph returns.

Changes

The sockets come from the model. spec_from_model adds to from_model what the model knows and the sockets did not: which fields are required (a field with a default is optional), the default itself in JSON form, dict[str, T] fields as dynamic namespaces (one socket per key, validated through T), and, per socket, whether the body receives plain Python or the engine's stored node. A socket is named after its field, whatever alias the field carries. What a body receives for a field the caller did not write depends on the field's shape: a scalar arrives at the model's default; a nested model, or a mapping of models, arrives as the members that were written — {} when the field itself was omitted — so a namelist a caller wrote three keywords into reaches the body as those three, not as the model's hundred defaults. A write may fill a nested model a few members at a time; the missing ones may come by link. (list[Model] items arrive as model instances; the rule stops at mappings.)

Two descriptions of one task, checked against each other. With input_model=, the model and the function signature both describe the inputs. The model is authoritative; the signature is still what the body is called with, so at decoration ModelContractError names any disagreement: a field with no parameter, a parameter with no field, an annotation that contradicts its field (annotations may repeat the model, never differ from it), *args/**kwargs, or a default written in the signature ("defaults live in the model — move it").

Two model shapes are refused at the same point. We refuse extra='allow' (at any depth) because a model that accepts undeclared fields would accept inputs that have no socket; declare every field. We refuse a list of engine objects, such as list[orm.StructureData], because a list is stored as one plain-data node and its members come back as plain data, not as the objects that went in; the fix is dict[str, orm.StructureData], which stores each member as its own node. A model that names itself in a field is refused too, naming the cycle, rather than exhausting the stack.

The checks travel with the task, not with its name. The validated callable is bound in the task's module under a name of its own and the executor records that name, so a task rebuilt from its stored form in another process — or a handle bound to a name other than the function's — still enforces its models. A model this process cannot rebuild (a __main__ model whose nested model carries a @model_validator under from __future__ import annotations fails pydantic's cloudpickle round-trip) is refused with ModelContractError telling the author to move it into an importable module, rather than silently running without checkpoint A.

A refusal can be read, not only printed. TaskInputValidationError carries task, model and errors (pydantic's loc/type/msg list); a value the socket layer refuses before any model runs raises SocketValueError with loc and type, so a consumer translating either into its own vocabulary never parses text.

Validation may change how a value is spelled, never what it says. '60' may become a Decimal; a validator that derives or rewrites an input is refused, because the body would run on a value that never reached storage. The written value is rendered to JSON before the model runs and compared afterwards, so in-place mutation, nested-instance edits and serializer tricks are all caught. The one permitted difference is the model's own config coercion: with str_strip_whitespace, storage keeps ' silicon ' and the body gets 'silicon'.

What a body receives is decided per socket from the field's type: plain Python where pydantic can rebuild the type from plain data (str, int, Decimal, an Enum), the engine's stored node otherwise. A union whose arms disagree is refused at decoration. The mechanism and its alternatives are discussed in the paired aiida-workgraph pull request (see Open decision).

Without either keyword nothing changes. A task declared the usual way takes exactly the path it took before.

Testing

tests/test_input_model.py, 201 collected tests; every checkpoint's test is paired with a control that fails without it.

• A. The same bad value is refused by all three write routes, and so is a constraint (gt=0) the socket layer cannot see; field validators are refused at the wiring in the ph/eps pair above with no ph built. Controls: the check stubbed out, add_task accepts the value; the same graph with collinear runs and hands the body the enum member. (reproduced)
• B and C. Model validators run at expansion on a copy and the body keeps its tagged originals; the leaf run edge refuses a missing, mistyped or undeclared output. (reproduced)
• Content invariance. Four ways to rewrite a value under validation — in-place list mutation, replacement, nested-instance edit, a PlainSerializer — are each refused; str_strip_whitespace is accepted as coercion while a hand-written validator doing the same is refused. (reproduced)
• What a body receives, per field kind and per member of a TypedDict; unions in both directions with Optional[int | Marker] as the control. (reproduced)
• What an unwritten field is worth. A nested model written with one member reaches the body as that one member; an omitted nested or mapping field arrives as {} and the body is still called; a scalar arrives at its default. Control: with default materialization restored the body sees every member. (reproduced)
• The checks survive a round trip. A spec dumped in one process and rebuilt in another still refuses a bad write; a handle bound to another name refuses like the decorator spelling. Control: at the base the rebuilt spec carries no model and accepts the write. (reproduced)
• Documented limits are pinned: a field validator does not fire at A on a linked field; a model validator waits for B, on nested models too; a mode='before' normalizer is not honoured at A; the direct socket write stores what set_inputs refuses; a rule's own KeyError escapes as a bare KeyError.

Full suite: 564 passed, 1 failed on this tip (363 passed, same failure, at the tip of #178). The failure is tests/test_socket.py::test_await_forbidden, which needs an async pytest plugin this interpreter lacks.

Open decision

The read-edge mechanism — body_receives, the per-socket mark saying whether a body gets plain Python or the engine's stored node, plus the removal of SerializationAdapter.to_python — sits in one commit and can be dropped without touching the rest. It is one of several ways to give a body the form its field declares; the alternatives and their costs are in the paired aiida-workgraph pull request.

[codecov]

Codecov Report

❌ Patch coverage is 92.14340% with 103 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.90%. Comparing base (8d85e61) to head (2fc6dc9).

Files with missing lines	Patch %	Lines
src/node_graph/input_model.py	89.81%	76 Missing ⚠️
tests/test_enum_literal_sockets.py	95.29%	12 Missing ⚠️
src/node_graph/link.py	92.18%	5 Missing ⚠️
src/node_graph/utils/struct_utils.py	94.87%	4 Missing ⚠️
src/node_graph/task.py	77.77%	2 Missing ⚠️
src/node_graph/utils/graph.py	92.59%	2 Missing ⚠️
src/node_graph/socket.py	95.00%	1 Missing ⚠️
src/node_graph/socket_spec.py	98.11%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #182      +/-   ##
==========================================
+ Coverage   89.68%   90.90%   +1.22%     
==========================================
  Files          81       84       +3     
  Lines        8984    11881    +2897     
==========================================
+ Hits         8057    10801    +2744     
- Misses        927     1080     +153     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.