Skip to content

kv_cache: report abandoned pushes as failed, not sent #679

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Oseltamivir wants to merge 2 commits into
base: main
Choose a base branch
from
+31 −3
Open

kv_cache: report abandoned pushes as failed, not sent #679

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ecf4974
kv_cache: report abandoned pushes as failed, not sent
Oseltamivir Aug 16, 2026
7190118
chore: empty commit to re-run checks
Oseltamivir Aug 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 31 additions & 3 deletions tpu_sync/core/kv_cache_manager_with_transfer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1531,6 +1531,13 @@ void KVCacheManagerWithTransfer::StartRead(
free_slots_.empty()) {
// Request larger than a slot, or staging pool exhausted: surface as a
// recv failure (the connector can recompute) rather than throwing.
LOG(ERROR) << "StartRead: cannot stage req_id=" << req_id
<< " uuid=" << uuid
<< ": unique_local_blocks=" << unique_local_bids.size()
<< " max_blocks=" << max_blocks_
<< " free_slots=" << free_slots_.size()
<< " num_slots=" << all_slots_.size()
<< ". Reported as a recv failure.";
failed_recving_.insert(req_id);
return;
}
Expand Down Expand Up @@ -2256,9 +2263,28 @@ void KVCacheManagerWithTransfer::StartPushInternal(
free_slots_.empty()) {
auto it = send_entries_.find(uuid);
if (it != send_entries_.end()) {
done_sending_.insert(it->second->req_id);
// The transfer is abandoned here. The protocol carries no
// producer->consumer failure message, so the consumer is not told and
// blocks until its own deadline; this is the only place the cause is
// known. Report failed rather than sent, mirroring what the pool
// reshard send completion path already does for a send-side failure.
LOG(ERROR) << "StartPushInternal: abandoning transfer req_id="
<< it->second->req_id << " uuid=" << uuid
<< ": src_blocks=" << src_block_ids.size()
<< " max_blocks=" << max_blocks_
<< " free_slots=" << free_slots_.size()
<< " num_slots=" << all_slots_.size()
<< ". The consumer is not notified and will block until its "
"deadline.";
failed_recving_.insert(it->second->req_id);
ReleaseEntrySlotLocked(it->second);
send_entries_.erase(it);
} else {
LOG(ERROR) << "StartPushInternal: abandoning transfer uuid=" << uuid
<< " with no send entry: src_blocks=" << src_block_ids.size()
<< " max_blocks=" << max_blocks_
<< " free_slots=" << free_slots_.size()
<< ". Nothing is reported for this request.";
}
return;
}
Expand Down Expand Up @@ -2306,7 +2332,8 @@ void KVCacheManagerWithTransfer::StartPushInternal(
absl::MutexLock lock(mu_);
auto it = send_entries_.find(uuid);
if (it != send_entries_.end()) {
done_sending_.insert(it->second->req_id);
// The D2H never ran: this request did not send.
failed_recving_.insert(it->second->req_id);
ReleaseEntrySlotLocked(it->second);
send_entries_.erase(it);
}
Expand Down Expand Up @@ -2348,7 +2375,8 @@ void KVCacheManagerWithTransfer::SendNextLayer(uint64_t uuid, size_t l) {
absl::MutexLock lock(mu_);
auto it = send_entries_.find(uuid);
if (it != send_entries_.end()) {
done_sending_.insert(it->second->req_id);
// The D2H copy failed: this request did not send.
failed_recving_.insert(it->second->req_id);
ReleaseEntrySlotLocked(it->second);
send_entries_.erase(it);
}
Expand Down