fix: inline retry for rate-limited tasks in async scheduler

[andreatgretel]

📋 Summary

Fixes the async engine producing 0 records under sustained API rate limiting. Rate-limited tasks were deferred to salvage, which retried them in bursts through the same congested throttle - climbing the error rate, tripping early shutdown, and dropping all records before checkpointing. This adds a per-task cooldown queue that releases scheduler permits while a task waits, paces retries through the AIMD signal, and falls through to salvage only as a last resort. Also addresses related throttle-manager edge cases surfaced during review (Retry-After clipping, mixed-failure streak handling).

🔗 Related Issue

Fixes #575

🔄 Changes

🐛 Fixed

• Cooldown queue for rate-limited tasks (async_scheduler.py#L315-L351, #L820-L863): On ModelRateLimitError, the worker pushes (Task, wakeup_at) onto _cooldown_pending and exits cleanly, releasing both submission and LLM-wait permits. The dispatch loop drains entries whose wakeup has passed back into the frontier; they re-acquire permits via the standard path and naturally pace through the throttle's acquire_async. Per-task budget is DEFAULT_RATE_LIMIT_RETRIES = 5; exhaustion falls through to the existing salvage path.
• Retry-After honored as a floor (throttle_manager.py#L266-L289): release_rate_limited no longer clips an explicit Retry-After header below max_cooldown_seconds. The cap only constrains synthetic exponential growth; when the server provides a hint, we never retry earlier than it asked.
• ModelRateLimitError.retry_after (errors.py#L67-L80): _raise_from_provider_error now propagates the provider's Retry-After value through the exception, so the cooldown queue can wait the server-mandated duration on the throttled-client path.
• consecutive_429s resets on non-rate-limit failure (throttle_manager.py#L341-L356): a sequence like 429 → 500 → 429 is treated as two fresh cascades instead of sustained throttling.
• Stale cooldown entries are pruned: drain drops entries whose row group has been checkpointed or whose row was dropped by a sibling column's terminal failure, so the all_done gate isn't held hostage by a cooldown nothing is waiting on.
• from_scratch tasks route to salvage: seed tasks are dispatched by _dispatch_seeds, not the frontier, so the cooldown queue's drain wouldn't bring them back. They skip the cooldown queue and go straight to salvage, which already knows how to re-dispatch seeds.

🔧 Changed

• Exponential cooldown backoff at floor (throttle_manager.py): When concurrency is at minimum (1) and consecutive 429s keep arriving, cooldown duration grows exponentially (base * backoff_factor ^ (consecutive_429s - 1), capped at max_cooldown_seconds).
• Saturation detection (telemetry-only): After saturation_threshold consecutive 429s at minimum concurrency, the domain is marked as saturated with a warning log. is_saturated() is a public hook reserved for future scoping; no current behavioral effect on the client.
• New ThrottleConfig fields (run_config.py): cooldown_backoff_factor (default 2.0), max_cooldown_seconds (default 30.0), saturation_threshold (default 5). Sensible defaults; no user action needed.

Why a cooldown queue, not inline retry?

The first iteration of this PR put the retry loop inline in the scheduler worker. Adversarial review surfaced two real regressions: (1) a worker holding the LLM-wait permit across every cooldown attempt could starve unrelated row groups under multi-model load, and (2) a ModelRateLimitError raised by a custom generator outside ThrottledModelClient (no acquire_async to enforce cooldown) became a tight retry loop with no real backoff. Both trace to the same architectural choice: retry-while-holding-permits, with cooldown enforcement implicit.

Replacing it with a per-task cooldown queue:

• Workers exit immediately on 429, releasing all permits — no cross-model starvation.
• Cooldown wait is explicit (a wakeup time on the queue), so unthrottled 429s are paced just like throttled ones.
• The structural difference from "salvage but earlier" is that salvage retries deferred tasks in bursts (the original failure mode), while the cooldown queue dispatches tasks as their own wakeup expires — natural pacing without re-saturating the throttle.
• Salvage is preserved as the third-line backstop for non-rate-limit retryables and for tasks that exhaust the cooldown budget.

Behavior comparison

Before (stalls): Rate limit → task deferred to salvage → salvage retries through congested throttle → more rate limits → cooldowns compound → stalls 15+ min → early shutdown → 0 records.

After (cooldown queue): Rate limit → worker exits, releases permits → task waits in cooldown queue → drain re-dispatches when wakeup expires → throttle's acquire_async paces remainder → task usually succeeds within 1-4 retries → no salvage needed → no error rate increase → no early shutdown.

A/B test results (Anonymizer biographies pipeline, gpt-oss-120b at max_parallel=4, early shutdown enabled)

	With fix	Without fix
Detection (225 tasks)	225 ok, 0 failed	219 ok, 1 failed, 5 skipped
Salvage rounds	0	3 (13 tasks deferred)
Rewrite (325 tasks)	325 ok, 0 failed	240 ok, early shutdown triggered
Final result	✅ All stages completed	❌ 0 records produced

The A/B comparison was run on the original inline-retry version. The cooldown queue is equivalent in the common case (1-4 retries succeed) and strictly better under sustained 429s and on multi-model pipelines. A user-reported reproduction of the same #575 failure mode (Anonymizer notebook 04, rewrite phase: early shutdown at 345s, 156 ok, scheduler exited with 1 unfinished row group → 0 records → DataDesignerGenerationError) confirms the original failure mode independent of our test setup.

🧪 Testing

• make test passes (1909 engine tests; 9 new tests for the cooldown queue + throttle regressions)
• Unit tests added covering: cooldown-queue recovery, exhaustion fallback to salvage, permit-release-during-cooldown (P1), unthrottled-429 pacing (P2), stale-entry pruning, from_scratch routing, Retry-After honored
• Throttle manager regression tests: Retry-After floor, default cooldown still capped, streak reset on release_failure, saturation cleared on non-rate-limit failure
• A/B tested against live build.nvidia.com endpoint with rate limiting active
• E2E tests — N/A, covered by live A/B testing

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs updated — N/A, no user-facing API changes

🔍 Attention Areas

⚠️ Reviewers: Please pay special attention to the following:

• Cooldown queue + drain logic: core architectural change. The drain handles three states (ready-to-redispatch, stale, still-waiting) and the _main_dispatch_loop / _drain_frontier exit conditions and wake bounds are timeout-aware via _wait_for_wake_or_cooldown.
• from_scratch carve-out: seed/from-scratch tasks bypass the cooldown queue because the frontier wouldn't re-emit them. Worth confirming this is the right tradeoff for plugin-provided from-scratch generators that hit rate limits.
• ModelRateLimitError constructor: now accepts a kw-only retry_after argument. Existing call sites that construct it positionally are unaffected, but worth a sanity check.
• Architectural journey: this PR was rewritten mid-review. Original commits shipped inline retry; subsequent commits addressed adversarial review findings (Retry-After, mixed-failure streaks) and then refactored to the cooldown queue. Git history has the full progression: f7a17cb6 (adversarial fixes) → f390e4ad (cooldown queue refactor) → beed73d0 (Codex follow-up).
• The second rate-limit failure path (custom generators like chunked_validate in Anonymizer with internal retry pools that convert rate limits to non-retryable errors before they reach the scheduler) is out of scope for this PR. That path is unchanged.

---

Description updated with AI

[andreatgretel]

Worth considering: this currently catches every ModelRateLimitError regardless of AIMD state. We could scope it to only fire when the throttle is at the concurrency floor (is_saturated() or current_limit == DEFAULT_MIN_LIMIT), and let salvage handle above-floor cases.

Argument for scoping it down: Above floor, AIMD has room to reduce concurrency. Tasks deferred to salvage will retry with lower concurrency and should recover naturally - that path isn't actually broken. The catastrophic stall (#575) only manifests at the floor where AIMD can't reduce further.

Argument for current behavior: Even above floor, salvage retries deferred tasks in a burst, which can re-saturate the throttle and force AIMD to reduce again from scratch. Inline retry naturally paces tasks through cooldowns at any concurrency.

Our A/B test (without fix, disable_early_shutdown=True) showed salvage actually recovered all 13 deferred tasks across 3 rounds — so above-floor salvage isn't catastrophically broken. The 0-records failure required floor + early shutdown enabled.

Scoped version would look like:

except ModelRateLimitError:
    if self._early_shutdown or _rl_attempt == DEFAULT_RATE_LIMIT_RETRIES:
        raise
    # Only retry inline when AIMD is stuck at floor; let salvage handle
    # above-floor cases where AIMD can still reduce concurrency.
    if not self._tm.is_saturated(provider, model_id, domain):
        raise
    ...

[andreatgretel]

Closing in favor of a focused fix for #575 (see issue comment #575 (comment)). Diagnostic A/Bs showed rate-limit handling doesn't change the user-visible failure — the actual trigger is the async early-shutdown gate firing on clustered ModelTimeoutErrors. New PR will scope down to: (1) exclude retryable errors from the gate, (2) WARN log on degraded provider performance, (3) partial-RG checkpoint as safety net. ~150 prod lines vs. ~921 here.