🐛 Deserialize objects in the body of a task.graph

[elinscott]

Problem

The write path auto-serialises primitive task inputs into orm.Float / orm.Int / orm.Str / orm.Bool for provenance, but there was no symmetric read path: SerializationAdapter.deserialize existed in the node-graph base API and was never implemented here. A @task.graph body whose signature declared nspin: int could therefore receive an orm.Int after a round-trip through AiiDA storage, with two failure modes:

• silently wrong logic: if nspin == 2: style comparisons against a node object
• hard crashes at stdlib or AiiDA boundaries: conv_thr = 1.0e-9 * nelec, QueryBuilder / orm.Dict interfaces.

The same asymmetry bites dataclass-typed sockets crossing a @task.graph boundary: field values are node-promoted, then coerce_structured_value rebuilds the dataclass via cls(**value) with the fields still wrapped, so a field declared int arrives as orm.Int and e.g. range(self.ntyp) raises TypeError (orm.Int has no __index__).

Change

Implement AiidaSerializationAdapter.deserialize as the symmetric counterpart to serialize, invoked just before a @task.graph body runs:

• For sockets whose declared type is a float / int / string / bool): unwrap orm.BaseType to its .value, so the body sees the primitive its signature declared.
• Container-typed sockets get the same treatment: values persisted as orm.Dict/orm.List are unwrapped to plain dict/list on deserialize, so a @task.graph body that declared dict/list receives the Python container rather than the ORM node.
• For dataclass instances: unwrap any orm.BaseType-wrapped field via dataclasses.replace (frozen-safe).
• Everything else (workgraph.any, Aiass structures) passes through unchanged

Compatibility

Requires the node-graph side actually calling the adapter on the read path: the materialize_graph deserialize-on-read hook and routing of structured-namespace values through adapter.deserialize (today _deserialize_inputs only recurses into dicts). Both are part of the node-graph dynamic-namespace/deserialize work targeting node-graph main and not yet in a released node-graph, hence draft.

[elinscott]

@edan-bainglass this PR ties in with our discussion yesterday about whether regular Python is guaranteed to work inside a @task and/or a @task.graph body. Turns out the code currently implies some level of guarantee inside a @task.graph, too. See e.g. docs/source/overview.rst:92 ("Parallel tasks")

@task.graph()
def parallel_add(x, N):
    results = {}
    # Launch N parallel tasks to add x with each index
    for i in range(N): # <- regular python acting on an input node
        results[f"add_{i}"] = add(x, i).result
    return results

or also docs/source/overview.rst:130 ("Dynamic workflow")

@task.graph()
def conditional_workflow(x):
    if x > 0: # <- regular python acting on an input nod
        return add(x, 10).result
    else:
        return multiply(x, 10).result

Obviously being able to write the code this way is a lot nicer than the stricter alternative where both the loop generator and the comparison need to themselves be tasks.

So I think what we should be aiming for is that graph inputs should be concrete values whose type always matches their annotation. Meanwhile, task outputs remain sockets. In this context, this PR patches some instances where this expectation is not being met.

[elinscott]

Found a casualty of the unwrap change: a graph body can no longer echo one of its own inputs as an output. e.g. for the very natural

@task.graph
def RefinementLoop(prev_result: MyOutputs, remaining_steps: int, ...) -> MyOutputs:
    if remaining_steps <= 0:
        return MyOutputs(result=prev_result)   # Invalid graph return payload
    ...

now fails node-graph's return-payload validation: the validator accepts BaseSocket | TaggedValue, and with containers unwrapped in the body, nothing tagged survives to the return. (Against main the echo passes, because the input arrives TaggedValue-wrapped.)

An ugly workaround is an artificial passthrough task i.e.

@task
def echo(x)
    return x

but obviously this is not what we want long-term.

We probably want to preserve enough input-origin tagging through the unwrap that using the variable as an output still satisfies the validator.

[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.}