From 67e12131dfea1a63256e900d8e76124fa51666dc Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sun, 16 Aug 2026 14:58:06 -0700 Subject: [PATCH] Build Rust extensions on demand in source checkouts (#34994) --- .github/actions/download-rust-ext/action.yml | 2 +- .github/workflows/_pr-test-rust-ext-build.yml | 8 +- .gitignore | 1 + docker/rocm.Dockerfile | 4 +- python/pyproject.toml | 1 + python/setup.py | 5 +- python/sglang/srt/entrypoints/http_server.py | 14 +- python/sglang/srt/environ.py | 3 + python/sglang/srt/grpc/__init__.py | 1 - python/sglang/srt/managers/rust_server.py | 6 +- .../inkling/image_processing_rust.py | 4 +- .../srt/multimodal/processors/inkling.py | 3 +- python/sglang/srt/rust_extensions/__init__.py | 5 + python/sglang/srt/rust_extensions/loader.py | 412 ++++++++++++++++++ python/sglang/test/test_utils.py | 13 +- rust/Cargo.lock | 37 +- rust/sglang-grpc/Cargo.toml | 2 +- rust/sglang-grpc/pyproject.toml | 6 +- rust/sglang-grpc/src/lib.rs | 2 +- rust/sglang-mm/Cargo.toml | 10 +- rust/sglang-mm/README.md | 13 +- rust/sglang-mm/bench/bench_parity.py | 2 +- rust/sglang-mm/pyproject.toml | 2 +- rust/sglang-mm/src/lib.rs | 4 +- rust/sglang-mm/tests/test_golden.py | 2 +- rust/sglang-mm/tests/test_hash_fetch.py | 2 +- rust/sglang-mm/tests/test_resize_parity.py | 4 +- rust/sglang-server/Cargo.toml | 2 +- rust/sglang-server/pyproject.toml | 2 +- .../src/api_server/openai/chat.rs | 2 + rust/sglang-server/src/lib.rs | 2 +- scripts/ci/cuda/ci_install_dependency.sh | 14 +- scripts/ci/musa/musa_install_dependency.sh | 3 +- scripts/ci/utils/stage_rust_ext_modules.sh | 18 +- test/registered/rust/test_rust_extension.py | 350 +++++++++++++++ .../unit/multimodal/rust/_mm_rust_utils.py | 6 +- .../unit/server_args/test_server_args.py | 18 +- .../registered/vlm/test_rust_native_mm_e2e.py | 2 +- .../vlm/test_rust_native_mm_mmmu.py | 2 +- 39 files changed, 880 insertions(+), 109 deletions(-) delete mode 100644 python/sglang/srt/grpc/__init__.py create mode 100644 python/sglang/srt/rust_extensions/__init__.py create mode 100644 python/sglang/srt/rust_extensions/loader.py create mode 100644 test/registered/rust/test_rust_extension.py diff --git a/.github/actions/download-rust-ext/action.yml b/.github/actions/download-rust-ext/action.yml index eb30291bc..78893e1ad 100644 --- a/.github/actions/download-rust-ext/action.yml +++ b/.github/actions/download-rust-ext/action.yml @@ -49,7 +49,7 @@ runs: if: steps.artifact.outcome != 'success' uses: actions/cache/restore@v4 with: - path: python/sglang/srt/*/_core*.so + path: python/sglang/srt/rust_extensions/_*.so key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }} # Job-wide, but only setup.py reads it, and only while building. diff --git a/.github/workflows/_pr-test-rust-ext-build.yml b/.github/workflows/_pr-test-rust-ext-build.yml index 4cd832866..20a9cfae9 100644 --- a/.github/workflows/_pr-test-rust-ext-build.yml +++ b/.github/workflows/_pr-test-rust-ext-build.yml @@ -99,7 +99,7 @@ jobs: id: cache uses: actions/cache/restore@v4 with: - path: python/sglang/srt/*/_core*.so + path: python/sglang/srt/rust_extensions/_*.so key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }} # On a miss: different hash = rust/setup.py moved; no entries = evicted. @@ -112,7 +112,7 @@ jobs: run: | if [ -n "${MATCHED_KEY}" ]; then echo "hit: ${MATCHED_KEY}" - ls -l python/sglang/srt/*/_core*.so + ls -l python/sglang/srt/rust_extensions/_*.so else echo "miss: ${PRIMARY_KEY}" echo "entries under ${KEY_PREFIX}- (created / ref / size / key):" @@ -134,7 +134,7 @@ jobs: uses: actions/upload-artifact@v4 with: name: ${{ inputs.artifact_name }} - # Archive holds /_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/ if-no-files-found: error retention-days: 1 @@ -226,7 +226,7 @@ jobs: - name: Save built modules uses: actions/cache/save@v4 with: - path: python/sglang/srt/*/_core*.so + path: python/sglang/srt/rust_extensions/_*.so key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }} - name: Upload extension modules diff --git a/.gitignore b/.gitignore index 28c1613ab..cd1c6325d 100644 --- a/.gitignore +++ b/.gitignore @@ -230,6 +230,7 @@ work_dirs/ # Rust lib Cargo.lock +!rust/Cargo.lock # Generated vision test fixtures (regenerate with: python scripts/generate_vision_golden.py) sgl-model-gateway/tests/fixtures/golden/ diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 3699788e3..b0123cef9 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -285,8 +285,8 @@ RUN pip install IPython \ && pip install pybind11 # 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 -# sgl-model-gateway. Must precede the sglang install. +# (sglang.srt.rust_extensions._multimodal) during the sglang pip install below +# and later by sgl-model-gateway. Must precede the sglang install. ENV PATH="/root/.cargo/bin:${PATH}" RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ && rustc --version && cargo --version diff --git a/python/pyproject.toml b/python/pyproject.toml index bf1f6602f..a773f66bc 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -77,6 +77,7 @@ dependencies = [ "tilelang==0.1.11", "timm==1.0.16", "tokenspeed_mla==0.1.8", + "tomli ; python_version < '3.11'", "torch==2.13.0", "torch_memory_saver>=0.0.9.post1", "torchaudio==2.11.0", diff --git a/python/setup.py b/python/setup.py index 0f81e61c8..f9676f358 100644 --- a/python/setup.py +++ b/python/setup.py @@ -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: 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. + "sglang.srt.rust_extensions._grpc". It is read directly from os.environ + instead of sglang.srt.environ, which is not importable until the package is + built. """ import json diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 3371fcf0a..59458cc26 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -2687,16 +2687,10 @@ def _start_native_grpc_server_for_runtime( template_manager, scheduler_info, ): - try: - from sglang.srt.entrypoints.grpc_bridge import RuntimeHandle - from sglang.srt.grpc import _core as grpc_native - except ImportError as e: - 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 + from sglang.srt.entrypoints.grpc_bridge import RuntimeHandle + from sglang.srt.rust_extensions import load_rust_extension + + grpc_native = load_rust_extension("sglang.srt.rust_extensions._grpc") runtime_handle = RuntimeHandle( tokenizer_manager=tokenizer_manager, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d0089ebe1..5e8799ce4 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1523,6 +1523,9 @@ class Envs: # Rust server # =================================================================== 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. SGLANG_MAX_BATCH_REQS_PER_HTTP_REQ = EnvInt(4096) diff --git a/python/sglang/srt/grpc/__init__.py b/python/sglang/srt/grpc/__init__.py deleted file mode 100644 index de1d8e32a..000000000 --- a/python/sglang/srt/grpc/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# SGLang gRPC module diff --git a/python/sglang/srt/managers/rust_server.py b/python/sglang/srt/managers/rust_server.py index 3275baeec..9a6c1d07e 100644 --- a/python/sglang/srt/managers/rust_server.py +++ b/python/sglang/srt/managers/rust_server.py @@ -37,7 +37,7 @@ if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig from sglang.srt.managers.io_struct import BatchTokenIDOutput 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 logger = logging.getLogger(__name__) @@ -367,7 +367,9 @@ class RustServer: The caller gates this (``SGLANG_RUST_SERVER`` + rank 0); this always 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. os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") diff --git a/python/sglang/srt/multimodal/inkling/image_processing_rust.py b/python/sglang/srt/multimodal/inkling/image_processing_rust.py index ec573ee76..55a439c48 100644 --- a/python/sglang/srt/multimodal/inkling/image_processing_rust.py +++ b/python/sglang/srt/multimodal/inkling/image_processing_rust.py @@ -10,8 +10,10 @@ from PIL import Image from transformers.image_processing_utils import BaseImageProcessor, BatchFeature 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.rust_extensions import load_rust_extension + +_rs = load_rust_extension("sglang.srt.rust_extensions._multimodal").inkling def _bits_to_bthwc( diff --git a/python/sglang/srt/multimodal/processors/inkling.py b/python/sglang/srt/multimodal/processors/inkling.py index 8750e2912..803e97322 100644 --- a/python/sglang/srt/multimodal/processors/inkling.py +++ b/python/sglang/srt/multimodal/processors/inkling.py @@ -133,7 +133,8 @@ class InklingMultimodalProcessor(SGLangBaseProcessor): logger.info("Using Rust-accelerated Inkling image processor") except ImportError: 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." ) image_processor = InklingImageProcessor(patch_size=patch_size) diff --git a/python/sglang/srt/rust_extensions/__init__.py b/python/sglang/srt/rust_extensions/__init__.py new file mode 100644 index 000000000..564d437bc --- /dev/null +++ b/python/sglang/srt/rust_extensions/__init__.py @@ -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"] diff --git a/python/sglang/srt/rust_extensions/loader.py b/python/sglang/srt/rust_extensions/loader.py new file mode 100644 index 000000000..654576dbf --- /dev/null +++ b/python/sglang/srt/rust_extensions/loader.py @@ -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 diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 1c65776bc..cdd8f8c55 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -219,17 +219,10 @@ def is_rust_server_built(): """Return whether the embedded Rust server extension (``SGLANG_RUST_SERVER``) is importable. - ``sglang/srt/server/`` is not in the source tree — it is produced by - ``setup.py build_rust --inplace``, so on a build without it ``find_spec`` - 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. + The ``sglang.srt.rust_extensions`` Python package is always present; the + private ``_server`` module exists only when the PyO3 extension was built. """ - try: - return importlib.util.find_spec("sglang.srt.server._core") is not None - except ModuleNotFoundError: - return False + return importlib.util.find_spec("sglang.srt.rust_extensions._server") is not None def _use_cached_default_models(model_repo: str): diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 0df4da1aa..93e1a560d 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -982,9 +982,9 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "dynamo-parsers" -version = "7.0.1" +version = "7.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97cce1ec70f4c4896ff9a4e0bb29cd603f465e599c94bf2085b71872a3a58c70" +checksum = "7a7f84812e7ce1d3dace1b8d69853cc6c49759cbd07dfbb59452aa928645b031" dependencies = [ "aho-corasick", "anyhow", @@ -1005,9 +1005,9 @@ dependencies = [ [[package]] name = "dynamo-protocols" -version = "5.1.0" +version = "5.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "696b6df4c817c16dc0a619d43f6af5ecc4f045a6b6be7ef516bfe502f977764b" +checksum = "97fd951c32c033f4f220b6087db91d7f7d475f1bf1f886c8e60798ac1cf4cc9a" dependencies = [ "async-openai", "derive_builder", @@ -1022,9 +1022,9 @@ dependencies = [ [[package]] name = "dynamo-renderer" -version = "5.0.0" +version = "5.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0423da7bec83734e7567b1c8895fd2e7a1a98256253a29fe128032ac01cefac1" +checksum = "dfe04753d666e3462e3eed4abcf2a4d05997cf2dea45858733f2bdea5b61f733" dependencies = [ "anyhow", "chrono", @@ -1040,9 +1040,9 @@ dependencies = [ [[package]] name = "dynamo-tokenizers" -version = "1.7.0" +version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a81f7b7693282f4693a3b4c60174e14260cf011ab493a7ca9b4deb2fa98ecee" +checksum = "821c767e896f1f225b6411a5ce41dba7f114c2c97db862a13fd962195665075e" dependencies = [ "aho-corasick", "anyhow", @@ -1162,21 +1162,24 @@ dependencies = [ [[package]] name = "fastokens" -version = "0.2.1" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8728655e193e0d08d7a95d63cf1fdb9b768d282cab0a112ecb006615bae9f067" +checksum = "35ff8cf975d8f772e7d4c3be439659b0b23b1ae1086ebee7204573807053ff7f" dependencies = [ "daachorse", "fancy-regex 0.17.0", "hf-hub", "icu_normalizer", + "libc", "memchr", "pcre2", "rayon", + "regex-syntax", "serde", "serde_json", "strum", "thiserror", + "ureq", ] [[package]] @@ -1633,7 +1636,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.5", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -2123,9 +2126,9 @@ dependencies = [ [[package]] name = "minijinja" -version = "2.21.0" +version = "2.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" +checksum = "86886cf6dbf4e614b19c9a1eec9775f021869d7eadde0fc73921a81b90c9b4c9" dependencies = [ "indexmap 2.14.0", "memo-map", @@ -2135,9 +2138,9 @@ dependencies = [ [[package]] name = "minijinja-contrib" -version = "2.21.0" +version = "2.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85342f6fac0be8ccd5bd00d9066be538f34f393f577b75d81b17c8398a6b43bb" +checksum = "bd3e5f077bc2379f0f7d911e7cfdd921114ed99fc884533dca502944cb355b11" dependencies = [ "minijinja", "serde", @@ -2940,7 +2943,7 @@ dependencies = [ "quinn-udp", "rustc-hash 2.1.3", "rustls", - "socket2 0.6.5", + "socket2 0.5.10", "thiserror", "tokio", "tracing", @@ -2978,7 +2981,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.5", + "socket2 0.5.10", "tracing", "windows-sys 0.61.2", ] diff --git a/rust/sglang-grpc/Cargo.toml b/rust/sglang-grpc/Cargo.toml index 37decfb4a..d26625c77 100644 --- a/rust/sglang-grpc/Cargo.toml +++ b/rust/sglang-grpc/Cargo.toml @@ -8,7 +8,7 @@ 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" +python-module = "sglang.srt.rust_extensions._grpc" # Always build optimized, even for an editable install. debug = false diff --git a/rust/sglang-grpc/pyproject.toml b/rust/sglang-grpc/pyproject.toml index f226c9e4b..20db0e376 100644 --- a/rust/sglang-grpc/pyproject.toml +++ b/rust/sglang-grpc/pyproject.toml @@ -1,7 +1,7 @@ # 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 +# (python/pyproject.toml, target sglang.srt.rust_extensions._grpc). 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"] @@ -14,4 +14,4 @@ description = "In-process Rust gRPC server for SGLang" requires-python = ">=3.10" [tool.maturin] -module-name = "_core" +module-name = "_grpc" diff --git a/rust/sglang-grpc/src/lib.rs b/rust/sglang-grpc/src/lib.rs index cec1ea721..e9d0cd607 100644 --- a/rust/sglang-grpc/src/lib.rs +++ b/rust/sglang-grpc/src/lib.rs @@ -257,7 +257,7 @@ fn start_server( } #[pymodule] -fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { +fn _grpc(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(start_server, m)?)?; m.add_class::()?; m.add_class::()?; diff --git a/rust/sglang-mm/Cargo.toml b/rust/sglang-mm/Cargo.toml index 16b3ae3e4..ab6ed7339 100644 --- a/rust/sglang-mm/Cargo.toml +++ b/rust/sglang-mm/Cargo.toml @@ -11,17 +11,17 @@ license.workspace = true # features: the wheel build needs the PyO3 bindings, which are NOT default (see # [features] below). [package.metadata.sglang] -python-module = "sglang.srt.multimodal._core" +python-module = "sglang.srt.rust_extensions._multimodal" debug = false features = ["python", "parallel"] [lib] # 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. +# multiple `lib_core` cdylibs). The importable Python module is `_multimodal`: the +# name comes from the `#[pymodule]` entry point and the setuptools-rust target, +# which renames the built artifact. 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. crate-type = ["cdylib", "rlib"] diff --git a/rust/sglang-mm/README.md b/rust/sglang-mm/README.md index b4d528694..a5bd958ec 100644 --- a/rust/sglang-mm/README.md +++ b/rust/sglang-mm/README.md @@ -6,9 +6,9 @@ GIL-released. Built two ways: -- **PyO3 extension** `sglang.srt.multimodal._core` (features `python,parallel`, - requested by the wheel build) via setuptools-rust when installing sglang — - used by Python processors and parity tests. +- **PyO3 extension** `sglang.srt.rust_extensions._multimodal` (features + `python,parallel`, requested by the wheel build) via setuptools-rust when + installing sglang — used by Python processors and parity tests. - **Pure-Rust `rlib`** (default features, i.e. neither) linked by `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 @@ -19,7 +19,7 @@ Built two ways: ``` 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 │ # trait + the carriers (Tensor, TokenLayout, ...) ├── driver.rs # model-independent request driver (fetch → @@ -115,7 +115,7 @@ un-stripped URL to `open()`). ## Python API ```python -from sglang.srt.multimodal._core import common, inkling +from sglang.srt.rust_extensions._multimodal import common, inkling # Common (model-agnostic) 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)?;`. -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`) diff --git a/rust/sglang-mm/bench/bench_parity.py b/rust/sglang-mm/bench/bench_parity.py index 78fab9371..584245992 100644 --- a/rust/sglang-mm/bench/bench_parity.py +++ b/rust/sglang-mm/bench/bench_parity.py @@ -5,7 +5,6 @@ import numpy as np import torch from PIL import Image -from sglang.srt.multimodal._core import inkling as _rs_inkling from sglang.srt.multimodal.inkling.image_processing import ( IMAGE_MEAN, IMAGE_STD, @@ -13,6 +12,7 @@ from sglang.srt.multimodal.inkling.image_processing import ( _encode_image_bytes, _fill_patches_numba, ) +from sglang.srt.rust_extensions._multimodal import inkling as _rs_inkling PS = 40 diff --git a/rust/sglang-mm/pyproject.toml b/rust/sglang-mm/pyproject.toml index 057b611e4..1ef9cdec3 100644 --- a/rust/sglang-mm/pyproject.toml +++ b/rust/sglang-mm/pyproject.toml @@ -9,4 +9,4 @@ description = "Rust-accelerated multimodal preprocessing for SGLang" requires-python = ">=3.10" [tool.maturin] -module-name = "_core" +module-name = "_multimodal" diff --git a/rust/sglang-mm/src/lib.rs b/rust/sglang-mm/src/lib.rs index 1865d879e..de6eb2ecc 100644 --- a/rust/sglang-mm/src/lib.rs +++ b/rust/sglang-mm/src/lib.rs @@ -1,7 +1,7 @@ //! sglang-mm: Rust-accelerated multimodal preprocessing for SGLang. //! //! 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. //! * Pure-Rust `rlib` (`default-features = false`), linked by `sglang-server`'s //! MM worker path — no pyo3 in that dependency graph. @@ -18,7 +18,7 @@ use pyo3::prelude::*; #[cfg(feature = "python")] #[pymodule] -fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { +fn _multimodal(m: &Bound<'_, PyModule>) -> PyResult<()> { common::register(m)?; inkling::register(m)?; qwen_vl::register(m)?; diff --git a/rust/sglang-mm/tests/test_golden.py b/rust/sglang-mm/tests/test_golden.py index 1ec5e6f72..5d4ad2965 100644 --- a/rust/sglang-mm/tests/test_golden.py +++ b/rust/sglang-mm/tests/test_golden.py @@ -4,7 +4,7 @@ import os import numpy as np 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( "INKLING_MM_GOLDEN_DIR", diff --git a/rust/sglang-mm/tests/test_hash_fetch.py b/rust/sglang-mm/tests/test_hash_fetch.py index 76d32d55a..dfa53f7d2 100644 --- a/rust/sglang-mm/tests/test_hash_fetch.py +++ b/rust/sglang-mm/tests/test_hash_fetch.py @@ -15,12 +15,12 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "bench")) from bench_parity import make_photo_like 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.image_processing_rust import ( InklingRustImageProcessor, ) from sglang.srt.multimodal.processors import inkling as prc +from sglang.srt.rust_extensions._multimodal import common as _rs_common def png_bytes(arr): diff --git a/rust/sglang-mm/tests/test_resize_parity.py b/rust/sglang-mm/tests/test_resize_parity.py index 407dede65..727fef31e 100644 --- a/rust/sglang-mm/tests/test_resize_parity.py +++ b/rust/sglang-mm/tests/test_resize_parity.py @@ -6,8 +6,8 @@ import numpy as np import pytest from PIL import Image -from sglang.srt.multimodal._core import common as _rs_common -from sglang.srt.multimodal._core import inkling as _rs_inkling +from sglang.srt.rust_extensions._multimodal import common as _rs_common +from sglang.srt.rust_extensions._multimodal import inkling as _rs_inkling def py_scaled_dims( diff --git a/rust/sglang-server/Cargo.toml b/rust/sglang-server/Cargo.toml index 0a5a227c2..b5ccddb09 100644 --- a/rust/sglang-server/Cargo.toml +++ b/rust/sglang-server/Cargo.toml @@ -7,7 +7,7 @@ license.workspace = true # Consumed by python/setup.py: registers this crate as a PyO3 extension module. [package.metadata.sglang] -python-module = "sglang.srt.server._core" +python-module = "sglang.srt.rust_extensions._server" # Always build optimized, even for an editable install. debug = false diff --git a/rust/sglang-server/pyproject.toml b/rust/sglang-server/pyproject.toml index 1d2c1c424..7e332b509 100644 --- a/rust/sglang-server/pyproject.toml +++ b/rust/sglang-server/pyproject.toml @@ -13,4 +13,4 @@ classifiers = [ dynamic = ["version"] [tool.maturin] -module-name = "_core" +module-name = "_server" diff --git a/rust/sglang-server/src/api_server/openai/chat.rs b/rust/sglang-server/src/api_server/openai/chat.rs index 54ee59911..df3ed1725 100644 --- a/rust/sglang-server/src/api_server/openai/chat.rs +++ b/rust/sglang-server/src/api_server/openai/chat.rs @@ -836,6 +836,7 @@ pub(super) fn chat_logprobs(extras: Option<&ChunkExtras>) -> ChatChoiceLogprobs bytes: Some(token.as_bytes().to_vec()), token, logprob, + token_id: u32::try_from(token_id).ok(), top_logprobs, }); } @@ -1020,6 +1021,7 @@ mod tests { let logprobs = chat_logprobs(Some(&extras)); let token = &logprobs.content.unwrap()[0]; assert_eq!(token.token, "x"); + assert_eq!(token.token_id, Some(7)); assert_eq!(token.top_logprobs.len(), 2); assert_eq!(token.top_logprobs[1].token, "y"); } diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index b2aab2c54..2743ed648 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -292,7 +292,7 @@ static LOG_GUARD: std::sync::OnceLock) -> PyResult<()> { +fn _server(m: &Bound<'_, PyModule>) -> PyResult<()> { // Initialize tracing once; ignore if already set by the host process. // Non-blocking writer: emitting threads (axum workers, egress, detok) only // enqueue; a dedicated thread does the stdout formatting-flush + syscall. diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index 2ff7279e9..3642bf85d 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -460,7 +460,7 @@ require_prebuilt_rust_exts() { return 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 # 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 @@ -469,13 +469,13 @@ require_prebuilt_rust_exts() { local suffix suffix=$(python3 -c 'import sysconfig; print(sysconfig.get_config_var("EXT_SUFFIX"))') local missing=() - local pkg - for pkg in server grpc multimodal; do - [ -f "python/sglang/srt/${pkg}/_core${suffix}" ] || missing+=("${pkg}") + local module + for module in server grpc multimodal; do + [ -f "python/sglang/srt/rust_extensions/_${module}${suffix}" ] || missing+=("${module}") done if [ ${#missing[@]} -gt 0 ]; then - echo "::warning::no prebuilt _core${suffix} for: ${missing[*]}; building from source" - ls -l python/sglang/srt/*/_core*.so 2>/dev/null || echo "(no extension modules at all)" + echo "::warning::no prebuilt Rust extension ${suffix} for: ${missing[*]}; building from source" + ls -l python/sglang/srt/rust_extensions/_*.so 2>/dev/null || echo "(no extension modules at all)" export SGLANG_BUILD_RUST_EXTS= mark_step_done "${FUNCNAME[0]}" 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. import importlib for mod in ("server", "grpc", "multimodal"): - name = f"sglang.srt.{mod}._core" + name = f"sglang.srt.rust_extensions._{mod}" try: importlib.import_module(name) except Exception as exc: diff --git a/scripts/ci/musa/musa_install_dependency.sh b/scripts/ci/musa/musa_install_dependency.sh index 0dfe942bf..1056cc9cc 100755 --- a/scripts/ci/musa/musa_install_dependency.sh +++ b/scripts/ci/musa/musa_install_dependency.sh @@ -74,7 +74,8 @@ else 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 # 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 diff --git a/scripts/ci/utils/stage_rust_ext_modules.sh b/scripts/ci/utils/stage_rust_ext_modules.sh index f0dd3a4c6..63110075f 100755 --- a/scripts/ci/utils/stage_rust_ext_modules.sh +++ b/scripts/ci/utils/stage_rust_ext_modules.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copy the built PyO3 extension modules into rust-ext-staging// 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 # 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. rm -rf rust-ext-staging 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="" -for pkg in server grpc multimodal; do - found=(python/sglang/srt/"${pkg}"/_core*.so) +mkdir -p rust-ext-staging/rust_extensions +for module in server grpc multimodal; do + found=(python/sglang/srt/rust_extensions/_"${module}"*.so) if [ ${#found[@]} -eq 0 ]; then - echo "::error::no extension module found for ${pkg}" + echo "::error::no extension module found for ${module}" exit 1 fi - suffixes=$(printf '%s\n' "${found[@]##*/_core}" | sort) + suffixes=$(printf '%s\n' "${found[@]##*/_${module}}" | sort) if [ -z "${expected_suffixes}" ]; then expected_suffixes="${suffixes}" 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}" exit 1 fi - mkdir -p "rust-ext-staging/${pkg}" - cp "${found[@]}" "rust-ext-staging/${pkg}/" + cp "${found[@]}" rust-ext-staging/rust_extensions/ built+=("${found[@]}") done max_allowed="${MAX_GLIBC:-}" diff --git a/test/registered/rust/test_rust_extension.py b/test/registered/rust/test_rust_extension.py new file mode 100644 index 000000000..dc3fe94dd --- /dev/null +++ b/test/registered/rust/test_rust_extension.py @@ -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() diff --git a/test/registered/unit/multimodal/rust/_mm_rust_utils.py b/test/registered/unit/multimodal/rust/_mm_rust_utils.py index 7fe1c9be3..0496eaba4 100644 --- a/test/registered/unit/multimodal/rust/_mm_rust_utils.py +++ b/test/registered/unit/multimodal/rust/_mm_rust_utils.py @@ -21,13 +21,13 @@ register_cpu_ci( 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 skip — the CPU suite builds it from source.""" try: - from sglang.srt.multimodal import _core + from sglang.srt.rust_extensions import _multimodal - return _core + return _multimodal except ImportError: if is_in_ci(): raise diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 2acef3594..84d472924 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -2119,8 +2119,6 @@ class TestGrpcServerArgs(CustomTestCase): 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 launch. This mocks the native extension and locks the kwarg set.""" - import sys - from sglang.srt.entrypoints import http_server fake_core = SimpleNamespace(start_server=MagicMock(return_value="handle")) @@ -2128,13 +2126,14 @@ class TestGrpcServerArgs(CustomTestCase): server_args = SimpleNamespace( host="127.0.0.1", grpc_port=50051, grpc_worker_threads=4 ) - with patch.dict( - sys.modules, - { - "sglang.srt.grpc": SimpleNamespace(_core=fake_core), - "sglang.srt.grpc._core": fake_core, - "sglang.srt.entrypoints.grpc_bridge": fake_bridge, - }, + with ( + patch( + "sglang.srt.rust_extensions.load_rust_extension", + return_value=fake_core, + ) as load_rust_extension, + patch.dict( + "sys.modules", {"sglang.srt.entrypoints.grpc_bridge": fake_bridge} + ), ): handle = http_server._start_native_grpc_server_for_runtime( server_args=server_args, @@ -2144,6 +2143,7 @@ class TestGrpcServerArgs(CustomTestCase): ) self.assertEqual(handle, "handle") + load_rust_extension.assert_called_once_with("sglang.srt.rust_extensions._grpc") _, kwargs = fake_core.start_server.call_args self.assertEqual( set(kwargs), {"host", "port", "runtime_handle", "worker_threads"} diff --git a/test/registered/vlm/test_rust_native_mm_e2e.py b/test/registered/vlm/test_rust_native_mm_e2e.py index 412bc8953..f321fdce0 100644 --- a/test/registered/vlm/test_rust_native_mm_e2e.py +++ b/test/registered/vlm/test_rust_native_mm_e2e.py @@ -48,7 +48,7 @@ def solid_image_data_url(fmt): @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)", ) class TestRustServerNativeMm(CustomTestCase): diff --git a/test/registered/vlm/test_rust_native_mm_mmmu.py b/test/registered/vlm/test_rust_native_mm_mmmu.py index 1f5f72eb2..04e18047a 100644 --- a/test/registered/vlm/test_rust_native_mm_mmmu.py +++ b/test/registered/vlm/test_rust_native_mm_mmmu.py @@ -98,7 +98,7 @@ class QwenGenerateVisionSampler(SamplerBase): @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)", ) class TestRustNativeMmMMMU(CustomTestCase):