馃悰 Engine fixes for Map-zone cloning, gather races, and per-task (de)serializers - #797
Draft
elinscott wants to merge 3 commits into
Draft
馃悰 Engine fixes for Map-zone cloning, gather races, and per-task (de)serializers#797elinscott wants to merge 3 commits into
elinscott wants to merge 3 commits into
Conversation
* ``@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>
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
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>
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
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
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.
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.
@taskdecorator: spec-leveldeserializers/serializers@tasknow acceptsdeserializers=/serializers=kwargs and stores them on the spec.PyFunctionTask.executefalls 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 -> Atomsauto-deserializerOnly meaningful for plain-Python (PyFunction) tasks;
CalcJob/WorkChain/calcfunctiontasks are unaffected (AiiDA owns their deserialization).2.
_patch_cloned_tasks: materialise dynamic-namespace sockets on cloned tasksMap-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 raisedAttributeErrorat submit time. Cloning now resolves the target throughGraph._resolve_or_create_input_socket, which materialises the child socket on demand.3.
DROPPED (see comment below)update_map_task_state: gather race and failure propagationTwo related failure modes when a Map zone's iterations finish close together:Race:are_childen_finishedcan returnTruefor the map zone before every per-prefixgather_itemclone has recorded its result inctx._task_results. The gather loop thenKeyErrors. The method now bails out cleanly while any clone is still in flight; the nextcontinue_workgraphtick re-enters once results are populated.Failure masking: a failed mapped iteration leaves itsgather_itemclone SKIPPED (on_task_failedskips downstream children), so a check keyed on FAILED alone never fires and the gather loopKeyErrors 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:KeyErrorbefore 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.