fix: avoid duplicate policy updates during scheduler cancellation

[samsja]

Summary

generate_batch() cancels the previous step's update_policy_task before starting the next step. If that cancellation lands after update_weights() has started but before ckpt_step is committed, the next maybe_update_policy() can replay the same checkpoint update.

That duplicates both the weight sync and the off-policy bookkeeping for a single checkpoint step.

This PR fixes the race by:

• moving the checkpoint application into a shared in-flight task
• reusing that task across callers instead of starting duplicate updates
• shielding the in-flight task from outer-task cancellation
• committing ckpt_step immediately after the weight update completes
• explicitly cancelling the in-flight update task during shutdown so stop() does not hang

Testing

• uv run pytest tests/unit/orchestrator/test_scheduler.py -q

---

Note

Medium Risk
Touches async concurrency/cancellation paths for policy checkpoint updates; mistakes could cause missed updates or scheduler stalls, though changes are localized and unit-tested.

Overview
Prevents duplicate checkpoint applications in Scheduler.maybe_update_policy() by introducing a shared, lock-guarded inflight_policy_update_task and awaiting it via asyncio.shield() so outer cancellation doesn’t restart the same weight sync/off-policy update.

Refactors policy update logic into _compute_next_ckpt_step() and _apply_policy_update(), commits ckpt_step immediately after update_weights() completes, and extends stop() to cancel any in-flight policy update task. Adds unit tests covering reuse of an in-flight update after cancellation and shutdown cancellation behavior.

^{Written by Cursor Bugbot for commit 1392836. This will update automatically on new commits. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}