Build Rust extensions on demand in source checkouts (#34994)

This commit is contained in:
Lianmin Zheng
2026-08-16 14:58:06 -07:00
committed by GitHub
parent 0e231d365a
commit 67e12131df
39 changed files with 880 additions and 109 deletions
+1 -1
View File
@@ -49,7 +49,7 @@ runs:
if: steps.artifact.outcome != 'success' if: steps.artifact.outcome != 'success'
uses: actions/cache/restore@v4 uses: actions/cache/restore@v4
with: with:
path: python/sglang/srt/*/_core*.so path: python/sglang/srt/rust_extensions/_*.so
key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }} key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }}
# Job-wide, but only setup.py reads it, and only while building. # Job-wide, but only setup.py reads it, and only while building.
@@ -99,7 +99,7 @@ jobs:
id: cache id: cache
uses: actions/cache/restore@v4 uses: actions/cache/restore@v4
with: with:
path: python/sglang/srt/*/_core*.so path: python/sglang/srt/rust_extensions/_*.so
key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }} key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }}
# On a miss: different hash = rust/setup.py moved; no entries = evicted. # On a miss: different hash = rust/setup.py moved; no entries = evicted.
@@ -112,7 +112,7 @@ jobs:
run: | run: |
if [ -n "${MATCHED_KEY}" ]; then if [ -n "${MATCHED_KEY}" ]; then
echo "hit: ${MATCHED_KEY}" echo "hit: ${MATCHED_KEY}"
ls -l python/sglang/srt/*/_core*.so ls -l python/sglang/srt/rust_extensions/_*.so
else else
echo "miss: ${PRIMARY_KEY}" echo "miss: ${PRIMARY_KEY}"
echo "entries under ${KEY_PREFIX}- (created / ref / size / key):" echo "entries under ${KEY_PREFIX}- (created / ref / size / key):"
@@ -134,7 +134,7 @@ jobs:
uses: actions/upload-artifact@v4 uses: actions/upload-artifact@v4
with: with:
name: ${{ inputs.artifact_name }} name: ${{ inputs.artifact_name }}
# Archive holds <pkg>/_core*.so, so it unpacks into python/sglang/srt/. # Archive holds rust_extensions/_*.so, so it unpacks into python/sglang/srt/.
path: rust-ext-staging/ path: rust-ext-staging/
if-no-files-found: error if-no-files-found: error
retention-days: 1 retention-days: 1
@@ -226,7 +226,7 @@ jobs:
- name: Save built modules - name: Save built modules
uses: actions/cache/save@v4 uses: actions/cache/save@v4
with: with:
path: python/sglang/srt/*/_core*.so path: python/sglang/srt/rust_extensions/_*.so
key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }} key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }}
- name: Upload extension modules - name: Upload extension modules
+1
View File
@@ -230,6 +230,7 @@ work_dirs/
# Rust lib # Rust lib
Cargo.lock Cargo.lock
!rust/Cargo.lock
# Generated vision test fixtures (regenerate with: python scripts/generate_vision_golden.py) # Generated vision test fixtures (regenerate with: python scripts/generate_vision_golden.py)
sgl-model-gateway/tests/fixtures/golden/ sgl-model-gateway/tests/fixtures/golden/
+2 -2
View File
@@ -285,8 +285,8 @@ RUN pip install IPython \
&& pip install pybind11 && pip install pybind11
# Rust toolchain — needed by setuptools-rust to build the sglang-mm extension # Rust toolchain — needed by setuptools-rust to build the sglang-mm extension
# (sglang.srt.multimodal._core) during the sglang pip install below, and later by # (sglang.srt.rust_extensions._multimodal) during the sglang pip install below
# sgl-model-gateway. Must precede the sglang install. # and later by sgl-model-gateway. Must precede the sglang install.
ENV PATH="/root/.cargo/bin:${PATH}" ENV PATH="/root/.cargo/bin:${PATH}"
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \
&& rustc --version && cargo --version && rustc --version && cargo --version
+1
View File
@@ -77,6 +77,7 @@ dependencies = [
"tilelang==0.1.11", "tilelang==0.1.11",
"timm==1.0.16", "timm==1.0.16",
"tokenspeed_mla==0.1.8", "tokenspeed_mla==0.1.8",
"tomli ; python_version < '3.11'",
"torch==2.13.0", "torch==2.13.0",
"torch_memory_saver>=0.0.9.post1", "torch_memory_saver>=0.0.9.post1",
"torchaudio==2.11.0", "torchaudio==2.11.0",
+3 -2
View File
@@ -20,8 +20,9 @@ Two filters can narrow the discovered set:
- SGLANG_BUILD_RUST_EXTS env var, applied at build time on top of the above: - 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 unset or "all" builds everything, "none" builds nothing, and a
comma-separated list matches substrings, e.g. "grpc" matches comma-separated list matches substrings, e.g. "grpc" matches
"sglang.srt.grpc._core". It is read directly from os.environ instead of "sglang.srt.rust_extensions._grpc". It is read directly from os.environ
sglang.srt.environ, which is not importable until the package is built. instead of sglang.srt.environ, which is not importable until the package is
built.
""" """
import json import json
+4 -10
View File
@@ -2687,16 +2687,10 @@ def _start_native_grpc_server_for_runtime(
template_manager, template_manager,
scheduler_info, scheduler_info,
): ):
try: from sglang.srt.entrypoints.grpc_bridge import RuntimeHandle
from sglang.srt.entrypoints.grpc_bridge import RuntimeHandle from sglang.srt.rust_extensions import load_rust_extension
from sglang.srt.grpc import _core as grpc_native
except ImportError as e: grpc_native = load_rust_extension("sglang.srt.rust_extensions._grpc")
raise RuntimeError(
"Native gRPC extension (sglang.srt.grpc._core) not found in this wheel, "
"but --grpc-port was set. The extension is built from "
"rust/sglang-grpc/ via setuptools-rust during wheel build. Either "
"install a wheel that includes the extension or unset --grpc-port."
) from e
runtime_handle = RuntimeHandle( runtime_handle = RuntimeHandle(
tokenizer_manager=tokenizer_manager, tokenizer_manager=tokenizer_manager,
+3
View File
@@ -1523,6 +1523,9 @@ class Envs:
# Rust server # Rust server
# =================================================================== # ===================================================================
SGLANG_RUST_SERVER = EnvBool(False) SGLANG_RUST_SERVER = EnvBool(False)
# Build a missing Rust extension from source (auto), require a bundled or
# cached extension (never), or rebuild the local cache entry (force).
SGLANG_RUST_BUILD_MODE = EnvStr("auto")
# Most batched requests one /generate HTTP call may expand into. # Most batched requests one /generate HTTP call may expand into.
SGLANG_MAX_BATCH_REQS_PER_HTTP_REQ = EnvInt(4096) SGLANG_MAX_BATCH_REQS_PER_HTTP_REQ = EnvInt(4096)
-1
View File
@@ -1 +0,0 @@
# SGLang gRPC module
+4 -2
View File
@@ -37,7 +37,7 @@ if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.managers.io_struct import BatchTokenIDOutput from sglang.srt.managers.io_struct import BatchTokenIDOutput
from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.server._core import Server from sglang.srt.rust_extensions._server import Server
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -367,7 +367,9 @@ class RustServer:
The caller gates this (``SGLANG_RUST_SERVER`` + rank 0); this always The caller gates this (``SGLANG_RUST_SERVER`` + rank 0); this always
creates. creates.
""" """
from sglang.srt.server._core import Server from sglang.srt.rust_extensions import load_rust_extension
Server = load_rust_extension("sglang.srt.rust_extensions._server").Server
# Force turn off HF tokenizers rayon's unpinned global thread pool. # Force turn off HF tokenizers rayon's unpinned global thread pool.
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
@@ -10,8 +10,10 @@ from PIL import Image
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.image_utils import ImageInput from transformers.image_utils import ImageInput
from sglang.srt.multimodal._core import inkling as _rs
from sglang.srt.multimodal.inkling.image_processing import _load_image_bytes from sglang.srt.multimodal.inkling.image_processing import _load_image_bytes
from sglang.srt.rust_extensions import load_rust_extension
_rs = load_rust_extension("sglang.srt.rust_extensions._multimodal").inkling
def _bits_to_bthwc( def _bits_to_bthwc(
@@ -133,7 +133,8 @@ class InklingMultimodalProcessor(SGLangBaseProcessor):
logger.info("Using Rust-accelerated Inkling image processor") logger.info("Using Rust-accelerated Inkling image processor")
except ImportError: except ImportError:
logger.warning( logger.warning(
"SGLANG_INKLING_RS_MM_PREPROCESS=1 but sglang.srt.multimodal._core is not available; " "SGLANG_INKLING_RS_MM_PREPROCESS=1 but "
"sglang.srt.rust_extensions._multimodal is not available; "
"falling back to the default image processor." "falling back to the default image processor."
) )
image_processor = InklingImageProcessor(patch_size=patch_size) image_processor = InklingImageProcessor(patch_size=patch_size)
@@ -0,0 +1,5 @@
"""Loading support for SGLang's optional Rust extension modules."""
from sglang.srt.rust_extensions.loader import RustBuildMode, load_rust_extension
__all__ = ["RustBuildMode", "load_rust_extension"]
+412
View File
@@ -0,0 +1,412 @@
"""Load bundled Rust extensions or build them from an SGLang source tree."""
from __future__ import annotations
import fcntl
import hashlib
import importlib
import importlib.util
import json
import logging
import os
import shutil
import struct
import subprocess
import sys
import sysconfig
import tempfile
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from types import ModuleType
from typing import Iterator, Literal
try:
import tomllib
except ModuleNotFoundError:
import tomli as tomllib # type: ignore[no-redef]
from sglang.srt.environ import envs
logger = logging.getLogger(__name__)
RustBuildMode = Literal["auto", "never", "force"]
_FINGERPRINT_VERSION = 1
_IGNORED_SOURCE_DIRECTORIES = frozenset(
{".git", ".mypy_cache", ".pytest_cache", "__pycache__", "target"}
)
_BUILD_ENVIRONMENT_VARIABLES = (
"CARGO_BUILD_TARGET",
"CARGO_ENCODED_RUSTFLAGS",
"RUSTFLAGS",
)
_RUST_WORKSPACE = Path(__file__).resolve().parents[4] / "rust"
@dataclass(frozen=True)
class _CrateSpec:
"""One extension crate, discovered from its Cargo manifest."""
package: str
library: str
python_module: str
workspace: Path
features: tuple[str, ...]
@dataclass(frozen=True)
class _BuildContext:
source_digest: str
fingerprint: str
target_fingerprint: str
def load_rust_extension(
python_module: str,
*,
mode: RustBuildMode | None = None,
cache_dir: Path | None = None,
workspace: Path | None = None,
) -> ModuleType:
"""Import a PyO3 extension, compiling it locally when permitted and needed.
The crate is discovered from the workspace under ``rust/``: the one whose
Cargo manifest declares ``[package.metadata.sglang] python-module`` equal
to ``python_module`` (the same metadata setup.py uses for wheel builds), so
new crates need no registration here.
``auto`` prefers a module bundled in the installed wheel, then a cached
local build, and finally Cargo. ``never`` permits the first two but never
invokes Cargo. ``force`` rebuilds from source and replaces the cache entry.
``mode`` defaults to ``SGLANG_RUST_BUILD_MODE``.
"""
if mode is None:
mode = envs.SGLANG_RUST_BUILD_MODE.get()
if mode not in ("auto", "never", "force"):
raise ValueError(
f"invalid Rust extension build mode {mode!r}; expected auto, never, or force"
)
if mode != "force":
module = _import_bundled_extension(python_module)
if module is not None:
return module
elif python_module in sys.modules:
raise RuntimeError(
f"cannot force-build {python_module} after it has been imported; "
"start a new Python process"
)
if workspace is None:
workspace = _RUST_WORKSPACE
crate = _discover_crate(workspace, python_module)
context = _build_context(crate)
cache_root = _cache_root(cache_dir)
extension_path = _cached_extension_path(cache_root, crate, context.fingerprint)
lock_path = (
cache_root / "locks" / f"{crate.package}-{context.target_fingerprint}.lock"
)
with _filesystem_lock(lock_path):
if mode != "force" and extension_path.is_file():
return _load_extension_from_path(crate.python_module, extension_path)
if mode == "never":
raise ModuleNotFoundError(
f"{crate.python_module} is not bundled or cached, and Rust extension "
"build mode is 'never'",
name=crate.python_module,
)
target_dir = cache_root / "targets" / context.target_fingerprint
artifact = _cargo_build(crate, target_dir)
if _source_digest(crate.workspace) != context.source_digest:
raise RuntimeError(
f"Rust sources under {crate.workspace} changed during the build; "
"the result was not cached"
)
_stage_atomically(artifact, extension_path)
return _load_extension_from_path(crate.python_module, extension_path)
def _import_bundled_extension(module_name: str) -> ModuleType | None:
try:
return importlib.import_module(module_name)
except ModuleNotFoundError as exc:
if exc.name == module_name:
return None
raise
def _discover_crate(workspace: Path, python_module: str) -> _CrateSpec:
workspace = Path(workspace).resolve()
workspace_manifest = workspace / "Cargo.toml"
lockfile = workspace / "Cargo.lock"
if not workspace_manifest.is_file():
raise FileNotFoundError(
f"Rust workspace for {python_module} was not found at {workspace}"
)
if not lockfile.is_file():
raise FileNotFoundError(
f"{lockfile} is required for reproducible `cargo build --locked` builds"
)
matches: list[_CrateSpec] = []
declared_modules: list[str] = []
for manifest in _source_files(workspace):
if manifest.name != "Cargo.toml":
continue
with manifest.open("rb") as file:
document = tomllib.load(file)
package = document.get("package")
if not isinstance(package, dict):
continue
sglang_metadata = package.get("metadata", {}).get("sglang", {})
declared_module = sglang_metadata.get("python-module")
if declared_module is None:
continue
declared_modules.append(declared_module)
if declared_module != python_module:
continue
package_name = package.get("name")
library = document.get("lib", {}).get("name")
if not package_name or not library:
raise ValueError(
f"{manifest} declares python-module {python_module!r} but must "
"also set `package.name` and `lib.name`"
)
matches.append(
_CrateSpec(
package=package_name,
library=library,
python_module=python_module,
workspace=workspace,
features=tuple(sglang_metadata.get("features", ())),
)
)
if not matches:
raise ModuleNotFoundError(
f"no Cargo package under {workspace} declares "
f'`[package.metadata.sglang] python-module = "{python_module}"`; '
f"declared modules: {sorted(declared_modules)}",
name=python_module,
)
if len(matches) > 1:
raise ValueError(
f"multiple Cargo packages under {workspace} declare python module "
f"{python_module!r}: {sorted(crate.package for crate in matches)}"
)
return matches[0]
def _build_context(crate: _CrateSpec) -> _BuildContext:
source_digest = _source_digest(crate.workspace)
toolchain = {
"cargo": _command_version(
"cargo", "--version", "--verbose", cwd=crate.workspace
),
"rustc": _command_version("rustc", "-vV", cwd=crate.workspace),
}
python_abi = {
"cache_tag": sys.implementation.cache_tag,
"ext_suffix": sysconfig.get_config_var("EXT_SUFFIX"),
"platform": sysconfig.get_platform(),
"pointer_bits": struct.calcsize("P") * 8,
"soabi": sysconfig.get_config_var("SOABI"),
"version": list(sys.version_info[:3]),
}
build_environment = {
name: os.environ.get(name) for name in _BUILD_ENVIRONMENT_VARIABLES
}
target_inputs = {
"build_environment": build_environment,
"python_abi": python_abi,
"toolchain": toolchain,
}
target_fingerprint = _json_digest(target_inputs)[:24]
fingerprint = _json_digest(
{
"fingerprint_version": _FINGERPRINT_VERSION,
"package": crate.package,
"library": crate.library,
"python_module": crate.python_module,
"source_digest": source_digest,
**target_inputs,
}
)
return _BuildContext(
source_digest=source_digest,
fingerprint=fingerprint,
target_fingerprint=target_fingerprint,
)
def _source_digest(workspace: Path) -> str:
digest = hashlib.sha256()
for path in _source_files(workspace):
relative_path = path.relative_to(workspace).as_posix().encode()
digest.update(len(relative_path).to_bytes(8, "big"))
digest.update(relative_path)
if path.is_symlink():
contents = os.readlink(path).encode()
else:
contents = path.read_bytes()
digest.update(len(contents).to_bytes(8, "big"))
digest.update(contents)
return digest.hexdigest()
def _source_files(workspace: Path) -> Iterator[Path]:
for root, directories, filenames in os.walk(workspace):
directories[:] = sorted(
name for name in directories if name not in _IGNORED_SOURCE_DIRECTORIES
)
root_path = Path(root)
for filename in sorted(filenames):
yield root_path / filename
def _command_version(command: str, *arguments: str, cwd: Path) -> str:
try:
result = subprocess.run(
[command, *arguments],
check=True,
capture_output=True,
text=True,
cwd=cwd,
)
except (OSError, subprocess.CalledProcessError) as exc:
raise RuntimeError(
f"failed to query the Rust toolchain with `{command} {' '.join(arguments)}`"
) from exc
return result.stdout.strip()
def _json_digest(value: object) -> str:
serialized = json.dumps(
value, sort_keys=True, separators=(",", ":"), ensure_ascii=True
).encode()
return hashlib.sha256(serialized).hexdigest()
def _cache_root(cache_dir: Path | None) -> Path:
if cache_dir is not None:
return Path(cache_dir).expanduser().resolve()
sglang_cache = envs.SGLANG_CACHE_DIR.get()
return Path(sglang_cache).expanduser().resolve() / "rust_extensions"
def _cached_extension_path(
cache_root: Path, crate: _CrateSpec, fingerprint: str
) -> Path:
extension_suffix = sysconfig.get_config_var("EXT_SUFFIX")
if not extension_suffix:
raise RuntimeError("Python did not report an EXT_SUFFIX for native extensions")
module_leaf = crate.python_module.rsplit(".", 1)[-1]
return (
cache_root
/ "artifacts"
/ crate.package
/ fingerprint
/ (module_leaf + extension_suffix)
)
@contextmanager
def _filesystem_lock(path: Path) -> Iterator[None]:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("a+b") as lock_file:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
def _cargo_build(crate: _CrateSpec, target_dir: Path) -> Path:
target_dir.mkdir(parents=True, exist_ok=True)
command = [
"cargo",
"build",
"--release",
"--locked",
"--package",
crate.package,
]
if crate.features:
command.extend(("--features", ",".join(crate.features)))
environment = os.environ.copy()
environment["CARGO_TARGET_DIR"] = os.fspath(target_dir)
environment["PYO3_PYTHON"] = sys.executable
logger.info("Building %s with `%s`", crate.python_module, " ".join(command))
try:
subprocess.run(command, cwd=crate.workspace, env=environment, check=True)
except (OSError, subprocess.CalledProcessError) as exc:
raise RuntimeError(f"failed to build {crate.python_module} with Cargo") from exc
release_dir = target_dir / "release"
if target := environment.get("CARGO_BUILD_TARGET"):
release_dir = target_dir / target / "release"
artifact = release_dir / _cargo_library_filename(crate.library)
if not artifact.is_file():
raise FileNotFoundError(
f"Cargo completed but did not produce the expected artifact {artifact}"
)
return artifact
def _cargo_library_filename(library: str) -> str:
if sys.platform == "win32":
return f"{library}.dll"
if sys.platform == "darwin":
return f"lib{library}.dylib"
return f"lib{library}.so"
def _stage_atomically(source: Path, destination: Path) -> None:
destination.parent.mkdir(parents=True, exist_ok=True)
descriptor, temporary_name = tempfile.mkstemp(
prefix=f".{destination.name}.", dir=destination.parent
)
temporary_path = Path(temporary_name)
try:
with (
os.fdopen(descriptor, "wb") as destination_file,
source.open("rb") as source_file,
):
shutil.copyfileobj(source_file, destination_file)
destination_file.flush()
os.fsync(destination_file.fileno())
temporary_path.chmod(0o755)
os.replace(temporary_path, destination)
directory_descriptor = os.open(destination.parent, os.O_RDONLY)
try:
os.fsync(directory_descriptor)
finally:
os.close(directory_descriptor)
finally:
temporary_path.unlink(missing_ok=True)
def _load_extension_from_path(module_name: str, path: Path) -> ModuleType:
loaded = sys.modules.get(module_name)
if loaded is not None:
return loaded
module_spec = importlib.util.spec_from_file_location(module_name, path)
if module_spec is None or module_spec.loader is None:
raise ImportError(
f"could not create an import spec for {module_name} at {path}"
)
module = importlib.util.module_from_spec(module_spec)
sys.modules[module_name] = module
try:
module_spec.loader.exec_module(module)
except BaseException:
sys.modules.pop(module_name, None)
raise
return module
+3 -10
View File
@@ -219,17 +219,10 @@ def is_rust_server_built():
"""Return whether the embedded Rust server extension (``SGLANG_RUST_SERVER``) """Return whether the embedded Rust server extension (``SGLANG_RUST_SERVER``)
is importable. is importable.
``sglang/srt/server/`` is not in the source tree — it is produced by The ``sglang.srt.rust_extensions`` Python package is always present; the
``setup.py build_rust --inplace``, so on a build without it ``find_spec`` private ``_server`` module exists only when the PyO3 extension was built.
raises ``ModuleNotFoundError`` for the missing *parent* package rather than
returning ``None`` for the missing leaf. Suites gate a rust-server subclass on
this at class-definition time, so letting that escape would fail the whole
module import instead of skipping the one class.
""" """
try: return importlib.util.find_spec("sglang.srt.rust_extensions._server") is not None
return importlib.util.find_spec("sglang.srt.server._core") is not None
except ModuleNotFoundError:
return False
def _use_cached_default_models(model_repo: str): def _use_cached_default_models(model_repo: str):
+20 -17
View File
@@ -982,9 +982,9 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
[[package]] [[package]]
name = "dynamo-parsers" name = "dynamo-parsers"
version = "7.0.1" version = "7.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "97cce1ec70f4c4896ff9a4e0bb29cd603f465e599c94bf2085b71872a3a58c70" checksum = "7a7f84812e7ce1d3dace1b8d69853cc6c49759cbd07dfbb59452aa928645b031"
dependencies = [ dependencies = [
"aho-corasick", "aho-corasick",
"anyhow", "anyhow",
@@ -1005,9 +1005,9 @@ dependencies = [
[[package]] [[package]]
name = "dynamo-protocols" name = "dynamo-protocols"
version = "5.1.0" version = "5.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "696b6df4c817c16dc0a619d43f6af5ecc4f045a6b6be7ef516bfe502f977764b" checksum = "97fd951c32c033f4f220b6087db91d7f7d475f1bf1f886c8e60798ac1cf4cc9a"
dependencies = [ dependencies = [
"async-openai", "async-openai",
"derive_builder", "derive_builder",
@@ -1022,9 +1022,9 @@ dependencies = [
[[package]] [[package]]
name = "dynamo-renderer" name = "dynamo-renderer"
version = "5.0.0" version = "5.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0423da7bec83734e7567b1c8895fd2e7a1a98256253a29fe128032ac01cefac1" checksum = "dfe04753d666e3462e3eed4abcf2a4d05997cf2dea45858733f2bdea5b61f733"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"chrono", "chrono",
@@ -1040,9 +1040,9 @@ dependencies = [
[[package]] [[package]]
name = "dynamo-tokenizers" name = "dynamo-tokenizers"
version = "1.7.0" version = "1.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a81f7b7693282f4693a3b4c60174e14260cf011ab493a7ca9b4deb2fa98ecee" checksum = "821c767e896f1f225b6411a5ce41dba7f114c2c97db862a13fd962195665075e"
dependencies = [ dependencies = [
"aho-corasick", "aho-corasick",
"anyhow", "anyhow",
@@ -1162,21 +1162,24 @@ dependencies = [
[[package]] [[package]]
name = "fastokens" name = "fastokens"
version = "0.2.1" version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8728655e193e0d08d7a95d63cf1fdb9b768d282cab0a112ecb006615bae9f067" checksum = "35ff8cf975d8f772e7d4c3be439659b0b23b1ae1086ebee7204573807053ff7f"
dependencies = [ dependencies = [
"daachorse", "daachorse",
"fancy-regex 0.17.0", "fancy-regex 0.17.0",
"hf-hub", "hf-hub",
"icu_normalizer", "icu_normalizer",
"libc",
"memchr", "memchr",
"pcre2", "pcre2",
"rayon", "rayon",
"regex-syntax",
"serde", "serde",
"serde_json", "serde_json",
"strum", "strum",
"thiserror", "thiserror",
"ureq",
] ]
[[package]] [[package]]
@@ -1633,7 +1636,7 @@ dependencies = [
"libc", "libc",
"percent-encoding", "percent-encoding",
"pin-project-lite", "pin-project-lite",
"socket2 0.6.5", "socket2 0.5.10",
"tokio", "tokio",
"tower-service", "tower-service",
"tracing", "tracing",
@@ -2123,9 +2126,9 @@ dependencies = [
[[package]] [[package]]
name = "minijinja" name = "minijinja"
version = "2.21.0" version = "2.24.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" checksum = "86886cf6dbf4e614b19c9a1eec9775f021869d7eadde0fc73921a81b90c9b4c9"
dependencies = [ dependencies = [
"indexmap 2.14.0", "indexmap 2.14.0",
"memo-map", "memo-map",
@@ -2135,9 +2138,9 @@ dependencies = [
[[package]] [[package]]
name = "minijinja-contrib" name = "minijinja-contrib"
version = "2.21.0" version = "2.24.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85342f6fac0be8ccd5bd00d9066be538f34f393f577b75d81b17c8398a6b43bb" checksum = "bd3e5f077bc2379f0f7d911e7cfdd921114ed99fc884533dca502944cb355b11"
dependencies = [ dependencies = [
"minijinja", "minijinja",
"serde", "serde",
@@ -2940,7 +2943,7 @@ dependencies = [
"quinn-udp", "quinn-udp",
"rustc-hash 2.1.3", "rustc-hash 2.1.3",
"rustls", "rustls",
"socket2 0.6.5", "socket2 0.5.10",
"thiserror", "thiserror",
"tokio", "tokio",
"tracing", "tracing",
@@ -2978,7 +2981,7 @@ dependencies = [
"cfg_aliases", "cfg_aliases",
"libc", "libc",
"once_cell", "once_cell",
"socket2 0.6.5", "socket2 0.5.10",
"tracing", "tracing",
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]
+1 -1
View File
@@ -8,7 +8,7 @@ license.workspace = true
# Consumed by python/setup.py: registers this crate as a PyO3 extension module # Consumed by python/setup.py: registers this crate as a PyO3 extension module
# of the main sglang wheel at the given import path. # of the main sglang wheel at the given import path.
[package.metadata.sglang] [package.metadata.sglang]
python-module = "sglang.srt.grpc._core" python-module = "sglang.srt.rust_extensions._grpc"
# Always build optimized, even for an editable install. # Always build optimized, even for an editable install.
debug = false debug = false
+3 -3
View File
@@ -1,7 +1,7 @@
# Standalone dev builds only (`maturin build` / `maturin develop`); the # Standalone dev builds only (`maturin build` / `maturin develop`); the
# extension ships to users inside the main sglang wheel via setuptools-rust # extension ships to users inside the main sglang wheel via setuptools-rust
# (python/pyproject.toml, target sglang.srt.grpc._core). Requires a repo # (python/pyproject.toml, target sglang.srt.rust_extensions._grpc). Requires a
# checkout (build.rs reads ../../proto), so this is not publishable as an # 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. # sdist. protoc is not required: build.rs falls back to a vendored binary.
[build-system] [build-system]
requires = ["maturin>=1.5,<2"] requires = ["maturin>=1.5,<2"]
@@ -14,4 +14,4 @@ description = "In-process Rust gRPC server for SGLang"
requires-python = ">=3.10" requires-python = ">=3.10"
[tool.maturin] [tool.maturin]
module-name = "_core" module-name = "_grpc"
+1 -1
View File
@@ -257,7 +257,7 @@ fn start_server(
} }
#[pymodule] #[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { fn _grpc(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(start_server, m)?)?; m.add_function(wrap_pyfunction!(start_server, m)?)?;
m.add_class::<GrpcServerHandle>()?; m.add_class::<GrpcServerHandle>()?;
m.add_class::<ChunkSendStatus>()?; m.add_class::<ChunkSendStatus>()?;
+5 -5
View File
@@ -11,17 +11,17 @@ license.workspace = true
# features: the wheel build needs the PyO3 bindings, which are NOT default (see # features: the wheel build needs the PyO3 bindings, which are NOT default (see
# [features] below). # [features] below).
[package.metadata.sglang] [package.metadata.sglang]
python-module = "sglang.srt.multimodal._core" python-module = "sglang.srt.rust_extensions._multimodal"
debug = false debug = false
features = ["python", "parallel"] features = ["python", "parallel"]
[lib] [lib]
# Unique per-crate artifact name (the shared workspace target/ dir cannot hold # Unique per-crate artifact name (the shared workspace target/ dir cannot hold
# two `lib_core` cdylibs). The importable Python module is still `_core`: the # multiple `lib_core` cdylibs). The importable Python module is `_multimodal`: the
# name comes from the `#[pymodule]` entry point and the setuptools-rust # name comes from the `#[pymodule]` entry point and the setuptools-rust target,
# `target` in python/pyproject.toml, which renames the built artifact. # which renames the built artifact.
name = "sglang_mm_core" name = "sglang_mm_core"
# cdylib: the PyO3 module (`sglang.srt.multimodal._core`). # cdylib: the PyO3 module (`sglang.srt.rust_extensions._multimodal`).
# rlib: pure-Rust core linked by sglang-server's native MM path. # rlib: pure-Rust core linked by sglang-server's native MM path.
crate-type = ["cdylib", "rlib"] crate-type = ["cdylib", "rlib"]
+7 -6
View File
@@ -6,9 +6,9 @@ GIL-released.
Built two ways: Built two ways:
- **PyO3 extension** `sglang.srt.multimodal._core` (features `python,parallel`, - **PyO3 extension** `sglang.srt.rust_extensions._multimodal` (features
requested by the wheel build) via setuptools-rust when installing sglang — `python,parallel`, requested by the wheel build) via setuptools-rust when
used by Python processors and parity tests. installing sglang — used by Python processors and parity tests.
- **Pure-Rust `rlib`** (default features, i.e. neither) linked by - **Pure-Rust `rlib`** (default features, i.e. neither) linked by
`sglang-server`'s MM worker path — that copy needs no pyo3, no libpython, and `sglang-server`'s MM worker path — that copy needs no pyo3, no libpython, and
no rayon: it spawns no threads and runs inline on the calling thread, because no rayon: it spawns no threads and runs inline on the calling thread, because
@@ -19,7 +19,7 @@ Built two ways:
``` ```
src/ src/
├── lib.rs # module root; PyO3 module (_core) feature-gated ├── lib.rs # module root; PyO3 module (_multimodal) feature-gated
├── pipeline.rs # the server-pipeline contract: MmFamilyProcessor ├── pipeline.rs # the server-pipeline contract: MmFamilyProcessor
│ # trait + the carriers (Tensor, TokenLayout, ...) │ # trait + the carriers (Tensor, TokenLayout, ...)
├── driver.rs # model-independent request driver (fetch → ├── driver.rs # model-independent request driver (fetch →
@@ -115,7 +115,7 @@ un-stripped URL to `open()`).
## Python API ## Python API
```python ```python
from sglang.srt.multimodal._core import common, inkling from sglang.srt.rust_extensions._multimodal import common, inkling
# Common (model-agnostic) # Common (model-agnostic)
common.resize_rgb(arr, out_w, out_h) common.resize_rgb(arr, out_w, out_h)
@@ -174,7 +174,8 @@ impl ImageProcessorSpec for MyModelProcessor {
4. Wire up in `src/lib.rs`: `mod my_model;` and `my_model::register(m)?;`. 4. Wire up in `src/lib.rs`: `mod my_model;` and `my_model::register(m)?;`.
5. Add Python processor class that calls `from sglang.srt.multimodal._core import my_model`. 5. Add Python processor class that calls
`from sglang.srt.rust_extensions._multimodal import my_model`.
## Available transform primitives (`common::transforms`) ## Available transform primitives (`common::transforms`)
+1 -1
View File
@@ -5,7 +5,6 @@ import numpy as np
import torch import torch
from PIL import Image from PIL import Image
from sglang.srt.multimodal._core import inkling as _rs_inkling
from sglang.srt.multimodal.inkling.image_processing import ( from sglang.srt.multimodal.inkling.image_processing import (
IMAGE_MEAN, IMAGE_MEAN,
IMAGE_STD, IMAGE_STD,
@@ -13,6 +12,7 @@ from sglang.srt.multimodal.inkling.image_processing import (
_encode_image_bytes, _encode_image_bytes,
_fill_patches_numba, _fill_patches_numba,
) )
from sglang.srt.rust_extensions._multimodal import inkling as _rs_inkling
PS = 40 PS = 40
+1 -1
View File
@@ -9,4 +9,4 @@ description = "Rust-accelerated multimodal preprocessing for SGLang"
requires-python = ">=3.10" requires-python = ">=3.10"
[tool.maturin] [tool.maturin]
module-name = "_core" module-name = "_multimodal"
+2 -2
View File
@@ -1,7 +1,7 @@
//! sglang-mm: Rust-accelerated multimodal preprocessing for SGLang. //! sglang-mm: Rust-accelerated multimodal preprocessing for SGLang.
//! //!
//! Built two ways: //! Built two ways:
//! * PyO3 extension `sglang.srt.multimodal._core` (feature `python`, default), //! * PyO3 extension `sglang.srt.rust_extensions._multimodal` (feature `python`),
//! used by Python processors (e.g. Inkling) and by parity tests. //! used by Python processors (e.g. Inkling) and by parity tests.
//! * Pure-Rust `rlib` (`default-features = false`), linked by `sglang-server`'s //! * Pure-Rust `rlib` (`default-features = false`), linked by `sglang-server`'s
//! MM worker path — no pyo3 in that dependency graph. //! MM worker path — no pyo3 in that dependency graph.
@@ -18,7 +18,7 @@ use pyo3::prelude::*;
#[cfg(feature = "python")] #[cfg(feature = "python")]
#[pymodule] #[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { fn _multimodal(m: &Bound<'_, PyModule>) -> PyResult<()> {
common::register(m)?; common::register(m)?;
inkling::register(m)?; inkling::register(m)?;
qwen_vl::register(m)?; qwen_vl::register(m)?;
+1 -1
View File
@@ -4,7 +4,7 @@ import os
import numpy as np import numpy as np
import pytest import pytest
from sglang.srt.multimodal._core import inkling as _rs_inkling from sglang.srt.rust_extensions._multimodal import inkling as _rs_inkling
GOLDEN_DIR = os.environ.get( GOLDEN_DIR = os.environ.get(
"INKLING_MM_GOLDEN_DIR", "INKLING_MM_GOLDEN_DIR",
+1 -1
View File
@@ -15,12 +15,12 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "bench"))
from bench_parity import make_photo_like from bench_parity import make_photo_like
from sglang.srt.managers.mm_utils import data_hash, hash_feature from sglang.srt.managers.mm_utils import data_hash, hash_feature
from sglang.srt.multimodal._core import common as _rs_common
from sglang.srt.multimodal.inkling import InklingProcessor from sglang.srt.multimodal.inkling import InklingProcessor
from sglang.srt.multimodal.inkling.image_processing_rust import ( from sglang.srt.multimodal.inkling.image_processing_rust import (
InklingRustImageProcessor, InklingRustImageProcessor,
) )
from sglang.srt.multimodal.processors import inkling as prc from sglang.srt.multimodal.processors import inkling as prc
from sglang.srt.rust_extensions._multimodal import common as _rs_common
def png_bytes(arr): def png_bytes(arr):
+2 -2
View File
@@ -6,8 +6,8 @@ import numpy as np
import pytest import pytest
from PIL import Image from PIL import Image
from sglang.srt.multimodal._core import common as _rs_common from sglang.srt.rust_extensions._multimodal import common as _rs_common
from sglang.srt.multimodal._core import inkling as _rs_inkling from sglang.srt.rust_extensions._multimodal import inkling as _rs_inkling
def py_scaled_dims( def py_scaled_dims(
+1 -1
View File
@@ -7,7 +7,7 @@ license.workspace = true
# Consumed by python/setup.py: registers this crate as a PyO3 extension module. # Consumed by python/setup.py: registers this crate as a PyO3 extension module.
[package.metadata.sglang] [package.metadata.sglang]
python-module = "sglang.srt.server._core" python-module = "sglang.srt.rust_extensions._server"
# Always build optimized, even for an editable install. # Always build optimized, even for an editable install.
debug = false debug = false
+1 -1
View File
@@ -13,4 +13,4 @@ classifiers = [
dynamic = ["version"] dynamic = ["version"]
[tool.maturin] [tool.maturin]
module-name = "_core" module-name = "_server"
@@ -836,6 +836,7 @@ pub(super) fn chat_logprobs(extras: Option<&ChunkExtras>) -> ChatChoiceLogprobs
bytes: Some(token.as_bytes().to_vec()), bytes: Some(token.as_bytes().to_vec()),
token, token,
logprob, logprob,
token_id: u32::try_from(token_id).ok(),
top_logprobs, top_logprobs,
}); });
} }
@@ -1020,6 +1021,7 @@ mod tests {
let logprobs = chat_logprobs(Some(&extras)); let logprobs = chat_logprobs(Some(&extras));
let token = &logprobs.content.unwrap()[0]; let token = &logprobs.content.unwrap()[0];
assert_eq!(token.token, "x"); assert_eq!(token.token, "x");
assert_eq!(token.token_id, Some(7));
assert_eq!(token.top_logprobs.len(), 2); assert_eq!(token.top_logprobs.len(), 2);
assert_eq!(token.top_logprobs[1].token, "y"); assert_eq!(token.top_logprobs[1].token, "y");
} }
+1 -1
View File
@@ -292,7 +292,7 @@ static LOG_GUARD: std::sync::OnceLock<tracing_appender::non_blocking::WorkerGuar
std::sync::OnceLock::new(); std::sync::OnceLock::new();
#[pymodule] #[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { fn _server(m: &Bound<'_, PyModule>) -> PyResult<()> {
// Initialize tracing once; ignore if already set by the host process. // Initialize tracing once; ignore if already set by the host process.
// Non-blocking writer: emitting threads (axum workers, egress, detok) only // Non-blocking writer: emitting threads (axum workers, egress, detok) only
// enqueue; a dedicated thread does the stdout formatting-flush + syscall. // enqueue; a dedicated thread does the stdout formatting-flush + syscall.
+7 -7
View File
@@ -460,7 +460,7 @@ require_prebuilt_rust_exts() {
return return
fi fi
# Exact EXT_SUFFIX rather than a _core*.so glob: no crate sets abi3, so a module # Exact EXT_SUFFIX rather than an _*.so glob: no crate sets abi3, so a module
# built for another minor version satisfies the glob while the import system # built for another minor version satisfies the glob while the import system
# ignores it, leaving is_rust_server_built() false and the Rust-server tests # ignores it, leaving is_rust_server_built() false and the Rust-server tests
# silently skipped. Stages have no setup-python, so the interpreter is whatever # silently skipped. Stages have no setup-python, so the interpreter is whatever
@@ -469,13 +469,13 @@ require_prebuilt_rust_exts() {
local suffix local suffix
suffix=$(python3 -c 'import sysconfig; print(sysconfig.get_config_var("EXT_SUFFIX"))') suffix=$(python3 -c 'import sysconfig; print(sysconfig.get_config_var("EXT_SUFFIX"))')
local missing=() local missing=()
local pkg local module
for pkg in server grpc multimodal; do for module in server grpc multimodal; do
[ -f "python/sglang/srt/${pkg}/_core${suffix}" ] || missing+=("${pkg}") [ -f "python/sglang/srt/rust_extensions/_${module}${suffix}" ] || missing+=("${module}")
done done
if [ ${#missing[@]} -gt 0 ]; then if [ ${#missing[@]} -gt 0 ]; then
echo "::warning::no prebuilt _core${suffix} for: ${missing[*]}; building from source" echo "::warning::no prebuilt Rust extension ${suffix} for: ${missing[*]}; building from source"
ls -l python/sglang/srt/*/_core*.so 2>/dev/null || echo "(no extension modules at all)" ls -l python/sglang/srt/rust_extensions/_*.so 2>/dev/null || echo "(no extension modules at all)"
export SGLANG_BUILD_RUST_EXTS= export SGLANG_BUILD_RUST_EXTS=
mark_step_done "${FUNCNAME[0]}" mark_step_done "${FUNCNAME[0]}"
return return
@@ -812,7 +812,7 @@ print(f"sglang resolves to {spec.origin}")
# so a .so that cannot load passes find_spec and only fails inside some suite. # so a .so that cannot load passes find_spec and only fails inside some suite.
import importlib import importlib
for mod in ("server", "grpc", "multimodal"): for mod in ("server", "grpc", "multimodal"):
name = f"sglang.srt.{mod}._core" name = f"sglang.srt.rust_extensions._{mod}"
try: try:
importlib.import_module(name) importlib.import_module(name)
except Exception as exc: except Exception as exc:
+2 -1
View File
@@ -74,7 +74,8 @@ else
rm -f "${REPO_ROOT}/python/pyproject.toml" && mv "${REPO_ROOT}/python/pyproject_other.toml" "${REPO_ROOT}/python/pyproject.toml" 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) # setuptools-rust builds the sglang-mm extension
# (sglang.srt.rust_extensions._multimodal)
# declared in pyproject_other.toml, so a Rust toolchain must be present like # 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. # 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 # Export PATH here because the pip install below runs in this same shell
+9 -9
View File
@@ -1,5 +1,5 @@
#!/bin/bash #!/bin/bash
# Copy the built PyO3 extension modules into rust-ext-staging/<pkg>/ for # Copy the built PyO3 extension modules into rust-ext-staging/rust_extensions/ for
# upload-artifact. Shared by both jobs of _pr-test-rust-ext-build.yml, so the # upload-artifact. Shared by both jobs of _pr-test-rust-ext-build.yml, so the
# archive layout and the module-count check cannot drift between them. # archive layout and the module-count check cannot drift between them.
# #
@@ -12,24 +12,24 @@ shopt -s nullglob
# module would silently shift the archive layout. # module would silently shift the archive layout.
rm -rf rust-ext-staging rm -rf rust-ext-staging
built=() built=()
# Same suffix set across pkgs, or one ABI's Rust-server tests silently skip. # Same suffix set across modules, or one ABI's Rust-server tests silently skip.
expected_suffixes="" expected_suffixes=""
for pkg in server grpc multimodal; do mkdir -p rust-ext-staging/rust_extensions
found=(python/sglang/srt/"${pkg}"/_core*.so) for module in server grpc multimodal; do
found=(python/sglang/srt/rust_extensions/_"${module}"*.so)
if [ ${#found[@]} -eq 0 ]; then if [ ${#found[@]} -eq 0 ]; then
echo "::error::no extension module found for ${pkg}" echo "::error::no extension module found for ${module}"
exit 1 exit 1
fi fi
suffixes=$(printf '%s\n' "${found[@]##*/_core}" | sort) suffixes=$(printf '%s\n' "${found[@]##*/_${module}}" | sort)
if [ -z "${expected_suffixes}" ]; then if [ -z "${expected_suffixes}" ]; then
expected_suffixes="${suffixes}" expected_suffixes="${suffixes}"
elif [ "${suffixes}" != "${expected_suffixes}" ]; then elif [ "${suffixes}" != "${expected_suffixes}" ]; then
echo "::error::extension modules for ${pkg} do not match server's interpreter set" echo "::error::extension modules for ${module} do not match server's interpreter set"
printf 'have:\n%s\nwant:\n%s\n' "${suffixes}" "${expected_suffixes}" printf 'have:\n%s\nwant:\n%s\n' "${suffixes}" "${expected_suffixes}"
exit 1 exit 1
fi fi
mkdir -p "rust-ext-staging/${pkg}" cp "${found[@]}" rust-ext-staging/rust_extensions/
cp "${found[@]}" "rust-ext-staging/${pkg}/"
built+=("${found[@]}") built+=("${found[@]}")
done done
max_allowed="${MAX_GLIBC:-}" max_allowed="${MAX_GLIBC:-}"
+350
View File
@@ -0,0 +1,350 @@
import importlib.abc
import importlib.machinery
import multiprocessing
import os
import subprocess
import sys
import threading
import time
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from types import ModuleType
from unittest import mock
from sglang.srt.rust_extensions import load_rust_extension
from sglang.srt.rust_extensions import loader as rust_extension
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _hold_filesystem_lock(path: str, ready, release) -> None:
with rust_extension._filesystem_lock(Path(path)):
ready.set()
release.wait(timeout=10)
class _FailingExtensionLoader(importlib.abc.Loader):
def create_module(self, spec):
return None
def exec_module(self, module):
raise RuntimeError("broken extension")
class TestRustExtension(CustomTestCase):
def _workspace(self, root: Path) -> Path:
workspace = root / "rust"
crate = workspace / "demo"
crate.mkdir(parents=True)
(workspace / "Cargo.toml").write_text(
'[workspace]\nmembers = ["demo"]\n', encoding="utf-8"
)
(workspace / "Cargo.lock").write_text(
"# generated lockfile\n", encoding="utf-8"
)
(crate / "Cargo.toml").write_text(
"""
[package]
name = "demo-extension"
version = "0.1.0"
[package.metadata.sglang]
python-module = "demo._core"
features = ["python"]
[lib]
name = "demo_extension"
crate-type = ["cdylib"]
""".strip() + "\n",
encoding="utf-8",
)
(crate / "lib.rs").write_text("fn input() {}\n", encoding="utf-8")
return workspace
def test_bundled_wheel_extension_never_touches_source_or_cargo(self):
bundled = ModuleType("demo._core")
with (
mock.patch.object(
rust_extension.importlib, "import_module", return_value=bundled
),
mock.patch.object(rust_extension, "_discover_crate") as discover,
mock.patch.object(rust_extension, "_build_context") as fingerprint,
mock.patch.object(rust_extension, "_cargo_build") as cargo_build,
):
self.assertIs(
load_rust_extension(
"demo._core", mode="auto", workspace=Path("/workspace/not-present")
),
bundled,
)
discover.assert_not_called()
fingerprint.assert_not_called()
cargo_build.assert_not_called()
def test_discovery_reads_crate_manifest_metadata(self):
with TemporaryDirectory() as directory:
workspace = self._workspace(Path(directory))
crate = rust_extension._discover_crate(workspace, "demo._core")
self.assertEqual(crate.package, "demo-extension")
self.assertEqual(crate.library, "demo_extension")
self.assertEqual(crate.python_module, "demo._core")
self.assertEqual(crate.features, ("python",))
with self.assertRaisesRegex(
ModuleNotFoundError, r"declared modules: \['demo\._core'\]"
):
rust_extension._discover_crate(workspace, "demo._missing")
def test_fingerprint_is_content_based_and_covers_build_inputs(self):
with TemporaryDirectory() as directory:
workspace = self._workspace(Path(directory))
crate = rust_extension._discover_crate(workspace, "demo._core")
with mock.patch.object(
rust_extension,
"_command_version",
side_effect=lambda command, *args, **kwargs: f"{command} 1.0",
):
first = rust_extension._build_context(crate)
source = workspace / "demo" / "lib.rs"
os.utime(source, (1, 1))
self.assertEqual(first, rust_extension._build_context(crate))
source.write_text("fn changed() {}\n", encoding="utf-8")
changed_source = rust_extension._build_context(crate)
self.assertNotEqual(first.fingerprint, changed_source.fingerprint)
with mock.patch.dict(os.environ, {"RUSTFLAGS": "-Ctarget-cpu=native"}):
changed_flags = rust_extension._build_context(crate)
self.assertNotEqual(
changed_source.fingerprint, changed_flags.fingerprint
)
self.assertNotEqual(
changed_source.target_fingerprint,
changed_flags.target_fingerprint,
)
def test_auto_builds_once_then_uses_cache(self):
with TemporaryDirectory() as directory:
root = Path(directory)
workspace = self._workspace(root)
artifact = root / "libdemo_extension.so"
artifact.write_bytes(b"extension")
context = rust_extension._BuildContext("source", "fingerprint", "target")
loaded = ModuleType("demo._core")
with (
mock.patch.object(
rust_extension, "_import_bundled_extension", return_value=None
),
mock.patch.object(
rust_extension, "_build_context", return_value=context
),
mock.patch.object(
rust_extension, "_source_digest", return_value="source"
),
mock.patch.object(
rust_extension, "_cargo_build", return_value=artifact
) as cargo_build,
mock.patch.object(
rust_extension,
"_load_extension_from_path",
return_value=loaded,
),
):
self.assertIs(
rust_extension.load_rust_extension(
"demo._core", workspace=workspace, cache_dir=root / "cache"
),
loaded,
)
self.assertIs(
rust_extension.load_rust_extension(
"demo._core", workspace=workspace, cache_dir=root / "cache"
),
loaded,
)
cargo_build.assert_called_once()
def test_never_rejects_missing_cache_without_building(self):
with TemporaryDirectory() as directory:
root = Path(directory)
workspace = self._workspace(root)
context = rust_extension._BuildContext("source", "fingerprint", "target")
with (
mock.patch.object(
rust_extension, "_import_bundled_extension", return_value=None
),
mock.patch.object(
rust_extension, "_build_context", return_value=context
),
mock.patch.object(rust_extension, "_cargo_build") as cargo_build,
):
with self.assertRaisesRegex(
ModuleNotFoundError, "build mode is 'never'"
):
rust_extension.load_rust_extension(
"demo._core",
mode="never",
workspace=workspace,
cache_dir=root / "cache",
)
cargo_build.assert_not_called()
def test_force_skips_bundled_import_and_rebuilds_cached_artifact(self):
with TemporaryDirectory() as directory:
root = Path(directory)
workspace = self._workspace(root)
crate = rust_extension._discover_crate(workspace, "demo._core")
artifact = root / "libdemo_extension.so"
artifact.write_bytes(b"new extension")
context = rust_extension._BuildContext("source", "fingerprint", "target")
cached = rust_extension._cached_extension_path(
root / "cache", crate, context.fingerprint
)
cached.parent.mkdir(parents=True)
cached.write_bytes(b"old extension")
with (
mock.patch.object(
rust_extension, "_import_bundled_extension"
) as bundled_import,
mock.patch.object(
rust_extension, "_build_context", return_value=context
),
mock.patch.object(
rust_extension, "_source_digest", return_value="source"
),
mock.patch.object(
rust_extension, "_cargo_build", return_value=artifact
) as cargo_build,
mock.patch.object(
rust_extension,
"_load_extension_from_path",
return_value=ModuleType("demo._core"),
),
):
rust_extension.load_rust_extension(
"demo._core",
mode="force",
workspace=workspace,
cache_dir=root / "cache",
)
bundled_import.assert_not_called()
cargo_build.assert_called_once()
self.assertEqual(cached.read_bytes(), b"new extension")
def test_cargo_build_uses_locked_release_and_declared_features(self):
with TemporaryDirectory() as directory:
root = Path(directory)
workspace = self._workspace(root)
crate = rust_extension._discover_crate(workspace, "demo._core")
target_dir = root / "target"
def run(command, *, cwd, env, check):
self.assertTrue(check)
self.assertEqual(cwd, crate.workspace)
self.assertEqual(env["PYO3_PYTHON"], sys.executable)
artifact = Path(env["CARGO_TARGET_DIR"]) / "release"
artifact.mkdir(parents=True)
(artifact / "libdemo_extension.so").write_bytes(b"extension")
return subprocess.CompletedProcess(command, 0)
with mock.patch.object(
rust_extension.subprocess, "run", side_effect=run
) as cargo:
artifact = rust_extension._cargo_build(crate, target_dir)
self.assertEqual(artifact, target_dir / "release/libdemo_extension.so")
self.assertEqual(
cargo.call_args.args[0],
[
"cargo",
"build",
"--release",
"--locked",
"--package",
"demo-extension",
"--features",
"python",
],
)
def test_filesystem_lock_serializes_processes(self):
with TemporaryDirectory() as directory:
lock_path = Path(directory) / "build.lock"
context = multiprocessing.get_context("fork")
ready = context.Event()
release = context.Event()
process = context.Process(
target=_hold_filesystem_lock,
args=(os.fspath(lock_path), ready, release),
)
process.start()
self.assertTrue(ready.wait(timeout=5))
acquired = threading.Event()
def acquire_in_parent():
with rust_extension._filesystem_lock(lock_path):
acquired.set()
thread = threading.Thread(target=acquire_in_parent)
thread.start()
try:
time.sleep(0.1)
self.assertFalse(acquired.is_set())
finally:
release.set()
process.join(timeout=5)
thread.join(timeout=5)
self.assertEqual(process.exitcode, 0)
self.assertTrue(acquired.is_set())
def test_failed_import_does_not_poison_sys_modules(self):
module_name = "demo._broken_core"
module_spec = importlib.machinery.ModuleSpec(
module_name, _FailingExtensionLoader()
)
with mock.patch.object(
rust_extension.importlib.util,
"spec_from_file_location",
return_value=module_spec,
):
with self.assertRaisesRegex(RuntimeError, "broken extension"):
rust_extension._load_extension_from_path(
module_name, Path("/cache/_broken_core.so")
)
self.assertNotIn(module_name, sys.modules)
def test_checked_in_crates_are_discovered_from_wheel_metadata(self):
for python_module, package, library, features in (
(
"sglang.srt.rust_extensions._server",
"sglang-server",
"sglang_server",
(),
),
(
"sglang.srt.rust_extensions._grpc",
"sglang-grpc",
"sglang_grpc_core",
(),
),
(
"sglang.srt.rust_extensions._multimodal",
"sglang-mm",
"sglang_mm_core",
("python", "parallel"),
),
):
crate = rust_extension._discover_crate(
rust_extension._RUST_WORKSPACE, python_module
)
self.assertEqual(crate.package, package)
self.assertEqual(crate.library, library)
self.assertEqual(crate.features, features)
if __name__ == "__main__":
unittest.main()
@@ -21,13 +21,13 @@ register_cpu_ci(
def load_core(): def load_core():
"""The Rust ``_core`` extension, or ``None`` (→ skip) when not built """The Rust ``_multimodal`` extension, or ``None`` (→ skip) when not built
locally. In CI a missing extension is a hard failure, never a silent locally. In CI a missing extension is a hard failure, never a silent
skip — the CPU suite builds it from source.""" skip — the CPU suite builds it from source."""
try: try:
from sglang.srt.multimodal import _core from sglang.srt.rust_extensions import _multimodal
return _core return _multimodal
except ImportError: except ImportError:
if is_in_ci(): if is_in_ci():
raise raise
@@ -2119,8 +2119,6 @@ class TestGrpcServerArgs(CustomTestCase):
arg-parsing tests above never call start_server, so a stray kwarg (e.g. arg-parsing tests above never call start_server, so a stray kwarg (e.g.
the removed max_prefill_tokens) would only surface as a TypeError at the removed max_prefill_tokens) would only surface as a TypeError at
launch. This mocks the native extension and locks the kwarg set.""" launch. This mocks the native extension and locks the kwarg set."""
import sys
from sglang.srt.entrypoints import http_server from sglang.srt.entrypoints import http_server
fake_core = SimpleNamespace(start_server=MagicMock(return_value="handle")) fake_core = SimpleNamespace(start_server=MagicMock(return_value="handle"))
@@ -2128,13 +2126,14 @@ class TestGrpcServerArgs(CustomTestCase):
server_args = SimpleNamespace( server_args = SimpleNamespace(
host="127.0.0.1", grpc_port=50051, grpc_worker_threads=4 host="127.0.0.1", grpc_port=50051, grpc_worker_threads=4
) )
with patch.dict( with (
sys.modules, patch(
{ "sglang.srt.rust_extensions.load_rust_extension",
"sglang.srt.grpc": SimpleNamespace(_core=fake_core), return_value=fake_core,
"sglang.srt.grpc._core": fake_core, ) as load_rust_extension,
"sglang.srt.entrypoints.grpc_bridge": fake_bridge, patch.dict(
}, "sys.modules", {"sglang.srt.entrypoints.grpc_bridge": fake_bridge}
),
): ):
handle = http_server._start_native_grpc_server_for_runtime( handle = http_server._start_native_grpc_server_for_runtime(
server_args=server_args, server_args=server_args,
@@ -2144,6 +2143,7 @@ class TestGrpcServerArgs(CustomTestCase):
) )
self.assertEqual(handle, "handle") self.assertEqual(handle, "handle")
load_rust_extension.assert_called_once_with("sglang.srt.rust_extensions._grpc")
_, kwargs = fake_core.start_server.call_args _, kwargs = fake_core.start_server.call_args
self.assertEqual( self.assertEqual(
set(kwargs), {"host", "port", "runtime_handle", "worker_threads"} set(kwargs), {"host", "port", "runtime_handle", "worker_threads"}
@@ -48,7 +48,7 @@ def solid_image_data_url(fmt):
@unittest.skipIf( @unittest.skipIf(
importlib.util.find_spec("sglang.srt.server._core") is None, importlib.util.find_spec("sglang.srt.rust_extensions._server") is None,
"sglang-server rust extension not installed (e.g. AMD suite)", "sglang-server rust extension not installed (e.g. AMD suite)",
) )
class TestRustServerNativeMm(CustomTestCase): class TestRustServerNativeMm(CustomTestCase):
@@ -98,7 +98,7 @@ class QwenGenerateVisionSampler(SamplerBase):
@unittest.skipIf( @unittest.skipIf(
importlib.util.find_spec("sglang.srt.server._core") is None, importlib.util.find_spec("sglang.srt.rust_extensions._server") is None,
"sglang-server rust extension not installed (e.g. AMD suite)", "sglang-server rust extension not installed (e.g. AMD suite)",
) )
class TestRustNativeMmMMMU(CustomTestCase): class TestRustNativeMmMMMU(CustomTestCase):