fix(orchestrator): drain after final train batch

[jooooonas]

Summary

Enter drain mode immediately after successfully sending the final configured training batch, rather than waiting for an unused overflow batch to complete.

Problem

The orchestrator currently starts draining only when:

step > max_steps

After sending the valid batch at step == max_steps, it advances the step but continues scheduling rollouts. Another complete TrainBatch must reach finalize_train_batch() before the existing overflow guard stops training.

That batch is discarded and cannot serve as a fallback. This wastes rollout work and makes shutdown depend on an unnecessary batch
completing. A slow or unhealthy rollout tail can therefore delay shutdown.

Change

After sender.send() successfully sends the final valid batch, the orchestrator now:

• enters drain mode;
• disables further train scheduling;
• cancels in-flight train rollouts;
• leaves eval rollouts untouched so final evaluations can finish.

The existing step > max_steps check remains as a defensive guard.

Draining starts only after a valid batch is sent. Empty or fully filtered final-batch candidates continue to be retried.

Tests

Added one CPU-only regression test verifying that a valid final batch:

• is sent;
• immediately starts draining;
• advances progress normally;
• still triggers final-step evaluation.

uv run pytest -q tests/unit/orchestrator -m "not gpu"
95 passed, 1 skipped

Ruff lint and formatting checks also pass.