Fix collector scoping and invocation validation

[JPPhoto]

Summary

Fixes collector iteration scope bugs for empty and partially populated collections. invokeai/app/services/shared/graph.py now preserves upstream iteration paths when materializing collectors.

Also replaces assert-based invocation return validation in invokeai/app/invocations/baseinvocation.py, which failed under optimized Python.

Note that empty groups now materialize and may run downstream consumers with [].

Related Issues / Discussions

Closes #9381

• Actually, Optimize iterator graph materialization #9355 incidentally fixed its exact repro. Fix graph stall for collectors downstream of empty iterators #9349 fixed the adjacent empty-iterator stall, but did not close Consumer of both an iterator item and a per-iteration collector output silently skips iterations whose collector is empty #9381!

Closes #9382

QA Instructions

New tests cover empty, non-empty, and mixed collector iterations.

Merge Plan

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[lstein]

Adversarial review at abb324a7d0, diffed against the parent commit a21d74b8bd. No correctness regressions found — the collector fix is real and load-bearing, and I could not break it. Three things I'd like addressed or acknowledged before merge, none of them a code blocker.

What I verified

Check	Result
tests/test_graph_execution_state.py, tests/test_node_graph.py, tests/test_nodes.py, tests/app	2350 passed, 8 skipped, 8 xfailed
ruff check / ruff format --check	clean
25 hand-built graph shapes, HEAD vs parent, output diffed	every difference is an empty branch that now materializes with correct values — no wrong values, no lost non-empty data, no stalls
Serialization round-trip at 5 different mid-execution points on the fixed shape	identical final results
Prepare-time cost (20 outer × 10 inner, chained collectors)	2.42 s vs 2.46 s — no regression

Only 2 of the 4 new tests are load-bearing. Reverting graph.py to the parent commit and rerunning the new tests: test_graph_chained_collectors_preserve_outer_iteration_scope[True] and [None] fail; [False] (fully-populated) and all of test_graph_consumer_with_direct_iterator_and_empty_collector_preserves_outer_iterations pass without the fix. I also ran the verbatim repro script from #9381 against the parent commit and it already prints the expected output — #9381 was fixed by earlier work, and this PR adds regression coverage for it rather than fixing it. Worth saying so in the PR body so a future bisect isn't misled. The real behavior delta is confined to empty / partially-empty collections; that's a good thing, it just means the blast radius is narrower than the title suggests.

Findings

1. Materializing empty groups can turn a completing session into a hard failure (needs a conscious call, + release note)

Chained-collector shape with the inner collection empty for one outer item, where the per-outer consumer is any node that rejects an empty collection — grid, blend, and "at least one image" nodes are common in real workflows:

• parent: consumer runs once, session completes, result [['b.0','b.1']]
• this PR: consumer runs twice, the second time with [] → ValueError → the whole session fails

This is the intended consequence of preserving empty iterations (same family as #9349), not a defect. But it is user-visible for existing saved workflows in two ways: outputs change shape ([['b.0','b.1']] → [[], ['b.0','b.1']]), and a collector that used to yield one instance now yields N, so its downstream node runs N times instead of once. Please call this out in the PR description and in the release notes.

2. Closes #9382 overstates it — the same drop still happens one level deeper

src[a,b] → outer_iter → mid_map (empty for 'a') → mid_iter → inner_map → inner_iter → body → collect_a → per_x → top_collect

On this PR, collect_a materializes at ragged depths (0,), (1,0), (1,1), but top_collect still materializes a single group at (1,). per_x at (0,) executes and its output reaches nothing, and outer iteration 0 is still missing from the final collection — end result identical to the parent commit, but now with an orphaned node execution. Same class of bug as #9382. Either fix it here or file a follow-up and soften the issue reference.

3. The recursion branch in _get_iterator_input_iteration_paths is dead code

graph.py:543-544 fires only when an iterator's collection port is fed by another IterateInvocation. That edge cannot exist: every IterateInvocationOutput field is Any/int, so _validate_iterator_input_type rejects it with InvalidEdgeError("Iterator input must be a collection") — I confirmed by trying to build one. Injecting a raise into that branch leaves all 165 graph tests passing. The visited parameter exists only to serve this branch. Suggest dropping both.

4. A load-bearing invariant is left implicit across two functions

In _get_collect_candidate_group_keys the prepared-node contribution is truncated to group_depth, but the _get_iterator_input_iteration_paths contribution is not. When group_depth == 0 the truncation injects () into the candidate set on every call, and the fix works only because _get_collect_iteration_mapping_groups later drops any key that is a strict prefix of another. Correct today, but a future change to that prefix filter would silently re-merge every group this PR just split, with no test failing at the point of the change. Either guard the truncated update with if group_depth: (it is a no-op at 0) or leave a comment.

Two minor things in the same function: if prepared_nodes: on line 560 is redundant — an empty generator contributes nothing — and _get_ordered_prepared_nodes_for_source sorts a list whose order is immediately discarded into a set, so _get_prepared_nodes_for_source would do.

5. baseinvocation.py: correct, but the stated motivation is only half-addressed

The rewrite is semantically identical to the two asserts — issubclass on a non-class still raises TypeError, still caught by except Exception, still rewritten to the same ValueError. But if the goal is to survive python -O, the same decorator still relies on assert at lines 694-695, and invocation_output() at 822-823, plus line 619; under -O those checks are still stripped silently. Also raise TypeError with no message inside a try whose except Exception rewrites it reads oddly — hoisting the check out of the try (leaving only issubclass's own TypeError to be caught) would be clearer. And the motivating -O path has no test; test_nodes_must_return_invocation_output passed before this change too, and CI doesn't run optimized.

Attacks that found nothing

For the record, these all came out clean: If-node branch pruning above a chained collector (skipped-state filtering holds); cross-product of two independent iterators with chained collectors inside; top-level (non-nested) chains; empty source collection; all-empty fanout; mixed item/collection ports on the same collector; an extra top-level item edge feeding a per-outer collector; triple-chained collectors; and KeyError on source_prepared_mapping[...] for a not-yet-prepared source, which is unreachable because prepare() walks topological order and every ancestor of a preparable node is already prepared.

Two pre-existing bugs found while attacking

Both are byte-identical on the parent commit, so they are not caused by this PR. Filed separately so they don't get lost:

• #9484 — a session deadlocks with is_complete() == False and no error whenever an If node's branch input comes from a chain containing an Iterate node. Minimal repro is six nodes; a real queue item would hang forever.
• #9485 — two independent iterators feeding a node whose collection output is iterated again prepares the inner body twice for the off-diagonal iteration paths, so the downstream collector gets every item twice for those branches.

[lstein]

Re-reviewed at 6f2c306f78 (rebased onto 45ad17b358). Thanks for turning these around quickly — findings 3, 4 and 5 are properly addressed, and I confirmed the new -O subprocess test is load-bearing (it fails when baseinvocation.py is reverted to main).

Unfortunately the new hardening commit introduces a regression that is worse than the bug it fixes, so I have to switch to requesting changes. Details below, along with a fix direction I've already validated.

Blocker: nested collector scopes are flattened for any graph 3+ levels deep

_get_collect_iteration_group_key now returns () — instead of path[:-1] — whenever the item edge's source has no iterators of its own. But iterator_graph() cuts all in-edges to collectors, so every node downstream of a collector has no iterators by construction:

get_node_iterators('collect2') = []
get_node_iterators('collect1') = []

That's the ordinary chained-collector case, not an exotic one, so the new () branch fires on every collector-of-a-collector. The result is that all remaining iteration levels collapse into a single group.

This needs no empty collections at all to trigger. Plain collect2 → collect1 → collect0 over three nested iterators, everything fully populated:

                       main / previous head                 this commit
collect2 (0,0)         ['a00','a01']                        ['a00','a01']        (unchanged)
collect2 (0,1)         ['a10','a11']                        ['a10','a11']
collect2 (1,0)         ['b00','b01']                        ['b00','b01']
collect2 (1,1)         ['b10','b11']                        ['b10','b11']

collect1 (0,)          [['a00','a01'],['a10','a11']]        -- gone --
collect1 (1,)          [['b00','b01'],['b10','b11']]        -- gone --
collect1 ()            --                                   [['a00','a01'],['a10','a11'],
                                                             ['b00','b01'],['b10','b11']]

collect0 ()            [[['a00','a01'],['a10','a11']],      [[['a00','a01'],['a10','a11'],
                        [['b00','b01'],['b10','b11']]]        ['b00','b01'],['b10','b11']]]

The outer dimension is silently flattened away: collect1 goes from two per-outer collections to one merged collection, and collect0's nesting depth drops by one. At four levels, three of the four collectors collapse to a single () group and the two outer dimensions are both lost.

In user terms this is the "iterate models × iterate prompts × iterate seeds, with a collector per level building a grid" shape. Today it produces one grid per model; after this commit it produces a single grid containing every model's images merged together. Downstream nodes that index or zip against the expected structure get the wrong data rather than an error.

Repro (run from repo root with PYTHONPATH=.)

from typing import Any

from invokeai.app.invocations.baseinvocation import (
    BaseInvocation, BaseInvocationOutput, InvocationContext, invocation, invocation_output,
)
from invokeai.app.invocations.fields import InputField, OutputField
from invokeai.app.services.shared.graph import (
    CollectInvocation, Graph, GraphExecutionState, IterateInvocation,
)
from tests.test_nodes import create_edge, run_session_with_mock_context


@invocation_output("n3_col_out")
class ColOut(BaseInvocationOutput):
    collection: list[Any] = OutputField(default=[])


@invocation_output("n3_val_out")
class ValOut(BaseInvocationOutput):
    value: Any = OutputField(default=None)


@invocation("n3_source", version="1.0.0")
class Source(BaseInvocation):
    collection: list[Any] = InputField(default=[])

    def invoke(self, context: InvocationContext) -> ColOut:
        return ColOut(collection=self.collection)


@invocation("n3_fanout", version="1.0.0")
class Fanout(BaseInvocation):
    value: Any = InputField(default=None)

    def invoke(self, context: InvocationContext) -> ColOut:
        return ColOut(collection=[f"{self.value}{i}" for i in range(2)])


@invocation("n3_ident", version="1.0.0")
class Ident(BaseInvocation):
    value: Any = InputField(default=None)

    def invoke(self, context: InvocationContext) -> ValOut:
        return ValOut(value=self.value)


def build(levels: int) -> Graph:
    g = Graph()
    g.add_node(Source(id="src", collection=["a", "b"]))
    g.add_node(IterateInvocation(id="iter0"))
    g.add_edge(create_edge("src", "collection", "iter0", "collection"))
    for lvl in range(1, levels):
        g.add_node(Fanout(id=f"fan{lvl}"))
        g.add_node(IterateInvocation(id=f"iter{lvl}"))
        g.add_edge(create_edge(f"iter{lvl - 1}", "item", f"fan{lvl}", "value"))
        g.add_edge(create_edge(f"fan{lvl}", "collection", f"iter{lvl}", "collection"))
    g.add_node(Ident(id="body"))
    g.add_edge(create_edge(f"iter{levels - 1}", "item", "body", "value"))

    prev = ("body", "value")
    for lvl in reversed(range(levels)):
        g.add_node(CollectInvocation(id=f"collect{lvl}"))
        g.add_edge(create_edge(prev[0], prev[1], f"collect{lvl}", "item"))
        prev = (f"collect{lvl}", "collection")
    return g


for levels in (2, 3, 4):
    print(f"=== {levels} levels, fully populated ===")
    session = GraphExecutionState(graph=build(levels))
    run_session_with_mock_context(session)
    for i in range(levels):
        for p in sorted(session.source_prepared_mapping[f"collect{i}"], key=session._get_iteration_path):
            print(f"  collect{i} path={session._get_iteration_path(p)} out={session.results[p].collection}")

Two levels is unaffected (path[:-1] and () coincide when the path has length 1), which is why nothing catches this — I ran the full suite at 6f2c306f78 and got 4164 passed / 0 failed, and 167 passed on the three graph test modules. Every existing chained-collector test is two levels deep.

The underlying issue with the approach

Peeling a level is the right operation; the previous code just peeled it off the wrong thing. path[:-1] peels off this path's own depth, which misbehaves only when sibling prepared nodes have ragged depths — which is exactly the empty-branch situation from my finding #2 (collect_a at (0,) next to (1,0) and (1,1)). Falling back to () fixes the ragged case by discarding scope entirely, which also discards it for the well-formed case.

Peeling one level off the deepest sibling handles both. I tried this locally:

-    def _get_collect_iteration_group_key(self, edge: Edge) -> tuple[int, ...]:
+    def _get_collect_iteration_group_key(self, edge: Edge, sibling_depth: Optional[int] = None) -> tuple[int, ...]:
         path = self._state._get_iteration_path(edge.source.node_id)
         if edge.destination.field == ITEM_FIELD:
-            source_node_id = self._state.prepared_source_mapping[edge.source.node_id]
-            if self._get_collect_source_iterator_ids(source_node_id):
-                return path[:-1]
-            # No active iterator means the path is inherited from a collector boundary; keep it global.
-            return ()
+            depth = len(path) if sibling_depth is None else sibling_depth
+            return path[: max(depth - 1, 0)]
         return path

with the call site in _get_collect_iteration_mapping_groups passing

sibling_depth = max(
    (len(self._state._get_iteration_path(prepared_id)) for prepared_id in prepared_nodes), default=0
)

Results: the 3- and 4-level outputs above become byte-identical to main (nesting preserved), and the ragged shape from finding #2 resolves correctly rather than by flattening — top_collect materializes at (0,) → [[]] and (1,) → [[0, 1], [10, 11]], so the previously-dropped outer iteration 0 comes back and keeps its scope. 166 of 167 graph tests pass; the only failure is test_graph_chained_collectors_preserve_ragged_empty_scope, whose assertion encodes the flattened shape:

assert sorted(...) == [()]
E   assert [(0,), (1,)] == [()]

I'd argue [(0,), (1,)] is the assertion you want there — one group per outer iteration, matching what every other nesting level does. Treat the above as a validated direction rather than a finished patch; I haven't stress-tested it beyond the suite and my own battery.

Please also add a regression test at 3 levels

Whatever the fix, the coverage gap is the real lesson here: a change to core group-key logic flattened every nested collector graph in the product and the entire test suite stayed green. A fully-populated 3-level collect → collect → collect assertion would have caught it immediately.

Everything else from round 1 is resolved

• Finding 3 (dead recursion branch) — removed, along with visited. ✅
• Finding 4 (implicit prefix-filter invariant) — and group_depth guard plus a comment; the duplicated iterator-id logic is now shared via _get_collect_source_iterator_ids. ✅
• Finding 5 (-O motivation untested) — the subprocess test is real: reverting baseinvocation.py to main makes it fail. ✅ The other five asserts in that module (lines 619, 694-695, 822-823) are still -O-strippable, but they're different checks and fine as a follow-up.
• Finding 1 (empty-group behavior change) — the note is in the PR body now. Please make sure it reaches the release notes too, since it changes output shape for existing saved workflows.
• Closes #9381 — the clarification in the PR body is exactly what I was after. 👍

Once the flattening is addressed I'm happy to re-review; the rest of the PR is in good shape.