From 1d1972139464a32931f9458fa7d22743b05f7527 Mon Sep 17 00:00:00 2001 From: Alex Nails Date: Mon, 18 May 2026 22:31:38 -0700 Subject: [PATCH] [gRPC] Native server: Rust crate (1/N) (#23506) --- rust/sglang-grpc/Cargo.toml | 3 +- rust/sglang-grpc/src/bridge.rs | 804 ++++++++++++++ rust/sglang-grpc/src/bridge/tests.rs | 9 + rust/sglang-grpc/src/lib.rs | 256 ++++- rust/sglang-grpc/src/server.rs | 1010 ++++++++++++++++++ rust/sglang-grpc/src/server/tests.rs | 75 ++ rust/sglang-grpc/src/tokenizers.rs | 127 +++ rust/sglang-grpc/src/utils/mod.rs | 8 + rust/sglang-grpc/src/utils/py_utils.rs | 70 ++ rust/sglang-grpc/src/utils/py_utils/tests.rs | 8 + rust/sglang-grpc/src/utils/request_utils.rs | 239 +++++ 11 files changed, 2569 insertions(+), 40 deletions(-) create mode 100644 rust/sglang-grpc/src/bridge.rs create mode 100644 rust/sglang-grpc/src/bridge/tests.rs create mode 100644 rust/sglang-grpc/src/server.rs create mode 100644 rust/sglang-grpc/src/server/tests.rs create mode 100644 rust/sglang-grpc/src/tokenizers.rs create mode 100644 rust/sglang-grpc/src/utils/mod.rs create mode 100644 rust/sglang-grpc/src/utils/py_utils.rs create mode 100644 rust/sglang-grpc/src/utils/py_utils/tests.rs create mode 100644 rust/sglang-grpc/src/utils/request_utils.rs diff --git a/rust/sglang-grpc/Cargo.toml b/rust/sglang-grpc/Cargo.toml index 26f666fee..d2236ed01 100644 --- a/rust/sglang-grpc/Cargo.toml +++ b/rust/sglang-grpc/Cargo.toml @@ -14,13 +14,12 @@ pyo3 = { version = "0.23", features = ["extension-module"] } tokio = { version = "1", features = ["full"] } tonic = { version = "0.12", features = ["gzip", "transport"] } prost = "0.13" -crossbeam-channel = "0.5" uuid = { version = "1", features = ["v4"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } serde_json = "1" tokenizers = { version = "0.21", default-features = false, features = ["onig"] } -tokio-stream = "0.1" +tokio-stream = { version = "0.1", features = ["net"] } async-stream = "0.3" [build-dependencies] diff --git a/rust/sglang-grpc/src/bridge.rs b/rust/sglang-grpc/src/bridge.rs new file mode 100644 index 000000000..75a926c01 --- /dev/null +++ b/rust/sglang-grpc/src/bridge.rs @@ -0,0 +1,804 @@ +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyDict}; +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex, MutexGuard}; +use tokio::runtime::Handle; +use tokio::sync::mpsc::error::TrySendError; +use tokio::sync::mpsc::{self, Receiver, Sender}; + +use crate::tokenizers::RustTokenizer; +use crate::utils::{json_map_to_pydict, py_value_to_json_string}; + +#[derive(Debug, Clone)] +pub enum ResponseChunk { + Data(ResponseData), + Finished(ResponseData), + Error(String), +} + +impl ResponseChunk { + fn is_terminal(&self) -> bool { + matches!(self, Self::Finished(_) | Self::Error(_)) + } +} + +#[derive(Debug, Clone)] +pub struct ResponseData { + pub text: Option, + pub output_ids: Option>, + pub embedding: Option>, + pub json_bytes: Option>, + pub meta_info: HashMap, +} + +pub const DEFAULT_RESPONSE_CHANNEL_CAPACITY: usize = 64; + +type BridgeStateRef = Arc>; + +#[derive(Default)] +struct BridgeState { + channels: HashMap>, + pending_sends: HashSet, + ready_callbacks: HashMap, + ready_signals: HashSet, + terminal_errors: HashMap, +} + +#[derive(Debug, Clone)] +pub enum TerminalError { + ChannelFull { rid: String }, + ClientDisconnected { rid: String }, + Aborted { rid: String }, +} + +impl TerminalError { + pub fn message(&self) -> String { + match self { + Self::ChannelFull { rid } => { + format!("gRPC response channel full for {rid}: client not consuming") + } + Self::ClientDisconnected { rid } => { + format!("gRPC client disconnected for request {rid}") + } + Self::Aborted { rid } => format!("Request aborted: {rid}"), + } + } +} + +#[pyclass(eq, eq_int)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ChunkSendStatus { + Ready, + Pending, + Closed, +} + +fn lock_or_recover<'a, T>(mutex: &'a Mutex, name: &'static str) -> MutexGuard<'a, T> { + mutex.lock().unwrap_or_else(|poisoned| { + tracing::warn!(mutex = name, "Recovering from poisoned gRPC bridge mutex"); + poisoned.into_inner() + }) +} + +/// Holds a reference to the Python RuntimeHandle and manages per-request channels. +pub struct PyBridge { + runtime_handle: PyObject, + state: BridgeStateRef, + rust_tokenizer: Option, + context_len: i32, + response_channel_capacity: usize, + tokio_handle: Handle, +} + +impl PyBridge { + pub fn new( + runtime_handle: PyObject, + rust_tokenizer: Option, + context_len: i32, + response_channel_capacity: usize, + tokio_handle: Handle, + ) -> Self { + debug_assert!( + response_channel_capacity > 0, + "response_channel_capacity must be normalized by start_server" + ); + Self { + runtime_handle, + state: Arc::new(Mutex::new(BridgeState::default())), + rust_tokenizer, + context_len, + response_channel_capacity, + tokio_handle, + } + } + + /// Access the Rust tokenizer (if available). + pub fn rust_tokenizer(&self) -> Option<&RustTokenizer> { + self.rust_tokenizer.as_ref() + } + + /// Return the model's context length. + pub fn context_len(&self) -> i32 { + self.context_len + } + + // ------------------------------------------------------------------ + // Channel + callback helpers + // ------------------------------------------------------------------ + + fn create_channel(&self, rid: &str) -> PyResult> { + let (sender, receiver) = mpsc::channel(self.response_channel_capacity); + let mut state = lock_or_recover(self.state.as_ref(), "state"); + if state.channels.contains_key(rid) { + return Err(PyRuntimeError::new_err(format!( + "Duplicate active gRPC request id: {}", + rid + ))); + } + state.channels.insert(rid.to_string(), sender); + state.terminal_errors.remove(rid); + state.ready_callbacks.remove(rid); + state.ready_signals.remove(rid); + state.pending_sends.remove(rid); + Ok(receiver) + } + + fn make_chunk_callback(&self, py: Python<'_>, rid: String) -> PyResult { + let callback = ChunkCallback { + rid, + state: self.state.clone(), + runtime_handle: self.runtime_handle.clone_ref(py), + tokio_handle: self.tokio_handle.clone(), + }; + let py_callback = Py::new(py, callback)?; + Ok(py_callback.into_any()) + } + + fn make_json_callback(&self, py: Python<'_>, rid: String) -> PyResult { + let callback = JsonChunkCallback { + rid, + state: self.state.clone(), + runtime_handle: self.runtime_handle.clone_ref(py), + tokio_handle: self.tokio_handle.clone(), + }; + let py_callback = Py::new(py, callback)?; + Ok(py_callback.into_any()) + } + + // ------------------------------------------------------------------ + // Consolidated request submission (generate / embed / classify) + // ------------------------------------------------------------------ + + /// Submit a generate or embed request by passing a pre-built dict to Python. + /// + /// `req_type` is "generate", "embed", or "classify". + /// `req_dict` contains fields matching GenerateReqInput or EmbeddingReqInput. + pub fn submit_request( + &self, + rid: &str, + req_type: &str, + req_dict: HashMap, + ) -> PyResult> { + let receiver = self.create_channel(rid)?; + let rid_owned = rid.to_string(); + + let result = Python::with_gil(|py| -> PyResult<()> { + let py_req_dict = json_map_to_pydict(py, &req_dict)?; + let callback = self.make_chunk_callback(py, rid_owned)?; + + let kwargs = PyDict::new(py); + kwargs.set_item("req_type", req_type)?; + kwargs.set_item("req_dict", py_req_dict)?; + kwargs.set_item("chunk_callback", callback)?; + + self.runtime_handle + .call_method(py, "submit_request", (), Some(&kwargs))?; + Ok(()) + }); + + match result { + Ok(()) => Ok(receiver), + Err(err) => { + self.remove_channel(rid); + Err(err) + } + } + } + + // ------------------------------------------------------------------ + // Abort + // ------------------------------------------------------------------ + + pub fn abort(&self, rid: &str, abort_all: bool) -> PyResult<()> { + if !abort_all && rid.trim().is_empty() { + return Err(PyValueError::new_err( + "Abort requires a non-empty rid unless abort_all is true", + )); + } + + let should_call_python = if abort_all { + let mut state = lock_or_recover(self.state.as_ref(), "state"); + let rids = state + .channels + .drain() + .map(|(rid, _)| rid) + .collect::>(); + let affected = rids.len(); + state.pending_sends.clear(); + state.ready_callbacks.clear(); + state.ready_signals.clear(); + for channel_rid in rids { + state.terminal_errors.insert( + channel_rid.clone(), + TerminalError::Aborted { rid: channel_rid }, + ); + } + tracing::debug!(affected, "gRPC abort_all cleared active response channels"); + true + } else { + let mut state = lock_or_recover(self.state.as_ref(), "state"); + let was_active = remove_channel_refs_locked(&mut state, rid); + if was_active { + state + .terminal_errors + .insert(rid.to_string(), TerminalError::Aborted { rid: rid.into() }); + } else { + tracing::debug!(rid, "Ignoring abort for inactive gRPC request id"); + } + was_active + }; + + if !should_call_python { + return Ok(()); + } + + Python::with_gil(|py| { + self.runtime_handle + .call_method1(py, "abort", (rid, abort_all))?; + Ok(()) + }) + } + + // ------------------------------------------------------------------ + // Info / control RPCs (synchronous, small data) + // ------------------------------------------------------------------ + + pub fn get_model_info(&self) -> PyResult { + Python::with_gil(|py| { + let result = self.runtime_handle.call_method0(py, "get_model_info")?; + result.extract::(py) + }) + } + + pub fn get_server_info(&self) -> PyResult { + Python::with_gil(|py| { + let result = self.runtime_handle.call_method0(py, "get_server_info")?; + result.extract::(py) + }) + } + + pub fn health_check(&self) -> PyResult { + Python::with_gil(|py| { + let result = self.runtime_handle.call_method0(py, "health_check")?; + result.extract::(py) + }) + } + + /// Tokenize via Python (fallback when Rust tokenizer unavailable). + pub fn tokenize_py(&self, text: &str, add_special_tokens: bool) -> PyResult { + Python::with_gil(|py| { + let result = + self.runtime_handle + .call_method1(py, "tokenize", (text, add_special_tokens))?; + result.extract::(py) + }) + } + + /// Detokenize via Python (fallback when Rust tokenizer unavailable). + pub fn detokenize_py(&self, tokens: Vec) -> PyResult { + Python::with_gil(|py| { + let result = self + .runtime_handle + .call_method1(py, "detokenize", (tokens,))?; + result.extract::(py) + }) + } + + pub fn list_models(&self) -> PyResult { + Python::with_gil(|py| { + let result = self.runtime_handle.call_method0(py, "list_models")?; + result.extract::(py) + }) + } + + fn submit_json(&self, rid: &str, call: F) -> PyResult> + where + F: for<'py> FnOnce(Python<'py>, &PyObject, PyObject) -> PyResult<()>, + { + // Closure args are: current Python token, RuntimeHandle, and the JSON chunk callback. + let receiver = self.create_channel(rid)?; + let rid_owned = rid.to_string(); + + let result = Python::with_gil(|py| -> PyResult<()> { + let callback = self.make_json_callback(py, rid_owned)?; + call(py, &self.runtime_handle, callback) + }); + + match result { + Ok(()) => Ok(receiver), + Err(err) => { + self.remove_channel(rid); + Err(err) + } + } + } + + pub fn submit_get_load( + &self, + rid: &str, + dp_rank: Option, + ) -> PyResult> { + self.submit_json(rid, move |py, runtime_handle, callback| { + runtime_handle.call_method1(py, "get_load", (callback, dp_rank))?; + Ok(()) + }) + } + + pub fn submit_flush_cache(&self, rid: &str) -> PyResult> { + self.submit_json(rid, |py, runtime_handle, callback| { + runtime_handle.call_method1(py, "flush_cache", (callback,))?; + Ok(()) + }) + } + + pub fn submit_pause_generation( + &self, + rid: &str, + mode: &str, + ) -> PyResult> { + self.submit_json(rid, move |py, runtime_handle, callback| { + runtime_handle.call_method1(py, "pause_generation", (mode, callback))?; + Ok(()) + }) + } + + pub fn submit_continue_generation(&self, rid: &str) -> PyResult> { + self.submit_json(rid, |py, runtime_handle, callback| { + runtime_handle.call_method1(py, "continue_generation", (callback,))?; + Ok(()) + }) + } + + pub fn submit_start_profile( + &self, + rid: &str, + output_dir: Option<&str>, + ) -> PyResult> { + self.submit_json(rid, move |py, runtime_handle, callback| { + runtime_handle.call_method1(py, "start_profile", (output_dir, callback))?; + Ok(()) + }) + } + + pub fn submit_stop_profile(&self, rid: &str) -> PyResult> { + self.submit_json(rid, |py, runtime_handle, callback| { + runtime_handle.call_method1(py, "stop_profile", (callback,))?; + Ok(()) + }) + } + + pub fn submit_update_weights( + &self, + rid: &str, + model_path: &str, + load_format: Option<&str>, + ) -> PyResult> { + self.submit_json(rid, move |py, runtime_handle, callback| { + runtime_handle.call_method1( + py, + "update_weights_from_disk", + (model_path, load_format, callback), + )?; + Ok(()) + }) + } + + // ------------------------------------------------------------------ + // OpenAI pass-through RPCs + // ------------------------------------------------------------------ + + pub fn submit_openai( + &self, + rid: &str, + method_name: &str, + json_body: &[u8], + trace_headers: &HashMap, + ) -> PyResult> { + self.submit_json(rid, move |py, runtime_handle, callback| { + let kwargs = PyDict::new(py); + let py_bytes = PyBytes::new(py, json_body); + kwargs.set_item("json_body", py_bytes)?; + if !trace_headers.is_empty() { + let py_trace_headers = PyDict::new(py); + for (key, value) in trace_headers { + py_trace_headers.set_item(key, value)?; + } + kwargs.set_item("trace_headers", py_trace_headers)?; + } + + kwargs.set_item("chunk_callback", callback)?; + + runtime_handle.call_method(py, method_name, (), Some(&kwargs))?; + Ok(()) + }) + } + + pub fn remove_channel(&self, rid: &str) { + let mut state = lock_or_recover(self.state.as_ref(), "state"); + remove_channel_refs_locked(&mut state, rid); + state.terminal_errors.remove(rid); + } + + pub fn take_terminal_error(&self, rid: &str) -> Option { + let mut state = lock_or_recover(self.state.as_ref(), "state"); + state.terminal_errors.remove(rid) + } +} + +fn close_channel_with_error( + py: Python<'_>, + rid: &str, + state: &BridgeStateRef, + runtime_handle: &PyObject, + error: TerminalError, +) { + let mut state = lock_or_recover(state.as_ref(), "state"); + remove_channel_refs_locked(&mut state, rid); + state.terminal_errors.insert(rid.to_string(), error); + drop(state); + let _ = runtime_handle.call_method1(py, "abort", (rid, false)); +} + +fn remove_channel_refs_locked(state: &mut BridgeState, rid: &str) -> bool { + let had_channel = state.channels.remove(rid).is_some(); + let had_pending = state.pending_sends.remove(rid); + let had_callback = state.ready_callbacks.remove(rid).is_some(); + let had_signal = state.ready_signals.remove(rid); + had_channel || had_pending || had_callback || had_signal +} + +fn remove_channel_refs(rid: &str, state: &BridgeStateRef) { + let mut state = lock_or_recover(state.as_ref(), "state"); + remove_channel_refs_locked(&mut state, rid); +} + +fn register_pending_send(rid: &str, state: &BridgeStateRef) -> bool { + let mut state = lock_or_recover(state.as_ref(), "state"); + state.pending_sends.insert(rid.to_string()) +} + +fn mark_send_ready(py: Python<'_>, rid: &str, state: &BridgeStateRef) -> Option { + let mut state = lock_or_recover(state.as_ref(), "state"); + state.pending_sends.remove(rid); + if let Some(callback) = state.ready_callbacks.get(rid) { + Some(callback.clone_ref(py)) + } else { + state.ready_signals.insert(rid.to_string()); + None + } +} + +fn notify_ready(py: Python<'_>, rid: &str, callback: PyObject) { + if let Err(err) = callback.call0(py) { + tracing::warn!(rid, "gRPC on_ready callback failed: {}", err); + } +} + +fn set_on_ready_for_rid( + py: Python<'_>, + rid: &str, + state: &BridgeStateRef, + on_ready: PyObject, +) -> PyResult<()> { + let should_notify = { + let mut state = lock_or_recover(state.as_ref(), "state"); + state + .ready_callbacks + .insert(rid.to_string(), on_ready.clone_ref(py)); + state.ready_signals.remove(rid) + }; + if should_notify { + on_ready.call0(py)?; + } + Ok(()) +} + +fn clear_on_ready_for_rid(rid: &str, state: &BridgeStateRef) { + // End notifications for this rid. Do not call set_on_ready again for the same rid. + let mut state = lock_or_recover(state.as_ref(), "state"); + state.ready_callbacks.remove(rid); + state.ready_signals.remove(rid); +} + +fn try_send_chunk( + py: Python<'_>, + rid: &str, + state: &BridgeStateRef, + runtime_handle: &PyObject, + tokio_handle: &Handle, + sender: &Sender, + msg: ResponseChunk, +) -> PyResult { + let terminal = msg.is_terminal(); + match sender.try_send(msg) { + Ok(()) => { + if terminal { + remove_channel_refs(rid, state); + } + Ok(ChunkSendStatus::Ready) + } + Err(TrySendError::Full(msg)) => { + if !register_pending_send(rid, state) { + tracing::warn!( + rid, + "gRPC bridge received another chunk before the parked chunk drained; closing stream" + ); + close_channel_with_error( + py, + rid, + state, + runtime_handle, + TerminalError::ChannelFull { rid: rid.into() }, + ); + return Ok(ChunkSendStatus::Closed); + } + + let rid_owned = rid.to_string(); + let state = state.clone(); + let runtime_handle = runtime_handle.clone_ref(py); + let sender = sender.clone(); + + tokio_handle.spawn(async move { + match sender.send(msg).await { + Ok(()) => { + if terminal { + // Terminal chunks end the producer contract; no further on_ready + // signal is fired after a parked Finished/Error drains. + remove_channel_refs(&rid_owned, &state); + return; + } + + Python::with_gil(|py| { + if let Some(callback) = mark_send_ready(py, &rid_owned, &state) { + notify_ready(py, &rid_owned, callback); + } + }); + } + Err(_) => { + Python::with_gil(|py| { + close_channel_with_error( + py, + &rid_owned, + &state, + &runtime_handle, + TerminalError::ClientDisconnected { + rid: rid_owned.clone(), + }, + ); + }); + } + } + }); + + Ok(ChunkSendStatus::Pending) + } + Err(TrySendError::Closed(_)) => { + close_channel_with_error( + py, + rid, + state, + runtime_handle, + TerminalError::ClientDisconnected { rid: rid.into() }, + ); + Ok(ChunkSendStatus::Closed) + } + } +} + +// Typed chunk callback for SGLang-native RPCs (dict-based chunks). +#[pyclass] +struct ChunkCallback { + rid: String, + state: BridgeStateRef, + runtime_handle: PyObject, + tokio_handle: Handle, +} + +#[pymethods] +impl ChunkCallback { + /// Register before producing chunks. If a parked chunk drained before registration, + /// Rust fires `on_ready` immediately so late registration cannot miss the edge. + fn set_on_ready(&self, py: Python<'_>, on_ready: PyObject) -> PyResult<()> { + set_on_ready_for_rid(py, &self.rid, &self.state, on_ready) + } + + fn clear_on_ready(&self) { + clear_on_ready_for_rid(&self.rid, &self.state); + } + + #[pyo3(signature = (chunk, finished=false, error=None))] + fn __call__( + &self, + chunk: &Bound<'_, PyDict>, + finished: bool, + error: Option, + ) -> PyResult { + let py = chunk.py(); + let state = lock_or_recover(self.state.as_ref(), "state"); + let sender = match state.channels.get(&self.rid) { + Some(s) => s.clone(), + None => return Ok(ChunkSendStatus::Closed), + }; + drop(state); + + if let Some(err_msg) = error { + return try_send_chunk( + py, + &self.rid, + &self.state, + &self.runtime_handle, + &self.tokio_handle, + &sender, + ResponseChunk::Error(err_msg), + ); + } + + let text: Option = chunk + .get_item("text")? + .and_then(|v| v.extract::().ok()); + + let output_ids: Option> = chunk + .get_item("output_ids")? + .and_then(|v| v.extract::>().ok()); + + let embedding: Option> = chunk + .get_item("embedding")? + .and_then(|v| v.extract::>().ok()); + + let meta_info = extract_meta_info(chunk); + + let data = ResponseData { + text, + output_ids, + embedding, + json_bytes: None, + meta_info, + }; + + let msg = if finished { + ResponseChunk::Finished(data) + } else { + ResponseChunk::Data(data) + }; + + try_send_chunk( + py, + &self.rid, + &self.state, + &self.runtime_handle, + &self.tokio_handle, + &sender, + msg, + ) + } +} + +// JSON chunk callback for OpenAI pass-through RPCs (raw bytes). +#[pyclass] +struct JsonChunkCallback { + rid: String, + state: BridgeStateRef, + runtime_handle: PyObject, + tokio_handle: Handle, +} + +#[pymethods] +impl JsonChunkCallback { + /// Register before producing chunks. If a parked chunk drained before registration, + /// Rust fires `on_ready` immediately so late registration cannot miss the edge. + fn set_on_ready(&self, py: Python<'_>, on_ready: PyObject) -> PyResult<()> { + set_on_ready_for_rid(py, &self.rid, &self.state, on_ready) + } + + fn clear_on_ready(&self) { + clear_on_ready_for_rid(&self.rid, &self.state); + } + + #[pyo3(signature = (chunk_bytes, finished=false, error=None, status_code=None))] + fn __call__( + &self, + chunk_bytes: &Bound<'_, pyo3::PyAny>, + finished: bool, + error: Option, + status_code: Option, + ) -> PyResult { + let py = chunk_bytes.py(); + let state = lock_or_recover(self.state.as_ref(), "state"); + let sender = match state.channels.get(&self.rid) { + Some(s) => s.clone(), + None => return Ok(ChunkSendStatus::Closed), + }; + drop(state); + + if let Some(err_msg) = error { + return try_send_chunk( + py, + &self.rid, + &self.state, + &self.runtime_handle, + &self.tokio_handle, + &sender, + ResponseChunk::Error(err_msg), + ); + } + + let bytes_data: Vec = if let Ok(b) = chunk_bytes.extract::>() { + b + } else if let Ok(s) = chunk_bytes.extract::() { + s.into_bytes() + } else { + vec![] + }; + + let mut meta_info = HashMap::new(); + if let Some(code) = status_code { + meta_info.insert("status_code".to_string(), code.to_string()); + } + + let data = ResponseData { + text: None, + output_ids: None, + embedding: None, + json_bytes: Some(bytes_data), + meta_info, + }; + + let msg = if finished { + ResponseChunk::Finished(data) + } else { + ResponseChunk::Data(data) + }; + + try_send_chunk( + py, + &self.rid, + &self.state, + &self.runtime_handle, + &self.tokio_handle, + &sender, + msg, + ) + } +} + +fn extract_meta_info(chunk: &Bound<'_, PyDict>) -> HashMap { + let mut meta = HashMap::new(); + if let Ok(Some(meta_obj)) = chunk.get_item("meta_info") + && let Ok(meta_dict) = meta_obj.downcast::() + { + for (k, v) in meta_dict.iter() { + // The proto schema is map; encode each Python value as JSON + // so clients can recover numbers, booleans, arrays, and objects losslessly. + if let Ok(key) = k.extract::() + && let Ok(val) = py_value_to_json_string(&v) + { + meta.insert(key, val); + } + } + } + meta +} + +#[cfg(test)] +mod tests; diff --git a/rust/sglang-grpc/src/bridge/tests.rs b/rust/sglang-grpc/src/bridge/tests.rs new file mode 100644 index 000000000..d8b1138a5 --- /dev/null +++ b/rust/sglang-grpc/src/bridge/tests.rs @@ -0,0 +1,9 @@ +use super::*; + +#[test] +fn terminal_error_messages_include_request_id() { + let error = TerminalError::ClientDisconnected { + rid: "rid".to_string(), + }; + assert!(error.message().contains("rid")); +} diff --git a/rust/sglang-grpc/src/lib.rs b/rust/sglang-grpc/src/lib.rs index e3a560517..c3909322c 100644 --- a/rust/sglang-grpc/src/lib.rs +++ b/rust/sglang-grpc/src/lib.rs @@ -1,21 +1,32 @@ -use pyo3::prelude::*; -use std::sync::Arc; -use tokio::sync::Notify; +pub mod bridge; +pub mod server; +pub mod tokenizers; +pub(crate) mod utils; pub mod proto { tonic::include_proto!("sglang.runtime.v1"); } -/// Handle returned by `start_server` — used to shut down the gRPC server. +use pyo3::prelude::*; +use std::net::{SocketAddr, TcpListener}; +use std::sync::Arc; +use tokio::sync::Notify; +use tokio::time::Duration; +use tracing_subscriber::EnvFilter; + +use bridge::{ChunkSendStatus, DEFAULT_RESPONSE_CHANNEL_CAPACITY, PyBridge}; +use tokenizers::RustTokenizer; + +/// Handle returned to Python that controls the running gRPC server. #[pyclass] -pub struct GrpcServerHandle { +struct GrpcServerHandle { shutdown: Arc, join_handle: Option>, } #[pymethods] impl GrpcServerHandle { - /// Signal the server to stop and wait for the background thread to exit. + /// Gracefully shut down the gRPC server. fn shutdown(&mut self) { self.shutdown.notify_one(); if let Some(handle) = self.join_handle.take() { @@ -23,51 +34,220 @@ impl GrpcServerHandle { } } - /// Returns `true` while the server thread is still running. + /// Check if the server thread is still running. fn is_alive(&self) -> bool { - self.join_handle - .as_ref() - .map_or(false, |h| !h.is_finished()) + self.join_handle.as_ref().is_some_and(|h| !h.is_finished()) } } -/// Start the gRPC server in a background thread. +struct TokenizerInfo { + tokenizer_path: Option, + tokenizer_mode: Option, + context_len: i32, +} + +/// Extract tokenizer path/mode and context_len from the Python RuntimeHandle (one-time GIL). /// -/// * `host` – bind address (e.g. "0.0.0.0") -/// * `port` – listen port -/// * `runtime_handle` – Python `RuntimeHandle` object (from `grpc_bridge.py`) +/// Missing `tokenizer_manager` indicates a misconfigured runtime handle and should surface at +/// startup. Sub-fields are best-effort because unsupported native tokenizer backends can still +/// fall back to Python tokenization. +fn try_get_attr( + py: Python<'_>, + obj: &PyObject, + attr: &'static str, + context: &'static str, +) -> Option { + obj.getattr(py, attr).map(Some).unwrap_or_else(|err| { + tracing::debug!("{}.{} is unavailable: {}", context, attr, err); + None + }) +} + +fn try_get_attr_str( + py: Python<'_>, + obj: &PyObject, + attr: &'static str, + context: &'static str, +) -> Option { + try_get_attr(py, obj, attr, context).and_then(|value| { + value.extract(py).map(Some).unwrap_or_else(|err| { + tracing::debug!("Could not extract {}.{} as string: {}", context, attr, err); + None + }) + }) +} + +fn try_get_attr_i32( + py: Python<'_>, + obj: &PyObject, + attr: &'static str, + context: &'static str, +) -> Option { + try_get_attr(py, obj, attr, context).and_then(|value| { + value.extract(py).map(Some).unwrap_or_else(|err| { + tracing::debug!("Could not extract {}.{} as i32: {}", context, attr, err); + None + }) + }) +} + +fn extract_tokenizer_info(runtime_handle: &PyObject) -> PyResult { + Python::with_gil(|py| { + let tm = runtime_handle + .getattr(py, "tokenizer_manager") + .map_err(|err| { + pyo3::exceptions::PyValueError::new_err(format!( + "runtime_handle.tokenizer_manager is required: {}", + err + )) + })?; + + let server_args = try_get_attr(py, &tm, "server_args", "tokenizer_manager"); + + let tokenizer_path = server_args + .as_ref() + .and_then(|args| try_get_attr_str(py, args, "tokenizer_path", "server_args")) + .or_else(|| { + server_args + .as_ref() + .and_then(|args| try_get_attr_str(py, args, "model_path", "server_args")) + }) + .or_else(|| try_get_attr_str(py, &tm, "model_path", "tokenizer_manager")); + if tokenizer_path.is_none() { + tracing::warn!("Could not extract tokenizer path; Rust tokenizer disabled"); + } + + let tokenizer_mode = server_args + .as_ref() + .and_then(|args| try_get_attr_str(py, args, "tokenizer_mode", "server_args")); + + let context_len = try_get_attr(py, &tm, "model_config", "tokenizer_manager") + .and_then(|model_config| { + try_get_attr_i32(py, &model_config, "context_len", "model_config") + }) + .unwrap_or_else(|| { + tracing::warn!("Could not extract model_config.context_len; defaulting to 0"); + 0 + }); + + Ok(TokenizerInfo { + tokenizer_path, + tokenizer_mode, + context_len, + }) + }) +} + +/// Start the gRPC server in a background thread with its own Tokio runtime. /// -/// Returns a `GrpcServerHandle` that can be used to shut the server down. +/// Args: +/// host: Bind address (e.g., "0.0.0.0") +/// port: Port number (e.g., 40000) +/// runtime_handle: Python RuntimeHandle object with submit_generate, submit_embed, abort, etc. +/// +/// Returns: +/// GrpcServerHandle that can be used to shut down the server. #[pyfunction] -fn start_server(host: String, port: u16, runtime_handle: PyObject) -> PyResult { - let _ = &runtime_handle; // Will be used in Phase 1 PR 2 +#[pyo3(signature = (host, port, runtime_handle, worker_threads=4, response_channel_capacity=64, response_timeout_secs=300))] +fn start_server( + host: String, + port: u16, + runtime_handle: PyObject, + worker_threads: usize, + response_channel_capacity: usize, + response_timeout_secs: u64, +) -> PyResult { + // Best-effort: embedding processes may initialize tracing themselves. + let _ = tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")), + ) + .try_init(); + + let addr: SocketAddr = format!("{}:{}", host, port) + .parse() + .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid address: {}", e)))?; + let worker_threads = worker_threads.max(1); + let response_channel_capacity = if response_channel_capacity == 0 { + tracing::warn!( + default = DEFAULT_RESPONSE_CHANNEL_CAPACITY, + "response_channel_capacity must be positive; using default" + ); + DEFAULT_RESPONSE_CHANNEL_CAPACITY + } else { + response_channel_capacity + }; + let response_timeout_secs = if response_timeout_secs == 0 { + tracing::warn!( + default = server::DEFAULT_RESPONSE_TIMEOUT_SECS, + "response_timeout_secs must be positive; using default" + ); + server::DEFAULT_RESPONSE_TIMEOUT_SECS + } else { + response_timeout_secs + }; + let response_timeout = Duration::from_secs(response_timeout_secs); + let listener = TcpListener::bind(addr).map_err(|e| { + pyo3::exceptions::PyRuntimeError::new_err(format!( + "Failed to bind gRPC server to {}: {}", + addr, e + )) + })?; + listener.set_nonblocking(true).map_err(|e| { + pyo3::exceptions::PyRuntimeError::new_err(format!( + "Failed to configure gRPC listener for {}: {}", + addr, e + )) + })?; + + let tokenizer_info = extract_tokenizer_info(&runtime_handle)?; + + let rust_tokenizer = tokenizer_info.tokenizer_path.as_deref().and_then(|p| { + RustTokenizer::from_tokenizer_path( + p, + tokenizer_info.tokenizer_mode.as_deref(), + tokenizer_info.context_len, + ) + }); + + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(worker_threads) + .enable_all() + .thread_name("sglang-grpc-tokio") + .build() + .map_err(|err| { + pyo3::exceptions::PyRuntimeError::new_err(format!( + "Failed to build Tokio runtime for gRPC server: {}", + err + )) + })?; + let tokio_handle = rt.handle().clone(); + + let bridge = Arc::new(PyBridge::new( + runtime_handle, + rust_tokenizer, + tokenizer_info.context_len, + response_channel_capacity, + tokio_handle, + )); let shutdown = Arc::new(Notify::new()); let shutdown_clone = shutdown.clone(); - - let addr_str = format!("{}:{}", host, port); - let addr: std::net::SocketAddr = addr_str - .parse() - .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Bad address: {e}")))?; + let bridge_clone = bridge.clone(); let join_handle = std::thread::Builder::new() - .name("grpc-server".into()) + .name("sglang-grpc".to_string()) .spawn(move || { - let rt = tokio::runtime::Builder::new_multi_thread() - .worker_threads(4) - .enable_all() - .build() - .expect("Failed to build Tokio runtime"); - - rt.block_on(async move { - tracing::info!("gRPC server listening on {}", addr); - // Server implementation will be added in PR 2. - // For now, just wait for shutdown signal. - shutdown_clone.notified().await; - tracing::info!("gRPC server shutting down"); - }); + if let Err(e) = rt.block_on(server::run_grpc_server( + listener, + bridge_clone, + shutdown_clone, + response_timeout, + )) { + tracing::error!("gRPC server exited with error: {}", e); + } }) .map_err(|e| { - pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to spawn thread: {e}")) + pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to spawn gRPC thread: {}", e)) })?; Ok(GrpcServerHandle { @@ -76,10 +256,10 @@ fn start_server(host: String, port: u16, runtime_handle: PyObject) -> PyResult) -> PyResult<()> { m.add_function(wrap_pyfunction!(start_server, m)?)?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/rust/sglang-grpc/src/server.rs b/rust/sglang-grpc/src/server.rs new file mode 100644 index 000000000..33f33c6c9 --- /dev/null +++ b/rust/sglang-grpc/src/server.rs @@ -0,0 +1,1010 @@ +use std::collections::HashMap; +use std::pin::Pin; +use std::sync::Arc; + +use pyo3::PyErr; +use pyo3::Python; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use tokio::sync::{Notify, mpsc::Receiver}; +use tokio::time::{Duration, timeout}; +use tokio_stream::Stream; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::{Request, Response, Status}; + +use crate::bridge::{PyBridge, ResponseChunk, TerminalError}; +use crate::proto; +use crate::utils::{ + build_classify_dict, build_embed_dict, build_generate_dict, build_text_embed_dict, + build_text_generate_dict, extract_model_path, +}; + +pub struct SglangServiceImpl { + pub bridge: Arc, + pub response_timeout: Duration, +} + +type StreamResult = Pin> + Send + 'static>>; +pub const DEFAULT_RESPONSE_TIMEOUT_SECS: u64 = 300; + +/// 64 MiB — leaves headroom for multimodal inputs and OpenAI JSON pass-through bodies, +/// well above tonic's 4 MiB decode default. +pub const DEFAULT_GRPC_MAX_MESSAGE_SIZE: usize = 64 * 1024 * 1024; + +/// Resolve the per-message size cap (bytes) applied to the Tonic encoder/decoder. +// +// TODO(grpc-args): promote SGLANG_TONIC_PAYLOAD to a proper `--grpc-max-message-size` +// server argument once the launcher PR (3/4) wires server args through. +fn resolve_max_message_size() -> usize { + match std::env::var("SGLANG_TONIC_PAYLOAD") { + Ok(raw) => match raw.parse::() { + Ok(n) if n > 0 => { + tracing::info!( + bytes = n, + "Using SGLANG_TONIC_PAYLOAD override for gRPC max message size" + ); + n + } + _ => { + tracing::warn!( + value = %raw, + default = DEFAULT_GRPC_MAX_MESSAGE_SIZE, + "Ignoring invalid SGLANG_TONIC_PAYLOAD; using default" + ); + DEFAULT_GRPC_MAX_MESSAGE_SIZE + } + }, + Err(_) => DEFAULT_GRPC_MAX_MESSAGE_SIZE, + } +} + +/// Classify a bridge `PyErr` into the right gRPC `Status`. +/// +/// `PyValueError` / `PyTypeError` mean the client sent bad input — surface as +/// `INVALID_ARGUMENT` so callers can distinguish them from server failures. +/// Everything else (typically `PyRuntimeError`, but also Python tracebacks +/// from inside the tokenizer manager) maps to `INTERNAL`. +fn pyerr_to_status(err: PyErr, context: &str) -> Status { + let is_client_error = Python::with_gil(|py| { + err.is_instance_of::(py) || err.is_instance_of::(py) + }); + let msg = format!("{}: {}", context, err); + if is_client_error { + Status::invalid_argument(msg) + } else { + Status::internal(msg) + } +} + +async fn recv_chunk_with_timeout( + receiver: &mut Receiver, + response_timeout: Duration, + timeout_message: impl FnOnce() -> String, +) -> Result, Status> { + timeout(response_timeout, receiver.recv()) + .await + .map_err(|_| Status::deadline_exceeded(timeout_message())) +} + +struct RequestAbortGuard { + bridge: Arc, + rid: String, + armed: bool, +} + +impl RequestAbortGuard { + fn new(bridge: Arc, rid: impl Into) -> Self { + Self { + bridge, + rid: rid.into(), + armed: true, + } + } + + fn disarm(&mut self) { + self.armed = false; + } + + fn abort_now(&mut self) { + if self.armed { + self.armed = false; + spawn_abort(self.bridge.clone(), self.rid.clone()); + } + } +} + +impl Drop for RequestAbortGuard { + fn drop(&mut self) { + if self.armed { + // Dropping a response stream means the client stopped consuming; propagate + // cancellation to Python without blocking the Tokio worker. + spawn_abort(self.bridge.clone(), self.rid.clone()); + } + } +} + +fn spawn_abort(bridge: Arc, rid: String) { + match tokio::runtime::Handle::try_current() { + Ok(handle) => { + let _ = handle.spawn_blocking(move || { + let _ = bridge.abort(&rid, false); + }); + } + Err(_) => { + tracing::warn!( + rid, + "Skipping gRPC request abort because no Tokio runtime is available" + ); + } + } +} + +async fn recv_terminal_chunk_for_request( + bridge: &Arc, + rid: &str, + receiver: &mut Receiver, + response_timeout: Duration, +) -> Result { + let mut abort_guard = RequestAbortGuard::new(bridge.clone(), rid.to_string()); + + match recv_chunk_with_timeout(receiver, response_timeout, || { + format!("Request timed out after {}s", response_timeout.as_secs()) + }) + .await + { + Ok(Some(ResponseChunk::Data(_))) => { + tracing::warn!( + rid, + "Unary gRPC response received non-terminal Data chunk; expected Finished" + ); + abort_guard.abort_now(); + Err(Status::internal( + "Unary response protocol violation: expected Finished, got Data", + )) + } + Ok(Some(chunk @ (ResponseChunk::Finished(_) | ResponseChunk::Error(_)))) => { + abort_guard.disarm(); + Ok(chunk) + } + Ok(None) => { + let (status, should_abort) = closed_stream_status(bridge, rid); + if should_abort { + abort_guard.abort_now(); + } else { + abort_guard.disarm(); + } + Err(status) + } + Err(status) => { + if status.code() == tonic::Code::DeadlineExceeded { + abort_guard.abort_now(); + } else { + abort_guard.disarm(); + } + Err(status) + } + } +} + +fn closed_stream_status(bridge: &Arc, rid: &str) -> (Status, bool) { + if let Some(error) = bridge.take_terminal_error(rid) { + (terminal_error_status(error), false) + } else { + ( + Status::internal("gRPC response stream closed before a terminal response"), + true, + ) + } +} + +fn terminal_error_status(error: TerminalError) -> Status { + let message = error.message(); + match error { + TerminalError::ChannelFull { .. } => Status::resource_exhausted(message), + TerminalError::ClientDisconnected { .. } | TerminalError::Aborted { .. } => { + Status::cancelled(message) + } + } +} + +fn openai_status_code(meta_info: &HashMap, default: i32) -> i32 { + meta_info + .get("status_code") + .and_then(|value| value.parse::().ok()) + .unwrap_or(default) +} + +#[tonic::async_trait] +impl proto::sglang_service_server::SglangService for SglangServiceImpl { + // --- SGLang-native RPCs: TextGenerate / Generate --- + + type TextGenerateStream = StreamResult; + + async fn text_generate( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let rid = req + .rid + .clone() + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let req_dict = build_text_generate_dict(&rid, &req); + + let mut receiver = self + .bridge + .submit_request(&rid, "generate", req_dict) + .map_err(|e| pyerr_to_status(e, "Failed to submit request"))?; + + let bridge = self.bridge.clone(); + let rid_clone = rid.clone(); + let response_timeout = self.response_timeout; + + let stream = async_stream::stream! { + let mut abort_guard = RequestAbortGuard::new(bridge.clone(), rid_clone.clone()); + loop { + match recv_chunk_with_timeout(&mut receiver, response_timeout, || "Stream chunk timed out".to_string()).await { + Ok(Some(ResponseChunk::Data(data))) => { + yield Ok(proto::TextGenerateResponse { + text: data.text.unwrap_or_default(), + meta_info: data.meta_info, + finished: false, + }); + } + Ok(Some(ResponseChunk::Finished(data))) => { + abort_guard.disarm(); + yield Ok(proto::TextGenerateResponse { + text: data.text.unwrap_or_default(), + meta_info: data.meta_info, + finished: true, + }); + break; + } + Ok(Some(ResponseChunk::Error(msg))) => { + abort_guard.disarm(); + yield Err(Status::internal(msg)); + break; + } + Ok(None) => { + let (status, should_abort) = closed_stream_status(&bridge, &rid_clone); + if should_abort { + abort_guard.abort_now(); + } else { + abort_guard.disarm(); + } + yield Err(status); + break; + } + Err(status) => { + abort_guard.abort_now(); + yield Err(status); + break; + } + } + } + }; + + Ok(Response::new(Box::pin(stream))) + } + + type GenerateStream = StreamResult; + + async fn generate( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let rid = req + .rid + .clone() + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let req_dict = build_generate_dict(&rid, &req); + + let mut receiver = self + .bridge + .submit_request(&rid, "generate", req_dict) + .map_err(|e| pyerr_to_status(e, "Failed to submit request"))?; + + let bridge = self.bridge.clone(); + let rid_clone = rid.clone(); + let response_timeout = self.response_timeout; + + let stream = async_stream::stream! { + let mut abort_guard = RequestAbortGuard::new(bridge.clone(), rid_clone.clone()); + loop { + match recv_chunk_with_timeout(&mut receiver, response_timeout, || "Stream chunk timed out".to_string()).await { + Ok(Some(ResponseChunk::Data(data))) => { + yield Ok(proto::GenerateResponse { + output_ids: data.output_ids.unwrap_or_default(), + meta_info: data.meta_info, + finished: false, + }); + } + Ok(Some(ResponseChunk::Finished(data))) => { + abort_guard.disarm(); + yield Ok(proto::GenerateResponse { + output_ids: data.output_ids.unwrap_or_default(), + meta_info: data.meta_info, + finished: true, + }); + break; + } + Ok(Some(ResponseChunk::Error(msg))) => { + abort_guard.disarm(); + yield Err(Status::internal(msg)); + break; + } + Ok(None) => { + let (status, should_abort) = closed_stream_status(&bridge, &rid_clone); + if should_abort { + abort_guard.abort_now(); + } else { + abort_guard.disarm(); + } + yield Err(status); + break; + } + Err(status) => { + abort_guard.abort_now(); + yield Err(status); + break; + } + } + } + }; + + Ok(Response::new(Box::pin(stream))) + } + + // --- SGLang-native RPCs: Embed (text / tokenized) --- + + async fn text_embed( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let rid = req + .rid + .clone() + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let req_dict = build_text_embed_dict(&rid, &req); + + let mut receiver = self + .bridge + .submit_request(&rid, "embed", req_dict) + .map_err(|e| pyerr_to_status(e, "Failed to submit request"))?; + + let chunk = recv_terminal_chunk_for_request( + &self.bridge, + &rid, + &mut receiver, + self.response_timeout, + ) + .await?; + + match chunk { + ResponseChunk::Data(data) | ResponseChunk::Finished(data) => { + Ok(Response::new(proto::TextEmbedResponse { + embedding: data.embedding.unwrap_or_default(), + meta_info: data.meta_info, + })) + } + ResponseChunk::Error(msg) => Err(Status::internal(msg)), + } + } + + async fn embed( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let rid = req + .rid + .clone() + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let req_dict = build_embed_dict(&rid, &req); + + let mut receiver = self + .bridge + .submit_request(&rid, "embed", req_dict) + .map_err(|e| pyerr_to_status(e, "Failed to submit request"))?; + + let chunk = recv_terminal_chunk_for_request( + &self.bridge, + &rid, + &mut receiver, + self.response_timeout, + ) + .await?; + + match chunk { + ResponseChunk::Data(data) | ResponseChunk::Finished(data) => { + Ok(Response::new(proto::EmbedResponse { + embedding: data.embedding.unwrap_or_default(), + meta_info: data.meta_info, + })) + } + ResponseChunk::Error(msg) => Err(Status::internal(msg)), + } + } + + // --- SGLang-native RPCs: Classify --- + + async fn classify( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + if req.text.is_empty() && req.input_ids.is_empty() { + return Err(Status::invalid_argument( + "Classify requires either text or input_ids", + )); + } + let rid = req + .rid + .clone() + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let req_dict = build_classify_dict(&rid, &req); + + let mut receiver = self + .bridge + .submit_request(&rid, "embed", req_dict) + .map_err(|e| pyerr_to_status(e, "Failed to submit request"))?; + + let chunk = recv_terminal_chunk_for_request( + &self.bridge, + &rid, + &mut receiver, + self.response_timeout, + ) + .await?; + + match chunk { + ResponseChunk::Data(data) | ResponseChunk::Finished(data) => { + Ok(Response::new(proto::ClassifyResponse { + embedding: data.embedding.unwrap_or_default(), + meta_info: data.meta_info, + })) + } + ResponseChunk::Error(msg) => Err(Status::internal(msg)), + } + } + + // --- SGLang-native RPCs: Tokenize / Detokenize (Rust-native with fallback) --- + + async fn tokenize( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let add_special = req.add_special_tokens.unwrap_or(true); + + // Try Rust-native tokenizer first (no GIL) + if let Some(tok) = self.bridge.rust_tokenizer() { + let tokens = tok + .encode(&req.text, add_special) + .map_err(Status::internal)?; + let count = tokens.len() as i32; + return Ok(Response::new(proto::TokenizeResponse { + tokens: tokens.iter().map(|&t| t as i32).collect(), + count, + max_model_len: self.bridge.context_len(), + input_text: req.text, + })); + } + + // Fallback to Python + let json_str = tokio::task::spawn_blocking({ + let bridge = self.bridge.clone(); + let text = req.text.clone(); + move || bridge.tokenize_py(&text, add_special) + }) + .await + .map_err(|e| Status::internal(format!("Task join error: {}", e)))? + .map_err(|e| pyerr_to_status(e, "Tokenize failed"))?; + + let v: serde_json::Value = serde_json::from_str(&json_str) + .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; + Ok(Response::new(proto::TokenizeResponse { + tokens: v["tokens"] + .as_array() + .map(|a| { + a.iter() + .filter_map(|x| x.as_i64().map(|n| n as i32)) + .collect() + }) + .unwrap_or_default(), + count: v["count"].as_i64().unwrap_or(0) as i32, + max_model_len: self.bridge.context_len(), + input_text: req.text, + })) + } + + async fn detokenize( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + if req.tokens.iter().any(|&token| token < 0) { + return Err(Status::invalid_argument( + "Detokenize tokens must be non-negative", + )); + } + + // Try Rust-native tokenizer first (no GIL) + if let Some(tok) = self.bridge.rust_tokenizer() { + let ids: Vec = req.tokens.iter().map(|&t| t as u32).collect(); + let text = tok.decode(&ids, true).map_err(Status::internal)?; + return Ok(Response::new(proto::DetokenizeResponse { text })); + } + + // Fallback to Python + let json_str = tokio::task::spawn_blocking({ + let bridge = self.bridge.clone(); + let tokens = req.tokens; + move || bridge.detokenize_py(tokens) + }) + .await + .map_err(|e| Status::internal(format!("Task join error: {}", e)))? + .map_err(|e| pyerr_to_status(e, "Detokenize failed"))?; + + let v: serde_json::Value = serde_json::from_str(&json_str) + .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; + Ok(Response::new(proto::DetokenizeResponse { + text: v["text"].as_str().unwrap_or("").to_string(), + })) + } + + // --- SGLang-native RPCs: Info / control --- + + async fn health_check( + &self, + _request: Request, + ) -> Result, Status> { + let healthy = tokio::task::spawn_blocking({ + let bridge = self.bridge.clone(); + move || bridge.health_check() + }) + .await + .map_err(|e| Status::internal(format!("Task join error: {}", e)))? + .map_err(|e| pyerr_to_status(e, "Health check failed"))?; + + Ok(Response::new(proto::HealthCheckResponse { healthy })) + } + + async fn get_model_info( + &self, + _request: Request, + ) -> Result, Status> { + let json_info = tokio::task::spawn_blocking({ + let bridge = self.bridge.clone(); + move || bridge.get_model_info() + }) + .await + .map_err(|e| Status::internal(format!("Task join error: {}", e)))? + .map_err(|e| pyerr_to_status(e, "Failed to get model info"))?; + + Ok(Response::new(proto::GetModelInfoResponse { + model_path: extract_model_path(&json_info), + json_info, + })) + } + + async fn get_server_info( + &self, + _request: Request, + ) -> Result, Status> { + let json_info = tokio::task::spawn_blocking({ + let bridge = self.bridge.clone(); + move || bridge.get_server_info() + }) + .await + .map_err(|e| Status::internal(format!("Task join error: {}", e)))? + .map_err(|e| pyerr_to_status(e, "Failed to get server info"))?; + + Ok(Response::new(proto::GetServerInfoResponse { json_info })) + } + + async fn list_models( + &self, + _request: Request, + ) -> Result, Status> { + let json_str = tokio::task::spawn_blocking({ + let bridge = self.bridge.clone(); + move || bridge.list_models() + }) + .await + .map_err(|e| Status::internal(format!("Task join error: {}", e)))? + .map_err(|e| pyerr_to_status(e, "Failed to list models"))?; + + let models_arr: Vec = serde_json::from_str(&json_str) + .map_err(|e| Status::internal(format!("Failed to parse models JSON: {}", e)))?; + + let models = models_arr + .iter() + .map(|m| proto::ModelCard { + id: m["id"].as_str().unwrap_or("").to_string(), + root: m["root"].as_str().unwrap_or("").to_string(), + parent: m.get("parent").and_then(|v| v.as_str()).map(String::from), + max_model_len: m + .get("max_model_len") + .and_then(|v| v.as_i64()) + .map(|n| n as i32), + }) + .collect(); + + Ok(Response::new(proto::ListModelsResponse { models })) + } + + async fn get_load( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let rid = uuid::Uuid::new_v4().to_string(); + let receiver = self + .bridge + .submit_get_load(&rid, req.dp_rank) + .map_err(|e| pyerr_to_status(e, "Failed to get load"))?; + + let json_info = + recv_json_response(&self.bridge, &rid, receiver, self.response_timeout).await?; + Ok(Response::new(proto::GetLoadResponse { json_info })) + } + + async fn abort( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + if !req.abort_all && req.rid.trim().is_empty() { + return Err(Status::invalid_argument( + "Abort requires a non-empty rid unless abort_all is true", + )); + } + if req.abort_all { + tracing::warn!( + "Received abort_all over gRPC; this endpoint must only be exposed to trusted clients" + ); + } + self.bridge + .abort(&req.rid, req.abort_all) + .map_err(|e| pyerr_to_status(e, "Failed to abort"))?; + + Ok(Response::new(proto::AbortResponse { success: true })) + } + + async fn flush_cache( + &self, + _request: Request, + ) -> Result, Status> { + let rid = uuid::Uuid::new_v4().to_string(); + let receiver = self + .bridge + .submit_flush_cache(&rid) + .map_err(|e| pyerr_to_status(e, "Failed to flush cache"))?; + + let json_str = + recv_json_response(&self.bridge, &rid, receiver, self.response_timeout).await?; + let v: serde_json::Value = serde_json::from_str(&json_str) + .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; + Ok(Response::new(proto::FlushCacheResponse { + success: v["success"].as_bool().unwrap_or(false), + message: v["message"].as_str().unwrap_or("").to_string(), + })) + } + + async fn pause_generation( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let rid = uuid::Uuid::new_v4().to_string(); + let receiver = self + .bridge + .submit_pause_generation(&rid, &req.mode) + .map_err(|e| pyerr_to_status(e, "Failed to pause generation"))?; + + let json_str = + recv_json_response(&self.bridge, &rid, receiver, self.response_timeout).await?; + let v: serde_json::Value = serde_json::from_str(&json_str) + .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; + Ok(Response::new(proto::PauseGenerationResponse { + message: v["message"].as_str().unwrap_or("").to_string(), + })) + } + + async fn continue_generation( + &self, + _request: Request, + ) -> Result, Status> { + let rid = uuid::Uuid::new_v4().to_string(); + let receiver = self + .bridge + .submit_continue_generation(&rid) + .map_err(|e| pyerr_to_status(e, "Failed to continue generation"))?; + + let json_str = + recv_json_response(&self.bridge, &rid, receiver, self.response_timeout).await?; + let v: serde_json::Value = serde_json::from_str(&json_str) + .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; + Ok(Response::new(proto::ContinueGenerationResponse { + message: v["message"].as_str().unwrap_or("").to_string(), + })) + } + + // --- OpenAI-compatible RPCs (JSON pass-through) --- + + type ChatCompleteStream = StreamResult; + + async fn chat_complete( + &self, + request: Request, + ) -> Result, Status> { + self.openai_streaming_rpc(request, "submit_openai_chat") + .await + } + + type CompleteStream = StreamResult; + + async fn complete( + &self, + request: Request, + ) -> Result, Status> { + self.openai_streaming_rpc(request, "submit_openai_complete") + .await + } + + async fn open_ai_embed( + &self, + request: Request, + ) -> Result, Status> { + self.openai_unary_rpc(request, "submit_openai_embed").await + } + + async fn open_ai_classify( + &self, + request: Request, + ) -> Result, Status> { + self.openai_unary_rpc(request, "submit_openai_classify") + .await + } + + async fn score( + &self, + request: Request, + ) -> Result, Status> { + self.openai_unary_rpc(request, "submit_openai_score").await + } + + async fn rerank( + &self, + request: Request, + ) -> Result, Status> { + self.openai_unary_rpc(request, "submit_openai_rerank").await + } + + // --- Admin RPCs --- + + async fn start_profile( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let rid = uuid::Uuid::new_v4().to_string(); + let receiver = self + .bridge + .submit_start_profile(&rid, req.output_dir.as_deref()) + .map_err(|e| pyerr_to_status(e, "Failed to start profile"))?; + + let json_str = + recv_json_response(&self.bridge, &rid, receiver, self.response_timeout).await?; + let v: serde_json::Value = serde_json::from_str(&json_str) + .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; + Ok(Response::new(proto::StartProfileResponse { + message: v["message"].as_str().unwrap_or("").to_string(), + })) + } + + async fn stop_profile( + &self, + _request: Request, + ) -> Result, Status> { + let rid = uuid::Uuid::new_v4().to_string(); + let receiver = self + .bridge + .submit_stop_profile(&rid) + .map_err(|e| pyerr_to_status(e, "Failed to stop profile"))?; + + let json_str = + recv_json_response(&self.bridge, &rid, receiver, self.response_timeout).await?; + let v: serde_json::Value = serde_json::from_str(&json_str) + .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; + Ok(Response::new(proto::StopProfileResponse { + message: v["message"].as_str().unwrap_or("").to_string(), + })) + } + + async fn update_weights_from_disk( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let rid = uuid::Uuid::new_v4().to_string(); + let receiver = self + .bridge + .submit_update_weights(&rid, &req.model_path, req.load_format.as_deref()) + .map_err(|e| pyerr_to_status(e, "Failed to update weights"))?; + + let json_str = + recv_json_response(&self.bridge, &rid, receiver, self.response_timeout).await?; + let v: serde_json::Value = serde_json::from_str(&json_str) + .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; + Ok(Response::new(proto::UpdateWeightsResponse { + success: v["success"].as_bool().unwrap_or(false), + message: v["message"].as_str().unwrap_or("").to_string(), + })) + } +} + +// Helper methods for OpenAI pass-through RPCs. +impl SglangServiceImpl { + async fn openai_streaming_rpc( + &self, + request: Request, + method_name: &str, + ) -> Result>, Status> { + let req = request.into_inner(); + let rid = uuid::Uuid::new_v4().to_string(); + + let mut receiver = self + .bridge + .submit_openai(&rid, method_name, &req.json_body, &req.trace_headers) + .map_err(|e| pyerr_to_status(e, "Failed to submit request"))?; + + let bridge = self.bridge.clone(); + let rid_clone = rid.clone(); + let response_timeout = self.response_timeout; + + let stream = async_stream::stream! { + let mut abort_guard = RequestAbortGuard::new(bridge.clone(), rid_clone.clone()); + loop { + match recv_chunk_with_timeout(&mut receiver, response_timeout, || "Stream chunk timed out".to_string()).await { + Ok(Some(ResponseChunk::Data(data))) => { + yield Ok(proto::OpenAiStreamChunk { + json_chunk: data.json_bytes.unwrap_or_default(), + finished: false, + }); + } + Ok(Some(ResponseChunk::Finished(data))) => { + let bytes = data.json_bytes.unwrap_or_default(); + abort_guard.disarm(); + yield Ok(proto::OpenAiStreamChunk { + json_chunk: bytes, + finished: true, + }); + break; + } + Ok(Some(ResponseChunk::Error(msg))) => { + abort_guard.disarm(); + yield Err(Status::internal(msg)); + break; + } + Ok(None) => { + let (status, should_abort) = closed_stream_status(&bridge, &rid_clone); + if should_abort { + abort_guard.abort_now(); + } else { + abort_guard.disarm(); + } + yield Err(status); + break; + } + Err(status) => { + abort_guard.abort_now(); + yield Err(status); + break; + } + } + } + }; + + Ok(Response::new(Box::pin(stream))) + } + + async fn openai_unary_rpc( + &self, + request: Request, + method_name: &str, + ) -> Result, Status> { + let req = request.into_inner(); + let rid = uuid::Uuid::new_v4().to_string(); + + let mut receiver = self + .bridge + .submit_openai(&rid, method_name, &req.json_body, &req.trace_headers) + .map_err(|e| pyerr_to_status(e, "Failed to submit request"))?; + + let chunk = recv_terminal_chunk_for_request( + &self.bridge, + &rid, + &mut receiver, + self.response_timeout, + ) + .await?; + + match chunk { + ResponseChunk::Data(data) | ResponseChunk::Finished(data) => { + Ok(Response::new(proto::OpenAiResponse { + json_body: data.json_bytes.unwrap_or_default(), + status_code: openai_status_code(&data.meta_info, 200), + })) + } + ResponseChunk::Error(msg) => { + let error_json = serde_json::json!({"error": {"message": msg}}); + Ok(Response::new(proto::OpenAiResponse { + json_body: error_json.to_string().into_bytes(), + status_code: 500, + })) + } + } + } +} + +/// Receive a single JSON response from the bridge channel. +async fn recv_json_response( + bridge: &Arc, + rid: &str, + mut receiver: Receiver, + response_timeout: Duration, +) -> Result { + let chunk = + recv_terminal_chunk_for_request(bridge, rid, &mut receiver, response_timeout).await?; + + match chunk { + ResponseChunk::Data(data) | ResponseChunk::Finished(data) => { + let bytes = data.json_bytes.unwrap_or_default(); + String::from_utf8(bytes) + .map_err(|e| Status::internal(format!("Invalid UTF-8 in response: {}", e))) + } + ResponseChunk::Error(msg) => Err(Status::internal(msg)), + } +} + +/// Start the Tonic gRPC server on the given address. +// +// TODO(grpc-auth): this listener is currently unauthenticated. Before exposing +// it in any default deploy path, gate it with the same API-key / admin-key +// checks the HTTP server applies (see issue tracking gRPC auth parity). +pub async fn run_grpc_server( + listener: std::net::TcpListener, + bridge: Arc, + shutdown: Arc, + response_timeout: Duration, +) -> Result<(), Box> { + let addr = listener.local_addr()?; + let listener = tokio::net::TcpListener::from_std(listener)?; + let service = SglangServiceImpl { + bridge, + response_timeout, + }; + + let max_message_size = resolve_max_message_size(); + let svc = proto::sglang_service_server::SglangServiceServer::new(service) + .max_decoding_message_size(max_message_size) + .max_encoding_message_size(max_message_size); + + tracing::info!("gRPC server listening on {}", addr); + + tonic::transport::Server::builder() + .add_service(svc) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move { + shutdown.notified().await; + tracing::info!("gRPC server shutting down"); + }) + .await?; + + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/rust/sglang-grpc/src/server/tests.rs b/rust/sglang-grpc/src/server/tests.rs new file mode 100644 index 000000000..11ecb6876 --- /dev/null +++ b/rust/sglang-grpc/src/server/tests.rs @@ -0,0 +1,75 @@ +use super::{ + DEFAULT_GRPC_MAX_MESSAGE_SIZE, openai_status_code, resolve_max_message_size, + terminal_error_status, +}; +use crate::bridge::TerminalError; +use std::collections::HashMap; +use tonic::Code; + +#[test] +fn openai_status_code_uses_forwarded_status_when_present() { + let meta_info = HashMap::from([(String::from("status_code"), String::from("429"))]); + assert_eq!(openai_status_code(&meta_info, 200), 429); +} + +#[test] +fn openai_status_code_falls_back_when_missing_or_invalid() { + assert_eq!(openai_status_code(&HashMap::new(), 200), 200); + + let meta_info = HashMap::from([(String::from("status_code"), String::from("not-an-int"))]); + assert_eq!(openai_status_code(&meta_info, 200), 200); +} + +#[test] +fn terminal_error_status_maps_channel_full_to_resource_exhausted() { + let status = terminal_error_status(TerminalError::ChannelFull { + rid: "rid".to_string(), + }); + + assert_eq!(status.code(), Code::ResourceExhausted); +} + +#[test] +fn terminal_error_status_maps_abort_to_cancelled() { + let status = terminal_error_status(TerminalError::Aborted { + rid: "rid".to_string(), + }); + + assert_eq!(status.code(), Code::Cancelled); +} + +// SAFETY: env vars are process-global; bundle all SGLANG_TONIC_PAYLOAD cases into one +// serial test so they don't race each other under `cargo test`'s default parallelism. +#[test] +fn resolve_max_message_size_honors_env_var() { + const VAR: &str = "SGLANG_TONIC_PAYLOAD"; + + // Unset → default. + // SAFETY: single-threaded test mutating process env (see note above). + unsafe { + std::env::remove_var(VAR); + } + assert_eq!(resolve_max_message_size(), DEFAULT_GRPC_MAX_MESSAGE_SIZE); + + // Valid override → honored verbatim. + unsafe { + std::env::set_var(VAR, "1048576"); + } + assert_eq!(resolve_max_message_size(), 1_048_576); + + // Invalid string → warn + fall back to default. + unsafe { + std::env::set_var(VAR, "not-a-number"); + } + assert_eq!(resolve_max_message_size(), DEFAULT_GRPC_MAX_MESSAGE_SIZE); + + // Zero → treated as invalid, fall back to default. + unsafe { + std::env::set_var(VAR, "0"); + } + assert_eq!(resolve_max_message_size(), DEFAULT_GRPC_MAX_MESSAGE_SIZE); + + unsafe { + std::env::remove_var(VAR); + } +} diff --git a/rust/sglang-grpc/src/tokenizers.rs b/rust/sglang-grpc/src/tokenizers.rs new file mode 100644 index 000000000..7dc406d3e --- /dev/null +++ b/rust/sglang-grpc/src/tokenizers.rs @@ -0,0 +1,127 @@ +use std::path::{Path, PathBuf}; + +use tokenizers::Tokenizer; + +trait TokenizerBackend: Send + Sync { + fn name(&self) -> &'static str; + fn encode(&self, text: &str, add_special_tokens: bool) -> Result, String>; + fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result; +} + +struct HuggingFaceTokenizerBackend { + inner: Tokenizer, +} + +impl HuggingFaceTokenizerBackend { + fn from_file(path: &Path) -> Result { + Tokenizer::from_file(path) + .map(|inner| Self { inner }) + .map_err(|e| format!("failed to load HuggingFace tokenizer: {}", e)) + } +} + +impl TokenizerBackend for HuggingFaceTokenizerBackend { + fn name(&self) -> &'static str { + "huggingface-tokenizers" + } + + fn encode(&self, text: &str, add_special_tokens: bool) -> Result, String> { + let encoding = self + .inner + .encode(text, add_special_tokens) + .map_err(|e| format!("Tokenization failed: {}", e))?; + Ok(encoding.get_ids().to_vec()) + } + + fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result { + self.inner + .decode(ids, skip_special_tokens) + .map_err(|e| format!("Detokenization failed: {}", e)) + } +} + +/// Rust-native tokenizer wrapper with pluggable backends. +/// +/// This mirrors Python's `get_tokenizer` shape: inspect the tokenizer path, +/// choose a backend, and fall back to Python for unsupported tokenizer families. +pub struct RustTokenizer { + backend: Box, +} + +impl RustTokenizer { + /// Load a native tokenizer from a tokenizer path or model directory. + /// Returns `None` if no supported Rust backend is available. + pub fn from_tokenizer_path( + tokenizer_path: &str, + tokenizer_mode: Option<&str>, + context_len: i32, + ) -> Option { + let path = Path::new(tokenizer_path); + let Some(tokenizer_json) = resolve_tokenizer_json(path) else { + tracing::info!( + "No native tokenizer candidates found at {:?}; Rust tokenizer disabled", + path + ); + return None; + }; + + if matches!(tokenizer_mode, Some("slow")) { + tracing::info!( + "Rust tokenizer disabled because tokenizer_mode=slow for {:?}", + path + ); + return None; + } + + match load_backend(&tokenizer_json) { + Ok(backend) => { + let tokenizer = Self { backend }; + tracing::info!( + "Rust tokenizer loaded via {} from {:?} (context_len={})", + tokenizer.backend_name(), + tokenizer_json, + context_len + ); + Some(tokenizer) + } + Err(e) => { + tracing::warn!( + "Failed to load Rust tokenizer from {:?}: {}. Falling back to Python.", + tokenizer_json, + e + ); + None + } + } + } + + pub fn backend_name(&self) -> &'static str { + self.backend.name() + } + + /// Tokenize text, returning token IDs. + pub fn encode(&self, text: &str, add_special_tokens: bool) -> Result, String> { + self.backend.encode(text, add_special_tokens) + } + + /// Decode token IDs back to text. + pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result { + self.backend.decode(ids, skip_special_tokens) + } +} + +fn load_backend(tokenizer_json: &Path) -> Result, String> { + // Add new native backend probes here. Unsupported formats should return an + // error so callers can fall back to Python without changing the public API. + HuggingFaceTokenizerBackend::from_file(tokenizer_json) + .map(|backend| Box::new(backend) as Box) +} + +fn resolve_tokenizer_json(path: &Path) -> Option { + let candidate = if path.is_file() { + path.to_path_buf() + } else { + path.join("tokenizer.json") + }; + candidate.exists().then_some(candidate) +} diff --git a/rust/sglang-grpc/src/utils/mod.rs b/rust/sglang-grpc/src/utils/mod.rs new file mode 100644 index 000000000..2ed54c964 --- /dev/null +++ b/rust/sglang-grpc/src/utils/mod.rs @@ -0,0 +1,8 @@ +mod py_utils; +mod request_utils; + +pub(crate) use py_utils::{json_map_to_pydict, py_value_to_json_string}; +pub(crate) use request_utils::{ + build_classify_dict, build_embed_dict, build_generate_dict, build_text_embed_dict, + build_text_generate_dict, extract_model_path, +}; diff --git a/rust/sglang-grpc/src/utils/py_utils.rs b/rust/sglang-grpc/src/utils/py_utils.rs new file mode 100644 index 000000000..2e99def65 --- /dev/null +++ b/rust/sglang-grpc/src/utils/py_utils.rs @@ -0,0 +1,70 @@ +use pyo3::prelude::*; +use pyo3::types::{PyAny, PyDict, PyList}; +use std::collections::HashMap; + +fn json_value_to_py<'py>(py: Python<'py>, v: &serde_json::Value) -> PyResult { + match v { + serde_json::Value::Null => Ok(py.None()), + serde_json::Value::Bool(b) => Ok(b.into_pyobject(py)?.to_owned().into_any().unbind()), + serde_json::Value::Number(n) => { + if let Some(i) = n.as_i64() { + Ok(i.into_pyobject(py)?.into_any().unbind()) + } else if let Some(f) = n.as_f64() { + Ok(f.into_pyobject(py)?.into_any().unbind()) + } else { + Ok(py.None()) + } + } + serde_json::Value::String(s) => Ok(s.into_pyobject(py)?.into_any().unbind()), + serde_json::Value::Array(arr) => { + let items: Vec = arr + .iter() + .map(|item| json_value_to_py(py, item)) + .collect::>()?; + let py_list = PyList::new(py, &items)?; + Ok(py_list.into_any().unbind()) + } + serde_json::Value::Object(map) => { + let py_dict = PyDict::new(py); + for (k, val) in map { + py_dict.set_item(k, json_value_to_py(py, val)?)?; + } + Ok(py_dict.into_any().unbind()) + } + } +} + +pub(crate) fn json_map_to_pydict<'py>( + py: Python<'py>, + map: &HashMap, +) -> PyResult> { + let py_dict = PyDict::new(py); + for (k, v) in map { + py_dict.set_item(k, json_value_to_py(py, v)?)?; + } + Ok(py_dict) +} + +pub(crate) fn py_value_to_json_string(value: &Bound<'_, PyAny>) -> PyResult { + // gRPC meta_info is a map. JSON-encode every value, + // including strings, so clients can decode the map uniformly. + match value + .py() + .import("json") + .and_then(|json| json.call_method1("dumps", (value,))) + .and_then(|json_str| json_str.extract::()) + { + Ok(s) => Ok(s), + Err(_) => { + let fallback = value.str()?.to_string(); + Ok(json_encode_string(&fallback)) + } + } +} + +fn json_encode_string(value: &str) -> String { + serde_json::Value::String(value.to_string()).to_string() +} + +#[cfg(test)] +mod tests; diff --git a/rust/sglang-grpc/src/utils/py_utils/tests.rs b/rust/sglang-grpc/src/utils/py_utils/tests.rs new file mode 100644 index 000000000..7d14f4dda --- /dev/null +++ b/rust/sglang-grpc/src/utils/py_utils/tests.rs @@ -0,0 +1,8 @@ +use super::*; + +#[test] +fn fallback_string_is_json_encoded() { + let encoded = json_encode_string(""); + let decoded: String = serde_json::from_str(&encoded).unwrap(); + assert_eq!(decoded, ""); +} diff --git a/rust/sglang-grpc/src/utils/request_utils.rs b/rust/sglang-grpc/src/utils/request_utils.rs new file mode 100644 index 000000000..98d562a80 --- /dev/null +++ b/rust/sglang-grpc/src/utils/request_utils.rs @@ -0,0 +1,239 @@ +use std::collections::HashMap; + +use crate::proto; + +/// Convert proto SamplingParams to a serde_json map (used as Python dict via PyO3). +fn sampling_params_to_map(params: &Option) -> serde_json::Value { + match params { + Some(p) => { + let mut map = serde_json::Map::new(); + if let Some(v) = p.temperature { + map.insert("temperature".into(), serde_json::json!(v)); + } + if let Some(v) = p.top_p { + map.insert("top_p".into(), serde_json::json!(v)); + } + if let Some(v) = p.top_k { + map.insert("top_k".into(), serde_json::json!(v)); + } + if let Some(v) = p.min_p { + map.insert("min_p".into(), serde_json::json!(v)); + } + if let Some(v) = p.frequency_penalty { + map.insert("frequency_penalty".into(), serde_json::json!(v)); + } + if let Some(v) = p.presence_penalty { + map.insert("presence_penalty".into(), serde_json::json!(v)); + } + if let Some(v) = p.repetition_penalty { + map.insert("repetition_penalty".into(), serde_json::json!(v)); + } + if let Some(v) = p.max_new_tokens { + map.insert("max_new_tokens".into(), serde_json::json!(v)); + } + if let Some(v) = p.min_new_tokens { + map.insert("min_new_tokens".into(), serde_json::json!(v)); + } + if !p.stop.is_empty() { + map.insert("stop".into(), serde_json::json!(p.stop)); + } + if !p.stop_token_ids.is_empty() { + map.insert("stop_token_ids".into(), serde_json::json!(p.stop_token_ids)); + } + if let Some(v) = p.ignore_eos { + map.insert("ignore_eos".into(), serde_json::json!(v)); + } + if let Some(v) = p.n { + map.insert("n".into(), serde_json::json!(v)); + } + if let Some(ref v) = p.json_schema { + map.insert("json_schema".into(), serde_json::json!(v)); + } + if let Some(ref v) = p.regex { + map.insert("regex".into(), serde_json::json!(v)); + } + serde_json::Value::Object(map) + } + None => serde_json::Value::Object(serde_json::Map::new()), + } +} + +fn trace_headers_to_json(headers: &HashMap) -> Option { + if headers.is_empty() { + None + } else { + Some(serde_json::json!(headers)) + } +} + +fn now_timestamp() -> f64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64() +} + +pub(crate) fn extract_model_path(json_info: &str) -> String { + match serde_json::from_str::(json_info) { + Ok(value) => value + .get("model_path") + .and_then(|v| v.as_str()) + .map(str::to_owned) + .unwrap_or_default(), + Err(err) => { + tracing::warn!("Failed to parse model info JSON: {}", err); + String::new() + } + } +} + +/// Build a request dict for GenerateReqInput from proto TextGenerateRequest fields. +pub(crate) fn build_text_generate_dict( + rid: &str, + req: &proto::TextGenerateRequest, +) -> HashMap { + let mut d = HashMap::new(); + d.insert("rid".into(), serde_json::json!(rid)); + d.insert("text".into(), serde_json::json!(req.text)); + d.insert( + "sampling_params".into(), + sampling_params_to_map(&req.sampling_params), + ); + d.insert( + "stream".into(), + serde_json::json!(req.stream.unwrap_or(false)), + ); + d.insert( + "return_logprob".into(), + serde_json::json!(req.return_logprob.unwrap_or(false)), + ); + d.insert( + "top_logprobs_num".into(), + serde_json::json!(req.top_logprobs_num.unwrap_or(0)), + ); + d.insert( + "logprob_start_len".into(), + serde_json::json!(req.logprob_start_len.unwrap_or(-1)), + ); + d.insert( + "return_text_in_logprobs".into(), + serde_json::json!(req.return_text_in_logprobs.unwrap_or(false)), + ); + if let Some(ref lp) = req.lora_path { + d.insert("lora_path".into(), serde_json::json!(lp)); + } + if let Some(ref rk) = req.routing_key { + d.insert("routing_key".into(), serde_json::json!(rk)); + } + if let Some(rank) = req.routed_dp_rank { + d.insert("routed_dp_rank".into(), serde_json::json!(rank)); + } + if let Some(trace) = trace_headers_to_json(&req.trace_headers) { + d.insert("external_trace_header".into(), trace); + } + d.insert("received_time".into(), serde_json::json!(now_timestamp())); + d +} + +/// Build a request dict for GenerateReqInput from proto GenerateRequest (tokenized). +pub(crate) fn build_generate_dict( + rid: &str, + req: &proto::GenerateRequest, +) -> HashMap { + let mut d = HashMap::new(); + d.insert("rid".into(), serde_json::json!(rid)); + d.insert("input_ids".into(), serde_json::json!(req.input_ids)); + d.insert( + "sampling_params".into(), + sampling_params_to_map(&req.sampling_params), + ); + d.insert( + "stream".into(), + serde_json::json!(req.stream.unwrap_or(false)), + ); + d.insert( + "return_logprob".into(), + serde_json::json!(req.return_logprob.unwrap_or(false)), + ); + d.insert( + "top_logprobs_num".into(), + serde_json::json!(req.top_logprobs_num.unwrap_or(0)), + ); + d.insert( + "logprob_start_len".into(), + serde_json::json!(req.logprob_start_len.unwrap_or(-1)), + ); + if let Some(ref lp) = req.lora_path { + d.insert("lora_path".into(), serde_json::json!(lp)); + } + if let Some(ref rk) = req.routing_key { + d.insert("routing_key".into(), serde_json::json!(rk)); + } + if let Some(rank) = req.routed_dp_rank { + d.insert("routed_dp_rank".into(), serde_json::json!(rank)); + } + if let Some(trace) = trace_headers_to_json(&req.trace_headers) { + d.insert("external_trace_header".into(), trace); + } + d.insert("received_time".into(), serde_json::json!(now_timestamp())); + d +} + +/// Build a request dict for EmbeddingReqInput from proto TextEmbedRequest. +pub(crate) fn build_text_embed_dict( + rid: &str, + req: &proto::TextEmbedRequest, +) -> HashMap { + let mut d = HashMap::new(); + d.insert("rid".into(), serde_json::json!(rid)); + d.insert("text".into(), serde_json::json!(req.text)); + if let Some(ref rk) = req.routing_key { + d.insert("routing_key".into(), serde_json::json!(rk)); + } + if let Some(trace) = trace_headers_to_json(&req.trace_headers) { + d.insert("external_trace_header".into(), trace); + } + d.insert("received_time".into(), serde_json::json!(now_timestamp())); + d +} + +/// Build a request dict for EmbeddingReqInput from proto EmbedRequest (tokenized). +pub(crate) fn build_embed_dict( + rid: &str, + req: &proto::EmbedRequest, +) -> HashMap { + let mut d = HashMap::new(); + d.insert("rid".into(), serde_json::json!(rid)); + d.insert("input_ids".into(), serde_json::json!(req.input_ids)); + if let Some(ref rk) = req.routing_key { + d.insert("routing_key".into(), serde_json::json!(rk)); + } + if let Some(trace) = trace_headers_to_json(&req.trace_headers) { + d.insert("external_trace_header".into(), trace); + } + d.insert("received_time".into(), serde_json::json!(now_timestamp())); + d +} + +/// Build a request dict for EmbeddingReqInput from proto ClassifyRequest. +pub(crate) fn build_classify_dict( + rid: &str, + req: &proto::ClassifyRequest, +) -> HashMap { + let mut d = HashMap::new(); + d.insert("rid".into(), serde_json::json!(rid)); + if !req.text.is_empty() { + d.insert("text".into(), serde_json::json!(req.text)); + } + if !req.input_ids.is_empty() { + d.insert("input_ids".into(), serde_json::json!(req.input_ids)); + } + if let Some(ref rk) = req.routing_key { + d.insert("routing_key".into(), serde_json::json!(rk)); + } + if let Some(trace) = trace_headers_to_json(&req.trace_headers) { + d.insert("external_trace_header".into(), trace); + } + d.insert("received_time".into(), serde_json::json!(now_timestamp())); + d +}