fix: unblock adaptive scheduler bootstrap

[eric-tramel]

📋 Summary

Fix the adaptive row-group scheduler bootstrap condition so queued LLM demand cannot dead-end admission before any endpoint work has started. The queued-demand block now only applies after an LLM wait lease exists, and row-group admission diagnostics include resource demand/availability state for future no-progress reports.

🔗 Related Issue

Fixes #740

🔄 Changes

• Allow adaptive row-group admission to seed model work when queued LLM demand exists but no LLM wait leases have started yet.
• Preserve queued LLM demand blocking once LLM wait capacity is actively leased.
• Add structured row-group admission diagnostics for queued resource demand, in-flight tasks, limits, leases, and availability.
• Extend scheduler tests for zero-lease bootstrap, equality edge cases, post-bootstrap blocking, target growth, and diagnostics.

🧪 Testing

• make test passes
• Unit tests added/updated
• E2E tests added/updated (not applicable)

Additional validation:

• make check-all-fix
• uv run pytest packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py -q
• Large mock-provider experiments with 1,048,576-row and 8,388,608-row logical bootstrap scenarios plus a complete 256-row x 8-model-column run.

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs updated (not applicable; scheduler behavior fix only)

[greptile-apps]

Greptile Summary

This PR fixes a deadlock in the adaptive row-group scheduler's bootstrap phase: when no LLM wait leases were active yet, queued LLM demand could block all new row-group admissions before any endpoint work had even started. The fix gates the queued-demand block on llm_leased > 0, allowing the first batch of row-groups through freely.

• The core logic change in _adaptive_row_group_block_reason replaces the old llm_available <= queued_llm and queued_total > 0 guard with llm_leased > 0 and llm_available <= queued_llm, correctly separating the cold-start (no leases) case from the actively-running (leases exist) case.
• _row_group_admission_diagnostics is extended with full resource-state dicts alongside the existing per-resource scalar fields, giving richer context for no-progress reports.
• New tests cover the zero-lease bootstrap path, equality edge cases, post-bootstrap blocking, target growth, and the updated diagnostics shape.

Confidence Score: 5/5

Safe to merge — the change is small and well-targeted with no regressions introduced.

The logic change is minimal: one new variable read and one guard condition replacing another. The old queued_total > 0 guard is correctly dropped because queued_llm > 0 is already implied by the inequality when llm_available > 0, and the llm_wait_saturated branch already handles the llm_available <= 0 case. The new llm_leased > 0 gate precisely captures the intended bootstrap window. Six new tests cover the zero-lease bootstrap, equality boundary, post-bootstrap blocking, target growth, and the full diagnostics shape.

No files require special attention.

Important Files Changed

Filename	Overview
packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py	Bootstrap fix: gates queued-LLM-demand blocking on llm_leased > 0; adds full resource-state dicts to admission diagnostics alongside pre-existing scalar fields.
packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py	Refactors existing stub setup into a shared helper; adds six new tests covering zero-lease bootstrap, post-bootstrap blocking, target growth/stability, and diagnostics shape.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[_adaptive_row_group_block_reason] --> B{deferred tasks?}
    B -->|yes| C[return deferred_tasks]
    B -->|no| D{next unadmitted RG exists?}
    D -->|no| E[return no_pending_row_groups]
    D -->|yes| F{row guard allows size?}
    F -->|no| G[return max_admitted_rows]
    F -->|yes| H{queued_total >= queue_guard?}
    H -->|yes| I[return queued_task_guardrail]
    H -->|no| J{llm_limit <= 0?}
    J -->|yes| K[return no_llm_wait_resource]
    J -->|no| L{llm_available <= 0?}
    L -->|yes| M[return llm_wait_saturated]
    L -->|no| N{llm_leased > 0 AND llm_available <= queued_llm?}
    N -->|yes| O[return queued_llm_demand]
    N -->|no| P[return None - admit row group]

Loading

_{Reviews (1): Last reviewed commit: "fix: unblock adaptive scheduler bootstra..." | Re-trigger Greptile}

[github-actions]

Review: PR #744 — fix: unblock adaptive scheduler bootstrap

Summary

Fixes a bootstrap deadlock in the adaptive row-group admission path in async_scheduler.py. Previously, _adaptive_row_group_block_reason() returned "queued_llm_demand" whenever llm_available <= queued_llm and queued_total > 0, which could be true before any LLM wait lease had ever been acquired. With no admitted row group there is nothing to consume queued demand, so admission stayed blocked forever — confirmed at scale on the 1M / 8M-row mock runs called out in the PR description.

The fix gates the queued-LLM-demand block on llm_leased > 0, so the saturation-style backpressure only engages once at least one LLM wait lease exists. The diagnostics dict emitted on row_group_admission_blocked / ..._target_changed is also broadened to include full resource limits/leases/availability and queued demand by resource, which is purely additive.

Scope: 1 production file (5 added / 1 changed lines in scheduler logic), 1 test file (~125 lines of new tests + a shared helper).

Findings

Correctness — fix is sound

packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py:946-969 — the new condition if llm_leased > 0 and llm_available <= queued_llm is the right shape:

• The old queued_total > 0 clause is dropped, but that's safe: if queued_llm == 0, then llm_available <= queued_llm requires llm_available <= 0, which is already caught by the earlier llm_wait_saturated branch (llm_available <= 0). So losing the queued_total > 0 check cannot newly mis-block.
• Bootstrap path: with llm_leased == 0, the function returns None, allowing _admit_row_groups to acquire a slot, which eventually drives a lease.
• Post-bootstrap: once llm_leased > 0, the original demand-saturation backpressure is preserved.

Observability — additive, low-risk

async_scheduler.py:979-1002 — diagnostics now include queued_demand_by_resource, in_flight_tasks, resource_limits, leased_resources, resources_available as dicts. The legacy scalar fields (llm_wait_limit, llm_wait_leased, llm_wait_available, queued_llm_wait_demand) are kept, so any downstream consumer reading the previous keys is unaffected. The redundancy is mild but acceptable for a hot-path diagnostic that is only emitted on block / target-change events.

Test coverage — strong

packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py:3796-3820 introduces _stub_row_group_admission_resource_views, which cleanly DRYs up the _fair_queue / _task_admission stubbing and is reused across the existing ..._prefers_llm_saturation test plus four new cases:

• ..._allows_zero_llm_lease_bootstrap — the regression scenario.
• ..._blocks_queued_llm_demand_after_bootstrap — confirms the gate re-engages once llm_leased >= 1.
• ..._target_grows_for_zero_llm_lease_bootstrap — exercises the integration with _maybe_update_adaptive_row_group_target and verifies the row_group_admission_target_changed event fires.
• ..._target_stays_blocked_after_llm_lease_bootstrap — the converse, asserting _row_group_admission_blocked_reasons["queued_llm_demand"] is incremented and no target growth occurs.
• ..._diagnostics_include_resource_state_for_events — pins the new diagnostics fields and uses an ExplodingRequestPressureProvider to also assert request pressure is not captured here. Nice negative assertion.

The pairing of the equality-edge case (llm_available == queued_llm) plus a target-growth and a target-blocked test gives good confidence that the corner is fully covered.

Style / conventions

• Pure stdlib + SimpleNamespace-based stubs match the surrounding test patterns in the file.
• Type annotations on the new helper are complete; from __future__ import annotations is already in place.
• Follows the project's "declare, don't orchestrate" boundary — the change stays inside the scheduler block-reason method, no leakage into config or interface.
• No new heavy imports, no new dependencies, no relative imports.

Minor suggestions (non-blocking)

1. Consider a one-line comment near if llm_leased > 0 and llm_available <= queued_llm: documenting that the llm_leased > 0 gate exists specifically to avoid a bootstrap deadlock when queued LLM demand is non-zero before any lease has been acquired. The Why is non-obvious to a future reader, and the PR title — not the code — currently carries it. (Project style guidance favors comments only for non-obvious why; this qualifies.)
2. The _stub_row_group_admission_resource_views helper hardcodes queued_by_group={} and only emits the llm_wait key when non-zero. That's fine for current consumers, but if a future test needs to stub another resource the helper will need a slight generalization. Not worth doing pre-emptively.
3. The _row_group_admission_diagnostics now carries both per-resource dicts and the legacy scalar llm_wait_* fields. If you want to slim this in a follow-up, consider deprecating the scalars once any downstream telemetry consumers are migrated. Not a blocker.

Risks

• Blast radius: scheduler-internal, gated by adaptive_row_group_admission=True. Default config behavior is unchanged for non-adaptive paths.
• Backward compatibility: diagnostics keys are additive; no removals.
• Concurrency: the change reads the same task/queue views the surrounding code already reads under the same conditions; no new synchronization concerns.
• Performance: one extra dict lookup per _adaptive_row_group_block_reason call, plus the broadened diagnostics dict that's only built on infrequent events. Negligible.

Verdict

LGTM. Small, well-scoped fix with a clear failure mode, an explanation that lines up with the code, and a thorough set of unit tests that pin both the bootstrap regression and the post-bootstrap behavior. The only nit worth acting on is a brief inline comment explaining why the llm_leased > 0 gate exists; everything else is optional polish.

[andreatnvidia]

Reviewed #744. No blocking findings.

Validated the GitHub/local diff, ran graphify, ruff on changed Python files, and the full changed scheduler test file (100 passed). Claude cross-review also found no actionable bugs or test gaps.