perf(loops): make For-loop scheduling linear in the item count

[lstein]

This PR aims to fix the sporadic pytest failures on test_loop_scheduler_overhead_is_linear.

Scheduling a For loop was quadratic in its collection size. This showed up as test_loop_scheduler_overhead_is_linear failing on py3.11 windows-cpu in #221 — a frontend-only PR that touches no Python. That test was barely passing on merit: on a quiet Linux machine the 300→1200 ratio it measures came out anywhere between 1.30 and 1.84 against its own 1.5 threshold.

The two quadratics

The loop collection was deep-copied per iteration. _ExecutionMaterializer.create_for_iteration did new_node.collection = copydeep(collection) for every iteration, so a 1200-item loop made ~1.44M deepcopy calls (n²/2), almost all of them re-copying items the loop had already passed.

The collection now moves from one prepared For node to the next: the caller drops its own reference before handing the list over, so exactly one node owns it at a time — which is already the design, since the previous node's collection was being cleared immediately afterwards anyway. The assignment uses object.__setattr__ because pydantic's validate_assignment would otherwise re-validate and rebuild the list, making the move O(n) again; that is the same idiom _prepare_iterate_inputs already uses for this field.

Source-node completion rescanned every prepared node. _mark_source_node_complete and mark_exec_node_skipped ran all(n in executed for n in prepared_ids) on every completion, and prepared_ids grows to n for a loop source. GraphExecutionState now keeps a per-source count of prepared nodes that have not executed, built lazily and maintained at register / execute / drop, so the check is a dict lookup. The invariant is counts[s] == |{e in source_prepared_mapping[s] : e not in executed}|.

Effect

Per-node scheduling cost for a trivial for loop, GC disabled so the numbers reflect scheduling rather than gen-2 collections:

	n=300	n=1200	n=2400	n=4800	n=9600
before	399	733	1169	—	—
after	344	338	338	335	337

Flat across a 32× range. Profiling confirms every hot frame's call count now scales exactly 1.00 per node. iterate was already linear and is unchanged at ~130µs/node. With the GC on there is still growth past n≈4800, but that is gen-2 collection scanning the accumulating execution graph, which is inherent to retaining executed nodes and results.

Tests

• test_loop_scheduler_overhead_is_linear measured with time.process_time. Windows samples process CPU on the ~15.6ms scheduler tick, which quantized the 300-item baseline down to a single tick (every number in the CI failure is an exact multiple of 1/64s) and made the ratio meaningless. It now measures wall clock and takes the fastest of several samples, since a busy host only ever makes a run slower. The ratio it now sees is 0.99 for both for and iterate, against a 1.5 threshold.
• test_graph_for_iteration_does_not_deep_copy_collection_twice pinned the old behaviour of two copies per iteration. It is replaced by test_graph_for_iteration_does_not_copy_collection_per_iteration, which drives a full loop and asserts one copy per iteration rather than one per collection entry per iteration — the thing that actually distinguishes linear from quadratic.
• test_graph_for_scheduling_keeps_prepared_completion_counts_consistent asserts the new counts agree with a full rescan before and after every scheduling step.

Verification

Full suite: 6630 passed, 0 failed. (Three tests/backend/pid CUDA bf16 tolerance tests fail identically on unmodified main and do not run on CI's CPU runners.)

An adversarial review of the diff attacked both changes — every mutation site of executed and the prepared-node mappings repo-wide, cache staleness across model_copy/model_validate/deepcopy, dropped-then-executed prepared ids, nested For and For-in-Iterate, error paths, and JSON round-trips — with a differential fuzz harness (102 graph/round-trip combinations) that produced byte-identical traces against HEAD. It found no correctness break, and three real issues that are fixed here:

• The handover was still O(n) because validate_assignment rebuilt the list — hence object.__setattr__.
• Every iteration now observes the same element objects where each previously got a deep copy. This is safe because loop bodies read the item through the edge machinery, which hands out a copydeep, and nothing in the scheduler mutates a collection entry in place — but it is now documented rather than accidental.
• _reset_runtime_caches() reset every other derived private attr but not the new counts. Inert today (its only caller runs right after pydantic nulls private attrs), but a future caller would have got a count that never reaches 0, so the source would never be marked complete and the session would never finish.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Jcp5TPVjPL469HhuWYT45m

[lstein]

@JPPhoto Were you working on the collection For-loop handling? This is a small change that converts quadratic behavior into O(n). Could you take a look at it?

[JPPhoto]

@JPPhoto Were you working on the collection For-loop handling? This is a small change that converts quadratic behavior into O(n). Could you take a look at it?

I'm working on refactoring everything, hold off on merging this and maybe I can just incorporate that into my refactor if timing works out - otherwise I'll do a proper review and figure out how this fits in.

[lstein]

Sure. This PR is directed at a flaky pytest that times out on Windows systems. About 1 in 4 PRs initially fail CI because of this.