Harden legacy UUID KV transfer leases and teardown

[wangyang59]

Summary

• add versioned, bounded batched renew/cancel operations for legacy UUID-based KV transfers
• expose the lease operations through the JAX C++ and Python surfaces
• bound producer retention, early-unknown retries, tombstones, and transfer completion
• make transport shutdown drain queued, active, delayed-PJRT, and pool-reshard callbacks safely
• give asynchronous pull responses stable duplicated-fd ownership to prevent descriptor ABA/reuse bugs
• add hardware-free lifecycle, protocol, race, deadline, cancellation, and descriptor-ownership coverage across kv_cache/, core/, transport/, and frameworks/jax/

Control protocol operations

The numeric names are the values carried in ControlRequestHeader::op:

• op2 — ACK (kOpAck): acknowledges that the consumer no longer needs the producer's legacy UUID registration, allowing producer-side state to be released. It remains accepted for wire compatibility.
• op3 — pull stream (kOpPullStream): a receiver-initiated pull request. num_blocks > 0 claims the UUID and starts the requested transfer; num_blocks == 0 is the existing legacy ACK encoding and remains supported.
• op4 — renew leases (kOpRenewLeases): renews a bounded batch of producer UUID leases while the consumer is still waiting to pull them, subject to the maximum retention deadline.
• op5 — cancel leases (kOpCancelLeases): cancels/releases a bounded batch of producer UUID registrations. Pre-registration cancellations are tombstoned so a racing late registration cannot resurrect the transfer.

Ops 4/5 use protocol version 1 in ep_idx, accept at most 4096 UUIDs per wire request, preserve input order in their responses, and return one status per UUID: 1 applied, 0 unknown, -1 terminal, -2 transferring, or -3 maximum retention reached.

Motivation

A decoder can wait long enough that its KV pull arrives after the producer UUID expires. The current controller lease API does not cover the legacy JAX UUID path, so that path needs explicit renew/cancel support and fail-fast lifecycle handling.

Compatibility

• existing op2 ACK handling and both forms of op3 are preserved
• current controller/hash lease APIs are unchanged
• new ops 4/5 use an explicit protocol version and bounded request batches
• mixed old/new binaries fail closed for the new operations; producers and consumers should roll together
• lease bindings are JAX-only in this PR; Torch parity can follow separately

Test coverage

• //tpu_sync/kv_cache:kv_cache_manager_lifecycle_test: verifies manager teardown waits for tracked callbacks
• //tpu_sync/core:kv_cache_manager_with_transfer_lifecycle_test: covers expiry/tombstones, early-unknown retry, transfer and shutdown draining, lease status ordering, cancel/pull races, protocol versioning/chunking, and absolute deadlines
• //tpu_sync/transport:block_transport_test: verifies cancellation rejects a new asynchronous push exactly once and drains cleanly
• //tpu_sync/transport/lib:raw_buffer_transport_test: verifies cancellation interrupts a blocked request and connection-pool shutdown does not steal a borrowed descriptor's ownership
• //tpu_sync/frameworks/jax:kv_cache_manager_wrapper_test: verifies ordered per-UUID status forwarding, exactly one control client, and empty-manager rejection

Validation

• rebased onto current main
• git diff --check
• Apple clang-format 17 using Google style
• Bazel target analysis completed locally for the kv_cache, core, transport, and JAX wrapper targets
• full build/test execution is left to CI because the managed Mac kills locally generated Bazel executables (protoc_minimal, exit 137); the raw-buffer target also encounters the repository's existing missing //third_party/xxhash package during local analysis

Rollout notes

Roll out producer and consumer images together, begin with shadow traffic disabled, and ramp gradually while monitoring lease renew status, retryable-unknown pulls, transfer failures, local-prefill fallback, and shutdown/drain errors.

[juncgu-google]

Thanks for your contribution, @wangyang59.
We are actively reviewing this PR. Since it encompasses fixes across multiple areas, it may take us some time to complete the review.

To help us move forward, could you please address the following:

1. explain what op2, op3, op4, op5 mean in the PR description?
2. And more tests and verify the affected layers / modules: e.g., kv_cache/, transport, frameworks.

[wangyang59]

Thanks for your contribution, @wangyang59. We are actively reviewing this PR. Since it encompasses fixes across multiple areas, it may take us some time to complete the review.

To help us move forward, could you please address the following:

@juncgu-google thanks a lot for your review.

1. explain what op2, op3, op4, op5 mean in the PR description?

the values were carried in ControlRequestHeader::op, and their descriptions were added in the PR summary

2. And more tests and verify the affected layers / modules: e.g., kv_cache/, transport, frameworks.

added hardware-free lifecycle, protocol, race, deadline, cancellation, and descriptor-ownership coverage across kv_cache/, core/, transport/, and frameworks/jax/