diff --git a/crates/engine_zmq_client/src/connector.rs b/crates/engine_zmq_client/src/connector.rs index 1f3b74cda..392878805 100644 --- a/crates/engine_zmq_client/src/connector.rs +++ b/crates/engine_zmq_client/src/connector.rs @@ -227,6 +227,15 @@ impl ClientInner

{ /// Retire in-flight request ids from a rank. Idempotent: an id already /// retired by a racing path is simply absent. + /// + /// When the retirement empties the rank, its stored queue counts are + /// zeroed too: engines only report load on output batches (there is no + /// heartbeat), and the last batch a rank ever sends is sampled before its + /// final request's finish is committed — so an idle rank's last snapshot + /// permanently shows that request still resident. This client routed + /// every request, so an empty in-flight set is ground truth that the + /// rank's queue is empty; the KV term is left as reported (cache pages + /// outlive requests). fn release<'a>(&self, engine_index: u32, request_ids: impl IntoIterator) { let mut inflight = self.inflight.lock(); let Some(ids) = inflight.get_mut(&engine_index) else { @@ -235,6 +244,12 @@ impl ClientInner

{ for request_id in request_ids { ids.remove(request_id); } + if ids.is_empty() { + if let Some(load) = self.load.lock().get_mut(&engine_index) { + load.num_running = 0; + load.num_waiting = 0; + } + } } /// Retire one request from the rank named by `engine_id`, if it has one. @@ -895,12 +910,15 @@ mod tests { .unwrap(); engine.recv_request().await.unwrap(); + // Stats ride a mid-stream chunk: the request is still in flight, so + // the snapshot is current and must be stored verbatim. (A terminal + // batch's snapshot predates its own finish; once the rank empties, + // the queue counts are zeroed — covered separately.) let outputs = EngineCoreOutputs::RequestBatch(RequestBatchOutputs { engine_index: 0, outputs: vec![EngineCoreOutput { request_id: "r".into(), new_token_ids: vec![7], - finish_reason: Some(EngineCoreFinishReason::Stop), ..Default::default() }], scheduler_stats: Some(Box::new(SchedulerStats { @@ -908,7 +926,6 @@ mod tests { num_waiting_reqs: 5, ..Default::default() })), - finished_requests: Some(std::collections::BTreeSet::from(["r".to_string()])), ..Default::default() }); engine @@ -916,8 +933,8 @@ mod tests { .await .unwrap(); - // Drain the stream so the batch is processed. - while stream.next().await.is_some() {} + let chunk = stream.next().await.expect("mid-stream chunk").unwrap(); + assert!(!chunk.finished()); let load = client.engine_load(0).expect("load recorded"); assert_eq!(load.num_running, 3); assert_eq!(load.num_waiting, 5); @@ -927,8 +944,10 @@ mod tests { async fn unpinned_requests_prefer_the_least_loaded_engine() { let (client, mut engines, _ns) = connect_ranks(2, false).await; - // Make engine 0 report load via a pinned warm-up request; engine 1 - // never reports and therefore scores zero. + // Make engine 0 report load via a pinned warm-up request that stays + // in flight (a finished-and-emptied rank has its queue counts zeroed, + // because a terminal batch's snapshot predates its own finish); + // engine 1 never reports and therefore scores zero. let mut stream = client.submit(request_for("warm", 0)).await.unwrap(); engines[0].recv_request().await.unwrap(); let outputs = EngineCoreOutputs::RequestBatch(RequestBatchOutputs { @@ -936,7 +955,6 @@ mod tests { outputs: vec![EngineCoreOutput { request_id: "warm".into(), new_token_ids: vec![7], - finish_reason: Some(EngineCoreFinishReason::Stop), ..Default::default() }], scheduler_stats: Some(Box::new(SchedulerStats { @@ -944,14 +962,14 @@ mod tests { num_waiting_reqs: 5, ..Default::default() })), - finished_requests: Some(std::collections::BTreeSet::from(["warm".to_string()])), ..Default::default() }); engines[0] .send_output(vec![Bytes::from(encode_msgpack(&outputs).unwrap())]) .await .unwrap(); - while stream.next().await.is_some() {} + let chunk = stream.next().await.expect("warm chunk").unwrap(); + assert!(!chunk.finished()); assert!(client.engine_load(0).is_some(), "warm-up load recorded"); // An unpinned request must now land on the idle engine 1, not @@ -1087,6 +1105,66 @@ mod tests { } } + #[tokio::test] + async fn an_emptied_rank_sheds_its_stale_queue_counts() { + // Engines report load only on output batches, and a terminal batch's + // snapshot is sampled before its own finish commits — so the last + // thing an idle rank ever says is "still busy". The client routed + // every request: once a rank's in-flight set empties, its stored + // queue counts must be zeroed or the rank is shunned forever. + let (client, mut engines, _ns) = connect_ranks(2, false).await; + + // Rank 0 runs one request to completion; its terminal batch carries + // the stale pre-commit snapshot (3 running, 5 waiting). + let mut stream = client.submit(request_for("last", 0)).await.unwrap(); + engines[0].recv_request().await.unwrap(); + let outputs = EngineCoreOutputs::RequestBatch(RequestBatchOutputs { + engine_index: 0, + outputs: vec![EngineCoreOutput { + request_id: "last".into(), + new_token_ids: vec![7], + finish_reason: Some(EngineCoreFinishReason::Stop), + ..Default::default() + }], + scheduler_stats: Some(Box::new(SchedulerStats { + num_running_reqs: 3, + num_waiting_reqs: 5, + ..Default::default() + })), + finished_requests: Some(std::collections::BTreeSet::from(["last".to_string()])), + ..Default::default() + }); + engines[0] + .send_output(vec![Bytes::from(encode_msgpack(&outputs).unwrap())]) + .await + .unwrap(); + while stream.next().await.is_some() {} + let load = client.engine_load(0).expect("snapshot stored"); + assert_eq!((load.num_running, load.num_waiting), (0, 0)); + + // Rank 1 holds one live request. The next unpinned request must go to + // the genuinely idle rank 0 — trusting the stale snapshot would send + // it behind rank 1's real work. + let _held = client.submit(request_for("held", 1)).await.unwrap(); + engines[1].recv().await.unwrap(); + let _next = client + .submit(EngineCoreRequest { + request_id: "next".into(), + prompt_token_ids: Some(vec![1, 2, 3]), + ..EngineCoreRequest::default() + }) + .await + .unwrap(); + match tokio::time::timeout(TIMEOUT, engines[0].recv()) + .await + .expect("request should route to the emptied rank") + .unwrap() + { + EngineInbound::Add(request) => assert_eq!(request.request_id, "next"), + other => panic!("expected Add on engine 0, got {other:?}"), + } + } + /// The `EngineId` of one connected rank, by index. fn engines_id(client: &EngineCoreClient, engine_index: u32) -> EngineId { client @@ -1208,7 +1286,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![]], output_token_logprobs_idx: vec![vec![]], - engine_index: 0, + ..Default::default() }; let done = BatchTokenIDOutSlim { rids: vec!["ts-1".into()], @@ -1219,7 +1297,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![]], output_token_logprobs_idx: vec![vec![]], - engine_index: 0, + ..Default::default() }; engine .send_output(vec![Bytes::from(encode_msgpack(&chunk).unwrap())]) diff --git a/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs b/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs index 2eec47a35..abccf9721 100644 --- a/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs +++ b/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs @@ -31,7 +31,7 @@ use crate::{ output::{BatchTokenIDOutSlim, TokenSpeedOutput}, request::{TokenSpeedRequestType, TokenizedGenerateReqInput}, }, - EngineBatch, EngineProtocol, + EngineBatch, EngineLoad, EngineProtocol, }, }; @@ -142,6 +142,14 @@ impl EngineProtocol for TokenSpeedProtocol { let payload = frames.first().map(AsRef::as_ref).unwrap_or_default(); let batch: BatchTokenIDOutSlim = decode_msgpack(payload)?; let engine_index = batch.engine_index; + // `kv_total_pages == 0` marks a pre-piggyback sender (or a snapshot + // the engine could not take): report no load rather than fabricating + // an empty-scheduler signal that least-loaded selection would trust. + let load = (batch.kv_total_pages > 0).then(|| EngineLoad { + num_running: batch.num_running, + num_waiting: batch.num_waiting, + kv_cache_usage: batch.kv_active_pages as f64 / batch.kv_total_pages as f64, + }); let outputs = batch.into_outputs()?; let finished_request_ids = outputs .iter() @@ -152,9 +160,7 @@ impl EngineProtocol for TokenSpeedProtocol { engine_index, outputs, finished_request_ids, - // The slim batch piggybacks no scheduler load, so DP selection - // scores TokenSpeed ranks on the gateway's own in-flight counts. - load: None, + load, wave: None, }) } @@ -175,6 +181,10 @@ mod tests { output_token_logprobs_val: vec![vec![], vec![]], output_token_logprobs_idx: vec![vec![], vec![]], engine_index: 1, + num_running: 2, + num_waiting: 5, + kv_active_pages: 100, + kv_total_pages: 400, } } @@ -184,10 +194,31 @@ mod tests { let decoded = TokenSpeedProtocol::decode_batch(&frames).unwrap(); assert_eq!(decoded.outputs.len(), 2); assert_eq!(decoded.finished_request_ids, vec!["b".to_string()]); - assert!(decoded.load.is_none()); // The batch names its producing DP rank; the connector routes in-flight // release and scoring by it. assert_eq!(decoded.engine_index, 1); + // The piggybacked snapshot surfaces as the engine-neutral load signal. + assert_eq!( + decoded.load, + Some(EngineLoad { + num_running: 2, + num_waiting: 5, + kv_cache_usage: 0.25, + }) + ); + } + + #[test] + fn decode_batch_reports_no_load_without_a_snapshot() { + // kv_total_pages == 0 marks a pre-piggyback sender (or a snapshot the + // engine could not take): no load, not a fabricated empty scheduler. + let batch = BatchTokenIDOutSlim { + kv_total_pages: 0, + ..slim_batch() + }; + let frames = vec![Bytes::from(encode_msgpack(&batch).unwrap())]; + let decoded = TokenSpeedProtocol::decode_batch(&frames).unwrap(); + assert!(decoded.load.is_none()); } #[test] diff --git a/crates/engine_zmq_client/src/protocol/tokenspeed/output.rs b/crates/engine_zmq_client/src/protocol/tokenspeed/output.rs index 3c6810cee..3248f6920 100644 --- a/crates/engine_zmq_client/src/protocol/tokenspeed/output.rs +++ b/crates/engine_zmq_client/src/protocol/tokenspeed/output.rs @@ -56,11 +56,24 @@ pub struct BatchTokenIDOutSlim { /// tail field: an older (pre-DP) sender emits 9 elements and this defaults /// to `0`, which is also the sole rank of a single-engine worker. pub engine_index: u32, + /// Piggybacked scheduler-load snapshot, sampled by the producing rank at + /// send time (the msgpack wire drops control replies, so the output batch + /// is the only in-band load channel). Appended tail fields: all default + /// to `0` from older senders, and `kv_total_pages == 0` means "no + /// snapshot" — the decoder then reports no load at all rather than a + /// fabricated zero. + pub num_running: u64, + /// Scheduler waiting-queue depth at send time. + pub num_waiting: u64, + /// KV pages held by running requests (the usage ratio's numerator). + pub kv_active_pages: u64, + /// Usable KV pages on the rank (the usage ratio's denominator). + pub kv_total_pages: u64, } impl Serialize for BatchTokenIDOutSlim { fn serialize(&self, serializer: S) -> std::result::Result { - let mut tuple = serializer.serialize_tuple(10)?; + let mut tuple = serializer.serialize_tuple(14)?; tuple.serialize_element(BATCH_TOKEN_ID_OUT_SLIM_TAG)?; tuple.serialize_element(&self.rids)?; tuple.serialize_element(&self.output_ids)?; @@ -71,6 +84,10 @@ impl Serialize for BatchTokenIDOutSlim { tuple.serialize_element(&self.output_token_logprobs_val)?; tuple.serialize_element(&self.output_token_logprobs_idx)?; tuple.serialize_element(&self.engine_index)?; + tuple.serialize_element(&self.num_running)?; + tuple.serialize_element(&self.num_waiting)?; + tuple.serialize_element(&self.kv_active_pages)?; + tuple.serialize_element(&self.kv_total_pages)?; tuple.end() } } @@ -103,6 +120,12 @@ impl<'de> Deserialize<'de> for BatchTokenIDOutSlim { // Appended by the DP wire revision: a 9-element batch from // an older sender means rank 0 (the only rank it can be). engine_index: seq.next_element::()?.unwrap_or(0), + // Appended by the load-piggyback revision; zeros from + // older senders decode as "no snapshot". + num_running: seq.next_element::()?.unwrap_or(0), + num_waiting: seq.next_element::()?.unwrap_or(0), + kv_active_pages: seq.next_element::()?.unwrap_or(0), + kv_total_pages: seq.next_element::()?.unwrap_or(0), }; drain_trailing(&mut seq)?; Ok(batch) @@ -187,9 +210,13 @@ impl BatchTokenIDOutSlim { cached_tokens, output_token_logprobs_val, output_token_logprobs_idx, - // Batch-level rank tag, not a per-request column; the caller reads - // it off the batch before splitting. + // Batch-level fields, not per-request columns; the caller reads + // them off the batch before splitting. engine_index: _, + num_running: _, + num_waiting: _, + kv_active_pages: _, + kv_total_pages: _, } = self; Ok(rids @@ -233,8 +260,16 @@ mod tests { /// A slim output batch captured from the Python msgspec encoder: rids /// ["vec-1"], output_ids [[10, 11]], finished_reasons ["length"], prompt 3 /// / completion 2 / cached 1, logprobs [[-0.5, -0.25]] over tokens - /// [[10, 11]], engine_index 1 (a DP sender's rank-1 batch). + /// [[10, 11]], engine_index 1, load snapshot (2 running, 5 waiting, + /// 100/400 KV pages). const PYTHON_OUTPUT_VECTOR: &str = + "9eb34261746368546f6b656e49444f7574536c696d91a57665632d3191920a0b91a66c656e\ + 6774689103910291019192cbbfe0000000000000cbbfd000000000000091920a0b01020564\ + cd0190"; + + /// The same batch as encoded by the engine_index-era sender: 10 elements, + /// no load snapshot. + const PYTHON_OUTPUT_VECTOR_PRE_LOAD: &str = "9ab34261746368546f6b656e49444f7574536c696d91a57665632d3191920a0b91a66c656e\ 6774689103910291019192cbbfe0000000000000cbbfd000000000000091920a0b01"; @@ -267,11 +302,15 @@ mod tests { output_token_logprobs_val: vec![vec![-0.5, -0.25]], output_token_logprobs_idx: vec![vec![10, 11]], engine_index: 1, + num_running: 2, + num_waiting: 5, + kv_active_pages: 100, + kv_total_pages: 400, } } /// The pinned cross-language vector — the exact bytes the engine sends — - /// decodes into the full 10-element (tag + 8 columns + rank) batch. + /// decodes into the full 14-element (tag + 8 columns + rank + load) batch. #[test] fn python_output_vector_decodes() { let decoded: BatchTokenIDOutSlim = decode_msgpack(&python_output_bytes()).unwrap(); @@ -298,21 +337,26 @@ mod tests { } #[test] - fn batch_output_serializes_as_tagged_ten_element_array() { + fn batch_output_serializes_as_tagged_fourteen_element_array() { let encoded = encode_msgpack(&vector_batch()).unwrap(); let Value::Array(array) = decode_value(&encoded).unwrap() else { panic!("expected positional array"); }; - assert_eq!(array.len(), 10); + assert_eq!(array.len(), 14); assert_eq!(array[0], Value::from(BATCH_TOKEN_ID_OUT_SLIM_TAG)); assert_eq!(array[1], Value::Array(vec![Value::from("vec-1")])); // rids assert_eq!(array[3], Value::Array(vec![Value::from("length")])); // finished_reasons assert_eq!(array[6], Value::Array(vec![Value::from(1)])); // cached_tokens assert_eq!(array[9], Value::from(1)); // engine_index + // Load snapshot tail: running, waiting, active pages, total pages. + assert_eq!(array[10], Value::from(2)); + assert_eq!(array[11], Value::from(5)); + assert_eq!(array[12], Value::from(100)); + assert_eq!(array[13], Value::from(400)); } - /// A pre-DP sender emits 9 elements; the missing tail decodes as rank 0, - /// the only rank a single-engine worker can be. + /// A pre-DP sender emits 9 elements; the missing tail decodes as rank 0 + /// (the only rank a single-engine worker can be) with no load snapshot. #[test] fn pre_dp_nine_element_batch_decodes_as_rank_zero() { let decoded: BatchTokenIDOutSlim = @@ -321,6 +365,28 @@ mod tests { decoded, BatchTokenIDOutSlim { engine_index: 0, + num_running: 0, + num_waiting: 0, + kv_active_pages: 0, + kv_total_pages: 0, + ..vector_batch() + } + ); + } + + /// An engine_index-era sender emits 10 elements; the missing load tail + /// decodes as the zero "no snapshot" defaults. + #[test] + fn pre_load_ten_element_batch_decodes_with_zero_snapshot() { + let decoded: BatchTokenIDOutSlim = + decode_msgpack(&vector_bytes(PYTHON_OUTPUT_VECTOR_PRE_LOAD)).unwrap(); + assert_eq!( + decoded, + BatchTokenIDOutSlim { + num_running: 0, + num_waiting: 0, + kv_active_pages: 0, + kv_total_pages: 0, ..vector_batch() } ); @@ -361,7 +427,7 @@ mod tests { cached_tokens: vec![0, 1], output_token_logprobs_val: vec![vec![-0.5], vec![-1.0, -2.0]], output_token_logprobs_idx: vec![vec![10], vec![20, 21]], - engine_index: 0, + ..Default::default() }; let outputs = batch.into_outputs().unwrap(); assert_eq!(outputs.len(), 2); @@ -389,7 +455,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![]], output_token_logprobs_idx: vec![vec![]], - engine_index: 0, + ..Default::default() }; let outputs = batch.into_outputs().unwrap(); assert!(outputs[0].output_logprobs_val.is_empty()); @@ -407,7 +473,7 @@ mod tests { cached_tokens: vec![0, 1], output_token_logprobs_val: vec![vec![], vec![]], output_token_logprobs_idx: vec![vec![], vec![]], - engine_index: 0, + ..Default::default() }; assert!(batch.into_outputs().is_err()); } @@ -424,7 +490,7 @@ mod tests { // Only one logprob column entry for two requests. output_token_logprobs_val: vec![vec![]], output_token_logprobs_idx: vec![vec![], vec![]], - engine_index: 0, + ..Default::default() }; assert!(batch.into_outputs().is_err()); } diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index 4b0dab095..fda6dec0c 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -2145,7 +2145,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![-0.5]], output_token_logprobs_idx: vec![vec![10]], - engine_index: 0, + ..Default::default() }; let done = BatchTokenIDOutSlim { rids: vec!["r1".into()], @@ -2156,7 +2156,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![-1.25]], output_token_logprobs_idx: vec![vec![11]], - engine_index: 0, + ..Default::default() }; output .send_frames(vec![bytes::Bytes::from(encode_msgpack(&chunk).unwrap())]) @@ -2810,7 +2810,7 @@ mod tests { cached_tokens: vec![0, 0], output_token_logprobs_val: vec![vec![], vec![]], output_token_logprobs_idx: vec![vec![], vec![]], - engine_index: 0, + ..Default::default() }; output .send_frames(vec![bytes::Bytes::from(encode_msgpack(&done).unwrap())])