fix: localize deferred row admission

[eric-tramel]

📋 Summary

Fixes adaptive row-group admission so deferred retry work only blocks new row groups when the next row group has no independent resource/column work to expose. This keeps healthy model resources busy while preserving row/frontier guardrails and adds deferred-admission diagnostics.

🔗 Related Issue

Fixes #742

🔄 Changes

• Localize deferred retry row-group blockers by request resource, scheduler resource, custom-model group, multi-output column flow, and graph dependency.
• Track the pending admission row group directly so deferred-pressure analysis does not rescan completed row-group prefixes.
• Add stable deferred-admission diagnostics for blocker scope, counts by column/resource, candidate columns, and independent candidate columns.
• Add regression coverage for unrelated resources, same-resource blocks, dependency blocks, multi-output sibling blocks, shared scheduler resources across request domains, custom-model fallback localization, row-guard precedence, independent local branches, and async healthy-resource progress during cooldown.

🧪 Testing

• uv run --group dev pytest -q packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py -k 'adaptive_row_group or deferred_cooldown' — 13 passed
• uv run --group dev pytest -q packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py — 104 passed
• Large mock-provider experiment — 2048 row groups, one cooling resource, three healthy resources; healthy_total=300, observed_max_target=101, no deferred_tasks blockers
• make check-all-fix
• make test — config 616 passed; engine 2245 passed; interface 955 passed, 1 skipped

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Self-reviewed and review-nuke personas run; environment capped parallel reviewer threads at six, so scope reviewer was run immediately after the six-persona pass
• Architecture docs updated (N/A — internal scheduler behavior/diagnostics only)

[github-actions]

Code Review: PR #743 — fix: localize deferred row admission

Summary

Reworks adaptive row-group admission so that deferred retry work (e.g. cooling
LLM resources hit by a 429) only blocks the next row group when no independent
work could be exposed by admitting it. Previously, any deferred task forced
an "all stop" on row-group admission via the deferred_tasks reason. The new
logic computes per-row-group whether the deferred set's resource/group/dependency
footprint actually overlaps with what the next row group would dispatch.

Additionally, the helper _next_unadmitted_row_group_size() is replaced with
_next_unadmitted_row_group(), which reads from a new field
_row_group_admission_pending that _admit_row_groups maintains as it walks
the row-group list. This avoids re-scanning all completed row-group prefixes
on each adaptive tick.

Touches one source file (async_scheduler.py, +215/-36) and one test file
(+492/-1). 11 new tests cover the localization edge cases and an async smoke
test that exercises the keep-healthy-resources-busy path.

Findings

Correctness

• _admit_row_groups does not reset _row_group_admission_pending between
  the row-guard re-check and _rg_states[rg_id] = ... — that's intentional
  (the row group is still "pending" until admitted). But there's a small
  diagnostic gap: between iterations, _row_group_admission_pending is None,
  so _adaptive_row_group_block_reason() returns "no_pending_row_groups"
  even though the admit task is about to pick up the next row group. This is a
  behavioral change vs. the pre-PR _next_unadmitted_row_group_size(), which
  scanned the full row-group list on demand. In practice this is fine (the
  admit task immediately sets the field on its next iteration), but it's worth
  noting that adaptive decisions and emitted diagnostics during this micro-gap
  will now key off "no pending" rather than the next eligible row group. Not a
  bug, but a behavior delta worth a one-line comment in
  _next_unadmitted_row_group() so future readers don't re-derive this.

• Cache key for _deferred_admission_cache is order-sensitive
  (tuple(self._deferred)). The salvage path (lines ~1347, ~1453) reassigns
  self._deferred with reordered contents, so a deferred set that's logically
  identical but reordered will miss the cache. Correctness is unaffected — it
  just recomputes — but if you wanted maximum cache hit rate you could use
  frozenset(self._deferred). Tasks are frozen dataclasses so they're
  hashable. Minor nit; current behavior is correct.

• Cache key does not include _rg_states or _tracker state. The cached
  analysis is keyed only by (deferred_tuple, (row_group, row_group_size)).
  The candidate-task generation in _row_group_admission_candidate_tasks
  depends only on the graph (immutable post-init), so the cache is sound for a
  given pending row group. Confirmed by inspection — no hidden state read.

• _walk_graph memoizes per-starting-column but doesn't share suffix
  results. A → B → C: walking from A computes {B,C}, walking from B later is
  a fresh walk producing {C}. For typical column counts (dozens to low
  hundreds) this is negligible, but worth mentioning if graphs ever scale
  large. Not a blocker.

Style / project conventions

• All new code uses absolute imports, modern type syntax, and
  from __future__ import annotations is already present. ✅
• _DeferredAdmissionAnalysis is a frozen dataclass with three fields and
  no methods — appropriate for a value type used only as a return value.
• Naming is consistent with the file's existing patterns (_localized_…,
  _transitive_…, _row_group_admission_…). ✅
• The new _localized_deferred_admission_keys builds string keys with
  request_resource: / scheduler_resource: / group: namespacing — clean
  and avoids collisions across categories.

Performance

• Memoization of transitive upstream/downstream sets is the right call;
  the analysis is invoked on every adaptive tick.
• _row_group_admission_candidate_tasks walks the full topological order
  per call but is memoized via the analysis cache when the deferred set
  is unchanged. Should not be a hot-path concern.
• Diagnostics dict construction in _deferred_admission_diagnostics is
  unconditional even when nothing observes it. It's only called from
  _row_group_admission_diagnostics, which is itself called on event
  emission and from _adaptive_row_group_block_reason indirectly. Fine.

Test coverage

• 11 new tests covering: unrelated resources don't block; same-resource
  blocks; downstream dependency blocks; multi-output sibling generators
  (shared generator id) propagate the block; shared scheduler resource across
  request domains (chat vs embedding) blocks; custom-model fallback (no
  resource key, only group) localizes; row-guard precedence over deferred;
  independent local branches stay admittable; async healthy-resource progress
  during cooldown.
• The async test
  test_scheduler_adaptive_row_group_admission_keeps_healthy_resource_exposed_during_deferred_cooldown
  is the most realistic — it actually runs the scheduler with a permanently
  cooling generator and asserts that the healthy generator makes ≥16 calls
  before a 5s timeout. Good integration-style coverage. The cleanup in the
  finally block (cancel pending, gather, then re-cancel run_task with
  swallow) is defensively written and looks correct.
• Test helpers _build_adaptive_model_resource_scheduler and
  _set_next_row_group_pending appropriately reduce boilerplate. The
  latter does poke at private state (_row_group_admission_pending), which
  is the natural way to exercise _deferred_admission_analysis in
  isolation — acceptable for unit tests on internal scheduler invariants.
• Tests directly mutate scheduler._deferred, scheduler._deferred_errors,
  scheduler._rg_states — these are intentional white-box probes. The
  alternative (running the full scheduler to drive a 429 into the deferred
  list) would be much heavier; the targeted approach is reasonable and
  matches the existing test style in this file.

Risks

• Behavioral change for callers that observed block_reason == "deferred_tasks" as a global signal. Any external diagnostics consumer
  that relied on that reason firing the moment any task was deferred will
  now see it fire less often. The new diagnostics under deferred_admission
  expose the richer state, but downstream dashboards or tests outside this PR
  may need to be aware.
• _row_group_admission_pending is now load-bearing for the adaptive block
  reason. If a future change adds another caller that admits row groups
  outside _admit_row_groups, that path will need to maintain this field too.
  A short docstring on the field would help guard against this.
• Backward compatibility with the existing scheduler event schema: the
  new deferred_admission key inside row_group_admission_diagnostics is
  additive — existing consumers won't break, new consumers can opt in.

Suggestions (non-blocking)

1. Add a one-line docstring or comment on _row_group_admission_pending
   explaining that _admit_row_groups is its sole writer and what None
   means (no row group currently being admitted).
2. Consider a brief module-level docstring or comment near
   _deferred_admission_analysis explaining the three exclusion rules
   (resource/group key match, transitive dependency on a blocked column,
   transitive reach into a blocked column when not resource-scoped). The
   logic is correct but takes a few reads to follow.
3. The _localized_deferred_admission_keys function is named "localized" but
   actually returns the set of keys used to localize — naming nit. Maybe
   _admission_locality_keys or _admission_scope_keys. Subjective; skip if
   it would churn callers.

Structural Impact

No pre-computed structural impact analysis was available
(/tmp/structural-impact-743.md not present). Manual assessment:

• Blast radius: Limited to data_designer.engine.dataset_builders.async_scheduler.
  No public API changes; only adds private methods, one private field, and one
  internal frozen dataclass (_DeferredAdmissionAnalysis). The pre-existing
  helper _next_unadmitted_row_group_size is replaced with
  _next_unadmitted_row_group, but both are private (_-prefixed). Tests in
  the same package are the only direct consumers.
• Import direction: Imports remain within data_designer.engine
  (scheduling.resources, scheduling.task_admission). No upward imports
  into interface or sideways to other packages. ✅
• Cross-package dependencies: None added.

Verdict

Approve with minor suggestions. The change is well-scoped, well-tested,
and addresses a real fairness problem (deferred work on one cooling resource
shouldn't stall healthy resources). Correctness analysis holds: the cache
key is sound, the dependency walk is correct, and the row-guard precedence
is preserved (and explicitly tested). The new _row_group_admission_pending
field tightens the diagnostic semantics in a subtle way that's worth a
one-line comment, but doesn't change correctness. No blockers.

[greptile-apps]

Greptile Summary

This PR updates adaptive row-group admission so deferred retries block less broadly. The main changes are:

• Localized deferred retry blockers by request resource, scheduler resource, custom group, multi-output flow, and graph dependency.
• Added pending-row-group tracking for deferred admission checks.
• Added deferred-admission diagnostics for blocker scope and candidate columns.
• Added scheduler tests for resource isolation, dependency blocking, row guards, and cooldown progress.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py	Adds localized deferred-admission analysis, pending-row-group tracking, and diagnostics for adaptive row-group admission.
packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py	Adds tests for localized deferred retry behavior across resources, dependencies, multi-output flows, row guards, and cooldown progress.

_{Reviews (7): Last reviewed commit: "Merge branch 'main' into etramel/fix/742..." | Re-trigger Greptile}

[andreatnvidia]

Reviewed #743. No blocking findings.

Validated the GitHub/local diff, ran graphify, ruff on changed Python files, and the full changed scheduler test file (104 passed). Claude cross-review also found no correctness blockers. Only note is to preserve the localized deferred-admission behavior when rebasing with the overlapping scheduler PRs.

[eric-tramel]

Updated to resolve merge conflicts.

[nabinchha]

Thanks for putting this together, @eric-tramel — this is a clean, well-scoped fix with genuinely thorough test coverage.

Summary

This localizes deferred-retry row-group admission so a cooling resource no longer stalls all new row groups: admission is only blocked when the next row group exposes no independent resource/column work. It also tracks the pending row group directly (instead of rescanning completed prefixes) and adds structured deferred_admission diagnostics. The implementation matches the stated intent in the PR description, and the regression suite maps cleanly onto the enumerated scenarios.

I validated locally: ruff check / ruff format --check clean on both files, the targeted subset (-k 'adaptive_row_group or deferred_cooldown') passes 17/17, and the full test_async_scheduler.py passes 108/108.

Findings

Suggestions — Take it or leave it

async_scheduler.py:1115-1148 — Redundant _schedulable_task recomputation on the diagnostics path

• What: _deferred_admission_diagnostics() builds deferred_items = tuple(self._schedulable_task(task) ...) and then calls _deferred_admission_analysis(*next_row_group), which rebuilds the exact same deferred_items internally (it isn't covered by the _deferred_admission_cache short-circuit unless the cache key already matches).
• Why: _row_group_admission_diagnostics() runs on every admit/block event, so each emission re-derives the schedulable form of every deferred task twice. It's diagnostics-only and the deferred set is usually small, so impact is minor — but it's easy to avoid.
• Suggestion: Nothing required; if you want to trim it, the analysis already returns candidate_columns/independent_candidate_columns, so the diagnostics method only really needs the resource/scheduler counts from deferred_items. Leaving as-is is also perfectly reasonable.

async_scheduler.py:1200-1206 — Transient no_pending_row_groups during _dispatch_seeds

• What: _row_group_admission_pending is cleared at line 1200 before the await self._dispatch_seeds(...) at 1206. If the main loop's _maybe_update_adaptive_row_group_target() interleaves during that await, _adaptive_row_group_block_reason() returns "no_pending_row_groups", which resets _row_group_admission_pressure_ticks to 0.
• Why: This can very slightly delay soft-target ramp-up right after a row group is admitted (the pressure counter restarts). It's benign — the next tick re-establishes pressure — and arguably correct since admission just progressed, but worth being aware of since the old scan-based _next_unadmitted_row_group_size() didn't have this transient.
• Suggestion: If ramp-up latency ever looks sluggish in traces, consider clearing _row_group_admission_pending only after _dispatch_seeds (or after the emit block). Not worth changing preemptively.

architecture/dataset-builders.md:29 — Optional doc note

• What: The PR marks architecture docs N/A, which is fair — line 108's "additive ramp-up" description stays accurate. The deferred-retry paragraph at line 29 still reads as "stay deferred and retry after cooldown."
• Why: The new localized admission behavior (healthy resources keep flowing during a peer's cooldown) is a nice enough property that a future reader of the architecture doc might benefit from one sentence.
• Suggestion: Optional — a half-sentence at line 29 noting deferred pressure is scoped to the affected resource/column flow. Fine to skip or defer.

What Looks Good

• The blocks = bool(deferred_items) and not independent_candidates heuristic is well-factored. Splitting resource-lane sharing (keys & deferred_keys), upstream dependency (_task_depends_on_any), and downstream reachability (_task_reaches_any, gated on resource-scoped candidates) into small, individually-named predicates makes the localization logic readable and each branch is exercised by a dedicated test.
• Caching is correctly keyed. _deferred_admission_cache keys on (tuple(self._deferred), (row_group, row_group_size)), and the transitive upstream/downstream caches key on column — all derived from immutable graph/deferred state, so no stale-cache hazard. The single event loop also means the _row_group_admission_pending read/write pattern is race-free.
• Test coverage is excellent and scenario-driven. Unrelated-resource, same-resource, dependency, multi-output sibling, cross-domain shared scheduler resource, custom-model fallback, row-guard precedence, independent local branch, and the async healthy-resource-during-cooldown integration test each isolate one behavior, and the generation_kind hook on the mock generator is a tidy way to exercise the cross-domain case.

Verdict

Ship it (with nits) — no blocking findings. Everything above is optional; the fix is correct, well-tested, and preserves the existing admission guardrails (row guard now correctly precedes the deferred check, as the new precedence test confirms).

---

This review was generated by an AI assistant.

[nabinchha]

Ship it (with nits) — no blocking findings. See my review comment for optional suggestions. Verified locally: ruff clean, 108/108 scheduler tests pass.