From 7fe82dd02e17615fb8679479c9c48ffeb5efc1aa Mon Sep 17 00:00:00 2001 From: Rain Jiang <96632942+rainj-me@users.noreply.github.com> Date: Thu, 23 Jul 2026 12:02:41 -0700 Subject: [PATCH] create rust workspace (#32014) --- .github/workflows/lint.yml | 8 + .github/workflows/pr-test.yml | 4 +- .pre-commit-config.yaml | 15 ++ python/pyproject.toml | 14 +- python/pyproject_other.toml | 14 +- python/setup.py | 181 ++++++++++++++++----- rust/Cargo.toml | 33 ++++ rust/{sglang-grpc => }/rust-toolchain.toml | 1 + rust/sglang-grpc/Cargo.toml | 48 +++--- rust/sglang-grpc/build.rs | 12 ++ rust/sglang-grpc/pyproject.toml | 17 ++ rust/sglang-grpc/src/bridge.rs | 59 +++---- rust/sglang-grpc/src/lib.rs | 14 +- rust/sglang-grpc/src/server.rs | 7 +- rust/sglang-grpc/src/utils/py_utils.rs | 4 +- rust/sglang-mm/.gitignore | 1 - rust/sglang-mm/Cargo.toml | 34 ++-- rust/sglang-mm/pyproject.toml | 4 +- rust/sglang-mm/src/common/mod.rs | 20 +-- rust/sglang-mm/src/common/resize.rs | 57 +++---- rust/sglang-mm/src/common/transforms.rs | 10 +- rust/sglang-mm/src/inkling/mod.rs | 92 ++++++----- rust/sglang-mm/src/registry.rs | 19 ++- scripts/ci/musa/musa_install_dependency.sh | 9 + scripts/ci/utils/install_rustup.sh | 133 +++++++++++---- 25 files changed, 550 insertions(+), 260 deletions(-) create mode 100644 rust/Cargo.toml rename rust/{sglang-grpc => }/rust-toolchain.toml (58%) create mode 100644 rust/sglang-grpc/pyproject.toml delete mode 100644 rust/sglang-mm/.gitignore diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 368186bb6..7cb53216a 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -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 diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 7b99b548e..f1b7e99f8 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -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 }} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 483f950c4..88444f67b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -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' diff --git a/python/pyproject.toml b/python/pyproject.toml index 9a6b3d930..818eb4882 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -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 diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index b6727f358..3264b161c 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -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" = [ diff --git a/python/setup.py b/python/setup.py index 6bd5ef9cc..735c846b6 100644 --- a/python/setup.py +++ b/python/setup.py @@ -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.._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() diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 000000000..af8dd531d --- /dev/null +++ b/rust/Cargo.toml @@ -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 diff --git a/rust/sglang-grpc/rust-toolchain.toml b/rust/rust-toolchain.toml similarity index 58% rename from rust/sglang-grpc/rust-toolchain.toml rename to rust/rust-toolchain.toml index 67d745842..2739b9054 100644 --- a/rust/sglang-grpc/rust-toolchain.toml +++ b/rust/rust-toolchain.toml @@ -1,3 +1,4 @@ [toolchain] channel = "1.90" profile = "minimal" +components = ["clippy", "rustfmt"] diff --git a/rust/sglang-grpc/Cargo.toml b/rust/sglang-grpc/Cargo.toml index d2236ed01..9a223d072 100644 --- a/rust/sglang-grpc/Cargo.toml +++ b/rust/sglang-grpc/Cargo.toml @@ -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 diff --git a/rust/sglang-grpc/build.rs b/rust/sglang-grpc/build.rs index 5ee1bb2d5..1079581a6 100644 --- a/rust/sglang-grpc/build.rs +++ b/rust/sglang-grpc/build.rs @@ -1,4 +1,16 @@ fn main() -> Result<(), Box> { + // 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() diff --git a/rust/sglang-grpc/pyproject.toml b/rust/sglang-grpc/pyproject.toml new file mode 100644 index 000000000..f226c9e4b --- /dev/null +++ b/rust/sglang-grpc/pyproject.toml @@ -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" diff --git a/rust/sglang-grpc/src/bridge.rs b/rust/sglang-grpc/src/bridge.rs index 75a926c01..efabe61b4 100644 --- a/rust/sglang-grpc/src/bridge.rs +++ b/rust/sglang-grpc/src/bridge.rs @@ -40,7 +40,7 @@ type BridgeStateRef = Arc>; struct BridgeState { channels: HashMap>, pending_sends: HashSet, - ready_callbacks: HashMap, + ready_callbacks: HashMap>, ready_signals: HashSet, terminal_errors: HashMap, } @@ -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, 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, state: BridgeStateRef, rust_tokenizer: Option, context_len: i32, @@ -93,7 +96,7 @@ pub struct PyBridge { impl PyBridge { pub fn new( - runtime_handle: PyObject, + runtime_handle: Py, rust_tokenizer: Option, 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 { + fn make_chunk_callback(&self, py: Python<'_>, rid: String) -> PyResult> { 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 { + fn make_json_callback(&self, py: Python<'_>, rid: String) -> PyResult> { 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 { - Python::with_gil(|py| { + Python::attach(|py| { let result = self.runtime_handle.call_method0(py, "get_model_info")?; result.extract::(py) }) } pub fn get_server_info(&self) -> PyResult { - Python::with_gil(|py| { + Python::attach(|py| { let result = self.runtime_handle.call_method0(py, "get_server_info")?; result.extract::(py) }) } pub fn health_check(&self) -> PyResult { - Python::with_gil(|py| { + Python::attach(|py| { let result = self.runtime_handle.call_method0(py, "health_check")?; result.extract::(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 { - 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) -> PyResult { - 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 { - Python::with_gil(|py| { + Python::attach(|py| { let result = self.runtime_handle.call_method0(py, "list_models")?; result.extract::(py) }) @@ -314,13 +317,13 @@ impl PyBridge { fn submit_json(&self, rid: &str, call: F) -> PyResult> where - F: for<'py> FnOnce(Python<'py>, &PyObject, PyObject) -> PyResult<()>, + F: for<'py> FnOnce(Python<'py>, &Py, Py) -> 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, 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 { +fn mark_send_ready(py: Python<'_>, rid: &str, state: &BridgeStateRef) -> Option> { let mut state = lock_or_recover(state.as_ref(), "state"); state.pending_sends.remove(rid); if let Some(callback) = state.ready_callbacks.get(rid) { @@ -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) { 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, ) -> 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, tokio_handle: &Handle, sender: &Sender, 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, 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) -> 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, 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) -> 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 { let mut meta = HashMap::new(); if let Ok(Some(meta_obj)) = chunk.get_item("meta_info") - && let Ok(meta_dict) = meta_obj.downcast::() + && let Ok(meta_dict) = meta_obj.cast::() { for (k, v) in meta_dict.iter() { // The proto schema is map; encode each Python value as JSON diff --git a/rust/sglang-grpc/src/lib.rs b/rust/sglang-grpc/src/lib.rs index c3909322c..cec1ea721 100644 --- a/rust/sglang-grpc/src/lib.rs +++ b/rust/sglang-grpc/src/lib.rs @@ -53,10 +53,10 @@ struct TokenizerInfo { /// fall back to Python tokenization. fn try_get_attr( py: Python<'_>, - obj: &PyObject, + obj: &Py, attr: &'static str, context: &'static str, -) -> Option { +) -> Option> { 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, attr: &'static str, context: &'static str, ) -> Option { @@ -79,7 +79,7 @@ fn try_get_attr_str( fn try_get_attr_i32( py: Python<'_>, - obj: &PyObject, + obj: &Py, attr: &'static str, context: &'static str, ) -> Option { @@ -91,8 +91,8 @@ fn try_get_attr_i32( }) } -fn extract_tokenizer_info(runtime_handle: &PyObject) -> PyResult { - Python::with_gil(|py| { +fn extract_tokenizer_info(runtime_handle: &Py) -> PyResult { + 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 fn start_server( host: String, port: u16, - runtime_handle: PyObject, + runtime_handle: Py, worker_threads: usize, response_channel_capacity: usize, response_timeout_secs: u64, diff --git a/rust/sglang-grpc/src/server.rs b/rust/sglang-grpc/src/server.rs index 33f33c6c9..88ce594e9 100644 --- a/rust/sglang-grpc/src/server.rs +++ b/rust/sglang-grpc/src/server.rs @@ -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::(py) || err.is_instance_of::(py) }); let msg = format!("{}: {}", context, err); @@ -125,9 +125,10 @@ impl Drop for RequestAbortGuard { fn spawn_abort(bridge: Arc, 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!( diff --git a/rust/sglang-grpc/src/utils/py_utils.rs b/rust/sglang-grpc/src/utils/py_utils.rs index 2e99def65..58eee635a 100644 --- a/rust/sglang-grpc/src/utils/py_utils.rs +++ b/rust/sglang-grpc/src/utils/py_utils.rs @@ -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 { +fn json_value_to_py<'py>(py: Python<'py>, v: &serde_json::Value) -> PyResult> { match v { serde_json::Value::Null => Ok(py.None()), serde_json::Value::Bool(b) => Ok(b.into_pyobject(py)?.to_owned().into_any().unbind()), @@ -17,7 +17,7 @@ fn json_value_to_py<'py>(py: Python<'py>, v: &serde_json::Value) -> PyResult Ok(s.into_pyobject(py)?.into_any().unbind()), serde_json::Value::Array(arr) => { - let items: Vec = arr + let items: Vec> = arr .iter() .map(|item| json_value_to_py(py, item)) .collect::>()?; diff --git a/rust/sglang-mm/.gitignore b/rust/sglang-mm/.gitignore deleted file mode 100644 index ea8c4bf7f..000000000 --- a/rust/sglang-mm/.gitignore +++ /dev/null @@ -1 +0,0 @@ -/target diff --git a/rust/sglang-mm/Cargo.toml b/rust/sglang-mm/Cargo.toml index cbbd39317..b8b8a5fff 100644 --- a/rust/sglang-mm/Cargo.toml +++ b/rust/sglang-mm/Cargo.toml @@ -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" diff --git a/rust/sglang-mm/pyproject.toml b/rust/sglang-mm/pyproject.toml index b9c3015cf..057b611e4 100644 --- a/rust/sglang-mm/pyproject.toml +++ b/rust/sglang-mm/pyproject.toml @@ -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" diff --git a/rust/sglang-mm/src/common/mod.rs b/rust/sglang-mm/src/common/mod.rs index 1c4a4a8d6..792971985 100644 --- a/rust/sglang-mm/src/common/mod.rs +++ b/rust/sglang-mm/src/common/mod.rs @@ -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 = 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, ) -> PyResult<(usize, usize, Bound<'py, PyArray1>)> { 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) -> 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> { 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)?)?; diff --git a/rust/sglang-mm/src/common/resize.rs b/rust/sglang-mm/src/common/resize.rs index 0c1d58287..5a9ebb88f 100644 --- a/rust/sglang-mm/src/common/resize.rs +++ b/rust/sglang-mm/src/common/resize.rs @@ -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 { 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 { +pub fn resize_lanczos_rgb(src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize) -> Vec { 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, - cap: Option, -) -> (usize, usize) { +pub fn scaled_dims(w: usize, h: usize, frac: Option, cap: Option) -> (usize, usize) { let Some(frac) = frac else { return (w, h); }; diff --git a/rust/sglang-mm/src/common/transforms.rs b/rust/sglang-mm/src/common/transforms.rs index 6112e2c86..acb770b6d 100644 --- a/rust/sglang-mm/src/common/transforms.rs +++ b/rust/sglang-mm/src/common/transforms.rs @@ -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, 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)) } diff --git a/rust/sglang-mm/src/inkling/mod.rs b/rust/sglang-mm/src/inkling/mod.rs index 7f66873f5..d953d0e8b 100644 --- a/rust/sglang-mm/src/inkling/mod.rs +++ b/rust/sglang-mm/src/inkling/mod.rs @@ -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); +/// [`Patches`] plus the image content hash. +type HashedPatches = (usize, usize, Vec, u64); +/// [`Patches`] with the patch data as a numpy array bound to `'py`. +type PyPatches<'py> = (usize, usize, Bound<'py, PyArray1>); +/// [`PyPatches`] plus the image content hash. +type PyHashedPatches<'py> = (usize, usize, Bound<'py, PyArray1>, 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 { 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>)> { 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, rescale_cap: Option, -) -> PyResult>)>> { +) -> PyResult>> { check_ps(patch_size)?; - let results: Vec), 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> = 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, rescale_cap: Option, -) -> PyResult>, u64)>> { +) -> PyResult>> { check_ps(patch_size)?; - let results: Vec, 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> = 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)?)?; diff --git a/rust/sglang-mm/src/registry.rs b/rust/sglang-mm/src/registry.rs index 2bcb12d36..411ffa5bd 100644 --- a/rust/sglang-mm/src/registry.rs +++ b/rust/sglang-mm/src/registry.rs @@ -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, 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], patch_size: usize, rescale_frac: Option, rescale_cap: Option, - ) -> Result, u64)>, String>; + ) -> Result, String>; } /// Global registry of available processors. @@ -28,6 +26,12 @@ pub struct ProcessorRegistry { specs: Vec>, } +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> { diff --git a/scripts/ci/musa/musa_install_dependency.sh b/scripts/ci/musa/musa_install_dependency.sh index 0171a4440..33f2b1717 100755 --- a/scripts/ci/musa/musa_install_dependency.sh +++ b/scripts/ci/musa/musa_install_dependency.sh @@ -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" diff --git a/scripts/ci/utils/install_rustup.sh b/scripts/ci/utils/install_rustup.sh index 478f48da2..660634d52 100755 --- a/scripts/ci/utils/install_rustup.sh +++ b/scripts/ci/utils/install_rustup.sh @@ -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 ` 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