Track reused indices in controller checkpoint ownership

[rawsh]

Releasing a partition and reusing its indices restores partition ownership but omits the reused indices from allocated_indexes, so a saved controller undercounts live allocations.
Track the entire allocation and add same- and cross-partition tests through real controller save/load.
Validation: 11 focused tests pass without skips, including 8 regressions that fail on the base; Ruff and compileall pass. Full-suite, candidate-image and coherent-snapshot qualification remain pending.

[0oshowero0]

[AI Review]

Summary: The one-line fix is correct and minimal: allocated_indexes now includes reused indexes, which restores the invariant allocated_indexes == union(partition_to_indexes.values()). It applies cleanly on current main (#181 and #182 don't touch PartitionIndexManager). A few things should be addressed before merge.

1. The impact is a metrics undercount, not a checkpoint/ownership bug

allocated_indexes is never used for allocation or ownership decisions. Its only consumers are the checkpoint round-trip and the tq_global_index_allocated_total gauge (via _build_metrics_snapshot). Partition ownership (partition_to_indexes) was always correct, so there was no risk of double allocation. The undercount also starts at runtime as soon as an index is reused; checkpointing only persists it. Suggest retitling to something like "Fix tq_global_index_allocated_total undercount after index reuse" so the risk is clear to reviewers.

2. Consider deriving the count instead of keeping redundant state

allocated_indexes duplicates the union of partition_to_indexes, and its three mutation sites must stay in sync; this bug is exactly that drift. A smaller long-term fix would drop the field, compute the gauge as sum(len(s) for s in partition_to_indexes.values()), and ignore the key in load_checkpoint. If older TQ versions need to load new checkpoints, keep writing the derived value for one release. Not a blocker if maintainers prefer the one-liner.

3. Checkpoints written before this fix stay wrong after upgrade

load_checkpoint restores the saved allocated_indexes verbatim. I wrote a checkpoint with the base code (partitions old=3 and retained=2, clear old, then insert 4 into a new partition) and loaded it with this PR: the partitions own 6 indexes, but the gauge still reports 3. Either option in (2), or rebuilding the set from partition_to_indexes on load, would fix this. The impact is limited to observability and goes away once those partitions are cleared.

4. Tests could be much smaller

• All 8 new cases fail on the base at the same assertion, before any checkpoint step, so the 6-way parametrization and the two Ray tests add runtime without adding distinct coverage.
• pickle.loads(pickle.dumps(manager)) doesn't exercise the real checkpoint path: save_checkpoint serializes a dict of fields, not the object. This test is also a PartitionIndexManager unit test, so it doesn't belong in TestTransferQueueControllerCheckpoint.
• In test_controller_checkpoint_preserves_reused_index_ownership, the ownership assertions already pass on the base; only the allocated_indexes assertion covers new behavior.
• Suggest one short PartitionIndexManager unit test that asserts the invariant after release + reuse. If an integration test is kept, assert _build_metrics_snapshot()["global_index_allocated"] rather than reading pickle internals.

5. Merge blockers

• The commit has no Signed-off-by line, so dco-check will fail (git commit --amend -s and force-push).
• ruff format will collapse the controller.get_metadata.remote(..., batch_size=4, partition_id=destination, mode="insert") call in the new test onto one line (it fits in exactly 120 columns), so the pre-commit hook will fail.

Verification: On the PR head, all 11 tests in the checkpoint test class pass. With only the controller.py change reverted, all 8 new cases fail. On main + this PR, tests/test_controller.py, tests/test_metrics.py, and tests/test_controller_data_partitions.py pass (106/106), and compileall and ruff check pass. The full pytest -q suite was not run.