server: non-blocking stream send — stop one slow client stalling the batch
All three engines emitted tokens with blocking_send on the single decode/coordinator OS thread. A streaming client that drains slower than generation fills its 64-slot channel, and blocking_send then blocks the whole thread: under continuous batching one slow consumer stalls every other running sequence (and in the serial TP/PP path it blocks admission of the next request too). The whole point of continuous batching is defeated. Fix: switch to try_send. engine.rs sets a client_stalled flag on Full/Closed, reaped by is_finished() next iteration; tp_engine/pp_engine emit_text returns bool and the decode loop breaks with finish_reason "error". When the sequence/request is dropped its sender drops too, closing the channel so the client receive loop ends rather than hanging. A slow client now only loses its own sequence, never the batch. Verified on dash5: gpt-oss FP8 TP=1 streaming via tp_engine still streams correctly (SSE chunks, coherent content, no hang); bench-gpt-oss TP=2 5.9ms TPOT unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -294,9 +294,12 @@ pub fn run_tp(
|
||||
|
||||
let mut decode_buf: Vec<u8> = Vec::new();
|
||||
let mut generated = 1usize;
|
||||
emit_text(&tokenizer, &req, next, &mut decode_buf);
|
||||
let mut stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
|
||||
|
||||
let finish = loop {
|
||||
if stalled {
|
||||
break "error";
|
||||
}
|
||||
if tokenizer.is_eos(next) {
|
||||
break "stop";
|
||||
}
|
||||
@@ -317,17 +320,17 @@ pub fn run_tp(
|
||||
next = pick(&logits, &req.sampling, &gen_ids);
|
||||
gen_ids.push(next);
|
||||
generated += 1;
|
||||
emit_text(&tokenizer, &req, next, &mut decode_buf);
|
||||
stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
|
||||
};
|
||||
|
||||
let tail = tokenizer.flush_decode_stream(&mut decode_buf);
|
||||
if !tail.is_empty() {
|
||||
let _ = req.sender.blocking_send(GenerateEvent::Token {
|
||||
let _ = req.sender.try_send(GenerateEvent::Token {
|
||||
id: next,
|
||||
text: tail,
|
||||
});
|
||||
}
|
||||
let _ = req.sender.blocking_send(GenerateEvent::Done {
|
||||
let _ = req.sender.try_send(GenerateEvent::Done {
|
||||
finish_reason: finish.to_string(),
|
||||
});
|
||||
|
||||
@@ -340,14 +343,19 @@ pub fn run_tp(
|
||||
}
|
||||
|
||||
/// Stream a token's decoded text to the client (EOS contributes no text).
|
||||
fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &mut Vec<u8>) {
|
||||
/// Returns false if the send would block (client too slow) or the client is
|
||||
/// gone — the caller stops generating so the serial coordinator thread is free
|
||||
/// to admit the next request instead of blocking on one slow consumer.
|
||||
fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &mut Vec<u8>) -> bool {
|
||||
if tokenizer.is_eos(token_id) {
|
||||
return;
|
||||
return true;
|
||||
}
|
||||
let text = tokenizer.decode_token_stream(token_id, buf);
|
||||
if !text.is_empty() {
|
||||
let _ = req
|
||||
return req
|
||||
.sender
|
||||
.blocking_send(GenerateEvent::Token { id: token_id, text });
|
||||
.try_send(GenerateEvent::Token { id: token_id, text })
|
||||
.is_ok();
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user