minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

create rust workspace (#32014)

Browse Source
This commit is contained in:
Rain Jiang
2026-07-23 12:02:41 -07:00
committed by GitHub
parent d0b9689805
commit 7fe82dd02e
25 changed files with 550 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/sglang-grpc/Cargo.toml
+26 -22
View File
@@ -1,38 +1,42 @@
[package]
name = "sglang-grpc"
version = "0.1.0"
edition = "2024"
description = "In-process Rust gRPC server for SGLang"
license = "Apache-2.0"
version.workspace = true
edition.workspace = true
license.workspace = true
# Consumed by python/setup.py: registers this crate as a PyO3 extension module
# of the main sglang wheel at the given import path.
[package.metadata.sglang]
python-module = "sglang.srt.grpc._core"
[lib]
name = "_core"
# Unique per-crate artifact name (the shared workspace target/ dir cannot hold
# two `lib_core` cdylibs). The importable Python module is still `_core`: the
# name comes from the `#[pymodule]` entry point and the setuptools-rust
# `target` in python/pyproject.toml, which renames the built artifact.
name = "sglang_grpc_core"
crate-type = ["cdylib"]
[dependencies]
pyo3 = { version = "0.23", features = ["extension-module"] }
tokio = { version = "1", features = ["full"] }
tonic = { version = "0.12", features = ["gzip", "transport"] }
async-stream = { workspace = true }
pyo3 = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true }
tokio-stream = { workspace = true }
tracing = { workspace = true }
tracing-subscriber = { workspace = true }
uuid = { workspace = true }
prost = "0.13"
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 = { version = "0.1", features = ["net"] }
async-stream = "0.3"
tonic = { version = "0.12", features = ["gzip", "transport"] }
[build-dependencies]
tonic-build = "0.12"
# Fallback protoc for machines without a system install (see build.rs); keeps
# `cargo clippy/check/build` and lint CI working without apt/brew protobuf.
protoc-bin-vendored = "3"
[features]
default = ["pyo3/extension-module"]
[profile.release]
opt-level = 2
lto = "thin"
strip = true
[profile.dev]
opt-level = 0
debug = 1
rust/sglang-grpc/build.rs
+12
View File
@@ -1,4 +1,16 @@
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Prefer an explicitly configured protoc; otherwise use the vendored
// binary so builds (including `cargo clippy` and the pre-commit hook on
// machines/CI runners without protobuf installed) are self-contained.
// protoc_bin_path() errs on platforms the vendored crate doesn't cover;
// those fall back to prost-build's own lookup of `protoc` on PATH.
if std::env::var_os("PROTOC").is_none()
&& let Ok(vendored) = protoc_bin_vendored::protoc_bin_path()
{
// SAFETY: build scripts are single-threaded at this point.
unsafe { std::env::set_var("PROTOC", vendored) };
}
let proto_path = "../../proto/sglang/runtime/v1/sglang.proto";
tonic_build::configure()
rust/sglang-grpc/pyproject.toml
+17
View File
@@ -0,0 +1,17 @@
# Standalone dev builds only (`maturin build` / `maturin develop`); the
# extension ships to users inside the main sglang wheel via setuptools-rust
# (python/pyproject.toml, target sglang.srt.grpc._core). Requires a repo
# checkout (build.rs reads ../../proto), so this is not publishable as an
# sdist. protoc is not required: build.rs falls back to a vendored binary.
[build-system]
requires = ["maturin>=1.5,<2"]
build-backend = "maturin"
[project]
name = "sglang-grpc"
dynamic = ["version"]
description = "In-process Rust gRPC server for SGLang"
requires-python = ">=3.10"
[tool.maturin]
module-name = "_core"
rust/sglang-grpc/rust-toolchain.toml
-3
View File
@@ -1,3 +0,0 @@
[toolchain]
channel = "1.90"
profile = "minimal"
rust/sglang-grpc/src/bridge.rs
+31 -28
View File
@@ -40,7 +40,7 @@ type BridgeStateRef = Arc<Mutex<BridgeState>>;
struct BridgeState {
channels: HashMap<String, Sender<ResponseChunk>>,
pending_sends: HashSet<String>,
ready_callbacks: HashMap<String, PyObject>,
ready_callbacks: HashMap<String, Py<PyAny>>,
ready_signals: HashSet<String>,
terminal_errors: HashMap<String, TerminalError>,
}
@@ -66,7 +66,10 @@ impl TerminalError {
}
}
#[pyclass(eq, eq_int)]
// skip_from_py_object: this enum is only returned to Python, never received
// from it, so it opts out of pyo3's (deprecated-by-default) FromPyObject
// derive for Clone pyclasses.
#[pyclass(eq, eq_int, skip_from_py_object)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChunkSendStatus {
Ready,
@@ -83,7 +86,7 @@ fn lock_or_recover<'a, T>(mutex: &'a Mutex<T>, name: &'static str) -> MutexGuard
/// Holds a reference to the Python RuntimeHandle and manages per-request channels.
pub struct PyBridge {
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
state: BridgeStateRef,
rust_tokenizer: Option<RustTokenizer>,
context_len: i32,
@@ -93,7 +96,7 @@ pub struct PyBridge {
impl PyBridge {
pub fn new(
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
rust_tokenizer: Option<RustTokenizer>,
context_len: i32,
response_channel_capacity: usize,
@@ -144,7 +147,7 @@ impl PyBridge {
Ok(receiver)
}
fn make_chunk_callback(&self, py: Python<'_>, rid: String) -> PyResult<PyObject> {
fn make_chunk_callback(&self, py: Python<'_>, rid: String) -> PyResult<Py<PyAny>> {
let callback = ChunkCallback {
rid,
state: self.state.clone(),
@@ -155,7 +158,7 @@ impl PyBridge {
Ok(py_callback.into_any())
}
fn make_json_callback(&self, py: Python<'_>, rid: String) -> PyResult<PyObject> {
fn make_json_callback(&self, py: Python<'_>, rid: String) -> PyResult<Py<PyAny>> {
let callback = JsonChunkCallback {
rid,
state: self.state.clone(),
@@ -183,7 +186,7 @@ impl PyBridge {
let receiver = self.create_channel(rid)?;
let rid_owned = rid.to_string();
let result = Python::with_gil(|py| -> PyResult<()> {
let result = Python::attach(|py| -> PyResult<()> {
let py_req_dict = json_map_to_pydict(py, &req_dict)?;
let callback = self.make_chunk_callback(py, rid_owned)?;
@@ -253,7 +256,7 @@ impl PyBridge {
return Ok(());
}
Python::with_gil(|py| {
Python::attach(|py| {
self.runtime_handle
.call_method1(py, "abort", (rid, abort_all))?;
Ok(())
@@ -265,21 +268,21 @@ impl PyBridge {
// ------------------------------------------------------------------
pub fn get_model_info(&self) -> PyResult<String> {
Python::with_gil(|py| {
Python::attach(|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| {
Python::attach(|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| {
Python::attach(|py| {
let result = self.runtime_handle.call_method0(py, "health_check")?;
result.extract::<bool>(py)
})
@@ -287,7 +290,7 @@ impl PyBridge {
/// 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| {
Python::attach(|py| {
let result =
self.runtime_handle
.call_method1(py, "tokenize", (text, add_special_tokens))?;
@@ -297,7 +300,7 @@ impl PyBridge {
/// Detokenize via Python (fallback when Rust tokenizer unavailable).
pub fn detokenize_py(&self, tokens: Vec<i32>) -> PyResult<String> {
Python::with_gil(|py| {
Python::attach(|py| {
let result = self
.runtime_handle
.call_method1(py, "detokenize", (tokens,))?;
@@ -306,7 +309,7 @@ impl PyBridge {
}
pub fn list_models(&self) -> PyResult<String> {
Python::with_gil(|py| {
Python::attach(|py| {
let result = self.runtime_handle.call_method0(py, "list_models")?;
result.extract::<String>(py)
})
@@ -314,13 +317,13 @@ impl PyBridge {
fn submit_json<F>(&self, rid: &str, call: F) -> PyResult<Receiver<ResponseChunk>>
where
F: for<'py> FnOnce(Python<'py>, &PyObject, PyObject) -> PyResult<()>,
F: for<'py> FnOnce(Python<'py>, &Py<PyAny>, Py<PyAny>) -> 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 result = Python::attach(|py| -> PyResult<()> {
let callback = self.make_json_callback(py, rid_owned)?;
call(py, &self.runtime_handle, callback)
});
@@ -450,7 +453,7 @@ fn close_channel_with_error(
py: Python<'_>,
rid: &str,
state: &BridgeStateRef,
runtime_handle: &PyObject,
runtime_handle: &Py<PyAny>,
error: TerminalError,
) {
let mut state = lock_or_recover(state.as_ref(), "state");
@@ -478,7 +481,7 @@ fn register_pending_send(rid: &str, state: &BridgeStateRef) -> bool {
state.pending_sends.insert(rid.to_string())
}
fn mark_send_ready(py: Python<'_>, rid: &str, state: &BridgeStateRef) -> Option<PyObject> {
fn mark_send_ready(py: Python<'_>, rid: &str, state: &BridgeStateRef) -> Option<Py<PyAny>> {
let mut state = lock_or_recover(state.as_ref(), "state");
state.pending_sends.remove(rid);
if let Some(callback) = state.ready_callbacks.get(rid) {
@@ -489,7 +492,7 @@ fn mark_send_ready(py: Python<'_>, rid: &str, state: &BridgeStateRef) -> Option<
}
}
fn notify_ready(py: Python<'_>, rid: &str, callback: PyObject) {
fn notify_ready(py: Python<'_>, rid: &str, callback: Py<PyAny>) {
if let Err(err) = callback.call0(py) {
tracing::warn!(rid, "gRPC on_ready callback failed: {}", err);
}
@@ -499,7 +502,7 @@ fn set_on_ready_for_rid(
py: Python<'_>,
rid: &str,
state: &BridgeStateRef,
on_ready: PyObject,
on_ready: Py<PyAny>,
) -> PyResult<()> {
let should_notify = {
let mut state = lock_or_recover(state.as_ref(), "state");
@@ -525,7 +528,7 @@ fn try_send_chunk(
py: Python<'_>,
rid: &str,
state: &BridgeStateRef,
runtime_handle: &PyObject,
runtime_handle: &Py<PyAny>,
tokio_handle: &Handle,
sender: &Sender<ResponseChunk>,
msg: ResponseChunk,
@@ -569,14 +572,14 @@ fn try_send_chunk(
return;
}
Python::with_gil(|py| {
Python::attach(|py| {
if let Some(callback) = mark_send_ready(py, &rid_owned, &state) {
notify_ready(py, &rid_owned, callback);
}
});
}
Err(_) => {
Python::with_gil(|py| {
Python::attach(|py| {
close_channel_with_error(
py,
&rid_owned,
@@ -611,7 +614,7 @@ fn try_send_chunk(
struct ChunkCallback {
rid: String,
state: BridgeStateRef,
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
tokio_handle: Handle,
}
@@ -619,7 +622,7 @@ struct ChunkCallback {
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<()> {
fn set_on_ready(&self, py: Python<'_>, on_ready: Py<PyAny>) -> PyResult<()> {
set_on_ready_for_rid(py, &self.rid, &self.state, on_ready)
}
@@ -699,7 +702,7 @@ impl ChunkCallback {
struct JsonChunkCallback {
rid: String,
state: BridgeStateRef,
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
tokio_handle: Handle,
}
@@ -707,7 +710,7 @@ struct JsonChunkCallback {
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<()> {
fn set_on_ready(&self, py: Python<'_>, on_ready: Py<PyAny>) -> PyResult<()> {
set_on_ready_for_rid(py, &self.rid, &self.state, on_ready)
}
@@ -785,7 +788,7 @@ impl JsonChunkCallback {
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>()
&& let Ok(meta_dict) = meta_obj.cast::<PyDict>()
{
for (k, v) in meta_dict.iter() {
// The proto schema is map<string, string>; encode each Python value as JSON
rust/sglang-grpc/src/lib.rs
+7 -7
View File
@@ -53,10 +53,10 @@ struct TokenizerInfo {
/// fall back to Python tokenization.
fn try_get_attr(
py: Python<'_>,
obj: &PyObject,
obj: &Py<PyAny>,
attr: &'static str,
context: &'static str,
) -> Option<PyObject> {
) -> Option<Py<PyAny>> {
obj.getattr(py, attr).map(Some).unwrap_or_else(|err| {
tracing::debug!("{}.{} is unavailable: {}", context, attr, err);
None
@@ -65,7 +65,7 @@ fn try_get_attr(
fn try_get_attr_str(
py: Python<'_>,
obj: &PyObject,
obj: &Py<PyAny>,
attr: &'static str,
context: &'static str,
) -> Option<String> {
@@ -79,7 +79,7 @@ fn try_get_attr_str(
fn try_get_attr_i32(
py: Python<'_>,
obj: &PyObject,
obj: &Py<PyAny>,
attr: &'static str,
context: &'static str,
) -> Option<i32> {
@@ -91,8 +91,8 @@ fn try_get_attr_i32(
})
}
fn extract_tokenizer_info(runtime_handle: &PyObject) -> PyResult<TokenizerInfo> {
Python::with_gil(|py| {
fn extract_tokenizer_info(runtime_handle: &Py<PyAny>) -> PyResult<TokenizerInfo> {
Python::attach(|py| {
let tm = runtime_handle
.getattr(py, "tokenizer_manager")
.map_err(|err| {
@@ -152,7 +152,7 @@ fn extract_tokenizer_info(runtime_handle: &PyObject) -> PyResult<TokenizerInfo>
fn start_server(
host: String,
port: u16,
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
worker_threads: usize,
response_channel_capacity: usize,
response_timeout_secs: u64,
rust/sglang-grpc/src/server.rs
+4 -3
View File
@@ -64,7 +64,7 @@ fn resolve_max_message_size() -> usize {
/// 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| {
let is_client_error = Python::attach(|py| {
err.is_instance_of::<PyValueError>(py) || err.is_instance_of::<PyTypeError>(py)
});
let msg = format!("{}: {}", context, err);
@@ -125,9 +125,10 @@ impl Drop for RequestAbortGuard {
fn spawn_abort(bridge: Arc<PyBridge>, rid: String) {
match tokio::runtime::Handle::try_current() {
Ok(handle) => {
let _ = handle.spawn_blocking(move || {
// Fire-and-forget: dropping the JoinHandle detaches the task.
drop(handle.spawn_blocking(move || {
let _ = bridge.abort(&rid, false);
});
}));
}
Err(_) => {
tracing::warn!(
rust/sglang-grpc/src/utils/py_utils.rs
+2 -2
View File
@@ -2,7 +2,7 @@ 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> {
fn json_value_to_py<'py>(py: Python<'py>, v: &serde_json::Value) -> PyResult<Py<PyAny>> {
match v {
serde_json::Value::Null => Ok(py.None()),
serde_json::Value::Bool(b) => Ok(b.into_pyobject(py)?.to_owned().into_any().unbind()),
@@ -17,7 +17,7 @@ fn json_value_to_py<'py>(py: Python<'py>, v: &serde_json::Value) -> PyResult<PyO
}
serde_json::Value::String(s) => Ok(s.into_pyobject(py)?.into_any().unbind()),
serde_json::Value::Array(arr) => {
let items: Vec<PyObject> = arr
let items: Vec<Py<PyAny>> = arr
.iter()
.map(|item| json_value_to_py(py, item))
.collect::<PyResult<_>>()?;