kv_cache: report abandoned pushes as failed, not sent - #679
Open
Oseltamivir wants to merge 2 commits into
Open
Conversation
StartPushInternal drops a request when it is larger than a slot or the staging pool is empty, and when a D2H cannot be issued or fails. All four paths insert the request into done_sending_, so the producer sees a transfer that never happened reported as sent, and the consumer -- which the protocol gives no failure message -- blocks until its own deadline and then times out naming no cause. Report these as failed_recving_ instead, mirroring the pool reshard send completion path, and log the slot accounting at the two exhaustion sites in StartPushInternal and StartRead, which is where the cause is known. The logs are LOG(ERROR) deliberately: absl holds the stderr threshold at ERROR until InitializeLog() is called, and nothing in the JAX path calls it, so LOG(INFO) diagnostics here would not be visible.
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
StartPushInternal(tpu_sync/core/kv_cache_manager_with_transfer.cc) drops a request in four places: larger than a slot, staging pool empty, D2H dispatch not ok, and D2H completion not ok. Each inserts intodone_sending_, so a transfer that did not happen is reported as sent. The protocol carries no producer to consumer failure message, so the consumer is not told either. It blocks until its own deadline and reports a timeout naming no cause. Nothing is logged at the exhaustion site, which is the only place the cause is known.This reports those four as
failed_recving_instead, mirroring what the pool reshard send completion path already does for a send-side failure, and addsLOG(ERROR)with the slot accounting (src_blocks,max_blocks_, free/total slots) at the two exhaustion sites, inStartPushInternaland inStartRead. It also removes an asymmetry: consumer-side exhaustion inStartReadalready reports viafailed_recving_, so today the same root cause is a clean failure from one side and a bare timeout from the other.LOG(ERROR)rather thanLOG(INFO)is deliberate. absl keeps the stderr threshold at ERROR untilInitializeLog()is called and nothing in the JAX path calls it, soLOG(INFO)diagnostics here are not visible in practice.No API change. I have no build environment for this repo, so it is not compiled or tested. Happy to adjust to whatever shape you prefer.