Skip to content

馃悰 Engine fixes for Map-zone cloning, gather races, and per-task (de)serializers - #797

Draft
elinscott wants to merge 3 commits into
aiidateam:mainfrom
elinscott:engine-dscf-fixes
Draft

馃悰 Engine fixes for Map-zone cloning, gather races, and per-task (de)serializers#797
elinscott wants to merge 3 commits into
aiidateam:mainfrom
elinscott:engine-dscf-fixes

Conversation

@elinscott

@elinscott elinscott commented Jul 2, 2026

Copy link
Copy Markdown
Collaborator

Summary

Three engine-level fixes uncovered by porting a production workflow with per-orbital Map zones. Each failure mode below is independent; they ship together because the reproducing workflow unearths all three.

1. @task decorator: spec-level deserializers / serializers

@task now accepts deserializers= / serializers= kwargs and stores them on the spec. PyFunctionTask.execute falls back to the spec-level values when no per-call kwarg is provided, in both the sync and async execution branches.

Motivation: provides a mechanism to opt out of the default StructureData -> Atoms auto-deserializer

@task(deserializers={"aiida.orm.StructureData": "my_plugin.utils.identity"})
def analyze(structure):
    ...

Only meaningful for plain-Python (PyFunction) tasks; CalcJob / WorkChain / calcfunction tasks are unaffected (AiiDA owns their deserialization).

2. _patch_cloned_tasks: materialise dynamic-namespace sockets on cloned tasks

Map-zone task cloning looked up link targets with a bare to_node.inputs[scoped_name]. A dotted path into a dynamic namespace (e.g. pseudos.O) does not exist on a fresh clone, so relinking raised AttributeError at submit time. Cloning now resolves the target through Graph._resolve_or_create_input_socket, which materialises the child socket on demand.

3. update_map_task_state: gather race and failure propagation DROPPED (see comment below)

Two related failure modes when a Map zone's iterations finish close together:

  • Race: are_childen_finished can return True for the map zone before every per-prefix gather_item clone has recorded its result in ctx._task_results. The gather loop then KeyErrors. The method now bails out cleanly while any clone is still in flight; the next continue_workgraph tick re-enters once results are populated.
  • Failure masking: a failed mapped iteration leaves its gather_item clone SKIPPED (on_task_failed skips downstream children), so a check keyed on FAILED alone never fires and the gather loop KeyErrors on the empty result, burying the original calculation failure under an unrelated trace. Gather clones that end SKIPPED or FAILED now propagate FAILED to the map task, preserving the real cause. Covered by a new test, test_map_zone_failed_iteration: KeyError before the fix, map task ends FAILED after.

Compatibility

Fix 2 requires Graph._resolve_or_create_input_socket, which is not yet in a released node-graph (introduced alongside the dynamic-namespace work targeting node-graph main). Draft until that release; fixes 1 and 3 have no new dependencies and can be split out if preferred.

* ``@task`` decorator: accept ``deserializers`` / ``serializers`` kwargs
  and stash them on the spec. ``PyFunctionTask.execute`` honours the
  spec-level values when no per-call kwarg is provided. Lets plugins
  opt out of the default ``StructureData -> Atoms`` auto-deserializer
  without per-call boilerplate at every CalcJob.

* ``_patch_cloned_tasks`` (Map zone task cloning): replace the bare
  ``to_node.inputs[scoped_name]`` lookup with the new
  ``Graph._resolve_or_create_input_socket`` helper so dotted paths into
  a dynamic namespace materialise the child socket on demand. Without
  this the cloned task's dynamic-namespace was empty and ``pseudos.O``
  raised AttributeError at submit time.

* ``update_map_task_state``: defensive guard against a race where
  ``are_childen_finished`` returned True for a Map zone while one of
  the per-prefix ``gather_item`` clones had not yet recorded its result
  in ``ctx._task_results``. Now bails out cleanly when any clone is
  still in-flight; the next ``continue_workgraph`` tick re-enters once
  the missing clones complete. Also handles FAILED mapped iterations
  by propagating failure to the parent instead of KeyError-ing on the
  missing gather output (which masked the original child failure with
  a confusing engine trace).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@elinscott elinscott changed the title 馃悰聼聬聸 Engine fixes for Map-zone cloning, gather races, and per-task (de)serializers 馃悰 Engine fixes for Map-zone cloning, gather races, and per-task (de)serializers Jul 2, 2026
A failed compute iteration marks its downstream gather_item clone
SKIPPED (on_task_failed skips the child_node set), not FAILED, so
update_map_task_state's "state == FAILED" check never fired and the
gather loop raised "KeyError: '<socket>'" on the pre-seeded empty
result dict, masking the original failure. Reproduced at
task_state.py:345 with a map where one iteration raises; the failed
iteration's gather clone was observed SKIPPED with an empty result
dict while sibling clones were FINISHED.

Treat SKIPPED (and FAILED) gather clones as failed iterations that
propagate FAILED to the map task. Also replace the dead
"name not in ctx._task_results" guard (copy_task pre-seeds an empty
dict, so it never fired) with an empty-result check that still covers
the finished-but-not-yet-recorded race. Add
test_map_zone_failed_iteration covering the failure path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@GeigerJ2
GeigerJ2 self-requested a review July 20, 2026 11:55
update_map_task_state changes (gather-race bail-out, FAILED/SKIPPED
propagation) and the test_map_zone_failed_iteration test are removed:

- both failure modes are covered by the update_map_task_state rewrite
  in aiidateam#776, which gathers directly from source
  tasks and fails the zone on failed/skipped iterations
- remaining changes are the spec-level (de)serializers and the
  dynamic-namespace clone relinking

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@elinscott

Copy link
Copy Markdown
Collaborator Author

Dropped fix for the race condition from this PR as it is separately addressed by #776

elinscott added a commit to elinscott/aiida-workgraph that referenced this pull request Aug 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant