sgl-router: experimental Rust HTTP router for SGLang worker pools (#25851)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-25 15:34:05 +08:00
parent aae04b1241
commit 6e8fe176be
131 changed files with 28623 additions and 55 deletions
@@ -0,0 +1,42 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::{Context, Result};
+use dynamo_tokenizers::{traits::DecodeResult, Tokenizer};
+use std::sync::Arc;
+
+pub fn load(path: &str) -> Result<Arc<Tokenizer>> {
+    Tokenizer::from_file(path)
+        .map(Arc::new)
+        .with_context(|| format!("load tokenizer from {path}"))
+}
+
+pub fn encode(t: &Tokenizer, text: &str) -> Result<Vec<u32>> {
+    let enc = t.encode(text).context("encode")?;
+    Ok(enc.token_ids().to_vec())
+}
+
+/// Decode token ids to a complete UTF-8 string.
+///
+/// Non-streaming callers (e.g. `/v1/detokenize`) get the full result either way:
+/// - `DecodeResult::Complete(s)` — the token sequence ends on a codepoint boundary.
+/// - `DecodeResult::Partial(s)` — the token sequence ends mid-codepoint; `s` ends
+///   in U+FFFD. We return `s` as-is so the client sees the closest-possible string.
+///
+/// Streaming callers should NOT use this; they should consume `DecodeResult`
+/// directly and withhold the trailing U+FFFD until the next decode produces a
+/// `Complete` result.
+pub fn decode_complete(t: &Tokenizer, ids: &[u32], skip_special: bool) -> Result<String> {
+    let res = t.decode(ids, skip_special).context("decode")?;
+    Ok(match res {
+        DecodeResult::Complete(s) => s,
+        DecodeResult::Partial(s) => {
+            tracing::debug!(
+                n_tokens = ids.len(),
+                trailing_bytes = s.len(),
+                "decode_complete: tokenizer returned Partial for non-streaming call"
+            );
+            s
+        }
+    })
+}
@@ -0,0 +1,199 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+pub mod adapter;
+
+use anyhow::Result;
+use dashmap::DashMap;
+use dynamo_tokenizers::Tokenizer;
+use std::sync::Arc;
+
+#[derive(Default)]
+pub struct TokenizerRegistry {
+    inner: DashMap<String, Arc<Tokenizer>>,
+}
+
+impl std::fmt::Debug for TokenizerRegistry {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("TokenizerRegistry")
+            .field("models", &self.ids())
+            .finish()
+    }
+}
+
+impl TokenizerRegistry {
+    pub fn load_from_config(cfg: &crate::config::Config) -> Result<Self> {
+        let me = TokenizerRegistry::default();
+        for m in &cfg.models {
+            let t = adapter::load(&m.tokenizer_path)?;
+            me.inner.insert(m.id.clone(), t);
+        }
+        Ok(me)
+    }
+
+    pub fn get(&self, model_id: &str) -> Option<Arc<Tokenizer>> {
+        self.inner.get(model_id).map(|r| Arc::clone(&*r))
+    }
+
+    pub fn ids(&self) -> Vec<String> {
+        self.inner.iter().map(|kv| kv.key().clone()).collect()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Arc;
+
+    use crate::config::PolicyKind;
+
+    fn cfg() -> crate::config::Config {
+        crate::config::Config {
+            server: crate::config::ServerConfig {
+                host: "0".into(),
+                port: 0,
+            },
+            observability: Default::default(),
+            models: vec![crate::config::ModelConfig {
+                id: "tiny".into(),
+                tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
+                policy: PolicyKind::RoundRobin,
+                circuit_breaker: None,
+                cache_aware: None,
+            }],
+            discovery: crate::config::DiscoveryConfig {
+                backend: crate::config::DiscoveryBackend::StaticUrls(
+                    crate::config::StaticUrlsDiscoveryConfig {
+                        urls: vec!["http://placeholder:0".into()],
+                    },
+                ),
+            },
+            proxy: crate::config::ProxyConfig::default(),
+            active_load: crate::config::ActiveLoadConfig::default(),
+        }
+    }
+
+    #[test]
+    fn loads_from_config() {
+        let r = TokenizerRegistry::load_from_config(&cfg()).unwrap();
+        assert!(r.get("tiny").is_some());
+        assert!(r.get("missing").is_none());
+    }
+
+    #[test]
+    fn shared_arc_per_model() {
+        let r = TokenizerRegistry::load_from_config(&cfg()).unwrap();
+        let a = r.get("tiny").unwrap();
+        let b = r.get("tiny").unwrap();
+        assert!(
+            Arc::ptr_eq(&a, &b),
+            "registry should return shared Arc, not clones"
+        );
+    }
+
+    #[test]
+    fn decode_complete_preserves_round_trip() {
+        let r = TokenizerRegistry::load_from_config(&cfg()).unwrap();
+        let t = r.get("tiny").unwrap();
+        let ids = adapter::encode(&t, "hello world").unwrap();
+        assert!(!ids.is_empty());
+        let text = adapter::decode_complete(&t, &ids, true).unwrap();
+        // tiny BPE fixture is byte-level and lossless for ASCII.
+        assert_eq!(text, "hello world");
+    }
+
+    /// Forces `decode_complete` through its `DecodeResult::Partial` branch.
+    ///
+    /// Strategy A: the fixture is a GPT-2 byte-level BPE. The 4-byte UTF-8
+    /// emoji `😀` (`\xF0\x9F\x98\x80`) encodes into 2 byte-level BPE tokens
+    /// with this fixture: `[47249, 222]`. Decoding just the first token
+    /// yields a leading-bytes-only prefix that the HF adapter passes through
+    /// `String::from_utf8_lossy`, producing a trailing U+FFFD. dynamo's
+    /// `DecodeResult::from_decoded` then classifies that as `Partial`.
+    /// Pinning the literal token id keeps the test deterministic — if the
+    /// fixture or upstream BPE merges ever shift, this fails loudly rather
+    /// than silently dropping back into `Complete` and losing coverage.
+    #[test]
+    fn decode_complete_returns_string_on_partial_utf8() {
+        let r = TokenizerRegistry::load_from_config(&cfg()).unwrap();
+        let t = r.get("tiny").unwrap();
+
+        // Sanity-check that the fixture still tokenises `😀` the way we
+        // expect; if upstream changes this we want a loud failure here.
+        let full = adapter::encode(&t, "😀").unwrap();
+        assert_eq!(
+            full,
+            vec![47249, 222],
+            "fixture tokenisation drift: '😀' no longer encodes to [47249, 222]"
+        );
+
+        // Feed only the first token — its bytes are the leading 3 of a
+        // 4-byte UTF-8 codepoint, which is incomplete.
+        let s = adapter::decode_complete(&t, &full[..1], false).unwrap();
+
+        // We pin the exact output: the lossy decoder folds the 3 leading
+        // bytes into a single U+FFFD. Anything else (empty string, Err, or
+        // the original bytes) would be a regression.
+        assert_eq!(s, "\u{FFFD}");
+    }
+
+    /// Concurrent encode against one shared `Arc<Tokenizer>`. Pins that the
+    /// registry's `Arc<Tokenizer>` is `Send + Sync` and that
+    /// `dynamo_tokenizers::Tokenizer::encode` can be called concurrently
+    /// without interior mutability hazards. A regression that wraps
+    /// `Tokenizer` in `RefCell` / `!Sync` data would fail to compile;
+    /// a regression that introduces non-thread-safe internal caches
+    /// would surface as one of the tasks returning wrong ids (caught by
+    /// the per-task assertion against the sequentially-computed
+    /// reference).
+    ///
+    /// Uses a multi-thread runtime + `JoinSet` so the 10 tasks really do
+    /// run in parallel on distinct worker threads — a single-thread
+    /// runtime wouldn't exercise the `Sync` contract.
+    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+    async fn tokenizer_supports_concurrent_encode() {
+        use tokio::task::JoinSet;
+
+        let r = TokenizerRegistry::load_from_config(&cfg()).unwrap();
+        let t = r.get("tiny").unwrap();
+
+        // Build the reference sequentially — what each task should return.
+        let inputs: Vec<String> = (0..10).map(|i| format!("hello {i}")).collect();
+        let expected: Vec<Vec<u32>> = inputs
+            .iter()
+            .map(|s| adapter::encode(&t, s).unwrap())
+            .collect();
+
+        let mut set = JoinSet::new();
+        for (i, text) in inputs.into_iter().enumerate() {
+            let shared = Arc::clone(&t);
+            set.spawn(async move {
+                let ids = adapter::encode(&shared, &text).expect("concurrent encode must not fail");
+                (i, ids)
+            });
+        }
+
+        let mut got: Vec<Option<Vec<u32>>> = vec![None; expected.len()];
+        while let Some(joined) = set.join_next().await {
+            let (i, ids) = joined.expect("task panicked");
+            got[i] = Some(ids);
+        }
+
+        for (i, ids) in got.into_iter().enumerate() {
+            let ids = ids.unwrap_or_else(|| panic!("task {i} did not record a result"));
+            assert_eq!(
+                ids, expected[i],
+                "concurrent encode produced wrong tokens for task {i}; \
+                 sign of a non-thread-safe internal cache regression"
+            );
+        }
+    }
+
+    #[test]
+    fn missing_file_errors() {
+        let mut c = cfg();
+        c.models[0].tokenizer_path = "/nonexistent.json".into();
+        let err = TokenizerRegistry::load_from_config(&c).unwrap_err();
+        assert!(err.to_string().to_lowercase().contains("tokenizer"));
+    }
+}