[Bugfix][Rust Frontend] Account for new requests in DP routing (#58956)
Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Reid committed
Sep 29, 2026 at 18:24 UTC
3f3fbe286f25ead10f793dea0c549a99c196ad11
1 file changed
+148
-7
rust/src/engine-core-client/src/client/state.rs
+148
-7
@@ -67,25 +67,40 @@ struct EngineRoutingState {
67
inflight: usize,
68
/// The latest real scheduler snapshot received from this engine, if any.
69
last_scheduler_stats: Option<EngineLoadSnapshot>,
70
+ /// Requests admitted since the last scheduler snapshot was received.
71
+ ///
72
+ /// Added to the snapshot so each admission raises the routing score before
73
+ /// the next snapshot arrives. Only tracked once a snapshot exists.
74
+ admitted_since_stats: usize,
75
}
76
77
impl EngineRoutingState {
78
/// Compute the routing score used to pick the least-loaded engine.
79
///
75
- /// Scheduler stats can raise the load estimate above the frontend-local
76
- /// view, but they should not lower it below requests this frontend has
77
- /// already admitted.
80
+ /// Scheduler stats, plus admissions not yet reflected in them, can raise
81
+ /// the load estimate above the frontend-local view, but they should not
82
+ /// lower it below requests this frontend has already admitted.
83
fn routing_score(&self) -> usize {
84
let Some(stats) = self.last_scheduler_stats else {
85
return self.inflight;
86
};
87
83
- self.inflight.max(stats.running + stats.waiting)
88
+ self.inflight.max(stats.running + stats.waiting + self.admitted_since_stats)
89
+ }
90
+
91
+ /// Record one request admitted to this engine.
92
+ fn record_admission(&mut self) {
93
+ self.inflight += 1;
94
+ if self.last_scheduler_stats.is_some() {
95
+ self.admitted_since_stats += 1;
96
+ }
97
}
98
86
- /// Replace the local routing view with a fresh real scheduler snapshot.
99
+ /// Replace the local routing view with a fresh real scheduler snapshot,
100
+ /// resetting the admissions counted on top of the previous one.
101
fn apply_scheduler_counts(&mut self, next: EngineLoadSnapshot) {
102
self.last_scheduler_stats = Some(next);
103
+ self.admitted_since_stats = 0;
104
}
105
}
106
@@ -121,7 +136,7 @@ impl RequestRegistry {
136
///
137
/// When `data_parallel_rank` is provided, the request is routed directly to
138
/// the engine at that rank index, bypassing load balancing. Otherwise
124
- /// the engine with the fewest in-flight requests is chosen.
139
+ /// the engine with the lowest routing score is chosen.
140
pub fn register(
141
&mut self,
142
request_id: String,
@@ -154,7 +169,7 @@ impl RequestRegistry {
169
.routing_per_engine
170
.get_mut(&engine_id)
171
.expect("request registry must track all known engines");
157
- state.inflight += 1;
172
+ state.record_admission();
173
174
Ok((engine_id, rx))
175
}
@@ -736,6 +751,7 @@ mod tests {
751
let state = EngineRoutingState {
752
inflight: 3,
753
last_scheduler_stats: None,
754
+ ..Default::default()
755
};
756
757
assert_eq!(state.routing_score(), 3);
@@ -749,6 +765,7 @@ mod tests {
765
waiting: 0,
766
running: 2,
767
}),
768
+ ..Default::default()
769
};
770
771
assert_eq!(state.routing_score(), 7);
@@ -762,6 +779,7 @@ mod tests {
779
waiting: 3,
780
running: 2,
781
}),
782
+ ..Default::default()
783
};
784
785
assert_eq!(state.routing_score(), 5);
@@ -795,6 +813,129 @@ mod tests {
813
assert_eq!(chosen, engine_1);
814
}
815
816
+ #[test]
817
+ fn registry_spreads_bursts_after_cancellation_before_stats_refresh() {
818
+ let mut distributions = Vec::new();
819
+ for counts in [[10, 12], [12, 12]] {
820
+ let mut registry = RequestRegistry::new(
821
+ &[0, 1].map(|rank| connected_engine(EngineId::from_engine_index(rank))),
822
+ );
823
+ let mut old_requests = Vec::new();
824
+ for (rank, running) in counts.into_iter().enumerate() {
825
+ for index in 0..running {
826
+ let request_id = format!("old-{rank}-{index}");
827
+ registry.register(request_id.clone(), None, Some(rank as u32)).unwrap();
828
+ old_requests.push(request_id);
829
+ }
830
+ registry.apply_scheduler_counts(
831
+ rank as u32,
832
+ EngineLoadSnapshot {
833
+ waiting: 0,
834
+ running,
835
+ },
836
+ );
837
+ }
838
+ drop(registry.abort_many(&old_requests, 0.0));
839
+
840
+ let mut distribution = [0; 2];
841
+ for index in 0..14 {
842
+ let (engine, _) = registry.register(format!("new-{index}"), None, None).unwrap();
843
+ distribution[engine.engine_index().unwrap() as usize] += 1;
844
+ }
845
+ distributions.push(distribution);
846
+ }
847
+ expect_test::expect!["[[8, 6], [7, 7]]"].assert_eq(&format!("{distributions:?}"));
848
+ }
849
+
850
+ #[test]
851
+ fn registry_preserves_tie_order_across_finished_requests_and_stats_refreshes() {
852
+ let ranks = [2, 5, 9];
853
+ let mut registry = RequestRegistry::new(
854
+ &ranks.map(|rank| connected_engine(EngineId::from_engine_index(rank))),
855
+ );
856
+ let mut chosen = Vec::new();
857
+ for index in 0..6 {
858
+ for rank in ranks {
859
+ registry.apply_scheduler_counts(
860
+ u32::from(rank),
861
+ EngineLoadSnapshot {
862
+ waiting: 0,
863
+ running: 0,
864
+ },
865
+ );
866
+ }
867
+ let request_id = format!("req-{index}");
868
+ let (engine, _) = registry.register(request_id.clone(), None, None).unwrap();
869
+ chosen.push(engine.engine_index().unwrap());
870
+ drop(registry.finish_many(&[request_id]));
871
+ }
872
+ expect_test::expect!["[2, 2, 2, 2, 2, 2]"].assert_eq(&format!("{chosen:?}"));
873
+ }
874
+
875
+ #[test]
876
+ fn registry_refreshes_estimates_without_erasing_inflight() {
877
+ let mut registry = RequestRegistry::new(
878
+ &[0, 1].map(|rank| connected_engine(EngineId::from_engine_index(rank))),
879
+ );
880
+ let snapshot = EngineLoadSnapshot {
881
+ waiting: 0,
882
+ running: 10,
883
+ };
884
+ registry.apply_scheduler_counts(0, snapshot);
885
+ registry.apply_scheduler_counts(
886
+ 1,
887
+ EngineLoadSnapshot {
888
+ waiting: 0,
889
+ running: 12,
890
+ },
891
+ );
892
+ for index in 0..4 {
893
+ registry.register(format!("pinned-{index}"), None, Some(0)).unwrap();
894
+ }
895
+
896
+ // Even unchanged counts replace the optimistic estimate for that engine.
897
+ registry.apply_scheduler_counts(0, snapshot);
898
+ let (after_refresh, _) = registry.register("after-refresh".into(), None, None).unwrap();
899
+ // An older, empty snapshot must not erase the five admitted requests.
900
+ for rank in [0, 1] {
901
+ registry.apply_scheduler_counts(
902
+ rank,
903
+ EngineLoadSnapshot {
904
+ waiting: 0,
905
+ running: 0,
906
+ },
907
+ );
908
+ }
909
+ let (after_empty, _) = registry.register("after-empty".into(), None, None).unwrap();
910
+ expect_test::expect!["(Some(0), Some(1))"].assert_eq(&format!(
911
+ "{:?}",
912
+ (after_refresh.engine_index(), after_empty.engine_index())
913
+ ));
914
+ }
915
+
916
+ #[test]
917
+ fn registry_explicit_admissions_raise_routing_score() {
918
+ let mut registry = RequestRegistry::new(
919
+ &[2, 5, 9].map(|rank| connected_engine(EngineId::from_engine_index(rank))),
920
+ );
921
+ for rank in [2, 5, 9] {
922
+ registry.apply_scheduler_counts(
923
+ rank,
924
+ EngineLoadSnapshot {
925
+ waiting: 0,
926
+ running: 10,
927
+ },
928
+ );
929
+ }
930
+ registry.register("pinned".into(), None, Some(9)).unwrap();
931
+ let mut chosen = Vec::new();
932
+ for index in 0..3 {
933
+ let (engine, _) = registry.register(format!("auto-{index}"), None, None).unwrap();
934
+ chosen.push(engine.engine_index().unwrap());
935
+ }
936
+ expect_test::expect!["[2, 5, 2]"].assert_eq(&format!("{chosen:?}"));
937
+ }
938
+
939
#[test]
940
fn register_with_data_parallel_rank_routes_to_specified_engine() {
941
let engine_0 = EngineId::from_engine_index(0);