feat(zmq): score TokenSpeed DP ranks on piggybacked scheduler load

crates/engine_zmq_client/src/connector.rs

@@ -227,6 +227,15 @@ impl<P: EngineProtocol> ClientInner<P> {
 
     /// 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<Item = &'a String>) {
         let mut inflight = self.inflight.lock();
         let Some(ids) = inflight.get_mut(&engine_index) else {
@@ -235,6 +244,12 @@ impl<P: EngineProtocol> ClientInner<P> {
         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,29 +910,31 @@ 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 {
                 num_running_reqs: 3,
                 num_waiting_reqs: 5,
                 ..Default::default()
             })),
-            finished_requests: Some(std::collections::BTreeSet::from(["r".to_string()])),
             ..Default::default()
         });
         engine
             .send_output(vec![Bytes::from(encode_msgpack(&outputs).unwrap())])
             .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,31 +944,32 @@ 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 {
             engine_index: 0,
             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 {
                 num_running_reqs: 3,
                 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())])

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]