Build Rust extensions on demand in source checkouts (#34994)
This commit is contained in:
@@ -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
|
||||||
|
|||||||
@@ -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/
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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.grpc import _core as grpc_native
|
from sglang.srt.rust_extensions import load_rust_extension
|
||||||
except ImportError as e:
|
|
||||||
raise RuntimeError(
|
grpc_native = load_rust_extension("sglang.srt.rust_extensions._grpc")
|
||||||
"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,
|
||||||
|
|||||||
@@ -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 +0,0 @@
|
|||||||
# SGLang gRPC module
|
|
||||||
@@ -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"]
|
||||||
@@ -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
|
||||||
@@ -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):
|
||||||
|
|||||||
Generated
+20
-17
@@ -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",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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"
|
||||||
|
|||||||
@@ -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>()?;
|
||||||
|
|||||||
@@ -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"]
|
||||||
|
|
||||||
|
|||||||
@@ -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`)
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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"
|
||||||
|
|||||||
@@ -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)?;
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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):
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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:-}"
|
||||||
|
|||||||
@@ -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):
|
||||||
|
|||||||
Reference in New Issue
Block a user