From 740d0e148dd058cc328a6920ff014d926d79f357 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:26:46 -0700 Subject: [PATCH] fix(router): build a hash ring for model-less requests Requests that name no model resolve to the unknown wildcard, where the router widens the candidate set to every worker but still looks the ring up under that ID. Nothing keys a model index entry unknown, so the lookup always missed and consistent-hash policies fell back to their load path for the whole endpoint. Keep a wildcard ring alongside the per-model ones: the union of every model's workers deduped by URL, sharing the single model's ring when the registry serves one model. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- model_gateway/src/worker/registry.rs | 176 ++++++++++++++++++++++++++- 1 file changed, 170 insertions(+), 6 deletions(-) diff --git a/model_gateway/src/worker/registry.rs b/model_gateway/src/worker/registry.rs index 525f535268..af978d9b6a 100644 --- a/model_gateway/src/worker/registry.rs +++ b/model_gateway/src/worker/registry.rs @@ -326,6 +326,10 @@ impl WorkerRegistry { /// this registry except [`Self::get_by_model`]. A caller holding a /// client-supplied name resolves it with [`Self::resolve_model_alias`] /// once at request entry and passes the canonical ID from there on. + /// + /// [`UNKNOWN_MODEL_ID`] returns the wildcard ring spanning every worker, + /// matching the candidate set a request that names no model is routed + /// against. pub fn get_hash_ring(&self, model_id: &str) -> Option> { self.hash_rings.get(model_id).map(|r| Arc::clone(&r)) } @@ -1430,12 +1434,59 @@ impl WorkerRegistry { /// Rebuild the hash ring for a model based on current workers in the model index. fn rebuild_hash_ring(&self, model_id: &str) { - if let Some(workers) = self.model_index.get(model_id) { - let ring = HashRing::new(workers.value().iter().map(|w| w.url())); - self.hash_rings.insert(model_id.to_string(), Arc::new(ring)); - } else { - // No workers for this model, remove the ring - self.hash_rings.remove(model_id); + let ring = self + .model_index + .get(model_id) + .map(|workers| Arc::new(HashRing::new(workers.value().iter().map(|w| w.url())))); + + match ring { + Some(ring) => { + self.hash_rings.insert(model_id.to_string(), ring); + } + None => { + // No workers for this model, remove the ring + self.hash_rings.remove(model_id); + } + } + + self.rebuild_wildcard_hash_ring(); + } + + /// Rebuild the ring stored under [`UNKNOWN_MODEL_ID`], which requests that + /// name no model are routed against. Those requests may land on any worker, + /// so the ring spans every model's workers, deduplicated by URL. + fn rebuild_wildcard_hash_ring(&self) { + let model_ids: Vec = self + .model_index + .iter() + .map(|entry| entry.key().clone()) + .collect(); + + match model_ids.as_slice() { + [] => { + self.hash_rings.remove(UNKNOWN_MODEL_ID); + } + // A single model already covers every worker, so share its ring + // instead of hashing the same URLs a second time. + [only] => { + let ring = self.hash_rings.get(only).map(|ring| Arc::clone(&ring)); + match ring { + Some(ring) => { + self.hash_rings.insert(UNKNOWN_MODEL_ID.to_string(), ring); + } + None => { + self.hash_rings.remove(UNKNOWN_MODEL_ID); + } + } + } + _ => { + let mut urls: HashSet = HashSet::new(); + for entry in self.model_index.iter() { + urls.extend(entry.value().iter().map(|w| w.url().to_string())); + } + self.hash_rings + .insert(UNKNOWN_MODEL_ID.to_string(), Arc::new(HashRing::new(urls))); + } } } @@ -3311,4 +3362,117 @@ mod tests { other => panic!("Expected Removed event, got: {other:?}"), } } + + fn worker_serving(url: &str, model_ids: &[&str]) -> Arc { + Arc::new( + BasicWorkerBuilder::new(url) + .models( + model_ids + .iter() + .map(|id| ModelCard::new(*id)) + .collect::>(), + ) + .health_config(no_health_check()) + .build(), + ) + } + + #[test] + fn test_wildcard_hash_ring_matches_the_only_model() { + let registry = WorkerRegistry::new(); + registry + .register(worker_serving("http://w1:8080", &["llama-3"])) + .unwrap(); + registry + .register(worker_serving("http://w2:8080", &["llama-3"])) + .unwrap(); + + let wildcard = registry + .get_hash_ring(UNKNOWN_MODEL_ID) + .expect("requests naming no model need a ring"); + let model_ring = registry.get_hash_ring("llama-3").expect("per-model ring"); + + assert_eq!(wildcard.worker_count(), 2); + for key in ["alpha", "beta", "gamma"] { + assert_eq!( + wildcard.find_healthy_url(key, |_| true), + model_ring.find_healthy_url(key, |_| true), + "wildcard and single-model rings must agree on {key}" + ); + } + } + + #[test] + fn test_wildcard_hash_ring_unions_models() { + let registry = WorkerRegistry::new(); + registry + .register(worker_serving("http://w1:8080", &["llama-3"])) + .unwrap(); + registry + .register(worker_serving("http://w2:8080", &["gpt-4"])) + .unwrap(); + + let wildcard = registry.get_hash_ring(UNKNOWN_MODEL_ID).expect("ring"); + assert_eq!(wildcard.worker_count(), 2); + assert_eq!( + wildcard.find_healthy_url("key", |url| url == "http://w1:8080"), + Some("http://w1:8080") + ); + assert_eq!( + wildcard.find_healthy_url("key", |url| url == "http://w2:8080"), + Some("http://w2:8080") + ); + + // Per-model rings stay scoped to their own workers. + assert_eq!( + registry + .get_hash_ring("llama-3") + .expect("ring") + .worker_count(), + 1 + ); + } + + #[test] + fn test_wildcard_hash_ring_weights_multi_model_worker_once() { + let registry = WorkerRegistry::new(); + registry + .register(worker_serving("http://w1:8080", &["llama-3", "gpt-4"])) + .unwrap(); + registry + .register(worker_serving("http://w2:8080", &["gpt-4"])) + .unwrap(); + + let wildcard = registry.get_hash_ring(UNKNOWN_MODEL_ID).expect("ring"); + assert_eq!( + wildcard.worker_count(), + 2, + "a worker serving two models must not take a double share of the ring" + ); + } + + #[test] + fn test_wildcard_hash_ring_follows_removals() { + let registry = WorkerRegistry::new(); + registry + .register(worker_serving("http://w1:8080", &["llama-3"])) + .unwrap(); + registry + .register(worker_serving("http://w2:8080", &["gpt-4"])) + .unwrap(); + + registry.remove_by_url("http://w2:8080"); + let wildcard = registry.get_hash_ring(UNKNOWN_MODEL_ID).expect("ring"); + assert_eq!(wildcard.worker_count(), 1); + assert_eq!( + wildcard.find_healthy_url("key", |_| true), + Some("http://w1:8080") + ); + + registry.remove_by_url("http://w1:8080"); + assert!( + registry.get_hash_ring(UNKNOWN_MODEL_ID).is_none(), + "an empty registry has no ring to route against" + ); + } }