[gRPC] Native server: Rust crate (1/N) (#23506)

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]

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<String>,
+    pub output_ids: Option<Vec<i32>>,
+    pub embedding: Option<Vec<f32>>,
+    pub json_bytes: Option<Vec<u8>>,
+    pub meta_info: HashMap<String, String>,
+}
+
+pub const DEFAULT_RESPONSE_CHANNEL_CAPACITY: usize = 64;
+
+type BridgeStateRef = Arc<Mutex<BridgeState>>;
+
+#[derive(Default)]
+struct BridgeState {
+    channels: HashMap<String, Sender<ResponseChunk>>,
+    pending_sends: HashSet<String>,
+    ready_callbacks: HashMap<String, PyObject>,
+    ready_signals: HashSet<String>,
+    terminal_errors: HashMap<String, TerminalError>,
+}
+
+#[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<T>, 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<RustTokenizer>,
+    context_len: i32,
+    response_channel_capacity: usize,
+    tokio_handle: Handle,
+}
+
+impl PyBridge {
+    pub fn new(
+        runtime_handle: PyObject,
+        rust_tokenizer: Option<RustTokenizer>,
+        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<Receiver<ResponseChunk>> {
+        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<PyObject> {
+        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<PyObject> {
+        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<String, serde_json::Value>,
+    ) -> PyResult<Receiver<ResponseChunk>> {
+        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::<Vec<_>>();
+            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<String> {
+        Python::with_gil(|py| {
+            let result = self.runtime_handle.call_method0(py, "get_model_info")?;
+            result.extract::<String>(py)
+        })
+    }
+
+    pub fn get_server_info(&self) -> PyResult<String> {
+        Python::with_gil(|py| {
+            let result = self.runtime_handle.call_method0(py, "get_server_info")?;
+            result.extract::<String>(py)
+        })
+    }
+
+    pub fn health_check(&self) -> PyResult<bool> {
+        Python::with_gil(|py| {
+            let result = self.runtime_handle.call_method0(py, "health_check")?;
+            result.extract::<bool>(py)
+        })
+    }
+
+    /// Tokenize via Python (fallback when Rust tokenizer unavailable).
+    pub fn tokenize_py(&self, text: &str, add_special_tokens: bool) -> PyResult<String> {
+        Python::with_gil(|py| {
+            let result =
+                self.runtime_handle
+                    .call_method1(py, "tokenize", (text, add_special_tokens))?;
+            result.extract::<String>(py)
+        })
+    }
+
+    /// Detokenize via Python (fallback when Rust tokenizer unavailable).
+    pub fn detokenize_py(&self, tokens: Vec<i32>) -> PyResult<String> {
+        Python::with_gil(|py| {
+            let result = self
+                .runtime_handle
+                .call_method1(py, "detokenize", (tokens,))?;
+            result.extract::<String>(py)
+        })
+    }
+
+    pub fn list_models(&self) -> PyResult<String> {
+        Python::with_gil(|py| {
+            let result = self.runtime_handle.call_method0(py, "list_models")?;
+            result.extract::<String>(py)
+        })
+    }
+
+    fn submit_json<F>(&self, rid: &str, call: F) -> PyResult<Receiver<ResponseChunk>>
+    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<i32>,
+    ) -> PyResult<Receiver<ResponseChunk>> {
+        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<Receiver<ResponseChunk>> {
+        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<Receiver<ResponseChunk>> {
+        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<Receiver<ResponseChunk>> {
+        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<Receiver<ResponseChunk>> {
+        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<Receiver<ResponseChunk>> {
+        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<Receiver<ResponseChunk>> {
+        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<String, String>,
+    ) -> PyResult<Receiver<ResponseChunk>> {
+        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<TerminalError> {
+        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<PyObject> {
+    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<ResponseChunk>,
+    msg: ResponseChunk,
+) -> PyResult<ChunkSendStatus> {
+    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<String>,
+    ) -> PyResult<ChunkSendStatus> {
+        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<String> = chunk
+            .get_item("text")?
+            .and_then(|v| v.extract::<String>().ok());
+
+        let output_ids: Option<Vec<i32>> = chunk
+            .get_item("output_ids")?
+            .and_then(|v| v.extract::<Vec<i32>>().ok());
+
+        let embedding: Option<Vec<f32>> = chunk
+            .get_item("embedding")?
+            .and_then(|v| v.extract::<Vec<f32>>().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<String>,
+        status_code: Option<i32>,
+    ) -> PyResult<ChunkSendStatus> {
+        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<u8> = if let Ok(b) = chunk_bytes.extract::<Vec<u8>>() {
+            b
+        } else if let Ok(s) = chunk_bytes.extract::<String>() {
+            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<String, String> {
+    let mut meta = HashMap::new();
+    if let Ok(Some(meta_obj)) = chunk.get_item("meta_info")
+        && let Ok(meta_dict) = meta_obj.downcast::<PyDict>()
+    {
+        for (k, v) in meta_dict.iter() {
+            // The proto schema is map<string, string>; encode each Python value as JSON
+            // so clients can recover numbers, booleans, arrays, and objects losslessly.
+            if let Ok(key) = k.extract::<String>()
+                && let Ok(val) = py_value_to_json_string(&v)
+            {
+                meta.insert(key, val);
+            }
+        }
+    }
+    meta
+}
+
+#[cfg(test)]
+mod tests;

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"));
+}

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<Notify>,
     join_handle: Option<std::thread::JoinHandle<()>>,
 }
 
 #[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<String>,
+    tokenizer_mode: Option<String>,
+    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<PyObject> {
+    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<String> {
+    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<i32> {
+    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<TokenizerInfo> {
+    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<GrpcServerHandle> {
-    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<GrpcServerHandle> {
+    // 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<G
     })
 }
 
-/// Python module exported by the Rust extension.
 #[pymodule]
 fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(start_server, m)?)?;
     m.add_class::<GrpcServerHandle>()?;
+    m.add_class::<ChunkSendStatus>()?;
     Ok(())
 }

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);
+    }
+}

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<Vec<u32>, String>;
+    fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result<String, String>;
+}
+
+struct HuggingFaceTokenizerBackend {
+    inner: Tokenizer,
+}
+
+impl HuggingFaceTokenizerBackend {
+    fn from_file(path: &Path) -> Result<Self, String> {
+        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<Vec<u32>, 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<String, String> {
+        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<dyn TokenizerBackend>,
+}
+
+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<Self> {
+        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<Vec<u32>, 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<String, String> {
+        self.backend.decode(ids, skip_special_tokens)
+    }
+}
+
+fn load_backend(tokenizer_json: &Path) -> Result<Box<dyn TokenizerBackend>, 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<dyn TokenizerBackend>)
+}
+
+fn resolve_tokenizer_json(path: &Path) -> Option<PathBuf> {
+    let candidate = if path.is_file() {
+        path.to_path_buf()
+    } else {
+        path.join("tokenizer.json")
+    };
+    candidate.exists().then_some(candidate)
+}

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,
+};

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<PyObject> {
+    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<PyObject> = arr
+                .iter()
+                .map(|item| json_value_to_py(py, item))
+                .collect::<PyResult<_>>()?;
+            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<String, serde_json::Value>,
+) -> PyResult<Bound<'py, PyDict>> {
+    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<String> {
+    // gRPC meta_info is a map<string, string>. 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::<String>())
+    {
+        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;

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("<Foo at 0x123>");
+    let decoded: String = serde_json::from_str(&encoded).unwrap();
+    assert_eq!(decoded, "<Foo at 0x123>");
+}

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<proto::SamplingParams>) -> 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<String, String>) -> Option<serde_json::Value> {
+    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::<serde_json::Value>(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<String, serde_json::Value> {
+    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<String, serde_json::Value> {
+    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<String, serde_json::Value> {
+    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<String, serde_json::Value> {
+    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<String, serde_json::Value> {
+    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
+}