Harden legacy UUID KV transfer leases and teardown - #723
Conversation
|
Thanks for your contribution, @wangyang59. To help us move forward, could you please address the following:
|
8138086 to
a8228ba
Compare
@juncgu-google thanks a lot for your review.
the values were carried in ControlRequestHeader::op, and their descriptions were added in the PR summary
added hardware-free lifecycle, protocol, race, deadline, cancellation, and descriptor-ownership coverage across kv_cache/, core/, transport/, and frameworks/jax/ |
| send_tombstones_.size()); | ||
| } | ||
|
|
||
| bool KVCacheManagerWithTransfer::TerminalizeSendEntryLocked( |
There was a problem hiding this comment.
Does TerminalizeSendEntryLocked handle both success and failed entries? If yes, then all of them are inserted into done_sending_ (L2686)?
(In L2724 (FinishSendLayer), failed entries are also handled by TerminalizeSendEntryLocked.)
| std::list<uint64_t> pending_ack_order_; | ||
| std::chrono::milliseconds send_tombstone_ttl_{std::chrono::minutes(5)}; | ||
| std::chrono::milliseconds pending_ack_ttl_{std::chrono::seconds(30)}; | ||
| size_t max_send_tombstones_ = 4096; |
There was a problem hiding this comment.
I'd suggest to set these five values (send_tombstone_ttl_, pending_ack_ttl_, max_send_tombstones_ , max_pending_acks_, and control_io_timeout_) through env var (e.g., RAIDEN_xxx)
| static constexpr uint32_t kOpCancelLeases = 5; | ||
| static constexpr uint32_t kLeaseProtocolVersion = 1; | ||
| static constexpr int32_t kControlRetryableUnknown = 1; | ||
| static constexpr size_t kMaxLeaseBatchSize = 4096; |
There was a problem hiding this comment.
same here, please allow users to set this value (kMaxLeaseBatchSize) through env var.
| if (req.num_blocks == 0 || | ||
| req.num_blocks > static_cast<uint64_t>(max_blocks_)) { | ||
| throw std::invalid_argument("pull stream num_blocks is out of range"); | ||
| } |
There was a problem hiding this comment.
With RAIDEN_DYNAMIC_HOST_STAGING=1, a transfer larger than one slot is legal (like L3651). Please update the logic here.
| if (tombstone != send_tombstones_.end()) { | ||
| if (tombstone->second.kind == SendTombstoneKind::kPreRegistrationCancel) { | ||
| InstallSendTombstoneLocked(uuid, "send cancelled before registration", | ||
| now, SendTombstoneKind::kCancelled); | ||
| done_sending_.insert(req_id); | ||
| (void)ConsumePendingAckLocked(uuid, now); | ||
| consumed_early_cancel = true; | ||
| } else { | ||
| throw std::invalid_argument( | ||
| "send UUID is quarantined by a previous terminal outcome"); | ||
| } |
There was a problem hiding this comment.
Does it mean if a UUID is used and has a tombstone, it will not be usable for 5 mins (send_tombstone_ttl_)?
Add versioned batched renew/cancel support for the legacy JAX path, bound early-pull retries and producer retention, and make transport teardown cancellation-safe across queued, active, and delayed callbacks. Add hardware-free lifecycle, protocol, race, and descriptor-ownership regression coverage.
6d9bfc4 to
cea8eda
Compare
|
Thanks a lot for the comments! Addressed them as following:
|
2 participants
Summary
kv_cache/,core/,transport/, andframeworks/jax/Control protocol operations
The numeric names are the values carried in
ControlRequestHeader::op: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.kOpPullStream): a receiver-initiated pull request.num_blocks > 0claims the UUID and starts the requested transfer;num_blocks == 0is the existing legacy ACK encoding and remains supported.kOpRenewLeases): renews a bounded batch of producer UUID leases while the consumer is still waiting to pull them, subject to the maximum retention deadline.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:1applied,0unknown,-1terminal,-2transferring, or-3maximum 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
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 rejectionValidation
maingit diff --checkkv_cache,core,transport, and JAX wrapper targetsprotoc_minimal, exit 137); the raw-buffer target also encounters the repository's existing missing//third_party/xxhashpackage during local analysisRollout 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.