Iteration scope is not preserved across chained collectors (collector→iterator→collector merges outer iterations)

[lstein]

Summary

When a collector's output is iterated again inside an enclosing iteration — i.e. a collector → iterator → collector chain nested under an outer iterator — the downstream collector materializes as a single merged group instead of one group per outer iteration. Items from all outer iterations are mixed into one collection.

This was found while verifying #9349, but it is not caused by that PR: the behavior is byte-identical on main (verified empirically at af08179a90). #9349 fixed per-outer grouping for a collector directly under an unbroken iterator scope; this issue is about scope not surviving a mid-chain collector boundary.

Repro

Graph shape (all under outer_iter over ["a", "b"]):

src → outer_iter → mapper → inner_iter → body → collect_A → downstream_iter → downstream_body → collect_B → per_outer → outer_collect

Script (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 (
    PromptCollectionTestInvocation,
    PromptCollectionTestInvocationOutput,
    PromptTestInvocation,
    create_edge,
    run_session_with_mock_context,
)


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


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

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


@invocation("issue_fanout", version="1.0.0")
class FanoutTestInvocation(BaseInvocation):
    prompt: str = InputField(default="")

    def invoke(self, context: InvocationContext) -> PromptCollectionTestInvocationOutput:
        return PromptCollectionTestInvocationOutput(collection=[f"{self.prompt}.0", f"{self.prompt}.1"])


g = Graph()
for n in [
    PromptCollectionTestInvocation(id="src", collection=["a", "b"]),
    IterateInvocation(id="outer_iter"),
    FanoutTestInvocation(id="mapper"),
    IterateInvocation(id="inner_iter"),
    PromptTestInvocation(id="body"),
    CollectInvocation(id="collect_A"),
    IterateInvocation(id="downstream_iter"),
    PromptTestInvocation(id="downstream_body"),
    CollectInvocation(id="collect_B"),
    AnyCollectionPassthrough(id="per_outer"),
    CollectInvocation(id="outer_collect"),
]:
    g.add_node(n)
for e in [
    ("src", "collection", "outer_iter", "collection"),
    ("outer_iter", "item", "mapper", "prompt"),
    ("mapper", "collection", "inner_iter", "collection"),
    ("inner_iter", "item", "body", "prompt"),
    ("body", "prompt", "collect_A", "item"),
    ("collect_A", "collection", "downstream_iter", "collection"),
    ("downstream_iter", "item", "downstream_body", "prompt"),
    ("downstream_body", "prompt", "collect_B", "item"),
    ("collect_B", "collection", "per_outer", "collection"),
    ("per_outer", "collection", "outer_collect", "item"),
]:
    g.add_edge(create_edge(*e))

session = GraphExecutionState(graph=g)
run_session_with_mock_context(session)
for src in ["collect_B", "outer_collect"]:
    for p in sorted(session.source_prepared_mapping.get(src, set())):
        print(src, session._get_iteration_path(p), session.results[p].collection)

Observed (identical on main and on #9349's branch)

collect_B () ['a.0', 'b.0', 'a.1', 'b.1']     # one merged group across both outer iterations
outer_collect () [['a.0', 'b.0', 'a.1', 'b.1']]

Expected

collect_B (0,) ['a.0', 'a.1']
collect_B (1,) ['b.0', 'b.1']
outer_collect () [['a.0', 'a.1'], ['b.0', 'b.1']]

Analysis

get_node_iterators / iterator_graph deliberately cut edges into collectors, so from collect_B's perspective outer_iter is not an enclosing iterator — the iteration scope is considered "closed" at collect_A. That is the right model for collect_A itself, but the re-iteration by downstream_iter happens per outer iteration (one downstream_iter fan-out per collect_A instance), so collect_B should be grouped per outer iteration too.

Since #9349, execution nodes downstream of a collector carry an explicit iteration path inherited through the collector (e.g. downstream_iter copies get paths like (0, 0)), so the information needed to group collect_B per-outer now exists on the prepared nodes — _get_collect_candidate_group_keys / _get_collect_iteration_group_key just don't use it because they consult source-graph iterator ancestry, which stops at the collector boundary.

The empty-collection variant of this same shape (inner collections all empty) stalls forever on main and, after #9349, completes with a single merged empty group — consistent with the merged behavior above, but worth keeping in mind when fixing this: both the empty and non-empty variants should come out per-outer.

Related: #9343 (nested collector iteration scope), #9349 (empty collector groups). Also see #9381 for a second scope-related edge case found in the same verification pass.

🤖 Generated with Claude Code