minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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
.github/workflows/lint.yml
+8
View File
@@ -23,6 +23,14 @@ jobs:
python -m pip install pre-commit
pre-commit install
# The clippy-rust-workspace pre-commit hook compiles the rust/ workspace
# (debug profile); cache cargo registry + target across runs.
- name: Rust cache (rust/ workspace)
uses: Swatinem/rust-cache@v2
with:
workspaces: rust
shared-key: "rust-workspace-lint"
- name: Run pre-commit checks
run: SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure
.github/workflows/pr-test.yml
+2 -2
View File
@@ -309,10 +309,10 @@ jobs:
timeout-minutes: 10
run: bash scripts/ci/utils/install_rust_protoc.sh
- name: Rust cache (sglang-grpc)
- name: Rust cache (rust/ workspace)
uses: Swatinem/rust-cache@v2
with:
workspaces: rust/sglang-grpc
workspaces: rust
shared-key: "sglang-grpc-cpu"
save-if: ${{ matrix.partition == 0 }}
.pre-commit-config.yaml
+15
View File
@@ -117,6 +117,21 @@ repos:
language: system
files: ^experimental/sgl-router/.*\.rs$
pass_filenames: false
- id: clippy-rust-workspace
name: clippy rust/ workspace (auto-fix)
# Two steps: apply machine-applicable fixes, then fail on anything left
# (combining --fix with -D warnings can abort the fix phase). protoc is
# not required: sglang-grpc's build.rs falls back to a vendored binary.
entry: bash -c 'cd rust && cargo clippy --workspace --fix --allow-dirty --allow-staged && cargo clippy --workspace -- -D warnings'
language: system
files: ^rust/.*\.rs$
pass_filenames: false
- id: rustfmt-rust-workspace
name: rustfmt rust/ workspace
entry: bash -c 'cd rust && cargo fmt'
language: system
files: ^rust/.*\.rs$
pass_filenames: false
- id: mint-broken-links
name: check docs_new links (Mintlify; opt-in, requires `mint` on PATH)
entry: bash -c 'cd docs_new && mint broken-links --check-anchors --check-redirects'
python/pyproject.toml
+4 -10
View File
@@ -225,16 +225,10 @@ git_describe_command = ["python3", "python/tools/get_version_tag.py"]
# Allow editable installs even when .git metadata is not available.
fallback_version = "0.0.0.dev0"
[[tool.setuptools-rust.ext-modules]]
target = "sglang.srt.grpc._core"
path = "../rust/sglang-grpc/Cargo.toml"
binding = "PyO3"
[[tool.setuptools-rust.ext-modules]]
target = "sglang.srt.multimodal._core"
path = "../rust/sglang-mm/Cargo.toml"
binding = "PyO3"
debug = false
# Rust extension modules are auto-discovered by setup.py from the cargo
# workspace in ../rust ([package.metadata.sglang] python-module in each crate).
# This CUDA pyproject builds all of them; platform variants restrict the set
# via [tool.sglang] rust-extensions (see pyproject_other.toml).
[tool.kernels.dependencies]
"kernels-community/sgl-flash-attn3" = 1
python/pyproject_other.toml
+8 -6
View File
@@ -192,12 +192,14 @@ dev_mps = ["sglang[all_mps]", "sglang[test]"]
[project.scripts]
sglang = "sglang.cli.main:main"
# Rust-accelerated multimodal preprocessing (sglang.srt.multimodal._core).
# grpc is intentionally omitted here (it needs proto/tonic); ROCm only builds mm.
[[tool.setuptools-rust.ext-modules]]
target = "sglang.srt.multimodal._core"
path = "../rust/sglang-mm/Cargo.toml"
binding = "PyO3"
# Rust extension modules are auto-discovered by setup.py from the cargo
# workspace in ../rust; this allowlist restricts which of them this platform
# builds (case-insensitive substrings of the target module). grpc is
# intentionally omitted here (it needs proto/tonic); ROCm only builds mm.
# NOTE: setup.py reads this with a line regex, not a TOML parser — keep the
# assignment on a single line.
[tool.sglang]
rust-extensions = ["multimodal"]
[tool.setuptools.package-data]
"sglang" = [
python/setup.py
+144 -37
View File
@@ -1,39 +1,152 @@
"""sglang build hooks.
SGLANG_BUILD_RUST_EXTS controls which Rust extensions are built:
- unset or "all": build every declared Rust extension (the default).
- "none": build no Rust extensions.
- comma-separated names: build only extensions whose target matches one of the
given (case-insensitive) substrings, e.g. "grpc" matches
"sglang.srt.grpc._core".
Rust extensions are auto-discovered from the cargo workspace in ../rust: every
crate whose Cargo.toml declares
This is a build-time environment variable, so it is read directly from
os.environ instead of sglang.srt.environ, which is not available until after the
package has been built.
[package.metadata.sglang]
python-module = "sglang.srt.<pkg>._core" # import path inside the wheel
debug = false # optional RustExtension knob
is built as a PyO3 extension module at that import path. Adding a new extension
crate therefore needs no pyproject changes — declare the metadata in the crate.
Two filters can narrow the discovered set:
- [tool.sglang] rust-extensions in the active pyproject.toml: a list of
case-insensitive substrings of the target module. Platform pyprojects use
this to build a subset (e.g. pyproject_other.toml builds only "multimodal";
grpc needs proto/tonic and is intentionally CUDA-only).
- SGLANG_BUILD_RUST_EXTS env var, applied at build time on top of the above:
unset or "all" builds everything, "none" builds nothing, and a
comma-separated list matches substrings, e.g. "grpc" matches
"sglang.srt.grpc._core". It is read directly from os.environ instead of
sglang.srt.environ, which is not importable until the package is built.
"""
import json
import os
import re
import subprocess
from pathlib import Path
from setuptools import setup
try:
from setuptools_rust import build_rust
from setuptools_rust import Binding, RustExtension, build_rust
except ModuleNotFoundError as exc:
if exc.name != "setuptools_rust":
raise
# Alternate platform pyprojects do not declare Rust extensions.
# Alternate platform pyprojects that build no Rust extensions do not
# install setuptools-rust.
build_rust = None
_BUILD_RUST_EXTS_ENV = "SGLANG_BUILD_RUST_EXTS"
_PYTHON_DIR = Path(__file__).resolve().parent
_RUST_WORKSPACE_DIR = _PYTHON_DIR.parent / "rust"
def _cargo_workspace_metadata():
"""The rust/ cargo workspace as JSON, straight from cargo's own parser."""
manifest_path = _RUST_WORKSPACE_DIR / "Cargo.toml"
if not manifest_path.is_file():
raise RuntimeError(
f"no cargo workspace at {manifest_path} (building outside a repo "
f"checkout?); set {_BUILD_RUST_EXTS_ENV}=none to build without "
"Rust extensions"
)
try:
out = subprocess.run(
[
"cargo",
"metadata",
"--format-version",
"1",
"--no-deps",
"--manifest-path",
str(manifest_path),
],
capture_output=True,
check=True,
text=True,
)
except FileNotFoundError as exc:
raise RuntimeError(
"cargo is required to discover the Rust extension modules in "
f"{_RUST_WORKSPACE_DIR} (and to build them); install a Rust "
f"toolchain, or set {_BUILD_RUST_EXTS_ENV}=none to build without "
"Rust extensions"
) from exc
except subprocess.CalledProcessError as exc:
raise RuntimeError(f"cargo metadata failed:\n{exc.stderr}") from exc
return json.loads(out.stdout)
def _match_by_substring(declared, tokens, source):
"""Match tokens as case-insensitive substrings of extension names."""
matched = set()
unmatched = []
for token in tokens:
hits = {ext.name for ext in declared if token in ext.name.lower()}
if hits:
matched |= hits
else:
unmatched.append(token)
if unmatched:
declared_names = sorted(ext.name for ext in declared)
raise ValueError(
f"{source} matched no discovered Rust extension for: {unmatched}; "
f"discovered extensions are {declared_names}"
)
return [ext for ext in declared if ext.name in matched]
def _discovered_rust_extensions():
"""One RustExtension per workspace crate declaring a python-module."""
extensions = []
for package in sorted(
_cargo_workspace_metadata()["packages"], key=lambda p: p["name"]
):
sglang_meta = (package["metadata"] or {}).get("sglang", {})
if "python-module" not in sglang_meta:
continue
extensions.append(
RustExtension(
target=sglang_meta["python-module"],
path=package["manifest_path"],
binding=Binding.PyO3,
debug=sglang_meta.get("debug"),
)
)
if not extensions:
raise RuntimeError(
f"no crate under {_RUST_WORKSPACE_DIR} declares "
"[package.metadata.sglang] python-module; set "
f"{_BUILD_RUST_EXTS_ENV}=none to build without Rust extensions"
)
return extensions
# Deliberately not a TOML parser (keeps setup.py stdlib-only): the allowlist
# must be written as a single line, e.g. rust-extensions = ["multimodal"].
_ALLOWLIST_RE = re.compile(r"^rust-extensions\s*=\s*\[([^\]]*)\]", re.MULTILINE)
def _pyproject_rust_extensions(declared):
"""Apply the active pyproject's [tool.sglang] rust-extensions allowlist."""
pyproject_text = (_PYTHON_DIR / "pyproject.toml").read_text(encoding="utf-8")
match = _ALLOWLIST_RE.search(pyproject_text)
if match is None:
return declared
tokens = re.findall(r'"([^"]*)"', match.group(1))
return _match_by_substring(
declared=declared,
tokens=[token.lower() for token in tokens],
source="[tool.sglang] rust-extensions",
)
def _selected_rust_extensions(declared):
"""Return the Rust extensions selected by SGLANG_BUILD_RUST_EXTS.
`ext.name` is the fully-qualified target (e.g. "sglang.srt.grpc._core") for
the string-target declarations in pyproject.toml, so comma-separated names
are matched as case-insensitive substrings of it.
"""
"""Apply the SGLANG_BUILD_RUST_EXTS build-time filter."""
declared = list(declared)
raw = os.environ.get(_BUILD_RUST_EXTS_ENV)
if raw is None:
@@ -52,23 +165,17 @@ def _selected_rust_extensions(declared):
f"{_BUILD_RUST_EXTS_ENV}={raw!r} has an empty item; unset it or use "
"'all', 'none', or a comma-separated list of extension names"
)
return _match_by_substring(
declared=declared, tokens=tokens, source=_BUILD_RUST_EXTS_ENV
)
matched = set()
unmatched = []
for token in tokens:
hits = {ext.name for ext in declared if token in ext.name.lower()}
if hits:
matched |= hits
else:
unmatched.append(token)
if unmatched:
declared_names = sorted(ext.name for ext in declared)
raise ValueError(
f"{_BUILD_RUST_EXTS_ENV} matched no declared Rust extension for: "
f"{unmatched}; declared extensions are {declared_names}"
)
return [ext for ext in declared if ext.name in matched]
def _declared_rust_extensions():
# "none" short-circuits discovery so builds without a ../rust checkout
# (e.g. from an sdist) still work.
if (os.environ.get(_BUILD_RUST_EXTS_ENV) or "").strip().lower() == "none":
return []
return _pyproject_rust_extensions(_discovered_rust_extensions())
if build_rust is not None:
@@ -84,9 +191,9 @@ if build_rust is not None:
return
super().run()
_cmdclass = {"build_rust": BuildRust}
setup(
cmdclass={"build_rust": BuildRust},
rust_extensions=_declared_rust_extensions(),
)
else:
_cmdclass = {}
setup(cmdclass=_cmdclass)
setup()
rust/Cargo.toml
+33
View File
@@ -0,0 +1,33 @@
[workspace]
resolver = "3"
members = ["sglang-grpc", "sglang-mm"]
[workspace.package]
version = "0.1.0"
edition = "2024"
license = "Apache-2.0"
# Single source of truth for dependency versions across workspace members.
# Members pull these in with `dep = { workspace = true }` (plus extra features
# where needed).
[workspace.dependencies]
async-stream = "0.3"
pyo3 = { version = "0.29.0", features = ["extension-module"] }
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tokio-stream = { version = "0.1", features = ["net"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
uuid = { version = "1", features = ["v4"] }
# Profiles only take effect at the workspace root; per-member profile sections
# are ignored by cargo. Per-package tweaks live in the override tables below.
[profile.release]
lto = true
strip = true
opt-level = 3
codegen-units = 1
[profile.dev]
opt-level = 0
debug = 1
rust/sglang-grpc/rust-toolchain.toml → rust/rust-toolchain.toml
+1
View File
@@ -1,3 +1,4 @@
[toolchain]
channel = "1.90"
profile = "minimal"
components = ["clippy", "rustfmt"]
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/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<_>>()?;
rust/sglang-mm/.gitignore
-1
View File
@@ -1 +0,0 @@
/target
rust/sglang-mm/Cargo.toml
+21 -13
View File
@@ -1,23 +1,31 @@
[package]
name = "sglang-mm"
version = "0.1.0"
edition = "2024"
description = "Rust-accelerated multimodal preprocessing 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. debug = false keeps
# editable installs on release builds (image preprocessing is perf-sensitive).
[package.metadata.sglang]
python-module = "sglang.srt.multimodal._core"
debug = false
[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_mm_core"
crate-type = ["cdylib"]
[dependencies]
pyo3 = { version = "0.23", features = ["extension-module"] }
numpy = "0.23"
rayon = "1.10"
pyo3 = { workspace = true }
base64 = "0.22"
blake3 = "1"
half = "2.4"
image = { version = "0.25", default-features = false, features = ["jpeg", "png"] }
blake3 = "1"
base64 = "0.22"
[profile.release]
lto = true
codegen-units = 1
numpy = "0.29"
rayon = "1.10"
rust/sglang-mm/pyproject.toml
+2 -2
View File
@@ -4,9 +4,9 @@ build-backend = "maturin"
[project]
name = "sglang-mm"
version = "0.1.0"
dynamic = ["version"]
description = "Rust-accelerated multimodal preprocessing for SGLang"
requires-python = ">=3.9"
requires-python = ">=3.10"
[tool.maturin]
module-name = "_core"
rust/sglang-mm/src/common/mod.rs
+9 -11
View File
@@ -7,7 +7,6 @@ use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
pub fn pool() -> &'static rayon::ThreadPool {
static POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
POOL.get_or_init(|| {
@@ -72,10 +71,9 @@ pub fn resize_rgb<'py>(
.as_slice()
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
.to_vec();
let out = py.allow_threads(move || {
pool().install(|| resize::resize_lanczos_rgb(&data, h, w, out_h, out_w))
});
Ok(out.into_pyarray_bound(py))
let out =
py.detach(move || pool().install(|| resize::resize_lanczos_rgb(&data, h, w, out_h, out_w)));
Ok(out.into_pyarray(py))
}
#[pyfunction]
@@ -95,14 +93,14 @@ pub fn image_decode_rgb<'py>(
data: Vec<u8>,
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u8>>)> {
let (rgb, h, w) = py
.allow_threads(move || decode_rgb(&data))
.detach(move || decode_rgb(&data))
.map_err(PyValueError::new_err)?;
Ok((h, w, rgb.into_pyarray_bound(py)))
Ok((h, w, rgb.into_pyarray(py)))
}
#[pyfunction]
pub fn data_hash(py: Python<'_>, data: Vec<u8>) -> u64 {
py.allow_threads(move || {
py.detach(move || {
let digest = blake3::hash(&data);
u64::from_be_bytes(digest.as_bytes()[..8].try_into().unwrap())
})
@@ -115,17 +113,17 @@ pub fn base64_decode<'py>(
) -> PyResult<Bound<'py, pyo3::types::PyBytes>> {
use base64::Engine;
let decoded = py
.allow_threads(|| {
.detach(|| {
base64::engine::general_purpose::STANDARD
.decode(encoded)
.map_err(|e| format!("base64 decode error: {e}"))
})
.map_err(PyValueError::new_err)?;
Ok(pyo3::types::PyBytes::new_bound(py, &decoded))
Ok(pyo3::types::PyBytes::new(py, &decoded))
}
pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
let m = PyModule::new_bound(parent.py(), "common")?;
let m = PyModule::new(parent.py(), "common")?;
m.add_function(wrap_pyfunction!(resize_rgb, &m)?)?;
m.add_function(wrap_pyfunction!(scaled_dims, &m)?)?;
m.add_function(wrap_pyfunction!(image_decode_rgb, &m)?)?;
rust/sglang-mm/src/common/resize.rs
+22 -35
View File
@@ -46,14 +46,14 @@ fn precompute_coeffs(in_size: usize, out_size: usize) -> Coeffs {
let count = (xmax - xmin) as usize;
let k = &mut kkf[xx * ksize..(xx + 1) * ksize];
let mut ww = 0.0f64;
for x in 0..count {
for (x, kv) in k[..count].iter_mut().enumerate() {
let w = lanczos((x as f64 + xmin as f64 - center + 0.5) * ss);
k[x] = w;
*kv = w;
ww += w;
}
if ww != 0.0 {
for x in 0..count {
k[x] /= ww;
for kv in k[..count].iter_mut() {
*kv /= ww;
}
}
bounds[xx] = (xmin as usize, count);
@@ -111,35 +111,27 @@ fn resample_horizontal(src: &[u8], h: usize, w: usize, out_w: usize, c: &Coeffs)
fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec<u8> {
let mut out = vec![0u8; out_h * w * 3];
out.par_chunks_mut(w * 3)
.enumerate()
.for_each(|(yy, row)| {
let (ymin, count) = c.bounds[yy];
let k = &c.kk[yy * c.ksize..yy * c.ksize + count];
for x in 0..w {
let mut s = [1i32 << (PRECISION_BITS - 1); 3];
for (y, &coef) in k.iter().enumerate() {
let p = ((ymin + y) * w + x) * 3;
s[0] += src[p] as i32 * coef;
s[1] += src[p + 1] as i32 * coef;
s[2] += src[p + 2] as i32 * coef;
}
let o = x * 3;
row[o] = clip8(s[0]);
row[o + 1] = clip8(s[1]);
row[o + 2] = clip8(s[2]);
out.par_chunks_mut(w * 3).enumerate().for_each(|(yy, row)| {
let (ymin, count) = c.bounds[yy];
let k = &c.kk[yy * c.ksize..yy * c.ksize + count];
for x in 0..w {
let mut s = [1i32 << (PRECISION_BITS - 1); 3];
for (y, &coef) in k.iter().enumerate() {
let p = ((ymin + y) * w + x) * 3;
s[0] += src[p] as i32 * coef;
s[1] += src[p + 1] as i32 * coef;
s[2] += src[p + 2] as i32 * coef;
}
});
let o = x * 3;
row[o] = clip8(s[0]);
row[o + 1] = clip8(s[1]);
row[o + 2] = clip8(s[2]);
}
});
out
}
pub fn resize_lanczos_rgb(
src: &[u8],
h: usize,
w: usize,
out_h: usize,
out_w: usize,
) -> Vec<u8> {
pub fn resize_lanczos_rgb(src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize) -> Vec<u8> {
let need_h = out_w != w;
let need_v = out_h != h;
if need_h && need_v {
@@ -158,12 +150,7 @@ pub fn resize_lanczos_rgb(
}
}
pub fn scaled_dims(
w: usize,
h: usize,
frac: Option<f64>,
cap: Option<i64>,
) -> (usize, usize) {
pub fn scaled_dims(w: usize, h: usize, frac: Option<f64>, cap: Option<i64>) -> (usize, usize) {
let Some(frac) = frac else {
return (w, h);
};
rust/sglang-mm/src/common/transforms.rs
+7 -3
View File
@@ -2,6 +2,10 @@
//!
//! Model-specific processors compose these to build their preprocessing
//! pipelines. All functions operate on flat RGB byte arrays (HWC layout).
//!
//! Not every primitive is wired into a compiled-in processor yet; they are
//! kept available for upcoming model integrations.
#![allow(dead_code)]
/// Normalize u8 RGB pixels to f32 in a single pass: `(pixel/255 - mean) / std`.
///
@@ -37,8 +41,8 @@ pub fn pad_to_grid(
grid_w: usize,
pad_value: &[f32],
) -> (Vec<f32>, usize, usize) {
let new_h = ((h + grid_h - 1) / grid_h) * grid_h;
let new_w = ((w + grid_w - 1) / grid_w) * grid_w;
let new_h = h.div_ceil(grid_h) * grid_h;
let new_w = w.div_ceil(grid_w) * grid_w;
let mut out = vec![0.0f32; new_h * new_w * channels];
// Fill with pad value
for i in 0..new_h * new_w {
@@ -89,5 +93,5 @@ pub fn extract_patches_hwc(
/// Compute the patch grid dimensions for a given image size and patch size.
#[inline]
pub fn patch_grid(h: usize, w: usize, patch_h: usize, patch_w: usize) -> (usize, usize) {
((h + patch_h - 1) / patch_h, (w + patch_w - 1) / patch_w)
(h.div_ceil(patch_h), w.div_ceil(patch_w))
}
rust/sglang-mm/src/inkling/mod.rs
+52 -40
View File
@@ -8,6 +8,15 @@ use rayon::prelude::*;
use crate::common;
/// `(height, width, patches_as_u16_bits)` for one decoded image.
type Patches = (usize, usize, Vec<u16>);
/// [`Patches`] plus the image content hash.
type HashedPatches = (usize, usize, Vec<u16>, u64);
/// [`Patches`] with the patch data as a numpy array bound to `'py`.
type PyPatches<'py> = (usize, usize, Bound<'py, PyArray1<u16>>);
/// [`PyPatches`] plus the image content hash.
type PyHashedPatches<'py> = (usize, usize, Bound<'py, PyArray1<u16>>, u64);
const MEAN: [f32; 3] = [
0.48145466f64 as f32,
0.4578275f64 as f32,
@@ -40,7 +49,7 @@ fn luts() -> &'static [[u16; 256]; 3] {
#[inline]
pub fn grid(h: usize, w: usize, ps: usize) -> (usize, usize) {
((h + ps - 1) / ps, w / ps + 1)
(h.div_ceil(ps), w / ps + 1)
}
fn patchify_into(arr: &[u8], h: usize, w: usize, ps: usize, out: &mut [u16]) {
@@ -95,7 +104,9 @@ fn patchify_alloc(arr: &[u8], h: usize, w: usize, ps: usize) -> Vec<u16> {
fn check_ps(ps: usize) -> PyResult<()> {
if ps == 0 {
return Err(PyValueError::new_err("patch_size must be greater than zero"));
return Err(PyValueError::new_err(
"patch_size must be greater than zero",
));
}
Ok(())
}
@@ -118,8 +129,8 @@ fn patchify_rgb<'py>(
.as_slice()
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
.to_vec();
let out = py.allow_threads(move || patchify_alloc(&data, h, w, patch_size));
Ok(out.into_pyarray_bound(py))
let out = py.detach(move || patchify_alloc(&data, h, w, patch_size));
Ok(out.into_pyarray(py))
}
#[pyfunction]
@@ -133,14 +144,14 @@ fn decode_patchify<'py>(
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u16>>)> {
check_ps(patch_size)?;
let (h, w, out) = py
.allow_threads(move || {
.detach(move || {
common::pool().install(|| {
let (rgb, h, w) = common::decode_rescale(&data, rescale_frac, rescale_cap)?;
Ok::<_, String>((h, w, patchify_alloc(&rgb, h, w, patch_size)))
})
})
.map_err(PyValueError::new_err)?;
Ok((h, w, out.into_pyarray_bound(py)))
Ok((h, w, out.into_pyarray(py)))
}
#[pyfunction]
@@ -151,25 +162,24 @@ fn decode_patchify_batch<'py>(
patch_size: usize,
rescale_frac: Option<f64>,
rescale_cap: Option<i64>,
) -> PyResult<Vec<(usize, usize, Bound<'py, PyArray1<u16>>)>> {
) -> PyResult<Vec<PyPatches<'py>>> {
check_ps(patch_size)?;
let results: Vec<Result<(usize, usize, Vec<u16>), String>> =
py.allow_threads(move || {
common::pool().install(|| {
datas
.par_iter()
.map(|data| {
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size)))
})
.collect()
})
});
let results: Vec<Result<Patches, String>> = py.detach(move || {
common::pool().install(|| {
datas
.par_iter()
.map(|data| {
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size)))
})
.collect()
})
});
results
.into_iter()
.map(|r| {
let (h, w, v) = r.map_err(PyValueError::new_err)?;
Ok((h, w, v.into_pyarray_bound(py)))
Ok((h, w, v.into_pyarray(py)))
})
.collect()
}
@@ -182,26 +192,25 @@ fn preprocess_images<'py>(
patch_size: usize,
rescale_frac: Option<f64>,
rescale_cap: Option<i64>,
) -> PyResult<Vec<(usize, usize, Bound<'py, PyArray1<u16>>, u64)>> {
) -> PyResult<Vec<PyHashedPatches<'py>>> {
check_ps(patch_size)?;
let results: Vec<Result<(usize, usize, Vec<u16>, u64), String>> =
py.allow_threads(move || {
common::pool().install(|| {
datas
.par_iter()
.map(|data| {
let hash = common::sha256_u64(data);
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash))
})
.collect()
})
});
let results: Vec<Result<HashedPatches, String>> = py.detach(move || {
common::pool().install(|| {
datas
.par_iter()
.map(|data| {
let hash = common::sha256_u64(data);
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash))
})
.collect()
})
});
results
.into_iter()
.map(|r| {
let (h, w, v, hash) = r.map_err(PyValueError::new_err)?;
Ok((h, w, v.into_pyarray_bound(py), hash))
Ok((h, w, v.into_pyarray(py), hash))
})
.collect()
}
@@ -237,7 +246,6 @@ impl crate::registry::ImageProcessorSpec for InklingProcessor {
}
}
#[pyfunction]
#[pyo3(signature = (arr, raw_bytes, patch_size, rescale_frac=None, rescale_cap=None))]
fn rescale_patchify_hash<'py>(
@@ -261,22 +269,26 @@ fn rescale_patchify_hash<'py>(
.as_slice()
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
.to_vec();
let (oh, ow, out) = py.allow_threads(move || {
let (oh, ow, out) = py.detach(move || {
common::pool().install(|| {
let (tw, th) = common::resize::scaled_dims(w, h, rescale_frac, rescale_cap);
let (rgb, h, w) = if (tw, th) != (w, h) {
(common::resize::resize_lanczos_rgb(&rgb, h, w, th, tw), th, tw)
(
common::resize::resize_lanczos_rgb(&rgb, h, w, th, tw),
th,
tw,
)
} else {
(rgb, h, w)
};
(h, w, patchify_alloc(&rgb, h, w, patch_size))
})
});
Ok((oh, ow, out.into_pyarray_bound(py), hash))
Ok((oh, ow, out.into_pyarray(py), hash))
}
pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
let m = PyModule::new_bound(parent.py(), "inkling")?;
let m = PyModule::new(parent.py(), "inkling")?;
m.add_function(wrap_pyfunction!(patchify_rgb, &m)?)?;
m.add_function(wrap_pyfunction!(decode_patchify, &m)?)?;
m.add_function(wrap_pyfunction!(decode_patchify_batch, &m)?)?;
rust/sglang-mm/src/registry.rs
+13 -6
View File
@@ -3,8 +3,8 @@
//! Each model implements `ImageProcessorSpec` and registers itself. The Python
//! layer looks up a processor by model name at init time.
use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;
/// `(height, width, patches_as_u16_bits, content_hash)` for one image.
pub type PreprocessedImage = (usize, usize, Vec<u16>, u64);
/// Trait that each model's image processor must implement.
pub trait ImageProcessorSpec: Send + Sync {
@@ -12,15 +12,13 @@ pub trait ImageProcessorSpec: Send + Sync {
fn name(&self) -> &'static str;
/// Process a batch of raw image bytes: decode + preprocess + hash.
///
/// Returns `(height, width, patches_as_u16_bits, content_hash)` per image.
fn preprocess_batch(
&self,
datas: &[Vec<u8>],
patch_size: usize,
rescale_frac: Option<f64>,
rescale_cap: Option<i64>,
) -> Result<Vec<(usize, usize, Vec<u16>, u64)>, String>;
) -> Result<Vec<PreprocessedImage>, String>;
}
/// Global registry of available processors.
@@ -28,6 +26,12 @@ pub struct ProcessorRegistry {
specs: Vec<Box<dyn ImageProcessorSpec>>,
}
impl Default for ProcessorRegistry {
fn default() -> Self {
Self::new()
}
}
impl ProcessorRegistry {
pub fn new() -> Self {
Self { specs: Vec::new() }
@@ -38,7 +42,10 @@ impl ProcessorRegistry {
}
pub fn lookup(&self, name: &str) -> Option<&dyn ImageProcessorSpec> {
self.specs.iter().find(|s| s.name() == name).map(|s| s.as_ref())
self.specs
.iter()
.find(|s| s.name() == name)
.map(|s| s.as_ref())
}
pub fn list_names(&self) -> Vec<&'static str> {
scripts/ci/musa/musa_install_dependency.sh
+9
View File
@@ -73,6 +73,15 @@ else
find "$REPO_ROOT" -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true
rm -f "${REPO_ROOT}/python/pyproject.toml" && mv "${REPO_ROOT}/python/pyproject_other.toml" "${REPO_ROOT}/python/pyproject.toml"
# setuptools-rust builds the sglang-mm extension (sglang.srt.multimodal._core)
# declared in pyproject_other.toml, so a Rust toolchain must be present like
# on the CUDA/AMD CI paths. Idempotent; installs per-user under $HOME/.cargo.
# Export PATH here because the pip install below runs in this same shell
# (install_rustup.sh's own export/GITHUB_PATH only reach subsequent steps).
bash "${REPO_ROOT}/scripts/ci/utils/install_rustup.sh"
export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:${PATH}"
cd "${REPO_ROOT}" && ${PIP_INSTALL} -v -e "python[dev_musa]" --user
cd "${REPO_ROOT}/sgl-kernel"
scripts/ci/utils/install_rustup.sh
+101 -32
View File
@@ -2,8 +2,22 @@
# Ensure a Rust toolchain (rustc/cargo) is installed for crates built from
# source, e.g. the native gRPC extension bundled into the sglang wheel via
# setuptools-rust. Minimum supported version is 1.85 (edition 2024).
#
# Also pre-installs the workspace-pinned toolchain from rust/rust-toolchain.toml
# (best-effort) so cargo commands run inside rust/ don't pay the rustup
# auto-install on first use.
set -euxo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RUST_WORKSPACE_DIR="${SCRIPT_DIR}/../../../rust"
# Channel pinned by rust/rust-toolchain.toml, used as the default toolchain:
# setuptools-rust wheel builds run cargo from python/ — outside the pin's
# cwd-based scope — so only the default toolchain makes them use the same
# rustc as the workspace. Falls back to stable if the pin can't be parsed.
PINNED_CHANNEL="$(sed -n 's/^channel *= *"\([^"]*\)".*/\1/p' "${RUST_WORKSPACE_DIR}/rust-toolchain.toml" 2>/dev/null || true)"
DEFAULT_CHANNEL="${PINNED_CHANNEL:-stable}"
# Make cargo/rustc visible to the rest of this shell and to subsequent
# GitHub Actions steps in the same job.
export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:${PATH}"
@@ -15,46 +29,101 @@ if [ -n "${GITHUB_PATH:-}" ]; then
echo "${CARGO_HOME:-$HOME/.cargo}/bin" >> "${GITHUB_PATH}" || true
fi
if command -v cargo >/dev/null 2>&1 && command -v rustc >/dev/null 2>&1; then
echo "rust already installed: $(rustc --version), $(cargo --version)"
exit 0
# An in-cluster HTTP mirror may be available (e.g. on NPU runners); export it
# up front so every rustup invocation below (self-heal and pinned-toolchain
# install included) goes through the mirror.
if [ -n "${RUSTUP_CACHE_URL:-}" ]; then
export RUSTUP_DIST_SERVER="${RUSTUP_CACHE_URL}/rustup"
export RUSTUP_UPDATE_ROOT="${RUSTUP_CACHE_URL}/rustup/rustup"
fi
echo "rust not found, installing via rustup..."
install_workspace_pinned_toolchain() {
# Pre-install the toolchain pinned by rust/rust-toolchain.toml: with no
# arguments and cwd inside rust/, `rustup toolchain install` (rustup >=
# 1.28) resolves channel/profile from the toolchain file; older rustups
# fall back to parsing the channel out of the file. Best-effort: the pin
# only governs cargo runs with cwd inside rust/ — setuptools-rust wheel
# builds run cargo from python/ and use the default toolchain, so a failure
# here must not fail the build (rustup auto-installs the pin on first use
# anyway).
if ! command -v rustup >/dev/null 2>&1; then
return 0
fi
local toolchain_file="${RUST_WORKSPACE_DIR}/rust-toolchain.toml"
if [ ! -f "${toolchain_file}" ]; then
return 0
fi
if (cd "${RUST_WORKSPACE_DIR}" && rustup toolchain install); then
return 0
fi
local channel
channel="$(sed -n 's/^channel *= *"\([^"]*\)".*/\1/p' "${toolchain_file}")"
if [ -n "${channel}" ] && rustup toolchain install --profile minimal "${channel}"; then
return 0
fi
echo "WARNING: could not pre-install the toolchain pinned by ${toolchain_file}; rustup will auto-install it on first cargo use inside rust/"
}
# rustup.rs requires curl — make sure it's present.
if ! command -v curl >/dev/null 2>&1; then
if command -v apt-get &> /dev/null; then
apt-get update || true
apt-get install -y --no-install-recommends curl ca-certificates
elif command -v yum &> /dev/null; then
yum install -y curl ca-certificates
if command -v cargo >/dev/null 2>&1 && command -v rustc >/dev/null 2>&1; then
# `command -v` only proves the rustup shims exist. A runner image can ship
# them with no default toolchain configured (rustup-init --default-toolchain
# none, or a removed toolchain), in which case every cargo/rustc call fails
# with "rustup could not choose a version of rustc to run". Check
# functionally and self-heal instead of trusting shim presence.
if rustc --version >/dev/null 2>&1; then
echo "rust already installed: $(rustc --version), $(cargo --version)"
elif command -v rustup >/dev/null 2>&1; then
echo "rustup shims present but no usable default toolchain; installing ${DEFAULT_CHANNEL} as default..."
# `rustup default <channel>` alone is not enough: the toolchain dir may
# exist but be corrupt (a partial install baked into a runner image
# fails later with "Missing manifest in toolchain '...'"), and rustup
# treats any existing dir as installed. Remove it and install fresh
# before selecting it as default.
rustup toolchain uninstall "${DEFAULT_CHANNEL}" || true
rustup toolchain install "${DEFAULT_CHANNEL}"
rustup default "${DEFAULT_CHANNEL}"
else
echo "ERROR: curl is required to install rustup, but no supported package manager was found"
echo "ERROR: cargo/rustc on PATH but non-functional and rustup is missing; remove the stale binaries and re-run"
exit 1
fi
else
echo "rust not found, installing via rustup..."
# rustup.rs requires curl — make sure it's present.
if ! command -v curl >/dev/null 2>&1; then
if command -v apt-get &> /dev/null; then
apt-get update || true
apt-get install -y --no-install-recommends curl ca-certificates
elif command -v yum &> /dev/null; then
yum install -y curl ca-certificates
else
echo "ERROR: curl is required to install rustup, but no supported package manager was found"
exit 1
fi
fi
if [ -n "${RUSTUP_CACHE_URL:-}" ]; then
# Mirror env vars (RUSTUP_DIST_SERVER/RUSTUP_UPDATE_ROOT) were exported
# at the top of this script.
case "$(uname -m)" in
x86_64) RUSTUP_ARCH="x86_64-unknown-linux-gnu" ;;
aarch64) RUSTUP_ARCH="aarch64-unknown-linux-gnu" ;;
*) echo "ERROR: unsupported arch $(uname -m)"; exit 1 ;;
esac
RUSTUP_TMP="$(mktemp -d)"
trap 'rm -rf "${RUSTUP_TMP}"' EXIT
curl --retry 3 --retry-delay 2 -sSfL \
"${RUSTUP_UPDATE_ROOT}/dist/${RUSTUP_ARCH}/rustup-init" \
-o "${RUSTUP_TMP}/rustup-init"
chmod +x "${RUSTUP_TMP}/rustup-init"
"${RUSTUP_TMP}/rustup-init" -y --no-modify-path --default-toolchain "${DEFAULT_CHANNEL}"
else
curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://sh.rustup.rs \
| sh -s -- -y --no-modify-path --default-toolchain "${DEFAULT_CHANNEL}"
fi
fi
if [ -n "${RUSTUP_CACHE_URL:-}" ]; then
# An in-cluster HTTP mirror is available (e.g. on NPU runners).
export RUSTUP_DIST_SERVER="${RUSTUP_CACHE_URL}/rustup"
export RUSTUP_UPDATE_ROOT="${RUSTUP_CACHE_URL}/rustup/rustup"
case "$(uname -m)" in
x86_64) RUSTUP_ARCH="x86_64-unknown-linux-gnu" ;;
aarch64) RUSTUP_ARCH="aarch64-unknown-linux-gnu" ;;
*) echo "ERROR: unsupported arch $(uname -m)"; exit 1 ;;
esac
RUSTUP_TMP="$(mktemp -d)"
trap 'rm -rf "${RUSTUP_TMP}"' EXIT
curl --retry 3 --retry-delay 2 -sSfL \
"${RUSTUP_UPDATE_ROOT}/dist/${RUSTUP_ARCH}/rustup-init" \
-o "${RUSTUP_TMP}/rustup-init"
chmod +x "${RUSTUP_TMP}/rustup-init"
"${RUSTUP_TMP}/rustup-init" -y --no-modify-path
else
curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://sh.rustup.rs \
| sh -s -- -y --no-modify-path
fi
install_workspace_pinned_toolchain
rustc --version
cargo --version