[mm] rust-server: native multimodal processing for Qwen VL (integrate sglang-mm, e2e) (#32365)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
co-authored by
Claude Fable 5
Cursor
parent
dea07b348b
commit
32e5d788bd
@@ -2260,7 +2260,10 @@ def unwrap_from_pickle(obj: Optional[object]) -> Optional[object]:
|
|||||||
return None
|
return None
|
||||||
if _USE_PICKLE_IPC:
|
if _USE_PICKLE_IPC:
|
||||||
return obj
|
return obj
|
||||||
assert isinstance(obj, PickleWrapper)
|
if not isinstance(obj, PickleWrapper):
|
||||||
|
# Already materialized: the embedded Rust server attaches in-process
|
||||||
|
# objects (native-MM `mm_inputs`) without a pickle hop.
|
||||||
|
return obj
|
||||||
return pickle.loads(obj.data)
|
return pickle.loads(obj.data)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -1336,11 +1336,11 @@ def _get_is_default_transport():
|
|||||||
global _is_default_tensor_transport
|
global _is_default_tensor_transport
|
||||||
if _is_default_tensor_transport is None:
|
if _is_default_tensor_transport is None:
|
||||||
from sglang.srt.managers.tokenizer_manager import (
|
from sglang.srt.managers.tokenizer_manager import (
|
||||||
_determine_tensor_transport_mode,
|
determine_tensor_transport_mode,
|
||||||
)
|
)
|
||||||
|
|
||||||
_is_default_tensor_transport = (
|
_is_default_tensor_transport = (
|
||||||
_determine_tensor_transport_mode(get_server_args()) == "default"
|
determine_tensor_transport_mode(get_server_args()) == "default"
|
||||||
)
|
)
|
||||||
return _is_default_tensor_transport
|
return _is_default_tensor_transport
|
||||||
|
|
||||||
|
|||||||
@@ -41,14 +41,9 @@ def import_processors(package_name: str, overwrite: bool = False):
|
|||||||
PROCESSOR_MAPPING[arch] = cls
|
PROCESSOR_MAPPING[arch] = cls
|
||||||
|
|
||||||
|
|
||||||
def get_mm_processor(
|
def get_mm_processor_cls(hf_config, server_args: ServerArgs, model_config=None):
|
||||||
hf_config,
|
"""The class :func:`get_mm_processor` would instantiate, or ``None`` when the
|
||||||
server_args: ServerArgs,
|
architecture has no registered processor."""
|
||||||
processor,
|
|
||||||
transport_mode,
|
|
||||||
model_config=None,
|
|
||||||
**kwargs,
|
|
||||||
) -> BaseMultimodalProcessor:
|
|
||||||
model_impl = str(getattr(server_args, "model_impl", "auto")).lower()
|
model_impl = str(getattr(server_args, "model_impl", "auto")).lower()
|
||||||
uses_transformers_backend = model_impl == "transformers"
|
uses_transformers_backend = model_impl == "transformers"
|
||||||
if model_impl == "auto" and model_config is not None:
|
if model_impl == "auto" and model_config is not None:
|
||||||
@@ -64,20 +59,30 @@ def get_mm_processor(
|
|||||||
if not uses_transformers_backend or getattr(
|
if not uses_transformers_backend or getattr(
|
||||||
processor_cls, "supports_transformers_backend", False
|
processor_cls, "supports_transformers_backend", False
|
||||||
):
|
):
|
||||||
return processor_cls(
|
return processor_cls
|
||||||
hf_config, server_args, processor, transport_mode, **kwargs
|
|
||||||
)
|
|
||||||
|
|
||||||
if uses_transformers_backend:
|
if uses_transformers_backend:
|
||||||
from sglang.srt.multimodal.processors.transformers_auto import (
|
from sglang.srt.multimodal.processors.transformers_auto import (
|
||||||
TransformersAutoMultimodalProcessor,
|
TransformersAutoMultimodalProcessor,
|
||||||
)
|
)
|
||||||
|
|
||||||
return TransformersAutoMultimodalProcessor(
|
return TransformersAutoMultimodalProcessor
|
||||||
hf_config, server_args, processor, transport_mode, **kwargs
|
|
||||||
)
|
|
||||||
|
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def get_mm_processor(
|
||||||
|
hf_config,
|
||||||
|
server_args: ServerArgs,
|
||||||
|
processor,
|
||||||
|
transport_mode,
|
||||||
|
model_config=None,
|
||||||
|
**kwargs,
|
||||||
|
) -> BaseMultimodalProcessor:
|
||||||
|
processor_cls = get_mm_processor_cls(hf_config, server_args, model_config)
|
||||||
|
if processor_cls is None:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"No processor registered for architecture: {hf_config.architectures}.\n"
|
f"No processor registered for architecture: {hf_config.architectures}.\n"
|
||||||
f"Registered architectures: {[model_cls.__name__ for model_cls in PROCESSOR_MAPPING.keys()]}"
|
f"Registered architectures: {[model_cls.__name__ for model_cls in PROCESSOR_MAPPING.keys()]}"
|
||||||
)
|
)
|
||||||
|
return processor_cls(hf_config, server_args, processor, transport_mode, **kwargs)
|
||||||
|
|||||||
@@ -10,14 +10,17 @@ scheduler holds an `Optional[RustServer]` and delegates to it.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import importlib
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
from array import array
|
from array import array
|
||||||
from itertools import chain
|
from itertools import chain
|
||||||
from typing import TYPE_CHECKING, Any, List, Optional, Tuple
|
from typing import TYPE_CHECKING, Any, Dict, FrozenSet, List, Optional, Tuple
|
||||||
|
|
||||||
import msgspec
|
import msgspec
|
||||||
|
|
||||||
|
from sglang.srt.environ import envs
|
||||||
|
from sglang.srt.managers.io_struct import TokenizedGenerateReqInput
|
||||||
from sglang.srt.managers.utils import (
|
from sglang.srt.managers.utils import (
|
||||||
MsgpackDecodeError,
|
MsgpackDecodeError,
|
||||||
compute_num_reserved_tokens,
|
compute_num_reserved_tokens,
|
||||||
@@ -31,13 +34,315 @@ from sglang.srt.utils.flatten import (
|
|||||||
from sglang.version import __version__
|
from sglang.version import __version__
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
|
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.server._core import Server
|
||||||
|
from sglang.srt.server_args import ServerArgs
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class NativeMmSpec(msgspec.Struct, frozen=True, kw_only=True):
|
||||||
|
"""Resolved parameters of the native Rust MM pipeline for one model,
|
||||||
|
consumed by the Rust worker pool (:meth:`rust_json`) and the drain
|
||||||
|
adapter (:meth:`NativeMmHost.build_native_mm`)."""
|
||||||
|
|
||||||
|
family: str
|
||||||
|
feature_shm: bool
|
||||||
|
image_token_id: int
|
||||||
|
patch_size: int
|
||||||
|
merge_size: int
|
||||||
|
temporal_patch_size: int
|
||||||
|
min_pixels: int
|
||||||
|
max_pixels: int
|
||||||
|
image_mean: Tuple[float, ...]
|
||||||
|
image_std: Tuple[float, ...]
|
||||||
|
# Which HF processor the Rust resize must reproduce bit-exactly, from
|
||||||
|
# `NativeMmHost.NATIVE_IMAGE_PROCESSORS`.
|
||||||
|
resample: str
|
||||||
|
vision_start_token_id: Optional[int]
|
||||||
|
vision_end_token_id: Optional[int]
|
||||||
|
video_token_id: Optional[int]
|
||||||
|
|
||||||
|
# Used by the drain adapter only; every other field goes to Rust.
|
||||||
|
DRAIN_ONLY = ("vision_start_token_id", "vision_end_token_id", "video_token_id")
|
||||||
|
|
||||||
|
@property
|
||||||
|
def feature_dim(self) -> int:
|
||||||
|
return 3 * self.temporal_patch_size * self.patch_size * self.patch_size
|
||||||
|
|
||||||
|
def rust_json(self) -> str:
|
||||||
|
"""The subset `sglang_mm::registry::pipeline_from_spec` parses."""
|
||||||
|
fields = (f for f in self.__struct_fields__ if f not in self.DRAIN_ONLY)
|
||||||
|
return msgspec.json.encode({f: getattr(self, f) for f in fields}).decode()
|
||||||
|
|
||||||
|
|
||||||
|
class NativeMmFamily(msgspec.Struct, frozen=True, kw_only=True):
|
||||||
|
"""The Python half of one Rust MM family (an arm of
|
||||||
|
`sglang_mm::registry::pipeline_from_spec`): which models it serves.
|
||||||
|
Supporting a new model family = one entry in :data:`NATIVE_MM_FAMILIES`
|
||||||
|
plus its Rust arm — the launch gate is data-driven."""
|
||||||
|
|
||||||
|
name: str
|
||||||
|
# The registered Python mm-processor the native pipeline replaces, as
|
||||||
|
# "module:Class". Compared by identity, so an
|
||||||
|
# SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE override still disables the native path.
|
||||||
|
mm_processor: str
|
||||||
|
# Model types whose image-only M-RoPE matches the family's fast path.
|
||||||
|
model_types: FrozenSet[str]
|
||||||
|
# HF image processors the native resize reproduces bit-exactly, each mapped
|
||||||
|
# to the `resample` the Rust pipeline must use (see `NativeMmSpec.resample`).
|
||||||
|
image_processors: Dict[str, str]
|
||||||
|
|
||||||
|
def serves(self, mm_processor_cls: Any, model_type: Optional[str]) -> bool:
|
||||||
|
module_name, _, class_name = self.mm_processor.partition(":")
|
||||||
|
cls = getattr(importlib.import_module(module_name), class_name)
|
||||||
|
return mm_processor_cls is cls and model_type in self.model_types
|
||||||
|
|
||||||
|
|
||||||
|
NATIVE_MM_FAMILIES: Tuple[NativeMmFamily, ...] = (
|
||||||
|
NativeMmFamily(
|
||||||
|
name="qwen_vl",
|
||||||
|
mm_processor="sglang.srt.multimodal.processors.qwen_vl:QwenVLImageProcessor",
|
||||||
|
model_types=frozenset(
|
||||||
|
(
|
||||||
|
"qwen2_vl",
|
||||||
|
"qwen2_5_vl",
|
||||||
|
"qwen3_vl",
|
||||||
|
"qwen3_vl_moe",
|
||||||
|
"qwen3_5",
|
||||||
|
"qwen3_5_moe",
|
||||||
|
)
|
||||||
|
),
|
||||||
|
image_processors={
|
||||||
|
"Qwen2VLImageProcessor": "aten_u8",
|
||||||
|
"Qwen2VLImageProcessorFast": "aten_u8",
|
||||||
|
"Qwen2VLImageProcessorPil": "pil",
|
||||||
|
},
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def native_mm_family_for(
|
||||||
|
mm_processor_cls: Any, model_type: Optional[str]
|
||||||
|
) -> Optional[NativeMmFamily]:
|
||||||
|
"""The declared family serving this model, or ``None`` — which
|
||||||
|
:meth:`RustServer.launch` turns into a hard error (no Python fallback)."""
|
||||||
|
return next(
|
||||||
|
(f for f in NATIVE_MM_FAMILIES if f.serves(mm_processor_cls, model_type)), None
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class NativeMmHost:
|
||||||
|
"""Builds and validates the native Rust MM pipeline for one model.
|
||||||
|
|
||||||
|
Construction registers the same ``mm_processor`` mapping the Python
|
||||||
|
TokenizerManager would build — not to process requests (the Rust worker pool
|
||||||
|
does that, GIL-free) but as the source of truth
|
||||||
|
:meth:`resolve_native_spec` resolves the pipeline parameters from. At drain
|
||||||
|
time :meth:`build_native_mm` wraps the Rust-produced buffers into the
|
||||||
|
scheduler's ``MultimodalProcessorOutput``.
|
||||||
|
|
||||||
|
There is no Python fallback: a model without a native spec fails at launch,
|
||||||
|
and inputs outside the pipeline's scope are rejected per request.
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Rust mm-worker threads when --mm-processor-worker-num is 0. They are
|
||||||
|
# GIL-free, so unlike the Python processor pool more than one always helps.
|
||||||
|
AUTO_MM_WORKERS = 8
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
server_args: ServerArgs,
|
||||||
|
model_config: ModelConfig,
|
||||||
|
processor: Any = None,
|
||||||
|
):
|
||||||
|
# Lazy: this class exists only for multimodal models under
|
||||||
|
# SGLANG_RUST_SERVER.
|
||||||
|
from sglang.srt.managers.multimodal_processor import import_processors
|
||||||
|
from sglang.srt.managers.tokenizer_manager import get_processor_wrapper
|
||||||
|
|
||||||
|
self.server_args = server_args
|
||||||
|
self.model_config = model_config
|
||||||
|
# Worker threads == max concurrently-processed mm requests.
|
||||||
|
self.mm_workers = server_args.mm_processor_worker_num or self.AUTO_MM_WORKERS
|
||||||
|
|
||||||
|
# The mapping the Python TokenizerManager builds in
|
||||||
|
# init_tokenizer_and_processor. The caller's already-loaded HF
|
||||||
|
# AutoProcessor is reused when available (identical construction args).
|
||||||
|
import_processors("sglang.srt.multimodal.processors")
|
||||||
|
if mm_process_pkg := envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.get():
|
||||||
|
import_processors(mm_process_pkg, overwrite=True)
|
||||||
|
self._processor = processor or get_processor_wrapper(self.server_args)
|
||||||
|
|
||||||
|
def resolve_native_spec(self) -> Optional[NativeMmSpec]:
|
||||||
|
"""The :class:`NativeMmSpec` for this model, or ``None`` when it has no
|
||||||
|
native pipeline (the launch gate turns that into a hard error).
|
||||||
|
|
||||||
|
Carries only resolved settings — patch geometry, pixel limits,
|
||||||
|
normalization, token ids — never the HF config, and is conservative by
|
||||||
|
design: an unrecognized knob disables the native path rather than being
|
||||||
|
approximated."""
|
||||||
|
from sglang.srt.managers.multimodal_processor import get_mm_processor_cls
|
||||||
|
|
||||||
|
hf_config = self.model_config.hf_config
|
||||||
|
mm_processor_cls = get_mm_processor_cls(
|
||||||
|
hf_config, self.server_args, model_config=self.model_config
|
||||||
|
)
|
||||||
|
family = native_mm_family_for(
|
||||||
|
mm_processor_cls, getattr(hf_config, "model_type", None)
|
||||||
|
)
|
||||||
|
if family is None:
|
||||||
|
return None
|
||||||
|
ip = getattr(self._processor, "image_processor", None)
|
||||||
|
resample = family.image_processors.get(type(ip).__name__)
|
||||||
|
if resample is None:
|
||||||
|
return None
|
||||||
|
# The native pipeline always resizes, rescales by 1/255 and normalizes;
|
||||||
|
# Rust's fused normalize constants assume that factor. Anything else
|
||||||
|
# would silently produce different features.
|
||||||
|
stages = ("do_resize", "do_rescale", "do_normalize")
|
||||||
|
if not all(getattr(ip, stage, True) for stage in stages):
|
||||||
|
return None
|
||||||
|
if getattr(ip, "rescale_factor", None) != 1 / 255:
|
||||||
|
return None
|
||||||
|
|
||||||
|
# `--mm-process-config {"image": {...}}`: only pixel-limit overrides are
|
||||||
|
# mirrored natively, anything else disables the pipeline.
|
||||||
|
image_overrides = dict(
|
||||||
|
(self.server_args.mm_process_config or {}).get("image", {})
|
||||||
|
)
|
||||||
|
if not set(image_overrides) <= {"min_pixels", "max_pixels"}:
|
||||||
|
return None
|
||||||
|
|
||||||
|
size = getattr(ip, "size", None) or {}
|
||||||
|
min_pixels = image_overrides.get(
|
||||||
|
"min_pixels", getattr(ip, "min_pixels", None) or size.get("shortest_edge")
|
||||||
|
)
|
||||||
|
max_pixels = image_overrides.get(
|
||||||
|
"max_pixels", getattr(ip, "max_pixels", None) or size.get("longest_edge")
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
spec = NativeMmSpec(
|
||||||
|
family=family.name,
|
||||||
|
feature_shm=self._use_feature_shm(),
|
||||||
|
image_token_id=hf_config.image_token_id,
|
||||||
|
patch_size=ip.patch_size,
|
||||||
|
merge_size=ip.merge_size,
|
||||||
|
temporal_patch_size=ip.temporal_patch_size,
|
||||||
|
min_pixels=int(min_pixels),
|
||||||
|
max_pixels=int(max_pixels),
|
||||||
|
image_mean=tuple(float(x) for x in ip.image_mean),
|
||||||
|
image_std=tuple(float(x) for x in ip.image_std),
|
||||||
|
resample=resample,
|
||||||
|
vision_start_token_id=getattr(hf_config, "vision_start_token_id", None),
|
||||||
|
vision_end_token_id=getattr(hf_config, "vision_end_token_id", None),
|
||||||
|
video_token_id=getattr(hf_config, "video_token_id", None),
|
||||||
|
)
|
||||||
|
except (AttributeError, TypeError): # missing/odd processor attrs
|
||||||
|
return None
|
||||||
|
logger.info("rust server: native MM pipeline enabled (family=%s)", family.name)
|
||||||
|
return spec
|
||||||
|
|
||||||
|
def _use_feature_shm(self) -> bool:
|
||||||
|
"""Whether to park feature buffers in POSIX shm rather than inline.
|
||||||
|
|
||||||
|
On exactly when the drained request is broadcast across TP ranks *and*
|
||||||
|
the receiver's ``unwrap_shm_features`` will materialize the stubs (its
|
||||||
|
gates: non-default tensor transport, no ``skip_tokenizer_init``).
|
||||||
|
|
||||||
|
Inline, the whole ~20 MB/image buffer rides ``broadcast_pyobj`` serially
|
||||||
|
on the scheduler loop, so ranks 1..n start the TP-sharded ViT ~30 ms
|
||||||
|
after rank 0 and every rank then stalls that long at the first
|
||||||
|
collective. With shm the broadcast carries a ~100-byte stub and all ranks
|
||||||
|
map in parallel — the transport the Python TokenizerManager already uses.
|
||||||
|
Single-rank serving stays inline, where shm would only add a copy.
|
||||||
|
"""
|
||||||
|
from sglang.srt.managers.tokenizer_manager import (
|
||||||
|
determine_tensor_transport_mode,
|
||||||
|
)
|
||||||
|
|
||||||
|
return (
|
||||||
|
self.server_args.tp_size > 1
|
||||||
|
and determine_tensor_transport_mode(self.server_args) != "default"
|
||||||
|
and not self.server_args.skip_tokenizer_init
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def build_native_mm(spec: NativeMmSpec, entry):
|
||||||
|
"""Drain-time adapter: wrap the Rust-produced buffers of one ``MmHandoff``
|
||||||
|
into the scheduler's ``MultimodalProcessorOutput``. Wrapping only — load,
|
||||||
|
resize, patchify, token expansion and M-RoPE all ran in Rust.
|
||||||
|
|
||||||
|
Runs on the scheduler loop, so it must stay copy-free *and* hash-free:
|
||||||
|
``take_mm``'s numpy arrays own the Rust buffers, ``torch.from_numpy`` just
|
||||||
|
views them, and each item's ``hash`` is worker-precomputed so
|
||||||
|
``set_pad_value`` skips ``hash_feature``. Any per-byte work here — memcpy,
|
||||||
|
sha256, tens of MB per image-heavy request — measurably inflates every
|
||||||
|
running request's inter-token latency."""
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.managers.mm_utils import ShmPointerMMData
|
||||||
|
from sglang.srt.managers.schedule_batch import (
|
||||||
|
Modality,
|
||||||
|
MultimodalDataItem,
|
||||||
|
MultimodalProcessorOutput,
|
||||||
|
)
|
||||||
|
|
||||||
|
shm_names = entry.shm_names
|
||||||
|
if shm_names is None:
|
||||||
|
features = torch.from_numpy(entry.features.reshape(-1, spec.feature_dim))
|
||||||
|
items = []
|
||||||
|
row = 0
|
||||||
|
for index, ((t, h, w), item_hash, offset) in enumerate(
|
||||||
|
zip(entry.grids, entry.hashes, entry.offsets)
|
||||||
|
):
|
||||||
|
n = t * h * w
|
||||||
|
if shm_names is None:
|
||||||
|
feature = features[row : row + n]
|
||||||
|
else:
|
||||||
|
# The worker parked this item's buffer in a named POSIX
|
||||||
|
# segment (see `_use_feature_shm`). Build the stub in its
|
||||||
|
# post-`__setstate__` form: rank 0 never pickle-roundtrips its
|
||||||
|
# own copy, and `materialize()` needs the mapped view.
|
||||||
|
# Ownership of the unlink moved here with `take_mm`.
|
||||||
|
feature = ShmPointerMMData.__new__(ShmPointerMMData)
|
||||||
|
feature.__setstate__(
|
||||||
|
{
|
||||||
|
"shm_name": shm_names[index],
|
||||||
|
"shape": (n, spec.feature_dim),
|
||||||
|
"dtype": torch.float32,
|
||||||
|
"precomputed_hash": item_hash,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
items.append(
|
||||||
|
MultimodalDataItem(
|
||||||
|
modality=Modality.IMAGE,
|
||||||
|
feature=feature,
|
||||||
|
hash=item_hash,
|
||||||
|
offsets=[tuple(offset)],
|
||||||
|
model_specific_data={
|
||||||
|
"image_grid_thw": torch.tensor([[t, h, w]], dtype=torch.long)
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
row += n
|
||||||
|
if envs.SGLANG_MM_PRECOMPUTE_HASH.get():
|
||||||
|
for item in items:
|
||||||
|
item.set_pad_value()
|
||||||
|
return MultimodalProcessorOutput(
|
||||||
|
mm_items=items,
|
||||||
|
im_token_id=spec.image_token_id,
|
||||||
|
im_start_id=spec.vision_start_token_id,
|
||||||
|
im_end_id=spec.vision_end_token_id,
|
||||||
|
video_token_id=spec.video_token_id,
|
||||||
|
mrope_positions=torch.from_numpy(entry.mrope.reshape(3, -1)),
|
||||||
|
mrope_position_delta=torch.tensor([[entry.mrope_delta]], dtype=torch.long),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class RustServer:
|
class RustServer:
|
||||||
"""Owns the embedded multi-threaded Rust server (``sglang_server.Server``).
|
"""Owns the embedded multi-threaded Rust server (``sglang_server.Server``).
|
||||||
|
|
||||||
@@ -45,8 +350,14 @@ class RustServer:
|
|||||||
all implemented as Rust threads in scheduler process.
|
all implemented as Rust threads in scheduler process.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, server: Server, max_per_poll: int = 256):
|
def __init__(
|
||||||
|
self,
|
||||||
|
server: Server,
|
||||||
|
mm_spec: Optional[NativeMmSpec] = None,
|
||||||
|
max_per_poll: int = 256,
|
||||||
|
):
|
||||||
self.server = server
|
self.server = server
|
||||||
|
self.mm_spec = mm_spec
|
||||||
self._max_per_poll = max_per_poll
|
self._max_per_poll = max_per_poll
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -83,7 +394,13 @@ class RustServer:
|
|||||||
if dp_rank is not None:
|
if dp_rank is not None:
|
||||||
http_addr = f"{server_args.host}:{server_args.port + dp_rank}"
|
http_addr = f"{server_args.host}:{server_args.port + dp_rank}"
|
||||||
|
|
||||||
launch_cores, server_cores = cls._partition_cores()
|
launch_cores, server_cores = cls._partition_cores(
|
||||||
|
mm_workers=(
|
||||||
|
(server_args.mm_processor_worker_num or NativeMmHost.AUTO_MM_WORKERS)
|
||||||
|
if scheduler.model_config.is_multimodal
|
||||||
|
else 0
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
server = Server(
|
server = Server(
|
||||||
# None -> run unpinned; the list carries the pinning decision.
|
# None -> run unpinned; the list carries the pinning decision.
|
||||||
@@ -92,6 +409,40 @@ class RustServer:
|
|||||||
server_args_json=cls._build_server_args(scheduler),
|
server_args_json=cls._build_server_args(scheduler),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Multimodal models must have a native Rust pipeline — there is no Python
|
||||||
|
# fallback.
|
||||||
|
mm_spec = None
|
||||||
|
if scheduler.model_config.is_multimodal:
|
||||||
|
# New threads inherit the spawning thread's affinity, and this launch
|
||||||
|
# thread still holds the full mask. Narrow it first so every MM thread
|
||||||
|
# created below (the processor's executors, the Rust MM workers) stays
|
||||||
|
# off the scheduler's reserved cores, where MM preprocessing would
|
||||||
|
# preempt the scheduler loop and inflate inter-token latency.
|
||||||
|
if server_cores is not None:
|
||||||
|
try:
|
||||||
|
os.sched_setaffinity(0, set(server_cores))
|
||||||
|
except OSError as e:
|
||||||
|
logger.warning(
|
||||||
|
"rust server: cannot confine mm threads to server cores: %s", e
|
||||||
|
)
|
||||||
|
mm_host = NativeMmHost(
|
||||||
|
server_args=server_args,
|
||||||
|
model_config=scheduler.model_config,
|
||||||
|
processor=scheduler.processor,
|
||||||
|
)
|
||||||
|
mm_spec = mm_host.resolve_native_spec()
|
||||||
|
if mm_spec is None:
|
||||||
|
supported = sorted(
|
||||||
|
set(chain.from_iterable(f.model_types for f in NATIVE_MM_FAMILIES))
|
||||||
|
)
|
||||||
|
raise RuntimeError(
|
||||||
|
"SGLANG_RUST_SERVER=1: no native Rust MM pipeline for "
|
||||||
|
f"model_type={scheduler.model_config.hf_config.model_type!r} "
|
||||||
|
f"(supported: {', '.join(supported)}; "
|
||||||
|
"images only). Unset SGLANG_RUST_SERVER to serve this model."
|
||||||
|
)
|
||||||
|
server.start_mm_workers(mm_spec.rust_json(), mm_host.mm_workers)
|
||||||
|
|
||||||
# Narrow the scheduler thread only after the server threads are launched.
|
# Narrow the scheduler thread only after the server threads are launched.
|
||||||
if launch_cores is not None:
|
if launch_cores is not None:
|
||||||
try:
|
try:
|
||||||
@@ -111,7 +462,7 @@ class RustServer:
|
|||||||
dp_note,
|
dp_note,
|
||||||
)
|
)
|
||||||
|
|
||||||
return cls(server)
|
return cls(server, mm_spec=mm_spec)
|
||||||
|
|
||||||
def wait_ingress(self, timeout_ms: int) -> None:
|
def wait_ingress(self, timeout_ms: int) -> None:
|
||||||
"""Block until a request is pushed into the in-process ring or the timeout
|
"""Block until a request is pushed into the in-process ring or the timeout
|
||||||
@@ -162,6 +513,13 @@ class RustServer:
|
|||||||
ids.frombytes(ids_view[pos : pos + nbytes])
|
ids.frombytes(ids_view[pos : pos + nbytes])
|
||||||
obj.input_ids = ids
|
obj.input_ids = ids
|
||||||
pos += nbytes
|
pos += nbytes
|
||||||
|
if self.mm_spec is not None and isinstance(obj, TokenizedGenerateReqInput):
|
||||||
|
# The buffers rode the Rust sidecar, parked before the ring push;
|
||||||
|
# wrapping them into tensors is the only Python step of the native
|
||||||
|
# path. `None` for a text-only request on a multimodal model.
|
||||||
|
native = self.server.take_mm(obj.rid)
|
||||||
|
if native is not None:
|
||||||
|
obj.mm_inputs = NativeMmHost.build_native_mm(self.mm_spec, native)
|
||||||
out.append(obj)
|
out.append(obj)
|
||||||
return out
|
return out
|
||||||
|
|
||||||
@@ -374,7 +732,9 @@ class RustServer:
|
|||||||
return msgspec.json.encode(server_args, enc_hook=str).decode("utf-8")
|
return msgspec.json.encode(server_args, enc_hook=str).decode("utf-8")
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _partition_cores() -> Tuple[Optional[List[int]], Optional[List[int]]]:
|
def _partition_cores(
|
||||||
|
mm_workers: int = 0,
|
||||||
|
) -> Tuple[Optional[List[int]], Optional[List[int]]]:
|
||||||
"""Split this rank's allowed cores into ``(launch_cores, server_cores)``.
|
"""Split this rank's allowed cores into ``(launch_cores, server_cores)``.
|
||||||
|
|
||||||
Pure computation — no affinity is changed here. Both sets are a subset
|
Pure computation — no affinity is changed here. Both sets are a subset
|
||||||
@@ -403,7 +763,17 @@ class RustServer:
|
|||||||
# effectively serial) and never take more than a quarter of the cores.
|
# effectively serial) and never take more than a quarter of the cores.
|
||||||
reserve = min(2, len(allowed) // 4)
|
reserve = min(2, len(allowed) // 4)
|
||||||
launch_cores = allowed[:reserve]
|
launch_cores = allowed[:reserve]
|
||||||
server_cores = allowed[reserve:]
|
# Bound the pool instead of taking the whole remainder: this rank's
|
||||||
|
# allowed cores are usually the entire NUMA node, shared with the sibling
|
||||||
|
# TP ranks' processes, so an unbounded mask lets MM preprocessing bursts
|
||||||
|
# preempt a sibling's CUDA-launch thread and inflate every rank's forward
|
||||||
|
# through the TP collectives. Measured on Qwen3.5-35B TP4 at one 720p
|
||||||
|
# image per request: ~20 ms of ViT wall time on the worst sibling, gone
|
||||||
|
# once bounded. The budget covers the CPU-hot threads (MM workers, plus
|
||||||
|
# the I/O-shaped tokenizer/ingress/egress/api ones that are rarely all hot
|
||||||
|
# at once) and leaves the rest of the node to the scheduler ranks.
|
||||||
|
pool_budget = max(8, mm_workers + 4)
|
||||||
|
server_cores = allowed[reserve : reserve + pool_budget]
|
||||||
logger.info(
|
logger.info(
|
||||||
"rust server cores=%s, scheduler launch cores=%s",
|
"rust server cores=%s, scheduler launch cores=%s",
|
||||||
server_cores,
|
server_cores,
|
||||||
|
|||||||
@@ -469,8 +469,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
import_processors("sglang.srt.multimodal.processors")
|
import_processors("sglang.srt.multimodal.processors")
|
||||||
if mm_process_pkg := envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.get():
|
if mm_process_pkg := envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.get():
|
||||||
import_processors(mm_process_pkg, overwrite=True)
|
import_processors(mm_process_pkg, overwrite=True)
|
||||||
_processor = _get_processor_wrapper(server_args)
|
_processor = get_processor_wrapper(server_args)
|
||||||
transport_mode = _determine_tensor_transport_mode(self.server_args)
|
transport_mode = determine_tensor_transport_mode(self.server_args)
|
||||||
|
|
||||||
# We want to parallelize the image pre-processing so we create an executor for it
|
# We want to parallelize the image pre-processing so we create an executor for it
|
||||||
# We create mm_processor for any skip_tokenizer_init to make sure we still encode
|
# We create mm_processor for any skip_tokenizer_init to make sure we still encode
|
||||||
@@ -3699,7 +3699,7 @@ async def print_exception_wrapper(func):
|
|||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
|
|
||||||
def _get_processor_wrapper(server_args):
|
def get_processor_wrapper(server_args):
|
||||||
try:
|
try:
|
||||||
processor = get_processor(
|
processor = get_processor(
|
||||||
server_args.tokenizer_path,
|
server_args.tokenizer_path,
|
||||||
@@ -3730,7 +3730,7 @@ def _get_processor_wrapper(server_args):
|
|||||||
return processor
|
return processor
|
||||||
|
|
||||||
|
|
||||||
def _determine_tensor_transport_mode(server_args: ServerArgs) -> TensorTransportMode:
|
def determine_tensor_transport_mode(server_args: ServerArgs) -> TensorTransportMode:
|
||||||
is_cross_node = server_args.dist_init_addr
|
is_cross_node = server_args.dist_init_addr
|
||||||
|
|
||||||
if is_cross_node:
|
if is_cross_node:
|
||||||
|
|||||||
Generated
+4
@@ -3690,6 +3690,7 @@ version = "0.1.0"
|
|||||||
dependencies = [
|
dependencies = [
|
||||||
"async-stream",
|
"async-stream",
|
||||||
"axum 0.8.9",
|
"axum 0.8.9",
|
||||||
|
"bytemuck",
|
||||||
"bytes",
|
"bytes",
|
||||||
"core_affinity",
|
"core_affinity",
|
||||||
"dynamo-parsers",
|
"dynamo-parsers",
|
||||||
@@ -3700,12 +3701,15 @@ dependencies = [
|
|||||||
"futures",
|
"futures",
|
||||||
"hf-hub",
|
"hf-hub",
|
||||||
"itertools",
|
"itertools",
|
||||||
|
"libc",
|
||||||
|
"numpy",
|
||||||
"pyo3",
|
"pyo3",
|
||||||
"regex-syntax",
|
"regex-syntax",
|
||||||
"rmp-serde",
|
"rmp-serde",
|
||||||
"rmpv",
|
"rmpv",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
|
"sglang-mm",
|
||||||
"socket2 0.6.5",
|
"socket2 0.6.5",
|
||||||
"thiserror",
|
"thiserror",
|
||||||
"tokio",
|
"tokio",
|
||||||
|
|||||||
@@ -15,49 +15,111 @@ use base64::Engine;
|
|||||||
/// payloads reject the request here).
|
/// payloads reject the request here).
|
||||||
pub const MAX_FETCH_BYTES: u64 = 64 << 20;
|
pub const MAX_FETCH_BYTES: u64 = 64 << 20;
|
||||||
|
|
||||||
|
/// Charge granularity of a streaming read: the most an in-flight source can
|
||||||
|
/// over-charge a shared [`ByteBudget`] by.
|
||||||
|
const CHUNK_BYTES: u64 = 256 << 10;
|
||||||
|
|
||||||
|
/// A byte allowance shared by every source of one request, charged *as they
|
||||||
|
/// stream*, so concurrent fetches stop at their combined size rather than each
|
||||||
|
/// stopping at [`MAX_FETCH_BYTES`].
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct ByteBudget(std::sync::atomic::AtomicU64);
|
||||||
|
|
||||||
|
impl ByteBudget {
|
||||||
|
pub fn new(total: u64) -> Self {
|
||||||
|
Self(std::sync::atomic::AtomicU64::new(total))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Claim `n` bytes, or `Err` once the allowance is spent.
|
||||||
|
fn claim(&self, n: u64) -> Result<(), ()> {
|
||||||
|
use std::sync::atomic::Ordering::{AcqRel, Acquire};
|
||||||
|
self.0
|
||||||
|
.fetch_update(AcqRel, Acquire, |left| left.checked_sub(n))
|
||||||
|
.map(|_| ())
|
||||||
|
.map_err(|_| ())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Give back bytes claimed for a chunk but not filled by the read.
|
||||||
|
fn release(&self, n: u64) {
|
||||||
|
self.0.fetch_add(n, std::sync::atomic::Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Resolve one string-typed image source into raw encoded-image bytes.
|
/// Resolve one string-typed image source into raw encoded-image bytes.
|
||||||
/// An `Err` rejects the request, matching the Python per-request
|
/// An `Err` rejects the request, matching the Python per-request
|
||||||
/// exception → 400.
|
/// exception → 400.
|
||||||
pub fn fetch_bytes(src: &str) -> Result<Vec<u8>, String> {
|
pub fn fetch_bytes(src: &str) -> Result<Vec<u8>, String> {
|
||||||
|
fetch_bytes_budgeted(src, &ByteBudget::new(MAX_FETCH_BYTES))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`fetch_bytes`] against a caller-owned allowance, for resolving several
|
||||||
|
/// sources under one whole-request bound. [`MAX_FETCH_BYTES`] still caps each.
|
||||||
|
pub fn fetch_bytes_budgeted(src: &str, budget: &ByteBudget) -> Result<Vec<u8>, String> {
|
||||||
if src.starts_with("http://") || src.starts_with("https://") {
|
if src.starts_with("http://") || src.starts_with("https://") {
|
||||||
return http_get(src);
|
return http_get(src, budget);
|
||||||
}
|
}
|
||||||
if let Some(path) = src.strip_prefix("file://") {
|
if let Some(path) = src.strip_prefix("file://") {
|
||||||
return read_file(path);
|
return read_file(path, budget);
|
||||||
}
|
}
|
||||||
if src.starts_with('/') {
|
if src.starts_with('/') {
|
||||||
return read_file(src);
|
return read_file(src, budget);
|
||||||
}
|
}
|
||||||
if let Some(rest) = src.strip_prefix("data:") {
|
if let Some(rest) = src.strip_prefix("data:") {
|
||||||
let encoded = rest
|
let encoded = rest
|
||||||
.split_once(',')
|
.split_once(',')
|
||||||
.ok_or_else(|| "media fetch: malformed data: URL".to_string())?
|
.ok_or_else(|| "media fetch: malformed data: URL".to_string())?
|
||||||
.1;
|
.1;
|
||||||
return b64(encoded);
|
return charge_decoded(b64(encoded)?, budget);
|
||||||
}
|
}
|
||||||
// Python treats any other string as bare base64.
|
// Python treats any other string as bare base64.
|
||||||
b64(src)
|
charge_decoded(b64(src)?, budget)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Base64 payloads are already resident in the request body — they cannot
|
||||||
|
/// amplify the way a download can, so they charge once decoded, not per chunk.
|
||||||
|
fn charge_decoded(decoded: Vec<u8>, budget: &ByteBudget) -> Result<Vec<u8>, String> {
|
||||||
|
budget
|
||||||
|
.claim(decoded.len() as u64)
|
||||||
|
.map_err(|()| over_budget("base64 payload"))?;
|
||||||
|
Ok(decoded)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn over_budget(what: &str) -> String {
|
||||||
|
format!("media fetch: {what}: exceeds the request media byte budget")
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Bounded read: never trusts metadata, so huge and non-regular files
|
/// Bounded read: never trusts metadata, so huge and non-regular files
|
||||||
/// (`/dev/zero`) hit the cap instead of exhausting memory.
|
/// (`/dev/zero`) hit the cap instead of exhausting memory.
|
||||||
fn read_file(path: &str) -> Result<Vec<u8>, String> {
|
fn read_file(path: &str, budget: &ByteBudget) -> Result<Vec<u8>, String> {
|
||||||
let file = std::fs::File::open(path).map_err(|e| format!("media fetch: {path}: {e}"))?;
|
let file = std::fs::File::open(path).map_err(|e| format!("media fetch: {path}: {e}"))?;
|
||||||
read_capped(file, path)
|
read_capped(file, path, budget)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn read_capped(reader: impl Read, what: &str) -> Result<Vec<u8>, String> {
|
/// Read to EOF, charging `budget` per chunk, so an oversized source stops
|
||||||
|
/// mid-stream instead of going fully resident first.
|
||||||
|
fn read_capped(mut reader: impl Read, what: &str, budget: &ByteBudget) -> Result<Vec<u8>, String> {
|
||||||
|
let too_big = || format!("media fetch: {what}: exceeds {MAX_FETCH_BYTES} bytes");
|
||||||
let mut buf = Vec::new();
|
let mut buf = Vec::new();
|
||||||
reader
|
loop {
|
||||||
.take(MAX_FETCH_BYTES + 1)
|
// `+ 1`: read one byte past the cap, so oversized is detected, not truncated.
|
||||||
|
let want = CHUNK_BYTES.min(MAX_FETCH_BYTES + 1 - buf.len() as u64);
|
||||||
|
if want == 0 {
|
||||||
|
return Err(too_big());
|
||||||
|
}
|
||||||
|
budget.claim(want).map_err(|()| over_budget(what))?;
|
||||||
|
let read = reader
|
||||||
|
.by_ref()
|
||||||
|
.take(want)
|
||||||
.read_to_end(&mut buf)
|
.read_to_end(&mut buf)
|
||||||
.map_err(|e| format!("media fetch: read {what}: {e}"))?;
|
.map_err(|e| format!("media fetch: read {what}: {e}"))? as u64;
|
||||||
|
budget.release(want - read);
|
||||||
if buf.len() as u64 > MAX_FETCH_BYTES {
|
if buf.len() as u64 > MAX_FETCH_BYTES {
|
||||||
return Err(format!(
|
return Err(too_big());
|
||||||
"media fetch: {what}: exceeds {MAX_FETCH_BYTES} bytes"
|
}
|
||||||
));
|
if read < want {
|
||||||
|
return Ok(buf); // short read == EOF
|
||||||
|
}
|
||||||
}
|
}
|
||||||
Ok(buf)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn b64(encoded: &str) -> Result<Vec<u8>, String> {
|
fn b64(encoded: &str) -> Result<Vec<u8>, String> {
|
||||||
@@ -163,7 +225,7 @@ fn in_ipv4_network(ip: std::net::Ipv4Addr, net: std::net::Ipv4Addr, bits: u32) -
|
|||||||
u32::from(ip) & mask == u32::from(net) & mask
|
u32::from(ip) & mask == u32::from(net) & mask
|
||||||
}
|
}
|
||||||
|
|
||||||
fn http_get(url: &str) -> Result<Vec<u8>, String> {
|
fn http_get(url: &str, budget: &ByteBudget) -> Result<Vec<u8>, String> {
|
||||||
// Python: `int(os.getenv("REQUEST_TIMEOUT", "3"))` seconds per image GET.
|
// Python: `int(os.getenv("REQUEST_TIMEOUT", "3"))` seconds per image GET.
|
||||||
let timeout = std::env::var("REQUEST_TIMEOUT")
|
let timeout = std::env::var("REQUEST_TIMEOUT")
|
||||||
.ok()
|
.ok()
|
||||||
@@ -178,7 +240,7 @@ fn http_get(url: &str) -> Result<Vec<u8>, String> {
|
|||||||
.timeout(std::time::Duration::from_secs(timeout))
|
.timeout(std::time::Duration::from_secs(timeout))
|
||||||
.call()
|
.call()
|
||||||
.map_err(|e| format!("media fetch: GET {url}: {e}"))?;
|
.map_err(|e| format!("media fetch: GET {url}: {e}"))?;
|
||||||
read_capped(resp.into_reader(), url)
|
read_capped(resp.into_reader(), url, budget)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -221,6 +283,38 @@ mod tests {
|
|||||||
assert!(err.contains("exceeds"), "{err}");
|
assert!(err.contains("exceeds"), "{err}");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// One budget spans sources: each fits alone, the set does not.
|
||||||
|
#[test]
|
||||||
|
fn shared_budget_spans_sources() {
|
||||||
|
let payload = base64::engine::general_purpose::STANDARD.encode([7u8; 4096]);
|
||||||
|
let budget = ByteBudget::new(6144);
|
||||||
|
assert_eq!(fetch_bytes_budgeted(&payload, &budget).unwrap().len(), 4096);
|
||||||
|
let err = fetch_bytes_budgeted(&payload, &budget).err().unwrap();
|
||||||
|
assert!(err.contains("request media byte budget"), "{err}");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Unused claims come back, so small sources fit in a budget their
|
||||||
|
/// worst-case sizes would have exhausted.
|
||||||
|
#[test]
|
||||||
|
fn short_reads_release_their_claim() {
|
||||||
|
let path = std::env::temp_dir().join(format!("sglang-budget-{}", std::process::id()));
|
||||||
|
std::fs::write(&path, [0u8; 1024]).unwrap();
|
||||||
|
let src = path.display().to_string();
|
||||||
|
let budget = ByteBudget::new(CHUNK_BYTES + 4096);
|
||||||
|
for _ in 0..4 {
|
||||||
|
assert_eq!(fetch_bytes_budgeted(&src, &budget).unwrap().len(), 1024);
|
||||||
|
}
|
||||||
|
std::fs::remove_file(&path).ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The per-source cap holds even under a larger shared budget.
|
||||||
|
#[test]
|
||||||
|
fn per_source_cap_survives_a_large_budget() {
|
||||||
|
let budget = ByteBudget::new(MAX_FETCH_BYTES * 4);
|
||||||
|
let err = fetch_bytes_budgeted("/dev/zero", &budget).err().unwrap();
|
||||||
|
assert!(err.contains(&format!("exceeds {MAX_FETCH_BYTES}")), "{err}");
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn host_parsing_strips_userinfo_and_path() {
|
fn host_parsing_strips_userinfo_and_path() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
|
|||||||
@@ -80,16 +80,32 @@ mod python {
|
|||||||
|
|
||||||
use super::{decode_rgb, resize};
|
use super::{decode_rgb, resize};
|
||||||
|
|
||||||
|
/// `resample` names the implementation to reproduce: `"pil_lanczos"` (the
|
||||||
|
/// inkling default), `"pil_bicubic"`, or `"aten_u8"` (torchvision's uint8
|
||||||
|
/// antialias bicubic). Exposed so the bit-exactness tests can cover each.
|
||||||
#[pyfunction]
|
#[pyfunction]
|
||||||
|
#[pyo3(signature = (arr, out_w, out_h, resample="pil_lanczos"))]
|
||||||
pub fn resize_rgb<'py>(
|
pub fn resize_rgb<'py>(
|
||||||
py: Python<'py>,
|
py: Python<'py>,
|
||||||
arr: PyReadonlyArray3<'py, u8>,
|
arr: PyReadonlyArray3<'py, u8>,
|
||||||
out_w: usize,
|
out_w: usize,
|
||||||
out_h: usize,
|
out_h: usize,
|
||||||
|
resample: &str,
|
||||||
) -> PyResult<Bound<'py, PyArray1<u8>>> {
|
) -> PyResult<Bound<'py, PyArray1<u8>>> {
|
||||||
if out_w == 0 || out_h == 0 {
|
if out_w == 0 || out_h == 0 {
|
||||||
return Err(PyValueError::new_err("output size must be positive"));
|
return Err(PyValueError::new_err("output size must be positive"));
|
||||||
}
|
}
|
||||||
|
let resample = match resample {
|
||||||
|
"pil_lanczos" => resize::Resample::Pil(resize::Filter::Lanczos),
|
||||||
|
"pil_bicubic" => resize::Resample::Pil(resize::Filter::Bicubic),
|
||||||
|
"aten_u8" => resize::Resample::AtenU8,
|
||||||
|
other => {
|
||||||
|
return Err(PyValueError::new_err(format!(
|
||||||
|
"unknown resample {other:?}; expected \"pil_lanczos\", \
|
||||||
|
\"pil_bicubic\" or \"aten_u8\""
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
};
|
||||||
let shape = arr.shape();
|
let shape = arr.shape();
|
||||||
let (h, w, c) = (shape[0], shape[1], shape[2]);
|
let (h, w, c) = (shape[0], shape[1], shape[2]);
|
||||||
if c != 3 {
|
if c != 3 {
|
||||||
@@ -101,7 +117,7 @@ mod python {
|
|||||||
.as_slice()
|
.as_slice()
|
||||||
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
|
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
|
||||||
.to_vec();
|
.to_vec();
|
||||||
let out = py.detach(move || resize::resize_lanczos_rgb(&data, h, w, out_h, out_w));
|
let out = py.detach(move || resize::resize_rgb(&data, h, w, out_h, out_w, resample));
|
||||||
Ok(out.into_pyarray(py))
|
Ok(out.into_pyarray(py))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -1,17 +1,52 @@
|
|||||||
use super::par;
|
use super::par;
|
||||||
|
|
||||||
const PRECISION_BITS: i32 = 32 - 8 - 2;
|
/// PIL's `PRECISION_BITS` for 8-bit images: weights quantized to i32.
|
||||||
|
const PIL_PRECISION_BITS: u32 = 32 - 8 - 2;
|
||||||
|
|
||||||
/// Resampling filters, bit-exact clones of PIL's kernels.
|
/// Resampling filters, bit-exact clones of PIL's kernels.
|
||||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||||
pub enum Filter {
|
pub enum Filter {
|
||||||
/// support 3.0 — PIL `LANCZOS`.
|
/// support 3.0 — PIL `LANCZOS`.
|
||||||
Lanczos,
|
Lanczos,
|
||||||
/// support 2.0, a = -0.5 — PIL `BICUBIC` (≈ torchvision antialiased
|
/// support 2.0, a = -0.5 — PIL `BICUBIC`.
|
||||||
/// bicubic, which the HF "fast" image processors use).
|
|
||||||
Bicubic,
|
Bicubic,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// A resampler reproduced bit-exactly. Both share PIL's geometry, kernels and
|
||||||
|
/// per-pass u8 rounding, and differ only in how the weights are quantized.
|
||||||
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||||
|
pub enum Resample {
|
||||||
|
/// PIL `Image.resize`, i32 weights.
|
||||||
|
Pil(Filter),
|
||||||
|
/// ATen's uint8 antialias bicubic — torchvision `resize(antialias=True)` on
|
||||||
|
/// a uint8 tensor. i16 weights, so it rounds unlike `Pil(Bicubic)`.
|
||||||
|
AtenU8,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Resample {
|
||||||
|
fn filter(self) -> Filter {
|
||||||
|
match self {
|
||||||
|
Resample::Pil(filter) => filter,
|
||||||
|
Resample::AtenU8 => Filter::Bicubic,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fixed-point precision for one axis's already-normalized weights. ATen
|
||||||
|
/// (`_compute_weights_precision`) takes the widest that stays inside i16.
|
||||||
|
fn precision(self, weights: &[f64]) -> u32 {
|
||||||
|
match self {
|
||||||
|
Resample::Pil(_) => PIL_PRECISION_BITS,
|
||||||
|
Resample::AtenU8 => {
|
||||||
|
let wmax = weights.iter().fold(0.0f64, |m, w| m.max(w.abs()));
|
||||||
|
(1..PIL_PRECISION_BITS)
|
||||||
|
.take_while(|&p| (0.5 + wmax * (1u64 << p) as f64) < (1 << 15) as f64)
|
||||||
|
.last()
|
||||||
|
.unwrap_or(1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl Filter {
|
impl Filter {
|
||||||
fn support(self) -> f64 {
|
fn support(self) -> f64 {
|
||||||
match self {
|
match self {
|
||||||
@@ -60,9 +95,11 @@ struct Coeffs {
|
|||||||
bounds: Vec<(usize, usize)>,
|
bounds: Vec<(usize, usize)>,
|
||||||
kk: Vec<i32>,
|
kk: Vec<i32>,
|
||||||
ksize: usize,
|
ksize: usize,
|
||||||
|
prec: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
fn precompute_coeffs(in_size: usize, out_size: usize, filter: Filter) -> Coeffs {
|
fn precompute_coeffs(in_size: usize, out_size: usize, resample: Resample) -> Coeffs {
|
||||||
|
let filter = resample.filter();
|
||||||
let scale = in_size as f64 / out_size as f64;
|
let scale = in_size as f64 / out_size as f64;
|
||||||
let filterscale = if scale < 1.0 { 1.0 } else { scale };
|
let filterscale = if scale < 1.0 { 1.0 } else { scale };
|
||||||
let support = filter.support() * filterscale;
|
let support = filter.support() * filterscale;
|
||||||
@@ -97,7 +134,8 @@ fn precompute_coeffs(in_size: usize, out_size: usize, filter: Filter) -> Coeffs
|
|||||||
bounds[xx] = (xmin as usize, count);
|
bounds[xx] = (xmin as usize, count);
|
||||||
}
|
}
|
||||||
|
|
||||||
let factor = (1i64 << PRECISION_BITS) as f64;
|
let prec = resample.precision(&kkf);
|
||||||
|
let factor = (1i64 << prec) as f64;
|
||||||
let kk = kkf
|
let kk = kkf
|
||||||
.iter()
|
.iter()
|
||||||
.map(|&v| {
|
.map(|&v| {
|
||||||
@@ -108,17 +146,22 @@ fn precompute_coeffs(in_size: usize, out_size: usize, filter: Filter) -> Coeffs
|
|||||||
}
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
Coeffs { bounds, kk, ksize }
|
Coeffs {
|
||||||
|
bounds,
|
||||||
|
kk,
|
||||||
|
ksize,
|
||||||
|
prec,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
#[inline]
|
||||||
fn clip8(v: i32) -> u8 {
|
fn clip8(v: i32, prec: u32) -> u8 {
|
||||||
if v >= 1 << (PRECISION_BITS + 8) {
|
if v >= 1 << (prec + 8) {
|
||||||
255
|
255
|
||||||
} else if v <= 0 {
|
} else if v <= 0 {
|
||||||
0
|
0
|
||||||
} else {
|
} else {
|
||||||
(v >> PRECISION_BITS) as u8
|
(v >> prec) as u8
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -129,7 +172,7 @@ fn resample_horizontal(src: &[u8], h: usize, w: usize, out_w: usize, c: &Coeffs)
|
|||||||
for xx in 0..out_w {
|
for xx in 0..out_w {
|
||||||
let (xmin, count) = c.bounds[xx];
|
let (xmin, count) = c.bounds[xx];
|
||||||
let k = &c.kk[xx * c.ksize..xx * c.ksize + count];
|
let k = &c.kk[xx * c.ksize..xx * c.ksize + count];
|
||||||
let mut s = [1i32 << (PRECISION_BITS - 1); 3];
|
let mut s = [1i32 << (c.prec - 1); 3];
|
||||||
for (x, &coef) in k.iter().enumerate() {
|
for (x, &coef) in k.iter().enumerate() {
|
||||||
let p = (xmin + x) * 3;
|
let p = (xmin + x) * 3;
|
||||||
s[0] += src_row[p] as i32 * coef;
|
s[0] += src_row[p] as i32 * coef;
|
||||||
@@ -137,9 +180,9 @@ fn resample_horizontal(src: &[u8], h: usize, w: usize, out_w: usize, c: &Coeffs)
|
|||||||
s[2] += src_row[p + 2] as i32 * coef;
|
s[2] += src_row[p + 2] as i32 * coef;
|
||||||
}
|
}
|
||||||
let o = xx * 3;
|
let o = xx * 3;
|
||||||
row[o] = clip8(s[0]);
|
row[o] = clip8(s[0], c.prec);
|
||||||
row[o + 1] = clip8(s[1]);
|
row[o + 1] = clip8(s[1], c.prec);
|
||||||
row[o + 2] = clip8(s[2]);
|
row[o + 2] = clip8(s[2], c.prec);
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
out
|
out
|
||||||
@@ -151,7 +194,7 @@ fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec<u8>
|
|||||||
let (ymin, count) = c.bounds[yy];
|
let (ymin, count) = c.bounds[yy];
|
||||||
let k = &c.kk[yy * c.ksize..yy * c.ksize + count];
|
let k = &c.kk[yy * c.ksize..yy * c.ksize + count];
|
||||||
for x in 0..w {
|
for x in 0..w {
|
||||||
let mut s = [1i32 << (PRECISION_BITS - 1); 3];
|
let mut s = [1i32 << (c.prec - 1); 3];
|
||||||
for (y, &coef) in k.iter().enumerate() {
|
for (y, &coef) in k.iter().enumerate() {
|
||||||
let p = ((ymin + y) * w + x) * 3;
|
let p = ((ymin + y) * w + x) * 3;
|
||||||
s[0] += src[p] as i32 * coef;
|
s[0] += src[p] as i32 * coef;
|
||||||
@@ -159,27 +202,27 @@ fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec<u8>
|
|||||||
s[2] += src[p + 2] as i32 * coef;
|
s[2] += src[p + 2] as i32 * coef;
|
||||||
}
|
}
|
||||||
let o = x * 3;
|
let o = x * 3;
|
||||||
row[o] = clip8(s[0]);
|
row[o] = clip8(s[0], c.prec);
|
||||||
row[o + 1] = clip8(s[1]);
|
row[o + 1] = clip8(s[1], c.prec);
|
||||||
row[o + 2] = clip8(s[2]);
|
row[o + 2] = clip8(s[2], c.prec);
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
out
|
out
|
||||||
}
|
}
|
||||||
|
|
||||||
/// PIL-exact separable resize of a flat HWC RGB buffer with the given filter.
|
/// Separable resize of a flat HWC RGB buffer, bit-exact against `resample`.
|
||||||
///
|
///
|
||||||
/// Enters the fan-out pool once for both passes; the per-row `for_chunks_mut`
|
/// Enters the fan-out pool once for both passes; the per-row `for_chunks_mut`
|
||||||
/// calls inside then reuse that entry rather than injecting a job per pass.
|
/// calls inside then reuse that entry rather than injecting a job per pass.
|
||||||
pub fn resize_rgb_filter(
|
pub fn resize_rgb(
|
||||||
src: &[u8],
|
src: &[u8],
|
||||||
h: usize,
|
h: usize,
|
||||||
w: usize,
|
w: usize,
|
||||||
out_h: usize,
|
out_h: usize,
|
||||||
out_w: usize,
|
out_w: usize,
|
||||||
filter: Filter,
|
resample: Resample,
|
||||||
) -> Vec<u8> {
|
) -> Vec<u8> {
|
||||||
par::in_pool(move || resize_passes(src, h, w, out_h, out_w, filter))
|
par::in_pool(move || resize_passes(src, h, w, out_h, out_w, resample))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn resize_passes(
|
fn resize_passes(
|
||||||
@@ -188,28 +231,23 @@ fn resize_passes(
|
|||||||
w: usize,
|
w: usize,
|
||||||
out_h: usize,
|
out_h: usize,
|
||||||
out_w: usize,
|
out_w: usize,
|
||||||
filter: Filter,
|
resample: Resample,
|
||||||
) -> Vec<u8> {
|
) -> Vec<u8> {
|
||||||
let need_h = out_w != w;
|
// Per-axis coefficients — and, under `AtenU8`, a per-axis precision.
|
||||||
let need_v = out_h != h;
|
let coeffs = |in_size, out_size| precompute_coeffs(in_size, out_size, resample);
|
||||||
if need_h && need_v {
|
match (out_w != w, out_h != h) {
|
||||||
let ch = precompute_coeffs(w, out_w, filter);
|
(true, true) => {
|
||||||
let tmp = resample_horizontal(src, h, w, out_w, &ch);
|
let tmp = resample_horizontal(src, h, w, out_w, &coeffs(w, out_w));
|
||||||
let cv = precompute_coeffs(h, out_h, filter);
|
resample_vertical(&tmp, out_w, out_h, &coeffs(h, out_h))
|
||||||
resample_vertical(&tmp, out_w, out_h, &cv)
|
}
|
||||||
} else if need_h {
|
(true, false) => resample_horizontal(src, h, w, out_w, &coeffs(w, out_w)),
|
||||||
let ch = precompute_coeffs(w, out_w, filter);
|
(false, true) => resample_vertical(src, w, out_h, &coeffs(h, out_h)),
|
||||||
resample_horizontal(src, h, w, out_w, &ch)
|
(false, false) => src.to_vec(),
|
||||||
} else if need_v {
|
|
||||||
let cv = precompute_coeffs(h, out_h, filter);
|
|
||||||
resample_vertical(src, w, out_h, &cv)
|
|
||||||
} else {
|
|
||||||
src.to_vec()
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn resize_lanczos_rgb(src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize) -> Vec<u8> {
|
pub fn resize_lanczos_rgb(src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize) -> Vec<u8> {
|
||||||
resize_rgb_filter(src, h, w, out_h, out_w, Filter::Lanczos)
|
resize_rgb(src, h, w, out_h, out_w, Resample::Pil(Filter::Lanczos))
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn scaled_dims(w: usize, h: usize, frac: Option<f64>, cap: Option<i64>) -> (usize, usize) {
|
pub fn scaled_dims(w: usize, h: usize, frac: Option<f64>, cap: Option<i64>) -> (usize, usize) {
|
||||||
|
|||||||
@@ -34,21 +34,60 @@ pub struct QwenVlSpec {
|
|||||||
pub max_pixels: usize,
|
pub max_pixels: usize,
|
||||||
pub image_mean: [f32; 3],
|
pub image_mean: [f32; 3],
|
||||||
pub image_std: [f32; 3],
|
pub image_std: [f32; 3],
|
||||||
|
#[serde(default)]
|
||||||
|
pub resample: Resampler,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The HF image processor the pipeline must match bit-exactly. Defaults to the
|
||||||
|
/// one a default server runs.
|
||||||
|
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Deserialize)]
|
||||||
|
#[serde(rename_all = "snake_case")]
|
||||||
|
pub enum Resampler {
|
||||||
|
/// `Qwen2VLImageProcessor` / `…Fast` — torchvision on a uint8 tensor.
|
||||||
|
#[default]
|
||||||
|
AtenU8,
|
||||||
|
/// `Qwen2VLImageProcessorPil`, behind `--disable-fast-image-processor`.
|
||||||
|
Pil,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<Resampler> for resize::Resample {
|
||||||
|
fn from(r: Resampler) -> Self {
|
||||||
|
match r {
|
||||||
|
Resampler::AtenU8 => resize::Resample::AtenU8,
|
||||||
|
Resampler::Pil => resize::Resample::Pil(resize::Filter::Bicubic),
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub struct QwenVlProcessor {
|
pub struct QwenVlProcessor {
|
||||||
spec: QwenVlSpec,
|
spec: QwenVlSpec,
|
||||||
/// Per-channel u8 → normalized-f32 lookup: `(v/255 - mean) / std`.
|
/// Per-channel u8 → normalized-f32 lookup; see [`normalize_lut`].
|
||||||
lut: [[f32; 256]; 3],
|
lut: [[f32; 256]; 3],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// `1 / rescale_factor`; `resolve_native_spec` rejects any other factor.
|
||||||
|
const INV_RESCALE: f32 = 255.0;
|
||||||
|
|
||||||
|
/// u8 → normalized f32, rounded as the mirrored processor rounds. The slow one
|
||||||
|
/// rescales then normalizes; the fast one folds the rescale into mean/std first
|
||||||
|
/// (`_fuse_mean_std_and_rescale_factor`), which differs on 128 of the 256 inputs.
|
||||||
|
fn normalize_lut(resample: Resampler, mean: f32, std: f32) -> [f32; 256] {
|
||||||
|
match resample {
|
||||||
|
Resampler::Pil => core::array::from_fn(|v| (v as f32 / INV_RESCALE - mean) / std),
|
||||||
|
Resampler::AtenU8 => {
|
||||||
|
let (mean, std) = (mean * INV_RESCALE, std * INV_RESCALE);
|
||||||
|
core::array::from_fn(|v| (v as f32 - mean) / std)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl QwenVlProcessor {
|
impl QwenVlProcessor {
|
||||||
pub fn new(spec: QwenVlSpec) -> Result<Self, String> {
|
pub fn new(spec: QwenVlSpec) -> Result<Self, String> {
|
||||||
if spec.patch_size == 0 || spec.merge_size == 0 || spec.temporal_patch_size == 0 {
|
if spec.patch_size == 0 || spec.merge_size == 0 || spec.temporal_patch_size == 0 {
|
||||||
return Err("qwen_vl spec: sizes must be positive".into());
|
return Err("qwen_vl spec: sizes must be positive".into());
|
||||||
}
|
}
|
||||||
let lut = core::array::from_fn(|c| {
|
let lut = core::array::from_fn(|c| {
|
||||||
core::array::from_fn(|v| (v as f32 / 255.0 - spec.image_mean[c]) / spec.image_std[c])
|
normalize_lut(spec.resample, spec.image_mean[c], spec.image_std[c])
|
||||||
});
|
});
|
||||||
Ok(Self { spec, lut })
|
Ok(Self { spec, lut })
|
||||||
}
|
}
|
||||||
@@ -127,7 +166,7 @@ impl MmFamilyProcessor for QwenVlProcessor {
|
|||||||
)?;
|
)?;
|
||||||
let resized;
|
let resized;
|
||||||
let data = if (th, tw) != (h, w) {
|
let data = if (th, tw) != (h, w) {
|
||||||
resized = resize::resize_rgb_filter(rgb, h, w, th, tw, resize::Filter::Bicubic);
|
resized = resize::resize_rgb(rgb, h, w, th, tw, self.spec.resample.into());
|
||||||
&resized
|
&resized
|
||||||
} else {
|
} else {
|
||||||
rgb.as_slice()
|
rgb.as_slice()
|
||||||
@@ -513,6 +552,22 @@ mod tests {
|
|||||||
max_pixels: 1 << 30,
|
max_pixels: 1 << 30,
|
||||||
image_mean: [0.0; 3],
|
image_mean: [0.0; 3],
|
||||||
image_std: [1.0; 3],
|
image_std: [1.0; 3],
|
||||||
|
resample: Resampler::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The fused and unfused normalize forms are not interchangeable: with
|
||||||
|
/// mean = std = 0.5 they disagree on 128 of the 256 u8 inputs, so picking
|
||||||
|
/// the wrong one silently costs bit-exactness with the HF processor.
|
||||||
|
#[test]
|
||||||
|
fn normalize_lut_differs_per_resampler() {
|
||||||
|
let pil = normalize_lut(Resampler::Pil, 0.5, 0.5);
|
||||||
|
let aten = normalize_lut(Resampler::AtenU8, 0.5, 0.5);
|
||||||
|
assert_eq!(pil.iter().zip(aten).filter(|(p, a)| *p != a).count(), 128);
|
||||||
|
// Both still span [-1, 1] — this is rounding, not a scale error.
|
||||||
|
for lut in [pil, aten] {
|
||||||
|
assert_eq!(lut[0], -1.0);
|
||||||
|
assert_eq!(lut[255], 1.0);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -34,15 +34,34 @@ def py_scaled_dims(
|
|||||||
return scale(width), scale(height)
|
return scale(width), scale(height)
|
||||||
|
|
||||||
|
|
||||||
def pil_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray:
|
def pil_resize(arr: np.ndarray, tw: int, th: int, filter=Image.Resampling.LANCZOS):
|
||||||
return np.array(
|
return np.array(Image.fromarray(arr).resize((tw, th), resample=filter), np.uint8)
|
||||||
Image.fromarray(arr).resize((tw, th), resample=Image.Resampling.LANCZOS),
|
|
||||||
dtype=np.uint8,
|
|
||||||
|
def tv_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray:
|
||||||
|
"""torchvision's uint8 antialias bicubic — ATen's fixed-point kernel."""
|
||||||
|
import torch
|
||||||
|
from torchvision.transforms.v2 import functional as F
|
||||||
|
|
||||||
|
tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
|
||||||
|
out = F.resize(
|
||||||
|
tensor, [th, tw], interpolation=F.InterpolationMode.BICUBIC, antialias=True
|
||||||
)
|
)
|
||||||
|
return out[0].permute(1, 2, 0).numpy()
|
||||||
|
|
||||||
|
|
||||||
def rs_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray:
|
# Every resampler the Rust resize claims, and its reference: `aten_u8` for a
|
||||||
return _rs_common.resize_rgb(arr, tw, th).reshape(th, tw, 3)
|
# default server, `pil_bicubic` for --disable-fast-image-processor, `pil_lanczos`
|
||||||
|
# for inkling.
|
||||||
|
REFERENCES = {
|
||||||
|
"pil_lanczos": lambda a, tw, th: pil_resize(a, tw, th, Image.Resampling.LANCZOS),
|
||||||
|
"pil_bicubic": lambda a, tw, th: pil_resize(a, tw, th, Image.Resampling.BICUBIC),
|
||||||
|
"aten_u8": tv_resize,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def rs_resize(arr, tw: int, th: int, resample: str = "pil_lanczos") -> np.ndarray:
|
||||||
|
return _rs_common.resize_rgb(arr, tw, th, resample).reshape(th, tw, 3)
|
||||||
|
|
||||||
|
|
||||||
CASES = [
|
CASES = [
|
||||||
@@ -58,13 +77,36 @@ CASES = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("resample", sorted(REFERENCES))
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"h,w,th,tw", CASES, ids=[f"{h}x{w}->{th}x{tw}" for h, w, th, tw in CASES]
|
"h,w,th,tw", CASES, ids=[f"{h}x{w}->{th}x{tw}" for h, w, th, tw in CASES]
|
||||||
)
|
)
|
||||||
def test_resize_bit_exact(h, w, th, tw):
|
def test_resize_bit_exact(h, w, th, tw, resample):
|
||||||
rng = np.random.default_rng(h * 10000 + w)
|
rng = np.random.default_rng(h * 10000 + w)
|
||||||
arr = rng.integers(0, 256, (h, w, 3), dtype=np.uint8)
|
arr = rng.integers(0, 256, (h, w, 3), dtype=np.uint8)
|
||||||
np.testing.assert_array_equal(rs_resize(arr, tw, th), pil_resize(arr, tw, th))
|
np.testing.assert_array_equal(
|
||||||
|
rs_resize(arr, tw, th, resample), REFERENCES[resample](arr, tw, th)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("resample", sorted(REFERENCES))
|
||||||
|
def test_resize_bit_exact_random_sweep(resample):
|
||||||
|
"""`aten_u8`'s weight precision varies with the scale factor, so the fixed
|
||||||
|
cases above are not enough coverage on their own."""
|
||||||
|
rng = np.random.default_rng(7)
|
||||||
|
for h, w, th, tw in rng.integers(1, 200, (40, 4)):
|
||||||
|
arr = rng.integers(0, 256, (h, w, 3), dtype=np.uint8)
|
||||||
|
np.testing.assert_array_equal(
|
||||||
|
rs_resize(arr, tw, th, resample),
|
||||||
|
REFERENCES[resample](arr, tw, th),
|
||||||
|
err_msg=f"{h}x{w}->{th}x{tw} under {resample}",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_unknown_resample_rejected():
|
||||||
|
arr = np.zeros((4, 4, 3), dtype=np.uint8)
|
||||||
|
with pytest.raises(ValueError, match="unknown resample"):
|
||||||
|
_rs_common.resize_rgb(arr, 2, 2, "nearest")
|
||||||
|
|
||||||
|
|
||||||
def test_scaled_dims_sweep():
|
def test_scaled_dims_sweep():
|
||||||
|
|||||||
@@ -43,14 +43,25 @@ dynamo-parsers = "7.0.1"
|
|||||||
dynamo-protocols = "5.1.0"
|
dynamo-protocols = "5.1.0"
|
||||||
dynamo-renderer = "5.0.0"
|
dynamo-renderer = "5.0.0"
|
||||||
flume = "0.12.0"
|
flume = "0.12.0"
|
||||||
|
# Safe POD slice casts (feature buffers viewed as bytes for the shm copy).
|
||||||
|
bytemuck = "1"
|
||||||
itertools = "0.14"
|
itertools = "0.14"
|
||||||
hf-hub = { version = "0.4", default-features = false }
|
hf-hub = { version = "0.4", default-features = false }
|
||||||
|
# POSIX shm for the MM feature fan-out (`mm::ShmSegment`).
|
||||||
|
libc = "0.2"
|
||||||
|
# Same major as the workspace pyo3: the zero-copy MM drain (`take_mm`) moves
|
||||||
|
# Rust vectors into numpy arrays.
|
||||||
|
numpy = "0.29.0"
|
||||||
rmp-serde = "1"
|
rmp-serde = "1"
|
||||||
rmpv = { version = "1", features = ["with-serde"] }
|
rmpv = { version = "1", features = ["with-serde"] }
|
||||||
# Pinned EXACTLY: this crate's accepted grammar defines the
|
# Pinned EXACTLY: this crate's accepted grammar defines the
|
||||||
# "anything Rust admits, Python can compile" invariant in `message::sampling`.
|
# "anything Rust admits, Python can compile" invariant in `message::sampling`.
|
||||||
# A minor bump can widen it and silently reopen a scheduler-killing hole.
|
# A minor bump can widen it and silently reopen a scheduler-killing hole.
|
||||||
regex-syntax = "=0.8.11"
|
regex-syntax = "=0.8.11"
|
||||||
|
# The pure-Rust core of the MM pipeline. `default-features = false` drops the
|
||||||
|
# pyo3 bindings so it links as a plain rlib; renamed so `use sglang_mm::…` reads
|
||||||
|
# naturally while the crate keeps its own artifact name in the shared target/.
|
||||||
|
sglang_mm = { package = "sglang-mm", path = "../sglang-mm", default-features = false }
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
# `Router::oneshot` for handler-level router tests.
|
# `Router::oneshot` for handler-level router tests.
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ mod log;
|
|||||||
mod native_api;
|
mod native_api;
|
||||||
mod openai;
|
mod openai;
|
||||||
mod pd_bootstrap;
|
mod pd_bootstrap;
|
||||||
|
mod prefetch;
|
||||||
mod submit;
|
mod submit;
|
||||||
|
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|||||||
@@ -154,7 +154,7 @@ async fn generate(
|
|||||||
let stream = body.stream;
|
let stream = body.stream;
|
||||||
// Fan `text`/`input_ids`/`sampling_params` (scalar or list) into per-request
|
// Fan `text`/`input_ids`/`sampling_params` (scalar or list) into per-request
|
||||||
// payloads. `is_batch` = list form → the response is a JSON array.
|
// payloads. `is_batch` = list form → the response is a JSON array.
|
||||||
let (payloads, is_batch) = match body.into_requests() {
|
let (mut payloads, is_batch) = match body.into_requests() {
|
||||||
Ok(v) => v,
|
Ok(v) => v,
|
||||||
// The error carries its own status (a bad batch is `Validation` → 400).
|
// The error carries its own status (a bad batch is `Validation` → 400).
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
@@ -162,6 +162,11 @@ async fn generate(
|
|||||||
return pre_submit_error(code, &e.to_string(), stream);
|
return pre_submit_error(code, &e.to_string(), stream);
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
// Media I/O (URL downloads, file reads) happens here, on the API runtime
|
||||||
|
// — never on the MM worker pool (see `prefetch`).
|
||||||
|
if let Err(e) = super::prefetch::prefetch_all(&mut payloads).await {
|
||||||
|
return pre_submit_error(StatusCode::BAD_REQUEST, &e, stream);
|
||||||
|
}
|
||||||
if !is_batch {
|
if !is_batch {
|
||||||
// `into_requests` guarantees exactly one payload for a non-batch body.
|
// `into_requests` guarantees exactly one payload for a non-batch body.
|
||||||
let payload = payloads
|
let payload = payloads
|
||||||
|
|||||||
@@ -0,0 +1,199 @@
|
|||||||
|
//! Resolve I/O-backed media sources on the API runtime, before MM dispatch.
|
||||||
|
//!
|
||||||
|
//! The MM worker pool is fixed, core-pinned CPU capacity: a slow image host — or
|
||||||
|
//! a file on a hanging network mount — must never occupy it, and a request's
|
||||||
|
//! images must download concurrently, not in `n * REQUEST_TIMEOUT`. URLs and
|
||||||
|
//! file paths resolve here through `sglang-mm`'s `fetch_bytes_budgeted` (one
|
||||||
|
//! owner for proxy/timeout/cap semantics) and ride out-of-band as
|
||||||
|
//! [`crate::message::MmData::prefetched`], which
|
||||||
|
//! [`crate::message::mm_payload::to_mm_input`] swaps back in.
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use bytes::Bytes;
|
||||||
|
use sglang_mm::common::fetch::{ByteBudget, fetch_bytes_budgeted};
|
||||||
|
use sglang_mm::driver::{MAX_ITEMS_PER_REQUEST, MAX_REQUEST_BYTES};
|
||||||
|
use tokio::sync::Semaphore;
|
||||||
|
|
||||||
|
use crate::message::mm_payload::{io_sources, item_count};
|
||||||
|
use crate::message::{GenerateRequest, MmData};
|
||||||
|
|
||||||
|
/// Global bound on concurrent media fetches across all in-flight requests;
|
||||||
|
/// excess acquisitions queue on the semaphore without holding a thread.
|
||||||
|
static PERMITS: Semaphore = Semaphore::const_new(32);
|
||||||
|
|
||||||
|
/// Fill [`MmData::prefetched`] for every request, all fetches across the batch
|
||||||
|
/// concurrent. Any failure rejects the call (a 400, as on the Python path).
|
||||||
|
///
|
||||||
|
/// The driver's budgets ([`MAX_ITEMS_PER_REQUEST`], [`MAX_REQUEST_BYTES`]) are
|
||||||
|
/// enforced *here* rather than in `sglang_mm::driver::process`, where 64 sources
|
||||||
|
/// of 64 MiB would already be resident. The driver keeps its own checks as the
|
||||||
|
/// backstop for callers without a prefetch layer.
|
||||||
|
pub async fn prefetch_all(requests: &mut [GenerateRequest]) -> Result<(), String> {
|
||||||
|
// The item budget rejects before a single byte is fetched.
|
||||||
|
let plan = |mm: &Option<Box<MmData>>| -> Result<Vec<String>, String> {
|
||||||
|
let Some(image_data) = mm.as_deref().and_then(|m| m.image_data.as_ref()) else {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
};
|
||||||
|
if item_count(image_data) > MAX_ITEMS_PER_REQUEST {
|
||||||
|
return Err(format!(
|
||||||
|
"multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(io_sources(image_data))
|
||||||
|
};
|
||||||
|
let plans = requests
|
||||||
|
.iter()
|
||||||
|
.map(|r| plan(&r.mm))
|
||||||
|
.collect::<Result<Vec<_>, String>>()?;
|
||||||
|
let fetches = plans
|
||||||
|
.into_iter()
|
||||||
|
.map(|sources| fetch_ordered(sources, MAX_REQUEST_BYTES));
|
||||||
|
let fetched = futures::future::try_join_all(fetches).await?;
|
||||||
|
for (req, bytes) in requests.iter_mut().zip(fetched) {
|
||||||
|
if !bytes.is_empty() {
|
||||||
|
req.mm.as_mut().expect("sources came from mm").prefetched = bytes;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve one request's sources concurrently (globally bounded), in order,
|
||||||
|
/// against one shared `total_bytes` allowance. Overflow rejects mid-download and
|
||||||
|
/// `try_join_all` drops the rest, so queued sources never start.
|
||||||
|
async fn fetch_ordered(sources: Vec<String>, total_bytes: u64) -> Result<Vec<Bytes>, String> {
|
||||||
|
let budget = Arc::new(ByteBudget::new(total_bytes));
|
||||||
|
futures::future::try_join_all(sources.into_iter().map(|src| {
|
||||||
|
let budget = Arc::clone(&budget);
|
||||||
|
async move {
|
||||||
|
let _permit = PERMITS.acquire().await.expect("semaphore never closed");
|
||||||
|
// Blocking I/O: parks a lazily-spawned blocking-pool thread, never
|
||||||
|
// an API worker. Those threads are pinned round-robin over the api
|
||||||
|
// core set (see `on_thread_start` in `runtime::start`) — off the
|
||||||
|
// CPU-bound stages, and mostly I/O-parked, so sharing is fine.
|
||||||
|
tokio::task::spawn_blocking(move || fetch_bytes_budgeted(&src, &budget))
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("media prefetch: {e}"))?
|
||||||
|
.map(Bytes::from)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use rmpv::Value;
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn serve(bodies: Vec<Vec<u8>>) -> std::net::SocketAddr {
|
||||||
|
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
std::thread::spawn(move || {
|
||||||
|
for body in bodies {
|
||||||
|
use std::io::{BufRead, Write};
|
||||||
|
let (stream, _) = listener.accept().unwrap();
|
||||||
|
let mut reader = std::io::BufReader::new(stream);
|
||||||
|
let mut line = String::new();
|
||||||
|
while reader.read_line(&mut line).unwrap() > 2 {
|
||||||
|
line.clear(); // headers until the blank line
|
||||||
|
}
|
||||||
|
let mut stream = reader.into_inner();
|
||||||
|
write!(
|
||||||
|
stream,
|
||||||
|
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n",
|
||||||
|
body.len()
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
stream.write_all(&body).unwrap();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
addr
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mm_request(image_data: Value) -> GenerateRequest {
|
||||||
|
GenerateRequest {
|
||||||
|
mm: Some(Box::new(MmData {
|
||||||
|
image_data: Some(image_data),
|
||||||
|
..Default::default()
|
||||||
|
})),
|
||||||
|
..Default::default()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// URLs and file paths resolve concurrently into `prefetched` in source
|
||||||
|
/// order; CPU-only sources and mm-free requests are untouched.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn resolves_io_sources() {
|
||||||
|
let addr = serve(vec![b"one".to_vec(), b"two".to_vec()]);
|
||||||
|
let path = std::env::temp_dir().join(format!("sglang-prefetch-{}", std::process::id()));
|
||||||
|
std::fs::write(&path, b"zzz").unwrap();
|
||||||
|
let mut requests = vec![
|
||||||
|
mm_request(Value::Array(vec![
|
||||||
|
Value::from(format!("http://{addr}/a.png")),
|
||||||
|
Value::from("data:image/png;base64,x"),
|
||||||
|
Value::from(format!("http://{addr}/b.png")),
|
||||||
|
Value::from(path.display().to_string()),
|
||||||
|
])),
|
||||||
|
GenerateRequest::default(),
|
||||||
|
];
|
||||||
|
prefetch_all(&mut requests).await.unwrap();
|
||||||
|
std::fs::remove_file(&path).ok();
|
||||||
|
let fetched = &requests[0].mm.as_ref().unwrap().prefetched;
|
||||||
|
// The one-shot server answers in accept order, so contents may swap
|
||||||
|
// between the two URLs; all three bodies must arrive.
|
||||||
|
let mut got: Vec<&[u8]> = fetched.iter().map(|b| b.as_ref()).collect();
|
||||||
|
got.sort();
|
||||||
|
assert_eq!(got, vec![b"one".as_ref(), b"two".as_ref(), b"zzz".as_ref()]);
|
||||||
|
assert!(requests[1].mm.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn failed_download_rejects() {
|
||||||
|
let mut requests = vec![mm_request(Value::from("http://127.0.0.1:1/nope.png"))];
|
||||||
|
let err = prefetch_all(&mut requests).await.err().unwrap();
|
||||||
|
assert!(err.contains("media fetch"), "{err}");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The item budget rejects before any source is touched: all of these would
|
||||||
|
/// fail to fetch, so a fetch error would prove fetching started.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn item_budget_rejects_before_fetching() {
|
||||||
|
let sources: Vec<Value> = (0..=MAX_ITEMS_PER_REQUEST)
|
||||||
|
.map(|i| Value::from(format!("/definitely/not/here-{i}.png")))
|
||||||
|
.collect();
|
||||||
|
let mut requests = vec![mm_request(Value::Array(sources))];
|
||||||
|
let err = prefetch_all(&mut requests).await.err().unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
err,
|
||||||
|
format!("multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items")
|
||||||
|
);
|
||||||
|
assert!(requests[0].mm.as_ref().unwrap().prefetched.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Sources legal alone but collectively over the limit are rejected while
|
||||||
|
/// downloading, not once every body is resident.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn byte_budget_is_shared_across_sources() {
|
||||||
|
let addr = serve(vec![vec![b'a'; 4096], vec![b'b'; 4096]]);
|
||||||
|
let sources = vec![
|
||||||
|
format!("http://{addr}/a.png"),
|
||||||
|
format!("http://{addr}/b.png"),
|
||||||
|
];
|
||||||
|
// Room for one body, not both.
|
||||||
|
let err = fetch_ordered(sources, 6144).await.err().unwrap();
|
||||||
|
assert!(err.contains("request media byte budget"), "{err}");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// ...and a fitting set still fetches: the budget never over-rejects.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn byte_budget_admits_a_fitting_request() {
|
||||||
|
let addr = serve(vec![vec![b'a'; 4096], vec![b'b'; 4096]]);
|
||||||
|
let sources = vec![
|
||||||
|
format!("http://{addr}/a.png"),
|
||||||
|
format!("http://{addr}/b.png"),
|
||||||
|
];
|
||||||
|
let fetched = fetch_ordered(sources, MAX_REQUEST_BYTES).await.unwrap();
|
||||||
|
assert_eq!(fetched.iter().map(|b| b.len()).sum::<usize>(), 8192);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -14,8 +14,6 @@
|
|||||||
|
|
||||||
use crate::error::Error;
|
use crate::error::Error;
|
||||||
|
|
||||||
// `Failed(Error)` carries the cause for observability even where it isn't read
|
|
||||||
// back yet; `EncodeDone` belongs to the deferred Encoder edge.
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub enum RequestState {
|
pub enum RequestState {
|
||||||
Received,
|
Received,
|
||||||
@@ -41,8 +39,8 @@ pub enum RequestState {
|
|||||||
/// Outcome of validation, selecting the ingress branch.
|
/// Outcome of validation, selecting the ingress branch.
|
||||||
#[derive(Debug, Clone, Copy)]
|
#[derive(Debug, Clone, Copy)]
|
||||||
pub enum ValidationOutcome {
|
pub enum ValidationOutcome {
|
||||||
/// Has multimodal inputs → Encoding. Deferred: no encoder yet.
|
/// Has multimodal inputs → Encoding, where an MM worker runs the native
|
||||||
#[allow(dead_code)]
|
/// pipeline and returns the final expanded `input_ids`.
|
||||||
HasMultimodal,
|
HasMultimodal,
|
||||||
/// Plain text → Tokenizing.
|
/// Plain text → Tokenizing.
|
||||||
NeedsTokenize,
|
NeedsTokenize,
|
||||||
@@ -52,7 +50,6 @@ pub enum ValidationOutcome {
|
|||||||
|
|
||||||
/// Events that drive transitions. Each variant maps 1:1 to an edge in the
|
/// Events that drive transitions. Each variant maps 1:1 to an edge in the
|
||||||
/// design's transition table.
|
/// design's transition table.
|
||||||
#[allow(dead_code)] // EncodeDone is the deferred Encoder edge.
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub enum Event {
|
pub enum Event {
|
||||||
// --- ingress ---
|
// --- ingress ---
|
||||||
@@ -125,7 +122,11 @@ impl RequestState {
|
|||||||
(Normalizing, Validated(HasMultimodal)) => Encoding,
|
(Normalizing, Validated(HasMultimodal)) => Encoding,
|
||||||
(Normalizing, Validated(NeedsTokenize)) => Tokenizing,
|
(Normalizing, Validated(NeedsTokenize)) => Tokenizing,
|
||||||
(Normalizing, Validated(AlreadyTokenized)) => PreSendValidating,
|
(Normalizing, Validated(AlreadyTokenized)) => PreSendValidating,
|
||||||
(Encoding, EncodeDone) => Tokenizing,
|
// The MM worker returns the *final* placeholder-expanded input_ids,
|
||||||
|
// so an encoded request skips the tokenizer pool — but not the
|
||||||
|
// pre-send checks: expanded image tokens count against the same
|
||||||
|
// input + max_new_tokens ceiling as tokenized text.
|
||||||
|
(Encoding, EncodeDone) => PreSendValidating,
|
||||||
// Every ingress branch funnels through the pre-send checks, so they
|
// Every ingress branch funnels through the pre-send checks, so they
|
||||||
// run exactly once per request no matter how it got its ids.
|
// run exactly once per request no matter how it got its ids.
|
||||||
(Tokenizing, TokenizeDone) => PreSendValidating,
|
(Tokenizing, TokenizeDone) => PreSendValidating,
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ mod error;
|
|||||||
mod fsm;
|
mod fsm;
|
||||||
mod ids;
|
mod ids;
|
||||||
mod message;
|
mod message;
|
||||||
|
mod mm;
|
||||||
mod ring;
|
mod ring;
|
||||||
mod runtime;
|
mod runtime;
|
||||||
mod tokenizer;
|
mod tokenizer;
|
||||||
@@ -31,6 +32,21 @@ use pyo3::types::PyBytes;
|
|||||||
|
|
||||||
use crate::runtime::{Runtime, RuntimeConfig};
|
use crate::runtime::{Runtime, RuntimeConfig};
|
||||||
|
|
||||||
|
/// One drained MM result (see [`Server::take_mm`]). Exactly one of
|
||||||
|
/// `features`/`shm_names` is `Some`: inline features for single-rank serving
|
||||||
|
/// (zero-copy into numpy), or one POSIX segment name per item when the scheduler
|
||||||
|
/// broadcasts across TP ranks and Python wraps each in a `ShmPointerMMData`.
|
||||||
|
#[pyclass(frozen, get_all)]
|
||||||
|
struct MmHandoff {
|
||||||
|
features: Option<Py<numpy::PyArray1<f32>>>,
|
||||||
|
shm_names: Option<Vec<String>>,
|
||||||
|
grids: Vec<(u32, u32, u32)>,
|
||||||
|
hashes: Vec<u64>,
|
||||||
|
offsets: Vec<(u32, u32)>,
|
||||||
|
mrope: Py<numpy::PyArray1<i64>>,
|
||||||
|
mrope_delta: i64,
|
||||||
|
}
|
||||||
|
|
||||||
/// Columnar ingress batch handed to Python by [`Server::recv_requests`].
|
/// Columnar ingress batch handed to Python by [`Server::recv_requests`].
|
||||||
/// `frozen`: immutable snapshot, so field access never contends on a borrow.
|
/// `frozen`: immutable snapshot, so field access never contends on a borrow.
|
||||||
#[pyclass(frozen, get_all)]
|
#[pyclass(frozen, get_all)]
|
||||||
@@ -199,6 +215,54 @@ impl Server {
|
|||||||
self.push_frame(py, crate::message::frame_egress_error(rid, message))
|
self.push_frame(py, crate::message::frame_egress_error(rid, message))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Spawn the MM worker pool for the pipeline in `spec_json` (built from the
|
||||||
|
/// resolved processor config; see `NativeMmHost.resolve_native_spec`).
|
||||||
|
/// Image-only requests are processed entirely in Rust and parked for
|
||||||
|
/// [`Server::take_mm`]; anything the pipeline cannot serve is rejected back to
|
||||||
|
/// the client — there is no Python fallback.
|
||||||
|
fn start_mm_workers(&self, spec_json: &str, workers: usize) -> PyResult<()> {
|
||||||
|
let ctx = mm::Context::new(
|
||||||
|
spec_json,
|
||||||
|
self.rt.tokenizer.clone(),
|
||||||
|
self.rt.mm_sidecar.clone(),
|
||||||
|
)
|
||||||
|
.map_err(PyErr::new::<pyo3::exceptions::PyValueError, _>)?;
|
||||||
|
self.rt.spawn_mm_pool(workers, std::sync::Arc::new(ctx));
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pop the MM result for `rid` — parked strictly before the request reached
|
||||||
|
/// the ingress ring — or `None` if there is none. The numeric buffers become
|
||||||
|
/// 1-D numpy arrays that take **ownership** of the Rust vectors, no copy.
|
||||||
|
///
|
||||||
|
/// Runs on the scheduler loop (`RustServer.drain`, under the GIL) between
|
||||||
|
/// decode steps, so any per-byte work here — memcpy or hashing, tens of MB
|
||||||
|
/// per image-heavy request — would stall every running request's ITL. Hence
|
||||||
|
/// the worker-precomputed `hashes`.
|
||||||
|
fn take_mm(&self, py: Python<'_>, rid: &str) -> Option<MmHandoff> {
|
||||||
|
use numpy::IntoPyArray;
|
||||||
|
|
||||||
|
let res = self.rt.mm_sidecar.take(rid)?;
|
||||||
|
let (features, shm_names) = match res.features {
|
||||||
|
mm::FeatureStore::Inline(v) => (Some(v.into_pyarray(py).unbind()), None),
|
||||||
|
// The segments — and the duty to unlink — move to Python here;
|
||||||
|
// `materialize()` unlinks after the post-broadcast clone on each rank.
|
||||||
|
mm::FeatureStore::Shm(segments) => (
|
||||||
|
None,
|
||||||
|
Some(segments.into_iter().map(|s| s.into_name()).collect()),
|
||||||
|
),
|
||||||
|
};
|
||||||
|
Some(MmHandoff {
|
||||||
|
features,
|
||||||
|
shm_names,
|
||||||
|
grids: res.grids.iter().map(|g| (g[0], g[1], g[2])).collect(),
|
||||||
|
hashes: res.hashes,
|
||||||
|
offsets: res.offsets,
|
||||||
|
mrope: res.mrope.into_pyarray(py).unbind(),
|
||||||
|
mrope_delta: res.mrope_delta,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
/// Signal all threads to stop (best effort).
|
/// Signal all threads to stop (best effort).
|
||||||
fn shutdown(&self) {
|
fn shutdown(&self) {
|
||||||
self.rt.request_shutdown();
|
self.rt.request_shutdown();
|
||||||
@@ -245,5 +309,6 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
|||||||
.try_init();
|
.try_init();
|
||||||
m.add_class::<Server>()?;
|
m.add_class::<Server>()?;
|
||||||
m.add_class::<IngressBatch>()?;
|
m.add_class::<IngressBatch>()?;
|
||||||
|
m.add_class::<MmHandoff>()?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,6 +11,7 @@
|
|||||||
mod egress;
|
mod egress;
|
||||||
mod finish_reason;
|
mod finish_reason;
|
||||||
mod io_struct;
|
mod io_struct;
|
||||||
|
pub mod mm_payload;
|
||||||
mod request;
|
mod request;
|
||||||
mod sampling;
|
mod sampling;
|
||||||
mod types;
|
mod types;
|
||||||
@@ -22,7 +23,10 @@ pub use egress::{
|
|||||||
};
|
};
|
||||||
pub use finish_reason::Matched;
|
pub use finish_reason::Matched;
|
||||||
pub(crate) use io_struct::{AbortReq, ControlRequest, GetInternalStateReq};
|
pub(crate) use io_struct::{AbortReq, ControlRequest, GetInternalStateReq};
|
||||||
pub use request::{GenerateBody, GenerateRequest, RequestKind};
|
pub use request::{GenerateBody, GenerateRequest, MmRequest, MmWorkItem, RequestKind};
|
||||||
|
// Constructed directly only by tests: `api_server::prefetch` fills its
|
||||||
|
// `prefetched` field, everything else gets it packed inside a `GenerateRequest`.
|
||||||
|
pub use request::MmData;
|
||||||
pub(crate) use sampling::{SamplingParams, SamplingParamsInput};
|
pub(crate) use sampling::{SamplingParams, SamplingParamsInput};
|
||||||
pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds};
|
pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds};
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,222 @@
|
|||||||
|
//! Convert a parked request's [`MmWorkItem`] into the typed [`MmInput`] the
|
||||||
|
//! `sglang-mm` driver consumes — an in-process handoff, nothing serialized.
|
||||||
|
//!
|
||||||
|
//! Every `Err` rejects the request back to the client; the message says whether
|
||||||
|
//! the input is malformed or merely outside the pipeline's scope (video/audio,
|
||||||
|
//! precomputed features, …).
|
||||||
|
|
||||||
|
use bytes::Bytes;
|
||||||
|
use rmpv::Value;
|
||||||
|
use sglang_mm::driver::{ImageSource, MmInput};
|
||||||
|
|
||||||
|
use super::request::MmWorkItem;
|
||||||
|
|
||||||
|
/// True for sources the API layer must resolve before MM dispatch: I/O — network
|
||||||
|
/// *or* disk, since a network mount can hang past any HTTP timeout — never runs
|
||||||
|
/// on the fixed MM worker pool (see `api_server::prefetch`). `data:` and bare
|
||||||
|
/// base64 are pure CPU and stay on the worker. Lives next to `collect_images` so
|
||||||
|
/// the prefetch walk and the parse walk cannot drift.
|
||||||
|
pub fn is_io_source(src: &str) -> bool {
|
||||||
|
src.starts_with("http://")
|
||||||
|
|| src.starts_with("https://")
|
||||||
|
|| src.starts_with("file://")
|
||||||
|
|| src.starts_with('/')
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The I/O-backed sources of an `image_data` value, in `collect_images` order.
|
||||||
|
pub fn io_sources(value: &Value) -> Vec<String> {
|
||||||
|
let mut out = Vec::new();
|
||||||
|
let mut walk = |value: &Value| {
|
||||||
|
if let Some(src) = value.as_str().filter(|s| is_io_source(s)) {
|
||||||
|
out.push(src.to_owned());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if let Value::Array(values) = value {
|
||||||
|
values.iter().for_each(&mut walk);
|
||||||
|
} else {
|
||||||
|
walk(value);
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
/// How many media items an `image_data` value contributes, walked the way
|
||||||
|
/// [`collect_images`] walks it, so the item budget can reject before fetching.
|
||||||
|
pub fn item_count(value: &Value) -> usize {
|
||||||
|
match value {
|
||||||
|
Value::Nil => 0,
|
||||||
|
Value::Array(values) => values.iter().map(item_count).sum(),
|
||||||
|
_ => 1,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// I/O-backed sources are swapped for their `work.prefetched` bytes (in
|
||||||
|
/// [`io_sources`] order); one left without an entry is an internal error here,
|
||||||
|
/// never a fetch.
|
||||||
|
pub fn to_mm_input(work: MmWorkItem) -> Result<MmInput, String> {
|
||||||
|
let present = |v: &Option<Value>| v.as_ref().is_some_and(value_present);
|
||||||
|
if present(&work.video_data) || present(&work.audio_data) {
|
||||||
|
return Err("unsupported modality: video/audio input".into());
|
||||||
|
}
|
||||||
|
let mut images = Vec::new();
|
||||||
|
if let Some(image_data) = &work.image_data {
|
||||||
|
collect_images(image_data, &mut work.prefetched.iter(), &mut images)?;
|
||||||
|
}
|
||||||
|
if images.is_empty() {
|
||||||
|
return Err("no raw image sources in mm input".into());
|
||||||
|
}
|
||||||
|
Ok(MmInput {
|
||||||
|
text: work.text,
|
||||||
|
input_ids: work.input_ids,
|
||||||
|
images,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn collect_images(
|
||||||
|
value: &Value,
|
||||||
|
prefetched: &mut std::slice::Iter<Bytes>,
|
||||||
|
out: &mut Vec<ImageSource>,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
match value {
|
||||||
|
Value::Nil => Ok(()),
|
||||||
|
Value::String(value) => {
|
||||||
|
let value = value
|
||||||
|
.as_str()
|
||||||
|
.ok_or_else(|| "non-utf8 image source".to_string())?;
|
||||||
|
if is_io_source(value) {
|
||||||
|
let bytes = prefetched
|
||||||
|
.next()
|
||||||
|
.ok_or_else(|| "I/O-backed image source was not prefetched".to_string())?;
|
||||||
|
out.push(ImageSource::Bytes(bytes.to_vec()));
|
||||||
|
} else {
|
||||||
|
out.push(ImageSource::String(value.to_owned()));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Value::Binary(value) => {
|
||||||
|
out.push(ImageSource::Bytes(value.clone()));
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Value::Array(values) => {
|
||||||
|
for value in values {
|
||||||
|
match value {
|
||||||
|
Value::String(_) | Value::Binary(_) | Value::Nil => {
|
||||||
|
collect_images(value, prefetched, out)?
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
return Err("unsupported image_data shape: nested/typed item".into());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
_ => Err("unsupported image_data shape".into()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rust mirror of Python `has_valid_data`: `nil` and (recursively) empty or
|
||||||
|
/// all-nil lists don't count as multimodal input. Shared with the ingress
|
||||||
|
/// `has_multimodal` check so routing and parsing cannot drift.
|
||||||
|
pub fn value_present(value: &Value) -> bool {
|
||||||
|
match value {
|
||||||
|
Value::Nil => false,
|
||||||
|
Value::Array(values) => values.iter().any(value_present),
|
||||||
|
_ => true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn image_work(image: Value) -> MmWorkItem {
|
||||||
|
MmWorkItem {
|
||||||
|
text: Some("prompt".into()),
|
||||||
|
image_data: Some(image),
|
||||||
|
..Default::default()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn converts_string_and_list_images() {
|
||||||
|
let one = to_mm_input(image_work(Value::from("data:image/png;base64,x"))).unwrap();
|
||||||
|
assert_eq!(one.images.len(), 1);
|
||||||
|
let many = to_mm_input(image_work(Value::Array(vec![
|
||||||
|
Value::from("a"),
|
||||||
|
Value::from("b"),
|
||||||
|
])))
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(many.images.len(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn unsupported_modalities_and_shapes_rejected() {
|
||||||
|
let video = MmWorkItem {
|
||||||
|
video_data: Some(Value::from("video.mp4")),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
assert!(to_mm_input(video).err().unwrap().contains("video/audio"));
|
||||||
|
|
||||||
|
let dict = Value::Map(vec![(Value::from("format"), Value::from("x"))]);
|
||||||
|
assert!(
|
||||||
|
to_mm_input(image_work(Value::Array(vec![dict])))
|
||||||
|
.err()
|
||||||
|
.unwrap()
|
||||||
|
.contains("image_data shape")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn empty_video_audio_lists_are_not_modalities() {
|
||||||
|
// Mirrors Python `has_valid_data`: nil / empty lists don't count.
|
||||||
|
let work = MmWorkItem {
|
||||||
|
input_ids: Some(vec![1]),
|
||||||
|
image_data: Some(Value::from("a")),
|
||||||
|
video_data: Some(Value::Array(vec![])),
|
||||||
|
audio_data: Some(Value::Array(vec![Value::Array(vec![])])),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
assert_eq!(to_mm_input(work).unwrap().images.len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// I/O-backed sources (URLs, file paths) take their prefetched bytes in walk
|
||||||
|
/// order; one left unfetched errors, so no I/O can reach an MM worker.
|
||||||
|
#[test]
|
||||||
|
fn io_sources_use_prefetched_bytes() {
|
||||||
|
let image = Value::Array(vec![
|
||||||
|
Value::from("http://a/x.png"),
|
||||||
|
Value::from("data:image/png;base64,x"),
|
||||||
|
Value::from("/mnt/nfs/y.png"),
|
||||||
|
]);
|
||||||
|
assert_eq!(io_sources(&image), vec!["http://a/x.png", "/mnt/nfs/y.png"]);
|
||||||
|
|
||||||
|
let mut work = image_work(image.clone());
|
||||||
|
work.prefetched = vec![Bytes::from_static(b"aa"), Bytes::from_static(b"bb")];
|
||||||
|
let input = to_mm_input(work).unwrap();
|
||||||
|
let as_bytes = |i: usize| match &input.images[i] {
|
||||||
|
ImageSource::Bytes(b) => b.as_slice(),
|
||||||
|
other => panic!("expected bytes, got {other:?}"),
|
||||||
|
};
|
||||||
|
assert_eq!(as_bytes(0), b"aa");
|
||||||
|
assert_eq!(as_bytes(2), b"bb");
|
||||||
|
assert!(matches!(&input.images[1], ImageSource::String(_)));
|
||||||
|
|
||||||
|
let err = to_mm_input(image_work(image)).err().unwrap();
|
||||||
|
assert!(err.contains("not prefetched"), "{err}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn image_free_work_rejected() {
|
||||||
|
assert!(
|
||||||
|
to_mm_input(image_work(Value::Nil))
|
||||||
|
.err()
|
||||||
|
.unwrap()
|
||||||
|
.contains("no raw image sources")
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
to_mm_input(MmWorkItem::default())
|
||||||
|
.err()
|
||||||
|
.unwrap()
|
||||||
|
.contains("no raw image sources")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -106,6 +106,23 @@ pub struct GenerateBody {
|
|||||||
pub routed_dp_rank: Option<i64>,
|
pub routed_dp_rank: Option<i64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub disagg_prefill_dp_rank: Option<i64>,
|
pub disagg_prefill_dp_rank: Option<i64>,
|
||||||
|
// Multimodal inputs, permissive `Value` so any shape Python's
|
||||||
|
// `GenerateReqInput` accepts (URL / base64 / list / list-of-lists) parses.
|
||||||
|
// `into_requests` fans them out per the Python
|
||||||
|
// `_normalize_{image,video,audio}_data` batch rules.
|
||||||
|
#[serde(default)]
|
||||||
|
pub image_data: Option<rmpv::Value>,
|
||||||
|
/// Caller-supplied per-item content hashes (hex) overriding the computed
|
||||||
|
/// ones, so an external router's keys align with the prefix cache. Single
|
||||||
|
/// requests only: Python declares the batched shapes but `__getitem__` never
|
||||||
|
/// forwards them, so a batch is rejected here rather than answered with
|
||||||
|
/// hashes it did not ask for.
|
||||||
|
#[serde(default)]
|
||||||
|
pub mm_hashes: Option<rmpv::Value>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub video_data: Option<rmpv::Value>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub audio_data: Option<rmpv::Value>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl GenerateBody {
|
impl GenerateBody {
|
||||||
@@ -135,6 +152,10 @@ impl GenerateBody {
|
|||||||
decode_tp_size,
|
decode_tp_size,
|
||||||
routed_dp_rank,
|
routed_dp_rank,
|
||||||
disagg_prefill_dp_rank,
|
disagg_prefill_dp_rank,
|
||||||
|
image_data,
|
||||||
|
video_data,
|
||||||
|
audio_data,
|
||||||
|
mm_hashes,
|
||||||
// Unported `GenerateReqInput` fields land here and are dropped, as they
|
// Unported `GenerateReqInput` fields land here and are dropped, as they
|
||||||
// are on the Python path.
|
// are on the Python path.
|
||||||
..
|
..
|
||||||
@@ -320,10 +341,26 @@ impl GenerateBody {
|
|||||||
let bootstrap_pair_keys =
|
let bootstrap_pair_keys =
|
||||||
flatten_column(fan_out(bootstrap_pair_key, n, "bootstrap_pair_key")?);
|
flatten_column(fan_out(bootstrap_pair_key, n, "bootstrap_pair_key")?);
|
||||||
let decode_tp_sizes = flatten_column(fan_out(decode_tp_size, n, "decode_tp_size")?);
|
let decode_tp_sizes = flatten_column(fan_out(decode_tp_size, n, "decode_tp_size")?);
|
||||||
|
// `mm_hashes` has no batch form: honoring it only here would give the two
|
||||||
|
// servers different prefix-cache keys for the same body. Reject instead of
|
||||||
|
// dropping it silently as Python does — the field exists to align a
|
||||||
|
// caller's keys, so ignoring it returns subtly wrong ones.
|
||||||
|
if is_batch && mm_value_present(&mm_hashes) {
|
||||||
|
return Err(Error::Validation(
|
||||||
|
"mm_hashes is not supported for batch requests; send one request per prompt".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
// Multimodal columns; see `split_mm_column` for the Python parity rules.
|
||||||
|
let images = split_mm_column(image_data, n, is_batch, MmBroadcast::WrapInList)
|
||||||
|
.map_err(|e| Error::Validation(format!("image_data: {e}")))?;
|
||||||
|
let videos = split_mm_column(video_data, n, is_batch, MmBroadcast::AsIs)
|
||||||
|
.map_err(|e| Error::Validation(format!("video_data: {e}")))?;
|
||||||
|
let audios = split_mm_column(audio_data, n, is_batch, MmBroadcast::AsIs)
|
||||||
|
.map_err(|e| Error::Validation(format!("audio_data: {e}")))?;
|
||||||
|
|
||||||
// Every column above is exactly `n` long, so zip them by value: each
|
// Every column above is exactly `n` long, so zip them by value: each
|
||||||
// request takes ownership of its cell, with no indexing or bounds checks.
|
// request takes ownership of its cell, with no indexing or bounds checks.
|
||||||
let requests = izip!(
|
let mut requests: Vec<GenerateRequest> = izip!(
|
||||||
rids,
|
rids,
|
||||||
texts,
|
texts,
|
||||||
id_lists,
|
id_lists,
|
||||||
@@ -338,6 +375,9 @@ impl GenerateBody {
|
|||||||
bootstrap_rooms,
|
bootstrap_rooms,
|
||||||
bootstrap_pair_keys,
|
bootstrap_pair_keys,
|
||||||
decode_tp_sizes,
|
decode_tp_sizes,
|
||||||
|
images,
|
||||||
|
videos,
|
||||||
|
audios,
|
||||||
)
|
)
|
||||||
.map(
|
.map(
|
||||||
|(
|
|(
|
||||||
@@ -355,6 +395,9 @@ impl GenerateBody {
|
|||||||
bootstrap_room,
|
bootstrap_room,
|
||||||
bootstrap_pair_key,
|
bootstrap_pair_key,
|
||||||
decode_tp_size,
|
decode_tp_size,
|
||||||
|
image_data,
|
||||||
|
video_data,
|
||||||
|
audio_data,
|
||||||
)| GenerateRequest {
|
)| GenerateRequest {
|
||||||
rid,
|
rid,
|
||||||
text,
|
text,
|
||||||
@@ -381,13 +424,126 @@ impl GenerateBody {
|
|||||||
decode_tp_size,
|
decode_tp_size,
|
||||||
routed_dp_rank,
|
routed_dp_rank,
|
||||||
disagg_prefill_dp_rank,
|
disagg_prefill_dp_rank,
|
||||||
|
mm: pack_mm(image_data, video_data, audio_data),
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
.collect();
|
.collect();
|
||||||
|
// Single requests only (batches rejected above). Malformed entries are
|
||||||
|
// dropped here and warned about in `mm::apply_caller_hashes`, never a 400.
|
||||||
|
if !is_batch
|
||||||
|
&& let (Some(rmpv::Value::Array(vals)), Some(req)) = (mm_hashes, requests.first_mut())
|
||||||
|
&& let Some(mm) = req.mm.as_deref_mut()
|
||||||
|
{
|
||||||
|
mm.mm_hashes = vals
|
||||||
|
.iter()
|
||||||
|
.filter_map(|v| v.as_str().map(str::to_owned))
|
||||||
|
.collect();
|
||||||
|
}
|
||||||
Ok((requests, is_batch))
|
Ok((requests, is_batch))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Box the per-item mm values, `None` when the item has none — the common
|
||||||
|
/// text-only case keeps `GenerateRequest` slim.
|
||||||
|
fn pack_mm(
|
||||||
|
image_data: Option<rmpv::Value>,
|
||||||
|
video_data: Option<rmpv::Value>,
|
||||||
|
audio_data: Option<rmpv::Value>,
|
||||||
|
) -> Option<Box<MmData>> {
|
||||||
|
if image_data.is_none() && video_data.is_none() && audio_data.is_none() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some(Box::new(MmData {
|
||||||
|
image_data,
|
||||||
|
video_data,
|
||||||
|
audio_data,
|
||||||
|
..Default::default()
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// How a scalar mm value broadcasts across a batch: images become a one-image
|
||||||
|
/// list per item (`[[img]] * num` in Python `_normalize_image_data`),
|
||||||
|
/// video/audio broadcast bare (`[v] * num` in `_normalize_video_data`).
|
||||||
|
#[derive(Clone, Copy)]
|
||||||
|
enum MmBroadcast {
|
||||||
|
WrapInList,
|
||||||
|
AsIs,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fan one mm field into per-item values, mirroring Python's
|
||||||
|
/// `_normalize_{image,video,audio}_data`:
|
||||||
|
/// * `None` / empty list → `None` for every item;
|
||||||
|
/// * single request → the raw value passes through (the processor wraps a
|
||||||
|
/// non-list into a one-element list);
|
||||||
|
/// * batch + scalar → broadcast to every item, per `MmBroadcast`;
|
||||||
|
/// * batch + list → per item, length must equal the batch size.
|
||||||
|
fn split_mm_column(
|
||||||
|
v: Option<rmpv::Value>,
|
||||||
|
n: usize,
|
||||||
|
is_batch: bool,
|
||||||
|
broadcast: MmBroadcast,
|
||||||
|
) -> Result<Vec<Option<rmpv::Value>>, String> {
|
||||||
|
let Some(v) = v else {
|
||||||
|
return Ok(vec![None; n]);
|
||||||
|
};
|
||||||
|
if v.is_nil() {
|
||||||
|
return Ok(vec![None; n]);
|
||||||
|
}
|
||||||
|
if !is_batch {
|
||||||
|
return Ok(vec![Some(v)]);
|
||||||
|
}
|
||||||
|
match v {
|
||||||
|
rmpv::Value::Array(items) if items.is_empty() => Ok(vec![None; n]),
|
||||||
|
rmpv::Value::Array(items) => {
|
||||||
|
if items.len() != n {
|
||||||
|
return Err(format!(
|
||||||
|
"list length {} does not match batch size {n}",
|
||||||
|
items.len()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(items.into_iter().map(Some).collect())
|
||||||
|
}
|
||||||
|
scalar => {
|
||||||
|
// A broadcast deep-clones once per prompt — same blow-up as
|
||||||
|
// sampling_params above, so bound the product before any clone.
|
||||||
|
check_broadcast_budget(scalar.heap_bytes(), n, "value").map_err(|e| e.to_string())?;
|
||||||
|
Ok(match broadcast {
|
||||||
|
MmBroadcast::WrapInList => vec![Some(rmpv::Value::Array(vec![scalar])); n],
|
||||||
|
MmBroadcast::AsIs => vec![Some(scalar); n],
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One request handed to the MM worker pool: the rid to correlate the result,
|
||||||
|
/// plus the owned inputs from [`GenerateRequest::take_mm_work`].
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct MmRequest {
|
||||||
|
pub rid: crate::ids::Rid,
|
||||||
|
pub work: MmWorkItem,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The parked request's fields the MM worker owns; converted to the driver input
|
||||||
|
/// by [`super::mm_payload::to_mm_input`].
|
||||||
|
#[derive(Debug, Default)]
|
||||||
|
pub struct MmWorkItem {
|
||||||
|
pub text: Option<String>,
|
||||||
|
pub input_ids: Option<Vec<i32>>,
|
||||||
|
pub image_data: Option<rmpv::Value>,
|
||||||
|
pub video_data: Option<rmpv::Value>,
|
||||||
|
pub audio_data: Option<rmpv::Value>,
|
||||||
|
/// See [`MmData::prefetched`].
|
||||||
|
pub prefetched: Vec<Bytes>,
|
||||||
|
/// See [`GenerateBody::mm_hashes`].
|
||||||
|
pub mm_hashes: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether an optional mm field counts as multimodal input, via the same
|
||||||
|
/// `value_present` the MM worker's payload parser uses.
|
||||||
|
fn mm_value_present(v: &Option<rmpv::Value>) -> bool {
|
||||||
|
v.as_ref().is_some_and(super::mm_payload::value_present)
|
||||||
|
}
|
||||||
|
|
||||||
/// Request variant — selects the ingress branch, scheduler wire message, and
|
/// Request variant — selects the ingress branch, scheduler wire message, and
|
||||||
/// egress shape. Each owns its body, so generate/control fields stay type-separate.
|
/// egress shape. Each owns its body, so generate/control fields stay type-separate.
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
@@ -475,6 +631,26 @@ pub struct GenerateRequest {
|
|||||||
/// so these are pure passthrough for the scheduler/LB protocol.
|
/// so these are pure passthrough for the scheduler/LB protocol.
|
||||||
pub routed_dp_rank: Option<i64>,
|
pub routed_dp_rank: Option<i64>,
|
||||||
pub disagg_prefill_dp_rank: Option<i64>,
|
pub disagg_prefill_dp_rank: Option<i64>,
|
||||||
|
/// Multimodal inputs, carried opaquely. Consumed by the Encoding stage,
|
||||||
|
/// which ships them to the MM worker pool; never read by the tokenizer or
|
||||||
|
/// serialized onto the scheduler header. Boxed so the common text-only
|
||||||
|
/// request doesn't grow every `Request` moved between stages.
|
||||||
|
pub mm: Option<Box<MmData>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The opaque multimodal fields of one request (see [`GenerateRequest::mm`]).
|
||||||
|
#[derive(Debug, Default)]
|
||||||
|
pub struct MmData {
|
||||||
|
pub image_data: Option<rmpv::Value>,
|
||||||
|
pub video_data: Option<rmpv::Value>,
|
||||||
|
pub audio_data: Option<rmpv::Value>,
|
||||||
|
/// Bytes of `image_data`'s I/O-backed sources, resolved by
|
||||||
|
/// `api_server::prefetch` in `mm_payload::io_sources` order so MM workers
|
||||||
|
/// never block on I/O. Out-of-band: the values above stay as the client
|
||||||
|
/// sent them.
|
||||||
|
pub prefetched: Vec<bytes::Bytes>,
|
||||||
|
/// See [`GenerateBody::mm_hashes`]; applied by the MM worker.
|
||||||
|
pub mm_hashes: Vec<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl GenerateRequest {
|
impl GenerateRequest {
|
||||||
@@ -483,11 +659,33 @@ impl GenerateRequest {
|
|||||||
self.input_ids.as_ref().is_some_and(|v| !v.is_empty())
|
self.input_ids.as_ref().is_some_and(|v| !v.is_empty())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Multimodal detection hook. Deferred (Encoder stubbed): always false until mm
|
/// True when the request carries a usable multimodal payload — the mirror of
|
||||||
/// fields are wired in.
|
/// Python `GenerateReqInput.contains_mm_input()`.
|
||||||
#[allow(dead_code)]
|
|
||||||
pub fn has_multimodal(&self) -> bool {
|
pub fn has_multimodal(&self) -> bool {
|
||||||
false
|
self.mm.as_ref().is_some_and(|mm| {
|
||||||
|
mm_value_present(&mm.image_data)
|
||||||
|
|| mm_value_present(&mm.video_data)
|
||||||
|
|| mm_value_present(&mm.audio_data)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Carve out the MM worker's inputs: `text` is cloned (the scheduler header
|
||||||
|
/// still needs it), `input_ids` is taken (the expanded ids replace it), and
|
||||||
|
/// the mm values move wholesale.
|
||||||
|
pub fn take_mm_work(&mut self) -> MmWorkItem {
|
||||||
|
let mut work = MmWorkItem {
|
||||||
|
text: self.text.clone(),
|
||||||
|
input_ids: self.input_ids.take(),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
if let Some(m) = self.mm.as_deref_mut() {
|
||||||
|
work.image_data = m.image_data.take();
|
||||||
|
work.video_data = m.video_data.take();
|
||||||
|
work.audio_data = m.audio_data.take();
|
||||||
|
work.prefetched = std::mem::take(&mut m.prefetched);
|
||||||
|
work.mm_hashes = std::mem::take(&mut m.mm_hashes);
|
||||||
|
}
|
||||||
|
work
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn encode_header(&self) -> Result<Bytes, Error> {
|
pub fn encode_header(&self) -> Result<Bytes, Error> {
|
||||||
@@ -539,6 +737,23 @@ impl<T: HeapBytes> HeapBytes for Option<T> {
|
|||||||
self.as_ref().map_or(0, HeapBytes::heap_bytes)
|
self.as_ref().map_or(0, HeapBytes::heap_bytes)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
impl HeapBytes for rmpv::Value {
|
||||||
|
fn heap_bytes(&self) -> usize {
|
||||||
|
use rmpv::Value;
|
||||||
|
const NODE: usize = std::mem::size_of::<rmpv::Value>();
|
||||||
|
match self {
|
||||||
|
Value::String(s) => s.as_bytes().len(),
|
||||||
|
Value::Binary(b) => b.len(),
|
||||||
|
Value::Ext(_, b) => b.len(),
|
||||||
|
Value::Array(items) => items.iter().map(|v| NODE + v.heap_bytes()).sum(),
|
||||||
|
Value::Map(entries) => entries
|
||||||
|
.iter()
|
||||||
|
.map(|(k, v)| 2 * NODE + k.heap_bytes() + v.heap_bytes())
|
||||||
|
.sum(),
|
||||||
|
_ => 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Collapse `fan_out`'s nullable-element output: outer `None` (field absent /
|
/// Collapse `fan_out`'s nullable-element output: outer `None` (field absent /
|
||||||
/// scalar broadcast of nothing) and inner `None` (an explicit `null` list
|
/// scalar broadcast of nothing) and inner `None` (an explicit `null` list
|
||||||
@@ -743,8 +958,9 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// The native `bench_serving` payload (a `GenerateReqInput` superset) parses:
|
/// The native `bench_serving` payload (a `GenerateReqInput` superset) parses:
|
||||||
/// its `lora_path`/`return_routed_experts`/`image_data` are accepted-but-ignored,
|
/// its `lora_path`/`return_routed_experts` are accepted-but-ignored and a
|
||||||
/// so `split` succeeds and drops them while the real fields survive.
|
/// `null` `image_data` means "no multimodal input", so `split` succeeds
|
||||||
|
/// while the real fields survive.
|
||||||
#[test]
|
#[test]
|
||||||
fn accepts_bench_serving_payload() {
|
fn accepts_bench_serving_payload() {
|
||||||
let (ps, is_batch) = requests(
|
let (ps, is_batch) = requests(
|
||||||
@@ -758,6 +974,111 @@ mod tests {
|
|||||||
assert_eq!(ps.len(), 1);
|
assert_eq!(ps.len(), 1);
|
||||||
assert_eq!(ps[0].text.as_deref(), Some("hi"));
|
assert_eq!(ps[0].text.as_deref(), Some("hi"));
|
||||||
assert!(ps[0].stream);
|
assert!(ps[0].stream);
|
||||||
|
assert!(!ps[0].has_multimodal());
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mm columns fan out per Python `_normalize_{image,video}_data`: a single
|
||||||
|
/// request passes the raw value through; a batch broadcasts a scalar image as
|
||||||
|
/// `[img]` per item, maps a list per item with matching lengths, and treats
|
||||||
|
/// `null`/`[]` as absent.
|
||||||
|
#[test]
|
||||||
|
fn split_mm_fanout_matches_python_normalize() {
|
||||||
|
let image_of = |p: &GenerateRequest| p.mm.as_ref().unwrap().image_data.clone().unwrap();
|
||||||
|
|
||||||
|
// Single request: raw value passes through untouched.
|
||||||
|
let (ps, _) = requests(r#"{"text": "a", "image_data": "http://x/i.jpg"}"#).unwrap();
|
||||||
|
assert_eq!(image_of(&ps[0]).as_str(), Some("http://x/i.jpg"));
|
||||||
|
assert!(ps[0].has_multimodal());
|
||||||
|
|
||||||
|
// Batch + scalar image: broadcast, wrapped as a one-image list per item.
|
||||||
|
let (ps, _) = requests(r#"{"text": ["a", "b"], "image_data": "u"}"#).unwrap();
|
||||||
|
for p in &ps {
|
||||||
|
assert_eq!(image_of(p).as_array().unwrap().len(), 1);
|
||||||
|
assert!(p.has_multimodal());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Batch + per-item list: element i goes to item i.
|
||||||
|
let (ps, _) = requests(r#"{"text": ["a", "b"], "image_data": ["u1", "u2"]}"#).unwrap();
|
||||||
|
assert_eq!(image_of(&ps[0]).as_str(), Some("u1"));
|
||||||
|
assert_eq!(image_of(&ps[1]).as_str(), Some("u2"));
|
||||||
|
|
||||||
|
// Batch + wrong-length list is a 400.
|
||||||
|
assert!(requests(r#"{"text": ["a", "b"], "image_data": ["u1"]}"#).is_err());
|
||||||
|
|
||||||
|
// null / [] mean "no multimodal input".
|
||||||
|
let (ps, _) = requests(r#"{"text": "a", "image_data": null}"#).unwrap();
|
||||||
|
assert!(!ps[0].has_multimodal());
|
||||||
|
let (ps, _) = requests(r#"{"text": "a", "image_data": []}"#).unwrap();
|
||||||
|
assert!(!ps[0].has_multimodal());
|
||||||
|
|
||||||
|
// Batch + scalar video: broadcast bare (not wrapped), per Python
|
||||||
|
// `_normalize_video_data`.
|
||||||
|
let (ps, _) = requests(r#"{"text": ["a", "b"], "video_data": "v"}"#).unwrap();
|
||||||
|
let video = ps[1].mm.as_ref().unwrap().video_data.clone().unwrap();
|
||||||
|
assert_eq!(video.as_str(), Some("v"));
|
||||||
|
assert!(ps[1].has_multimodal());
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A scalar broadcast is budget-checked before the deep clones (16 MiB ×
|
||||||
|
/// 4096 prompts would be 64 GiB and an abort); per-item lists clone nothing
|
||||||
|
/// and are never charged.
|
||||||
|
#[test]
|
||||||
|
fn oversized_mm_broadcast_rejected() {
|
||||||
|
let big = rmpv::Value::from("x".repeat(MAX_BROADCAST_CLONE_BYTES / 2 + 1));
|
||||||
|
let err = split_mm_column(Some(big.clone()), 2, true, MmBroadcast::WrapInList)
|
||||||
|
.err()
|
||||||
|
.unwrap();
|
||||||
|
assert!(err.contains("broadcast"), "{err}");
|
||||||
|
// A per-item list of the same total size moves, not clones: accepted.
|
||||||
|
let list = rmpv::Value::Array(vec![big, rmpv::Value::from("y")]);
|
||||||
|
assert!(split_mm_column(Some(list), 2, true, MmBroadcast::WrapInList).is_ok());
|
||||||
|
// Small scalars broadcast fine.
|
||||||
|
let small = rmpv::Value::from("u1");
|
||||||
|
assert!(split_mm_column(Some(small), 2, true, MmBroadcast::AsIs).is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `mm_hashes` rides only on single requests (Python `__getitem__`
|
||||||
|
/// parity: batches drop it) and moves into the work item.
|
||||||
|
#[test]
|
||||||
|
fn mm_hashes_single_only() {
|
||||||
|
let (mut ps, _) =
|
||||||
|
requests(r#"{"text": "a", "image_data": "u", "mm_hashes": ["a1b2", "0xff"]}"#).unwrap();
|
||||||
|
assert_eq!(ps[0].mm.as_ref().unwrap().mm_hashes, vec!["a1b2", "0xff"]);
|
||||||
|
assert_eq!(ps[0].take_mm_work().mm_hashes, vec!["a1b2", "0xff"]);
|
||||||
|
assert!(ps[0].mm.as_ref().unwrap().mm_hashes.is_empty());
|
||||||
|
|
||||||
|
// A batch cannot carry hashes (Python drops them), so it is rejected...
|
||||||
|
for body in [
|
||||||
|
r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": [["x"], ["y"]]}"#,
|
||||||
|
r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": ["x", "y"]}"#,
|
||||||
|
] {
|
||||||
|
let err = requests(body).err().unwrap();
|
||||||
|
assert!(matches!(err, Error::Validation(_)), "{body}: {err:?}");
|
||||||
|
}
|
||||||
|
// ...while an absent or empty field is not a payload and must still pass.
|
||||||
|
for body in [
|
||||||
|
r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": null}"#,
|
||||||
|
r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": []}"#,
|
||||||
|
] {
|
||||||
|
assert!(requests(body).is_ok(), "{body}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `take_mm_work` clones `text` (the scheduler header still needs it) and
|
||||||
|
/// moves everything the worker owns out of the request.
|
||||||
|
#[test]
|
||||||
|
fn mm_work_item_takes_owned_fields() {
|
||||||
|
let (mut ps, _) =
|
||||||
|
requests(r#"{"text": "hi", "image_data": ["u1", "u2"], "audio_data": "a"}"#).unwrap();
|
||||||
|
let work = ps[0].take_mm_work();
|
||||||
|
assert_eq!(work.text.as_deref(), Some("hi"));
|
||||||
|
assert!(work.input_ids.is_none());
|
||||||
|
assert_eq!(work.image_data.unwrap().as_array().unwrap().len(), 2);
|
||||||
|
assert!(work.video_data.is_none());
|
||||||
|
assert_eq!(work.audio_data.unwrap().as_str(), Some("a"));
|
||||||
|
// Moved out, not cloned; `text` survives for the header.
|
||||||
|
assert!(ps[0].mm.as_ref().unwrap().image_data.is_none());
|
||||||
|
assert_eq!(ps[0].text.as_deref(), Some("hi"));
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The body limit is disabled, so an unbounded batch turns a small body into an
|
/// The body limit is disabled, so an unbounded batch turns a small body into an
|
||||||
|
|||||||
@@ -0,0 +1,413 @@
|
|||||||
|
//! Multimodal worker pool.
|
||||||
|
//!
|
||||||
|
//! Rust threads drain requests parked in `Encoding` and run the `sglang-mm`
|
||||||
|
//! pipeline registered by `Server.start_mm_workers` (decode → preprocess →
|
||||||
|
//! placeholder expansion → M-RoPE, GIL-free). Each worker parks the result
|
||||||
|
//! buffers in the rid-keyed [`Sidecar`] and returns only the expanded ids;
|
||||||
|
//! Python attaches the buffers at drain time (`Server.take_mm`). Inputs the
|
||||||
|
//! pipeline cannot serve are rejected to the client — no Python fallback.
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
|
use crate::message::MmRequest;
|
||||||
|
use crate::runtime::Runnable;
|
||||||
|
use crate::tokenizer::TextTokenizer;
|
||||||
|
use crate::tokenizer_manager::TmEvent;
|
||||||
|
|
||||||
|
/// A named POSIX shared-memory segment owning its name: dropped → unlinked.
|
||||||
|
///
|
||||||
|
/// Written by an MM worker so the TP broadcast carries a ~100-byte
|
||||||
|
/// `ShmPointerMMData` stub instead of the ~20 MB feature tensor, and every
|
||||||
|
/// rank maps it in parallel. Python's `materialize()` unlinks after cloning;
|
||||||
|
/// this `Drop` covers the paths where the buffers never reach Python (aborted
|
||||||
|
/// while parked, late result purged).
|
||||||
|
pub struct ShmSegment {
|
||||||
|
name: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ShmSegment {
|
||||||
|
/// Create `/dev/shm/{name}` holding exactly `bytes`. No leading slash —
|
||||||
|
/// the name must suit Python's `SharedMemory(name=…)` (shm_open adds one).
|
||||||
|
pub fn create(name: String, bytes: &[u8]) -> Result<Self, String> {
|
||||||
|
let c_name = std::ffi::CString::new(format!("/{name}"))
|
||||||
|
.map_err(|_| "shm name contains NUL".to_string())?;
|
||||||
|
// SAFETY: plain POSIX calls on a name we own; every handle created
|
||||||
|
// below is closed/unmapped on all paths.
|
||||||
|
unsafe {
|
||||||
|
let fd = libc::shm_open(
|
||||||
|
c_name.as_ptr(),
|
||||||
|
libc::O_CREAT | libc::O_EXCL | libc::O_RDWR,
|
||||||
|
0o600,
|
||||||
|
);
|
||||||
|
if fd < 0 {
|
||||||
|
return Err(format!(
|
||||||
|
"shm_open({name}): {}",
|
||||||
|
std::io::Error::last_os_error()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let segment = Self { name }; // unlink from here on any failure
|
||||||
|
if libc::ftruncate(fd, bytes.len() as libc::off_t) != 0 {
|
||||||
|
let e = std::io::Error::last_os_error();
|
||||||
|
libc::close(fd);
|
||||||
|
return Err(format!("ftruncate({}): {e}", segment.name));
|
||||||
|
}
|
||||||
|
let ptr = libc::mmap(
|
||||||
|
std::ptr::null_mut(),
|
||||||
|
bytes.len(),
|
||||||
|
libc::PROT_WRITE,
|
||||||
|
libc::MAP_SHARED,
|
||||||
|
fd,
|
||||||
|
0,
|
||||||
|
);
|
||||||
|
libc::close(fd);
|
||||||
|
if ptr == libc::MAP_FAILED {
|
||||||
|
return Err(format!(
|
||||||
|
"mmap({}): {}",
|
||||||
|
segment.name,
|
||||||
|
std::io::Error::last_os_error()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
std::ptr::copy_nonoverlapping(bytes.as_ptr(), ptr.cast::<u8>(), bytes.len());
|
||||||
|
libc::munmap(ptr, bytes.len());
|
||||||
|
Ok(segment)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Hand the segment — and the duty to unlink — to the caller (Python, at
|
||||||
|
/// drain time).
|
||||||
|
pub fn into_name(self) -> String {
|
||||||
|
std::mem::take(&mut std::mem::ManuallyDrop::new(self).name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for ShmSegment {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
if let Ok(c_name) = std::ffi::CString::new(format!("/{}", self.name)) {
|
||||||
|
// SAFETY: unlinking a name we created; ENOENT (already unlinked
|
||||||
|
// by Python's materialize) is fine to ignore.
|
||||||
|
unsafe { libc::shm_unlink(c_name.as_ptr()) };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Unique segment names: the pid separates server restarts (a crash can leak
|
||||||
|
/// segments under the old pid), the counter separates results within one.
|
||||||
|
fn shm_name(item: usize) -> String {
|
||||||
|
static COUNTER: AtomicU64 = AtomicU64::new(0);
|
||||||
|
let n = COUNTER.fetch_add(1, Ordering::Relaxed);
|
||||||
|
format!("sglmm-{}-{n}-{item}", std::process::id())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Python parity: caller hashes override the computed ones so an external
|
||||||
|
/// router's keys align with the prefix cache. A length mismatch or malformed
|
||||||
|
/// entry warns and keeps the computed hash — never blocks the request.
|
||||||
|
fn apply_caller_hashes(hashes: &mut [u64], caller: &[String]) {
|
||||||
|
if caller.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if caller.len() != hashes.len() {
|
||||||
|
tracing::warn!(
|
||||||
|
caller = caller.len(),
|
||||||
|
items = hashes.len(),
|
||||||
|
"mm_hashes length != mm item count; ignoring caller hashes"
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
for (hash, entry) in hashes.iter_mut().zip(caller) {
|
||||||
|
match parse_caller_hash(entry) {
|
||||||
|
Some(v) => *hash = v,
|
||||||
|
None => tracing::warn!(%entry, "malformed mm_hashes entry; keeping computed hash"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Hex of any width, as Python's `int(hex_hash, 16)` takes it (a full SHA-256
|
||||||
|
/// being the common case), keeping the low 64 bits — only the low 30 are
|
||||||
|
/// observable, through `_compute_pad_value`.
|
||||||
|
fn parse_caller_hash(entry: &str) -> Option<u64> {
|
||||||
|
let hex = entry.strip_prefix("0x").unwrap_or(entry);
|
||||||
|
if hex.is_empty() || !hex.bytes().all(|b| b.is_ascii_hexdigit()) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
u64::from_str_radix(&hex[hex.len().saturating_sub(16)..], 16).ok()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One parked result: the buffers the drain-time Python adapter needs (the
|
||||||
|
/// expanded `input_ids` travel separately, via `TmEvent::MmEncoded`). The qwen
|
||||||
|
/// drain shape (`sglang_mm::qwen_vl::pack_drain`); generalizes to a
|
||||||
|
/// named-tensor handoff once a family needs a different one.
|
||||||
|
pub struct MmSidecarEntry {
|
||||||
|
pub features: FeatureStore,
|
||||||
|
pub grids: Vec<[u32; 3]>,
|
||||||
|
pub hashes: Vec<u64>,
|
||||||
|
pub offsets: Vec<(u32, u32)>,
|
||||||
|
pub mrope: Vec<i64>,
|
||||||
|
pub mrope_delta: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Where a result's feature buffers live between worker and drain.
|
||||||
|
pub enum FeatureStore {
|
||||||
|
/// In-process; the drain wraps them zero-copy. Single-rank serving, or the
|
||||||
|
/// shm fallback. Under TP the whole buffer would ride `broadcast_pyobj`.
|
||||||
|
Inline(Vec<f32>),
|
||||||
|
/// One POSIX segment per item, written by the worker; only the names cross
|
||||||
|
/// ranks. See [`ShmSegment`].
|
||||||
|
Shm(Vec<ShmSegment>),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Results parked between a worker's `MmEncoded` and the scheduler drain, keyed
|
||||||
|
/// by rid. Owns the lifecycle so entries never leak: [`park`](Self::park)
|
||||||
|
/// strictly before `MmEncoded`, [`take`](Self::take) at the drain,
|
||||||
|
/// [`purge`](Self::purge) for requests that die while parked.
|
||||||
|
#[derive(Clone, Default)]
|
||||||
|
pub struct Sidecar(Arc<Mutex<HashMap<String, MmSidecarEntry>>>);
|
||||||
|
|
||||||
|
impl Sidecar {
|
||||||
|
pub fn park(&self, rid: String, entry: MmSidecarEntry) {
|
||||||
|
self.0.lock().unwrap().insert(rid, entry);
|
||||||
|
}
|
||||||
|
pub fn take(&self, rid: &str) -> Option<MmSidecarEntry> {
|
||||||
|
self.0.lock().unwrap().remove(rid)
|
||||||
|
}
|
||||||
|
pub fn purge(&self, rid: &str) {
|
||||||
|
self.0.lock().unwrap().remove(rid);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Shared state of the mm path, built once at `start_mm_workers`.
|
||||||
|
pub struct Context {
|
||||||
|
pub family: Box<dyn sglang_mm::pipeline::MmFamilyProcessor>,
|
||||||
|
/// `None` under `skip_tokenizer_init` (requests must carry `input_ids`).
|
||||||
|
pub tokenizer: Option<Arc<dyn TextTokenizer>>,
|
||||||
|
pub sidecar: Sidecar,
|
||||||
|
/// Park feature buffers in POSIX shm. Set by the Python launcher
|
||||||
|
/// (`NativeMmHost._use_feature_shm`) exactly when the scheduler broadcasts
|
||||||
|
/// across TP ranks and will unwrap `ShmPointerMMData`.
|
||||||
|
pub feature_shm: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Context {
|
||||||
|
pub fn new(
|
||||||
|
spec_json: &str,
|
||||||
|
tokenizer: Option<Arc<dyn TextTokenizer>>,
|
||||||
|
sidecar: Sidecar,
|
||||||
|
) -> Result<Self, String> {
|
||||||
|
let feature_shm = serde_json::from_str::<serde_json::Value>(spec_json)
|
||||||
|
.ok()
|
||||||
|
.and_then(|v| v.get("feature_shm").and_then(|b| b.as_bool()))
|
||||||
|
.unwrap_or(false);
|
||||||
|
Ok(Self {
|
||||||
|
family: sglang_mm::registry::pipeline_from_spec(spec_json)?,
|
||||||
|
tokenizer,
|
||||||
|
sidecar,
|
||||||
|
feature_shm,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run the pipeline for one request. `Ok` returns the final expanded ids, the
|
||||||
|
/// buffers already parked; `Err` rejects the request back to the client.
|
||||||
|
fn process(
|
||||||
|
ctx: &Context,
|
||||||
|
rid: &crate::ids::Rid,
|
||||||
|
mut work: crate::message::MmWorkItem,
|
||||||
|
) -> Result<Vec<i32>, String> {
|
||||||
|
let caller_hashes = std::mem::take(&mut work.mm_hashes);
|
||||||
|
let input = crate::message::mm_payload::to_mm_input(work)?;
|
||||||
|
let output = sglang_mm::driver::process(ctx.family.as_ref(), input, |text| {
|
||||||
|
let tokenizer = ctx.tokenizer.as_ref().ok_or_else(|| {
|
||||||
|
"skip_tokenizer_init is set: multimodal text prompts require input_ids".to_string()
|
||||||
|
})?;
|
||||||
|
tokenizer.encode(text).map_err(|error| error.to_string())
|
||||||
|
})?;
|
||||||
|
let mut drain = sglang_mm::qwen_vl::pack_drain(output)?;
|
||||||
|
apply_caller_hashes(&mut drain.hashes, &caller_hashes);
|
||||||
|
let features = if ctx.feature_shm {
|
||||||
|
park_features_in_shm(&drain.features, &drain.grids)
|
||||||
|
} else {
|
||||||
|
FeatureStore::Inline(drain.features)
|
||||||
|
};
|
||||||
|
ctx.sidecar.park(
|
||||||
|
rid.as_str().to_owned(),
|
||||||
|
MmSidecarEntry {
|
||||||
|
features,
|
||||||
|
grids: drain.grids,
|
||||||
|
hashes: drain.hashes,
|
||||||
|
offsets: drain.offsets,
|
||||||
|
mrope: drain.mrope,
|
||||||
|
mrope_delta: drain.mrope_delta,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
Ok(drain.input_ids)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Split the flat feature buffer per item (`t*h*w` rows per grid) and park each
|
||||||
|
/// slice in its own segment. Any shm failure (`/dev/shm` full, odd shape) falls
|
||||||
|
/// back to inline, as Python's `_wrap_shm_or_inline` does: degrade to the slow
|
||||||
|
/// path, never fail the request.
|
||||||
|
fn park_features_in_shm(features: &[f32], grids: &[[u32; 3]]) -> FeatureStore {
|
||||||
|
let total_rows: usize = grids
|
||||||
|
.iter()
|
||||||
|
.map(|g| g[0] as usize * g[1] as usize * g[2] as usize)
|
||||||
|
.sum();
|
||||||
|
if total_rows == 0 || !features.len().is_multiple_of(total_rows) {
|
||||||
|
return FeatureStore::Inline(features.to_vec());
|
||||||
|
}
|
||||||
|
let dim = features.len() / total_rows;
|
||||||
|
let mut segments = Vec::with_capacity(grids.len());
|
||||||
|
let mut row = 0usize;
|
||||||
|
for (item, grid) in grids.iter().enumerate() {
|
||||||
|
let rows = grid[0] as usize * grid[1] as usize * grid[2] as usize;
|
||||||
|
let slice = &features[row * dim..(row + rows) * dim];
|
||||||
|
row += rows;
|
||||||
|
match ShmSegment::create(shm_name(item), bytemuck::cast_slice(slice)) {
|
||||||
|
Ok(segment) => segments.push(segment),
|
||||||
|
Err(error) => {
|
||||||
|
tracing::warn!(%error, "mm: shm feature transport failed; falling back to inline");
|
||||||
|
return FeatureStore::Inline(features.to_vec());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
FeatureStore::Shm(segments)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One MM worker, spawned via `Runtime::spawn_mm_pool` (which owns the
|
||||||
|
/// pinning policy for this pool — see its docs).
|
||||||
|
pub struct MmWorker {
|
||||||
|
rx: flume::Receiver<MmRequest>,
|
||||||
|
tm: flume::Sender<TmEvent>,
|
||||||
|
ctx: Arc<Context>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MmWorker {
|
||||||
|
pub fn new(
|
||||||
|
rx: flume::Receiver<MmRequest>,
|
||||||
|
tm: flume::Sender<TmEvent>,
|
||||||
|
ctx: Arc<Context>,
|
||||||
|
) -> Self {
|
||||||
|
Self { rx, tm, ctx }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Runnable for MmWorker {
|
||||||
|
/// Drain until the mm channel closes (tm-ingress drops its sender on
|
||||||
|
/// shutdown). One request at a time, so the pool size bounds MM
|
||||||
|
/// concurrency; an error rejects the request back to the client.
|
||||||
|
fn run(self) {
|
||||||
|
while let Ok(req) = self.rx.recv() {
|
||||||
|
let rid = req.rid;
|
||||||
|
let event = match process(&self.ctx, &rid, req.work) {
|
||||||
|
Ok(input_ids) => {
|
||||||
|
tracing::debug!(%rid, tokens = input_ids.len(), "mm: processed");
|
||||||
|
TmEvent::MmEncoded { rid, input_ids }
|
||||||
|
}
|
||||||
|
Err(message) => {
|
||||||
|
tracing::warn!(%rid, %message, "mm processing rejected");
|
||||||
|
TmEvent::MmFailed { rid, message }
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if self.tm.send(event).is_err() {
|
||||||
|
return; // tm-ingress gone: shutdown
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
/// Caller hashes override computed ones; mismatched lengths and malformed
|
||||||
|
/// entries fall back per item, never reject (Python parity).
|
||||||
|
#[test]
|
||||||
|
fn caller_hashes_override_with_fallback() {
|
||||||
|
let mut hashes = vec![1, 2, 3];
|
||||||
|
apply_caller_hashes(&mut hashes, &[]);
|
||||||
|
assert_eq!(hashes, [1, 2, 3]);
|
||||||
|
|
||||||
|
apply_caller_hashes(&mut hashes, &["ff".into()]); // length mismatch
|
||||||
|
assert_eq!(hashes, [1, 2, 3]);
|
||||||
|
|
||||||
|
apply_caller_hashes(&mut hashes, &["ff".into(), "not-hex".into(), "0x10".into()]);
|
||||||
|
assert_eq!(hashes, [0xff, 2, 0x10]);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A full SHA-256 (what routers send) keeps its low 64 bits rather than
|
||||||
|
/// falling back, so the pad value matches Python's wide `int`.
|
||||||
|
#[test]
|
||||||
|
fn caller_hashes_accept_arbitrary_width() {
|
||||||
|
let sha256 = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
|
||||||
|
let mut hashes = vec![1];
|
||||||
|
apply_caller_hashes(&mut hashes, &[sha256.into()]);
|
||||||
|
assert_eq!(hashes, [0xa495991b7852b855]);
|
||||||
|
assert_eq!(hashes[0] % (1 << 30), 944_945_237); // int(sha256, 16) % (1 << 30)
|
||||||
|
|
||||||
|
// Width alone is never malformed; a non-hex digit still is.
|
||||||
|
assert_eq!(parse_caller_hash(&"f".repeat(64)), Some(u64::MAX));
|
||||||
|
assert_eq!(parse_caller_hash("0x"), None);
|
||||||
|
assert_eq!(parse_caller_hash(""), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn shm_path(name: &str) -> std::path::PathBuf {
|
||||||
|
std::path::Path::new("/dev/shm").join(name)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The segment holds exactly the written bytes and dropping it unlinks —
|
||||||
|
/// the leak guard for results purged before Python takes them.
|
||||||
|
#[test]
|
||||||
|
fn segment_roundtrip_and_drop_unlinks() {
|
||||||
|
let name = shm_name(0);
|
||||||
|
let payload: Vec<u8> = (0..255u8).collect();
|
||||||
|
let segment = ShmSegment::create(name.clone(), &payload).unwrap();
|
||||||
|
assert_eq!(std::fs::read(shm_path(&name)).unwrap(), payload);
|
||||||
|
drop(segment);
|
||||||
|
assert!(!shm_path(&name).exists(), "drop must unlink");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `into_name` transfers the unlink duty to the caller (Python's
|
||||||
|
/// `materialize()`), so the segment must survive the handoff.
|
||||||
|
#[test]
|
||||||
|
fn into_name_disarms_the_unlink() {
|
||||||
|
let segment = ShmSegment::create(shm_name(0), &[1, 2, 3]).unwrap();
|
||||||
|
let name = segment.into_name();
|
||||||
|
assert!(shm_path(&name).exists(), "handoff must not unlink");
|
||||||
|
// manual cleanup for the test
|
||||||
|
let c = std::ffi::CString::new(format!("/{name}")).unwrap();
|
||||||
|
unsafe { libc::shm_unlink(c.as_ptr()) };
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-item slicing follows the grid row counts, so Python's
|
||||||
|
/// `(rows, feature_dim)` reshape of a segment sees only its own item.
|
||||||
|
#[test]
|
||||||
|
fn park_splits_features_by_grid() {
|
||||||
|
// Two items: grids (1,2,2)=4 rows and (1,1,2)=2 rows, dim=3.
|
||||||
|
let features: Vec<f32> = (0..18).map(|i| i as f32).collect();
|
||||||
|
let grids = [[1, 2, 2], [1, 1, 2]];
|
||||||
|
let FeatureStore::Shm(segments) = park_features_in_shm(&features, &grids) else {
|
||||||
|
panic!("expected shm store");
|
||||||
|
};
|
||||||
|
assert_eq!(segments.len(), 2);
|
||||||
|
let read = |seg: &ShmSegment| -> Vec<u8> { std::fs::read(shm_path(&seg.name)).unwrap() };
|
||||||
|
assert_eq!(
|
||||||
|
read(&segments[0]),
|
||||||
|
bytemuck::cast_slice::<f32, u8>(&features[..12])
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
read(&segments[1]),
|
||||||
|
bytemuck::cast_slice::<f32, u8>(&features[12..])
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A degenerate shape must degrade to inline, never a shm-side panic.
|
||||||
|
#[test]
|
||||||
|
fn shape_surprise_falls_back_inline() {
|
||||||
|
let features = vec![0.0f32; 7]; // not divisible by 2 rows
|
||||||
|
let grids = [[1, 1, 2]];
|
||||||
|
assert!(matches!(
|
||||||
|
park_features_in_shm(&features, &grids),
|
||||||
|
FeatureStore::Inline(_)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -8,6 +8,8 @@
|
|||||||
//! * Detokenizer — M pinned OS threads / shards (CPU bound), core set C
|
//! * Detokenizer — M pinned OS threads / shards (CPU bound), core set C
|
||||||
//! * TM ingress — 1 thread driving the ingress FSM
|
//! * TM ingress — 1 thread driving the ingress FSM
|
||||||
//! * TM egress — 1 thread draining the egress ring → detok shards
|
//! * TM egress — 1 thread draining the egress ring → detok shards
|
||||||
|
//! * MM workers — K unpinned OS threads, spawned late via
|
||||||
|
//! [`Runtime::spawn_mm_pool`] (multimodal models only)
|
||||||
//!
|
//!
|
||||||
//! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling
|
//! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling
|
||||||
//! axum's worker threads.
|
//! axum's worker threads.
|
||||||
@@ -38,6 +40,18 @@ pub use runnable::Runnable;
|
|||||||
pub struct Runtime {
|
pub struct Runtime {
|
||||||
pub ingress: IngressConsumer,
|
pub ingress: IngressConsumer,
|
||||||
pub egress: EgressProducer,
|
pub egress: EgressProducer,
|
||||||
|
/// Requests parked in `Encoding`, drained by the MM worker pool
|
||||||
|
/// (`Server.start_mm_workers`). Stays empty for non-multimodal models —
|
||||||
|
/// ingress never routes to it.
|
||||||
|
pub mm: flume::Receiver<crate::message::MmRequest>,
|
||||||
|
/// Back-channel for the MM workers' `MmEncoded` / `MmFailed` into tm-ingress.
|
||||||
|
pub tm: flume::Sender<TmEvent>,
|
||||||
|
/// The loaded tokenizer, shared with the MM worker path (`None` under
|
||||||
|
/// `skip_tokenizer_init`).
|
||||||
|
pub tokenizer: Option<Arc<dyn tokenizer::TextTokenizer>>,
|
||||||
|
/// MM results parked between a worker's `MmEncoded` and the scheduler drain
|
||||||
|
/// (`Server.take_mm`).
|
||||||
|
pub mm_sidecar: crate::mm::Sidecar,
|
||||||
/// Worker join handles, joined by `request_shutdown` / `Drop`.
|
/// Worker join handles, joined by `request_shutdown` / `Drop`.
|
||||||
threads: Mutex<Vec<JoinHandle<()>>>,
|
threads: Mutex<Vec<JoinHandle<()>>>,
|
||||||
/// The single shutdown sender.
|
/// The single shutdown sender.
|
||||||
@@ -49,6 +63,21 @@ pub struct Runtime {
|
|||||||
const SHUTDOWN_JOIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
|
const SHUTDOWN_JOIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
|
||||||
|
|
||||||
impl Runtime {
|
impl Runtime {
|
||||||
|
/// Spawn `workers` `mm-worker-{i}` threads into the shutdown join set —
|
||||||
|
/// late, once Python has built the mm spec (`Server::start_mm_workers`).
|
||||||
|
///
|
||||||
|
/// Deliberately unpinned: the threads inherit the launch thread's affinity,
|
||||||
|
/// already narrowed by `RustServer.launch` to the server cores, so bursty
|
||||||
|
/// MM preprocessing floats over that whole set (rather than owning cores
|
||||||
|
/// that idle between bursts) and never preempts the scheduler's reserved
|
||||||
|
/// cores.
|
||||||
|
pub fn spawn_mm_pool(&self, workers: usize, ctx: Arc<crate::mm::Context>) {
|
||||||
|
let mut threads = self.threads.lock().unwrap();
|
||||||
|
spawn_pool("mm-worker", None, workers.max(1), &mut threads, |_| {
|
||||||
|
crate::mm::MmWorker::new(self.mm.clone(), self.tm.clone(), ctx.clone())
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
/// Stop the runtime and join every worker thread (with a bounded wait).
|
/// Stop the runtime and join every worker thread (with a bounded wait).
|
||||||
///
|
///
|
||||||
/// Dropping `shutdown_tx` wakes the tm-ingress/tm-egress selectors (which
|
/// Dropping `shutdown_tx` wakes the tm-ingress/tm-egress selectors (which
|
||||||
@@ -107,6 +136,10 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
|||||||
let (tm_tx, tm_rx) = flume::bounded::<TmEvent>(cfg.rust_server_args.channel_cap);
|
let (tm_tx, tm_rx) = flume::bounded::<TmEvent>(cfg.rust_server_args.channel_cap);
|
||||||
let (tok_tx, tok_rx) =
|
let (tok_tx, tok_rx) =
|
||||||
flume::bounded::<crate::message::Request>(cfg.rust_server_args.channel_cap);
|
flume::bounded::<crate::message::Request>(cfg.rust_server_args.channel_cap);
|
||||||
|
// Encoding → MM worker pool. Bounded like the other stage edges so a slow
|
||||||
|
// pool back-pressures instead of buffering unboundedly.
|
||||||
|
let (mm_tx, mm_rx) =
|
||||||
|
flume::bounded::<crate::message::MmRequest>(cfg.rust_server_args.channel_cap);
|
||||||
let detokenizer_worker_num = cfg.server_args.detokenizer_worker_num;
|
let detokenizer_worker_num = cfg.server_args.detokenizer_worker_num;
|
||||||
let mut detok_tx = Vec::with_capacity(detokenizer_worker_num);
|
let mut detok_tx = Vec::with_capacity(detokenizer_worker_num);
|
||||||
let mut detok_rx = Vec::with_capacity(detokenizer_worker_num);
|
let mut detok_rx = Vec::with_capacity(detokenizer_worker_num);
|
||||||
@@ -140,6 +173,14 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
|||||||
cfg.server_args.revision.as_deref(),
|
cfg.server_args.revision.as_deref(),
|
||||||
skip_tokenizer_init,
|
skip_tokenizer_init,
|
||||||
)?;
|
)?;
|
||||||
|
// The `TextTokenizer` view of it, shared by the tokenizer pool and the MM
|
||||||
|
// worker path (which encodes the placeholder-expanded prompt itself).
|
||||||
|
let text_tokenizer: Option<Arc<dyn tokenizer::TextTokenizer>> = dyn_tokenizer
|
||||||
|
.as_ref()
|
||||||
|
.map(|t| Arc::new(tokenizer::DynamoTokenizer::new(t.clone())) as _);
|
||||||
|
|
||||||
|
// Shared: MM workers park, the Python drain pops, tm-ingress purges.
|
||||||
|
let mm_sidecar: crate::mm::Sidecar = Default::default();
|
||||||
|
|
||||||
// --- Detokenizer shards (pinned, CPU bound) ---
|
// --- Detokenizer shards (pinned, CPU bound) ---
|
||||||
{
|
{
|
||||||
@@ -168,10 +209,9 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
|||||||
// --- Tokenizer pool (pinned, CPU bound) ---
|
// --- Tokenizer pool (pinned, CPU bound) ---
|
||||||
// Only spawned when a real tokenizer is loaded; under `skip_tokenizer_init`
|
// Only spawned when a real tokenizer is loaded; under `skip_tokenizer_init`
|
||||||
// there is none and ingress never routes to the pool, so we skip it.
|
// there is none and ingress never routes to the pool, so we skip it.
|
||||||
if let Some(t) = &dyn_tokenizer {
|
if let Some(tokenizer) = &text_tokenizer {
|
||||||
// Reuse the single loaded tokenizer (shared with the detok shards).
|
// Reuse the single loaded tokenizer (shared with the detok shards).
|
||||||
let tokenizer: Arc<dyn tokenizer::TextTokenizer> =
|
let tokenizer = tokenizer.clone();
|
||||||
Arc::new(tokenizer::DynamoTokenizer::new(t.clone()));
|
|
||||||
let tok_cores = plan.as_ref().map(|p| p.tok.clone());
|
let tok_cores = plan.as_ref().map(|p| p.tok.clone());
|
||||||
// Workers share the MPMC inbox (`tok_rx`) and the read-only backend, so
|
// Workers share the MPMC inbox (`tok_rx`) and the read-only backend, so
|
||||||
// each gets a cheap clone of both.
|
// each gets a cheap clone of both.
|
||||||
@@ -220,6 +260,11 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
|||||||
.map(|c| vec![c]);
|
.map(|c| vec![c]);
|
||||||
let limits = tokenizer_manager::Limits::try_from(&*cfg.server_args)
|
let limits = tokenizer_manager::Limits::try_from(&*cfg.server_args)
|
||||||
.map_err(|e| format!("ingress limits: {e}"))?;
|
.map_err(|e| format!("ingress limits: {e}"))?;
|
||||||
|
let mm = tokenizer_manager::Mm {
|
||||||
|
enabled: cfg.server_args.model_is_multimodal(),
|
||||||
|
tx: mm_tx,
|
||||||
|
sidecar: mm_sidecar.clone(),
|
||||||
|
};
|
||||||
let mut parts = Some((tm_rx, ingress_tx)); // moved into the single worker
|
let mut parts = Some((tm_rx, ingress_tx)); // moved into the single worker
|
||||||
let shutdown_rx = shutdown_rx.clone();
|
let shutdown_rx = shutdown_rx.clone();
|
||||||
spawn_pool("tm-ingress", cores, 1, &mut threads, |_| {
|
spawn_pool("tm-ingress", cores, 1, &mut threads, |_| {
|
||||||
@@ -230,6 +275,7 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
|||||||
senders.clone(),
|
senders.clone(),
|
||||||
ingress_tx,
|
ingress_tx,
|
||||||
limits.clone(),
|
limits.clone(),
|
||||||
|
mm.clone(),
|
||||||
shutdown_rx.clone(),
|
shutdown_rx.clone(),
|
||||||
)
|
)
|
||||||
});
|
});
|
||||||
@@ -282,6 +328,10 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
|||||||
Ok(Runtime {
|
Ok(Runtime {
|
||||||
ingress: ingress_rx,
|
ingress: ingress_rx,
|
||||||
egress: egress_tx,
|
egress: egress_tx,
|
||||||
|
mm: mm_rx,
|
||||||
|
tm: tm_tx,
|
||||||
|
tokenizer: text_tokenizer,
|
||||||
|
mm_sidecar,
|
||||||
threads: Mutex::new(threads),
|
threads: Mutex::new(threads),
|
||||||
shutdown_tx: Mutex::new(Some(shutdown_tx)),
|
shutdown_tx: Mutex::new(Some(shutdown_tx)),
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -158,6 +158,11 @@ pub struct ModelConfig {
|
|||||||
/// boot ([`ServerArgs::validate_mandatory`]).
|
/// boot ([`ServerArgs::validate_mandatory`]).
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub vocab_size: Option<u64>,
|
pub vocab_size: Option<u64>,
|
||||||
|
/// Whether the model accepts multimodal inputs. Gates the MM Encoding branch
|
||||||
|
/// in tm-ingress; `false` silently ignores mm fields, as the Python
|
||||||
|
/// `TokenizerManager` does with `mm_processor is None`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub is_multimodal: bool,
|
||||||
/// Resolved default sampling parameters, stamped by
|
/// Resolved default sampling parameters, stamped by
|
||||||
/// `RustServer._build_server_args` from Python's
|
/// `RustServer._build_server_args` from Python's
|
||||||
/// `ModelConfig.get_default_sampling_params()`. Already gated on
|
/// `ModelConfig.get_default_sampling_params()`. Already gated on
|
||||||
@@ -259,6 +264,12 @@ impl ServerArgs {
|
|||||||
self.disaggregation_mode == "prefill"
|
self.disaggregation_mode == "prefill"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Whether the served model is multimodal, from the scheduler's dump. See
|
||||||
|
/// [`ModelConfig::is_multimodal`].
|
||||||
|
pub fn model_is_multimodal(&self) -> bool {
|
||||||
|
self.model_config.is_multimodal
|
||||||
|
}
|
||||||
|
|
||||||
/// Bind address `host:port`. `host` is expected to be an IP — the result is
|
/// Bind address `host:port`. `host` is expected to be an IP — the result is
|
||||||
/// parsed as a `SocketAddr`, so a bare IPv6 host gets bracketed.
|
/// parsed as a `SocketAddr`, so a bare IPv6 host gets bracketed.
|
||||||
pub fn bind(&self) -> String {
|
pub fn bind(&self) -> String {
|
||||||
|
|||||||
@@ -14,7 +14,7 @@ mod egress;
|
|||||||
mod ingress;
|
mod ingress;
|
||||||
|
|
||||||
pub use egress::{ActivityCounter, Egress};
|
pub use egress::{ActivityCounter, Egress};
|
||||||
pub use ingress::{Ingress, Limits};
|
pub use ingress::{Ingress, Limits, Mm};
|
||||||
|
|
||||||
use crate::ids::Rid;
|
use crate::ids::Rid;
|
||||||
use crate::message::{DetokMsg, Request};
|
use crate::message::{DetokMsg, Request};
|
||||||
@@ -36,6 +36,13 @@ pub enum TmEvent {
|
|||||||
/// A request back from the tokenizer pool: `PreSendValidating` (ids filled) on success,
|
/// A request back from the tokenizer pool: `PreSendValidating` (ids filled) on success,
|
||||||
/// or `Failed` on a tokenize error. `drive` handles both.
|
/// or `Failed` on a tokenize error. `drive` handles both.
|
||||||
Tokenized(Request),
|
Tokenized(Request),
|
||||||
|
/// An MM worker finished a request parked in `Encoding`: `input_ids` are the
|
||||||
|
/// final placeholder-expanded prompt ids. The buffers ride the rid-keyed
|
||||||
|
/// sidecar (`Server.take_mm`), not this event.
|
||||||
|
MmEncoded { rid: Rid, input_ids: Vec<i32> },
|
||||||
|
/// An MM worker rejected a request parked in `Encoding` (bad media URL,
|
||||||
|
/// unsupported modality, preprocess error, …).
|
||||||
|
MmFailed { rid: Rid, message: String },
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Producer-side handles, cloned into every stage that needs to emit.
|
/// Producer-side handles, cloned into every stage that needs to emit.
|
||||||
|
|||||||
@@ -17,14 +17,17 @@
|
|||||||
//! The egress edges (Streaming/Finalizing/Completed) are driven on the egress
|
//! The egress edges (Streaming/Finalizing/Completed) are driven on the egress
|
||||||
//! side (see `egress` + `detokenizer`).
|
//! side (see `egress` + `detokenizer`).
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use bytes::Bytes;
|
use bytes::Bytes;
|
||||||
|
|
||||||
use crate::error::Error;
|
use crate::error::Error;
|
||||||
use crate::fsm::{Event, RequestState, ValidationOutcome};
|
use crate::fsm::{Event, RequestState, ValidationOutcome};
|
||||||
|
use crate::ids::Rid;
|
||||||
|
|
||||||
use crate::message::{
|
use crate::message::{
|
||||||
AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, Request,
|
AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, MmRequest,
|
||||||
RequestKind,
|
Request, RequestKind,
|
||||||
};
|
};
|
||||||
use crate::ring::IngressProducer;
|
use crate::ring::IngressProducer;
|
||||||
use crate::runtime::{Runnable, ServerArgs};
|
use crate::runtime::{Runnable, ServerArgs};
|
||||||
@@ -41,9 +44,29 @@ pub struct Ingress {
|
|||||||
senders: Senders,
|
senders: Senders,
|
||||||
ingress: IngressProducer,
|
ingress: IngressProducer,
|
||||||
limits: Limits,
|
limits: Limits,
|
||||||
|
mm: Mm,
|
||||||
|
/// Requests parked in `Encoding` while an MM worker processes their media;
|
||||||
|
/// resumed by `MmEncoded` / `MmFailed`. Only this thread touches it, so no
|
||||||
|
/// lock.
|
||||||
|
pending_mm: HashMap<Rid, Request>,
|
||||||
shutdown: flume::Receiver<()>,
|
shutdown: flume::Receiver<()>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The ingress side of the MM path.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct Mm {
|
||||||
|
/// Whether the model is multimodal. When false, mm fields are silently
|
||||||
|
/// ignored, as the Python `TokenizerManager` does with `mm_processor is
|
||||||
|
/// None`.
|
||||||
|
pub enabled: bool,
|
||||||
|
/// → MM worker pool (spawned via `Server.start_mm_workers`).
|
||||||
|
pub tx: flume::Sender<MmRequest>,
|
||||||
|
/// Results sidecar. Purged here when a late result arrives for a request
|
||||||
|
/// that is no longer parked; otherwise it would leak, since only the
|
||||||
|
/// scheduler drain pops entries.
|
||||||
|
pub sidecar: crate::mm::Sidecar,
|
||||||
|
}
|
||||||
|
|
||||||
/// Longest client-supplied rid accepted. It keys the detok table and travels on
|
/// Longest client-supplied rid accepted. It keys the detok table and travels on
|
||||||
/// every chunk, so its length is a recurring cost; Python mints 32-byte uuid hex.
|
/// every chunk, so its length is a recurring cost; Python mints 32-byte uuid hex.
|
||||||
const MAX_RID_LEN: usize = 128;
|
const MAX_RID_LEN: usize = 128;
|
||||||
@@ -104,6 +127,7 @@ impl Ingress {
|
|||||||
senders: Senders,
|
senders: Senders,
|
||||||
ingress: IngressProducer,
|
ingress: IngressProducer,
|
||||||
limits: Limits,
|
limits: Limits,
|
||||||
|
mm: Mm,
|
||||||
shutdown: flume::Receiver<()>,
|
shutdown: flume::Receiver<()>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
Self {
|
Self {
|
||||||
@@ -112,6 +136,8 @@ impl Ingress {
|
|||||||
senders,
|
senders,
|
||||||
ingress,
|
ingress,
|
||||||
limits,
|
limits,
|
||||||
|
mm,
|
||||||
|
pending_mm: HashMap::new(),
|
||||||
shutdown,
|
shutdown,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -124,7 +150,7 @@ enum Lane {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl Runnable for Ingress {
|
impl Runnable for Ingress {
|
||||||
fn run(self) {
|
fn run(mut self) {
|
||||||
loop {
|
loop {
|
||||||
// Select, not a drain-then-block: an abort arriving while the inbox is
|
// Select, not a drain-then-block: an abort arriving while the inbox is
|
||||||
// idle must still be handled at once.
|
// idle must still be handled at once.
|
||||||
@@ -139,6 +165,12 @@ impl Runnable for Ingress {
|
|||||||
Some(Lane::Event(TmEvent::Ingress(req) | TmEvent::Tokenized(req))) => {
|
Some(Lane::Event(TmEvent::Ingress(req) | TmEvent::Tokenized(req))) => {
|
||||||
self.drive(req)
|
self.drive(req)
|
||||||
}
|
}
|
||||||
|
Some(Lane::Event(TmEvent::MmEncoded { rid, input_ids })) => {
|
||||||
|
self.on_mm_encoded(rid, input_ids)
|
||||||
|
}
|
||||||
|
Some(Lane::Event(TmEvent::MmFailed { rid, message })) => {
|
||||||
|
self.on_mm_failed(rid, message)
|
||||||
|
}
|
||||||
None => {
|
None => {
|
||||||
// Shutdown, or the inbox closed. Drain whatever is still queued
|
// Shutdown, or the inbox closed. Drain whatever is still queued
|
||||||
// on the abort lane first: those requests are in flight on the
|
// on the abort lane first: those requests are in flight on the
|
||||||
@@ -168,6 +200,9 @@ impl Ingress {
|
|||||||
if err.http_status() == 500 {
|
if err.http_status() == 500 {
|
||||||
tracing::error!(rid = %req.rid, error = %err, "ingress rejected request");
|
tracing::error!(rid = %req.rid, error = %err, "ingress rejected request");
|
||||||
}
|
}
|
||||||
|
// A rejected request never reaches the scheduler drain, so purge any
|
||||||
|
// parked MM result (no-op for the common non-mm request).
|
||||||
|
self.mm.sidecar.purge(req.rid.as_str());
|
||||||
let _ = req.state.apply(Event::Error(err.clone()));
|
let _ = req.state.apply(Event::Error(err.clone()));
|
||||||
let _ = req.sink.try_send(EgressItem::Error(err)); // client may be gone
|
let _ = req.sink.try_send(EgressItem::Error(err)); // client may be gone
|
||||||
if registered {
|
if registered {
|
||||||
@@ -178,11 +213,12 @@ impl Ingress {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Drive a request through its ingress states until it terminates (failed or
|
/// Drive a request through its ingress states until it terminates (failed or
|
||||||
/// pushed to the ring) or is handed to the tokenizer pool (re-entering as a
|
/// pushed to the ring), is handed to the tokenizer pool (re-entering as a
|
||||||
/// `Tokenized` event). Each arm acts and advances the FSM; the loop
|
/// `Tokenized` event), or is parked in `pending_mm` awaiting an MM worker
|
||||||
/// re-dispatches. The arms are the design table's states, `Failed` the single
|
/// (re-entering via `MmEncoded` / `MmFailed`). Each arm acts and advances
|
||||||
/// reject path.
|
/// the FSM; the loop re-dispatches. The arms are the design table's states,
|
||||||
fn drive(&self, mut req: Request) {
|
/// `Failed` the single reject path.
|
||||||
|
fn drive(&mut self, mut req: Request) {
|
||||||
// Flipped once `register_detok` succeeds; `fail` must not deregister before
|
// Flipped once `register_detok` succeeds; `fail` must not deregister before
|
||||||
// that (see `fail`). A pool return re-enters `drive` already registered.
|
// that (see `fail`). A pool return re-enters `drive` already registered.
|
||||||
let mut registered = !matches!(req.state, RequestState::Received);
|
let mut registered = !matches!(req.state, RequestState::Received);
|
||||||
@@ -235,6 +271,13 @@ impl Ingress {
|
|||||||
.normalize(self.limits.skip_tokenizer_init, self.limits.vocab_size)
|
.normalize(self.limits.skip_tokenizer_init, self.limits.vocab_size)
|
||||||
{
|
{
|
||||||
Err(e) => Err(e),
|
Err(e) => Err(e),
|
||||||
|
// The native pipeline produces the final input_ids,
|
||||||
|
// so it wins even over a pre-tokenized prompt (which
|
||||||
|
// still needs placeholder expansion) — the same
|
||||||
|
// precedence as the Python TokenizerManager.
|
||||||
|
Ok(()) if self.mm.enabled && g.has_multimodal() => {
|
||||||
|
Ok(ValidationOutcome::HasMultimodal)
|
||||||
|
}
|
||||||
// Client ids skip the pool; text goes to the tokenizer.
|
// Client ids skip the pool; text goes to the tokenizer.
|
||||||
Ok(()) if g.already_tokenized() => {
|
Ok(()) if g.already_tokenized() => {
|
||||||
Ok(ValidationOutcome::AlreadyTokenized)
|
Ok(ValidationOutcome::AlreadyTokenized)
|
||||||
@@ -252,6 +295,40 @@ impl Ingress {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
// Hand off to the MM worker pool and park the request; it
|
||||||
|
// re-enters via `MmEncoded` (→ PreSendValidating) or `MmFailed`
|
||||||
|
// (→ reject). Doesn't loop.
|
||||||
|
RequestState::Encoding => {
|
||||||
|
let work = {
|
||||||
|
let RequestKind::Generate(g) = &mut req.kind else {
|
||||||
|
self.fail(
|
||||||
|
&mut req,
|
||||||
|
Error::Internal("non-generate request in Encoding".into()),
|
||||||
|
registered,
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
g.take_mm_work()
|
||||||
|
};
|
||||||
|
let msg = MmRequest {
|
||||||
|
rid: req.rid.clone(),
|
||||||
|
work,
|
||||||
|
};
|
||||||
|
// Full = the pool can't keep up, so back-pressure like a full
|
||||||
|
// ingress ring. Disconnected = pool gone.
|
||||||
|
if let Err(e) = self.mm.tx.try_send(msg) {
|
||||||
|
let err = match e {
|
||||||
|
flume::TrySendError::Full(_) => Error::QueueFull,
|
||||||
|
flume::TrySendError::Disconnected(_) => {
|
||||||
|
Error::Internal("mm worker pool gone".into())
|
||||||
|
}
|
||||||
|
};
|
||||||
|
self.fail(&mut req, err, registered);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
self.pending_mm.insert(req.rid.clone(), req);
|
||||||
|
return;
|
||||||
|
}
|
||||||
// Hand off to the tokenizer pool; it returns the request as a
|
// Hand off to the tokenizer pool; it returns the request as a
|
||||||
// `Tokenized` event (PreSendValidating, or Failed on error).
|
// `Tokenized` event (PreSendValidating, or Failed on error).
|
||||||
// Doesn't loop.
|
// Doesn't loop.
|
||||||
@@ -392,6 +469,34 @@ impl Ingress {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// An MM worker finished a parked request: fill in the final expanded
|
||||||
|
/// `input_ids`, advance `Encoding → PreSendValidating`, and resume driving
|
||||||
|
/// (pre-send checks → ring). No pending entry means the request was already
|
||||||
|
/// rejected or aborted, so the result is dropped.
|
||||||
|
fn on_mm_encoded(&mut self, rid: Rid, input_ids: Vec<i32>) {
|
||||||
|
let Some(mut req) = self.pending_mm.remove(&rid) else {
|
||||||
|
tracing::debug!(rid = %rid, "mm result for unknown/finished request; dropped");
|
||||||
|
// It will never reach the scheduler drain, so purge or leak.
|
||||||
|
self.mm.sidecar.purge(rid.as_str());
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
if let RequestKind::Generate(g) = &mut req.kind {
|
||||||
|
g.input_ids = Some(input_ids);
|
||||||
|
}
|
||||||
|
let _ = req.state.apply(Event::EncodeDone); // Encoding → PreSendValidating
|
||||||
|
self.drive(req);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// An MM worker failed a parked request (bad URL, processor error): reject it
|
||||||
|
/// back to the client, as Python turns a per-request exception into a 400.
|
||||||
|
fn on_mm_failed(&mut self, rid: Rid, message: String) {
|
||||||
|
let Some(mut req) = self.pending_mm.remove(&rid) else {
|
||||||
|
tracing::debug!(rid = %rid, "mm failure for unknown/finished request; dropped");
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
self.fail(&mut req, Error::Encode(message), true); // parked ⇒ registered
|
||||||
|
}
|
||||||
|
|
||||||
/// Client disconnected (or a detok terminal): deregister the sink, then push an
|
/// Client disconnected (or a detok terminal): deregister the sink, then push an
|
||||||
/// `AbortReq(rid)` so the scheduler stops generating for it.
|
/// `AbortReq(rid)` so the scheduler stops generating for it.
|
||||||
///
|
///
|
||||||
@@ -400,8 +505,15 @@ impl Ingress {
|
|||||||
/// That wastes GPU work until the request finishes on its own, but it cannot be
|
/// That wastes GPU work until the request finishes on its own, but it cannot be
|
||||||
/// misdelivered — the rid is unique to this request for the process's lifetime
|
/// misdelivered — the rid is unique to this request for the process's lifetime
|
||||||
/// ([`Rid::from_client`]), so no later request can ever answer to it.
|
/// ([`Rid::from_client`]), so no later request can ever answer to it.
|
||||||
fn on_abort(&self, source: AbortSource) {
|
///
|
||||||
|
/// A request parked in `pending_mm` is cancelled here, so the worker's late
|
||||||
|
/// result lands in `on_mm_encoded`'s no-entry branch and purges the sidecar —
|
||||||
|
/// no generation runs for output nobody will read.
|
||||||
|
fn on_abort(&mut self, source: AbortSource) {
|
||||||
let rid = source.rid().clone();
|
let rid = source.rid().clone();
|
||||||
|
if self.pending_mm.remove(&rid).is_some() {
|
||||||
|
tracing::debug!(rid = %rid, "abort cancelled request parked for MM");
|
||||||
|
}
|
||||||
let _ = self
|
let _ = self
|
||||||
.senders
|
.senders
|
||||||
.detok_for(&rid)
|
.detok_for(&rid)
|
||||||
@@ -616,12 +728,14 @@ mod tests {
|
|||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
|
|
||||||
/// An `Ingress` plus its detok-shard receiver, ring consumer (keep alive —
|
/// An `Ingress` plus its detok-shard receiver, ring consumer (keep alive —
|
||||||
/// dropping it closes the ring → false QueueFull), and tm inbox sender.
|
/// dropping it closes the ring → false QueueFull), tm inbox sender, and the
|
||||||
|
/// mm-pool receiver (keep alive — dropping it makes mm submits fail).
|
||||||
fn make_ingress() -> (
|
fn make_ingress() -> (
|
||||||
Ingress,
|
Ingress,
|
||||||
flume::Receiver<DetokMsg>,
|
flume::Receiver<DetokMsg>,
|
||||||
IngressConsumer,
|
IngressConsumer,
|
||||||
flume::Sender<TmEvent>,
|
flume::Sender<TmEvent>,
|
||||||
|
flume::Receiver<MmRequest>,
|
||||||
) {
|
) {
|
||||||
make_ingress_with(test_limits())
|
make_ingress_with(test_limits())
|
||||||
}
|
}
|
||||||
@@ -633,6 +747,7 @@ mod tests {
|
|||||||
flume::Receiver<DetokMsg>,
|
flume::Receiver<DetokMsg>,
|
||||||
IngressConsumer,
|
IngressConsumer,
|
||||||
flume::Sender<TmEvent>,
|
flume::Sender<TmEvent>,
|
||||||
|
flume::Receiver<MmRequest>,
|
||||||
) {
|
) {
|
||||||
make_ingress_inner(test_limits(), abort_rx)
|
make_ingress_inner(test_limits(), abort_rx)
|
||||||
}
|
}
|
||||||
@@ -644,6 +759,7 @@ mod tests {
|
|||||||
flume::Receiver<DetokMsg>,
|
flume::Receiver<DetokMsg>,
|
||||||
IngressConsumer,
|
IngressConsumer,
|
||||||
flume::Sender<TmEvent>,
|
flume::Sender<TmEvent>,
|
||||||
|
flume::Receiver<MmRequest>,
|
||||||
) {
|
) {
|
||||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||||
std::mem::forget(abort_tx); // keep the lane open; tests end by dropping tm_tx
|
std::mem::forget(abort_tx); // keep the lane open; tests end by dropping tm_tx
|
||||||
@@ -658,6 +774,7 @@ mod tests {
|
|||||||
flume::Receiver<DetokMsg>,
|
flume::Receiver<DetokMsg>,
|
||||||
IngressConsumer,
|
IngressConsumer,
|
||||||
flume::Sender<TmEvent>,
|
flume::Sender<TmEvent>,
|
||||||
|
flume::Receiver<MmRequest>,
|
||||||
) {
|
) {
|
||||||
let (tok_tx, _tok_rx) = flume::unbounded();
|
let (tok_tx, _tok_rx) = flume::unbounded();
|
||||||
let (detok_tx, detok_rx) = flume::unbounded();
|
let (detok_tx, detok_rx) = flume::unbounded();
|
||||||
@@ -669,12 +786,30 @@ mod tests {
|
|||||||
};
|
};
|
||||||
let (ingress_producer, consumer) = ingress_ring(16);
|
let (ingress_producer, consumer) = ingress_ring(16);
|
||||||
let (tm_tx, tm_rx) = flume::unbounded();
|
let (tm_tx, tm_rx) = flume::unbounded();
|
||||||
|
let (mm_tx, mm_rx) = flume::unbounded();
|
||||||
// Keep the shutdown sender alive (leak) so its branch never fires — tests
|
// Keep the shutdown sender alive (leak) so its branch never fires — tests
|
||||||
// end `run` by dropping `tm_tx`, not by shutdown.
|
// end `run` by dropping `tm_tx`, not by shutdown.
|
||||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||||
std::mem::forget(sd_tx);
|
std::mem::forget(sd_tx);
|
||||||
let ingress = Ingress::new(tm_rx, abort_rx, senders, ingress_producer, limits, sd_rx);
|
let ingress = Ingress::new(
|
||||||
(ingress, detok_rx, consumer, tm_tx)
|
tm_rx,
|
||||||
|
abort_rx,
|
||||||
|
senders,
|
||||||
|
ingress_producer,
|
||||||
|
limits,
|
||||||
|
test_mm(mm_tx, true),
|
||||||
|
sd_rx,
|
||||||
|
);
|
||||||
|
(ingress, detok_rx, consumer, tm_tx, mm_rx)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// An [`Mm`] over `tx` with a fresh sidecar.
|
||||||
|
fn test_mm(tx: flume::Sender<MmRequest>, enabled: bool) -> Mm {
|
||||||
|
Mm {
|
||||||
|
enabled,
|
||||||
|
tx,
|
||||||
|
sidecar: Default::default(),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Both abort sources do the same two things: drop the detok entry so no
|
/// Both abort sources do the same two things: drop the detok entry so no
|
||||||
@@ -696,7 +831,7 @@ mod tests {
|
|||||||
let (ingress_producer, consumer) = ingress_ring(16);
|
let (ingress_producer, consumer) = ingress_ring(16);
|
||||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||||
std::mem::forget(sd_tx);
|
std::mem::forget(sd_tx);
|
||||||
let ingress = Ingress::new(
|
let mut ingress = Ingress::new(
|
||||||
flume::unbounded().1,
|
flume::unbounded().1,
|
||||||
flume::unbounded().1,
|
flume::unbounded().1,
|
||||||
Senders {
|
Senders {
|
||||||
@@ -707,6 +842,7 @@ mod tests {
|
|||||||
},
|
},
|
||||||
ingress_producer,
|
ingress_producer,
|
||||||
test_limits(),
|
test_limits(),
|
||||||
|
test_mm(flume::unbounded().0, true),
|
||||||
sd_rx,
|
sd_rx,
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -960,7 +1096,7 @@ mod tests {
|
|||||||
/// to the ring, after registration — so it must be deregistered, not leaked.
|
/// to the ring, after registration — so it must be deregistered, not leaked.
|
||||||
#[test]
|
#[test]
|
||||||
fn over_context_request_deregisters_and_never_reaches_the_ring() {
|
fn over_context_request_deregisters_and_never_reaches_the_ring() {
|
||||||
let (ingress, detok_rx, consumer, _tm_tx) = make_ingress_with(Limits {
|
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress_with(Limits {
|
||||||
context_len: 4,
|
context_len: 4,
|
||||||
..test_limits()
|
..test_limits()
|
||||||
});
|
});
|
||||||
@@ -993,7 +1129,7 @@ mod tests {
|
|||||||
/// pins. Nothing may reach the scheduler ring.
|
/// pins. Nothing may reach the scheduler ring.
|
||||||
#[test]
|
#[test]
|
||||||
fn detokenize_flows_register_then_decode_and_skips_the_ring() {
|
fn detokenize_flows_register_then_decode_and_skips_the_ring() {
|
||||||
let (ingress, detok_rx, consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
let (tx, mut rx) = mpsc::channel(8);
|
let (tx, mut rx) = mpsc::channel(8);
|
||||||
ingress.drive(Request {
|
ingress.drive(Request {
|
||||||
rid: "41".into(),
|
rid: "41".into(),
|
||||||
@@ -1028,7 +1164,7 @@ mod tests {
|
|||||||
/// leak and no decode job to drop).
|
/// leak and no decode job to drop).
|
||||||
#[test]
|
#[test]
|
||||||
fn detokenize_negative_ids_reject_before_registration() {
|
fn detokenize_negative_ids_reject_before_registration() {
|
||||||
let (ingress, detok_rx, consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
let (tx, mut rx) = mpsc::channel(8);
|
let (tx, mut rx) = mpsc::channel(8);
|
||||||
ingress.drive(Request {
|
ingress.drive(Request {
|
||||||
rid: "43".into(),
|
rid: "43".into(),
|
||||||
@@ -1071,7 +1207,15 @@ mod tests {
|
|||||||
let (_tm_tx, tm_rx) = flume::unbounded();
|
let (_tm_tx, tm_rx) = flume::unbounded();
|
||||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||||
std::mem::forget(sd_tx);
|
std::mem::forget(sd_tx);
|
||||||
let ingress = Ingress::new(tm_rx, abort_rx, senders, producer, test_limits(), sd_rx);
|
let mut ingress = Ingress::new(
|
||||||
|
tm_rx,
|
||||||
|
abort_rx,
|
||||||
|
senders,
|
||||||
|
producer,
|
||||||
|
test_limits(),
|
||||||
|
test_mm(flume::unbounded().0, true),
|
||||||
|
sd_rx,
|
||||||
|
);
|
||||||
|
|
||||||
ingress.on_abort(AbortSource::Guard("pushed".into()));
|
ingress.on_abort(AbortSource::Guard("pushed".into()));
|
||||||
ingress.on_abort(AbortSource::Guard("dropped".into()));
|
ingress.on_abort(AbortSource::Guard("dropped".into()));
|
||||||
@@ -1108,7 +1252,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn pre_registration_failure_does_not_deregister() {
|
fn pre_registration_failure_does_not_deregister() {
|
||||||
// Rejected inside `validate` (out-of-vocab id), which runs before registration.
|
// Rejected inside `validate` (out-of-vocab id), which runs before registration.
|
||||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
let mut req = generate_req(41, SamplingParams::default());
|
let mut req = generate_req(41, SamplingParams::default());
|
||||||
if let RequestKind::Generate(g) = &mut req.kind {
|
if let RequestKind::Generate(g) = &mut req.kind {
|
||||||
g.input_ids = Some(vec![2_000_000_000]);
|
g.input_ids = Some(vec![2_000_000_000]);
|
||||||
@@ -1121,7 +1265,7 @@ mod tests {
|
|||||||
);
|
);
|
||||||
|
|
||||||
// A post-registration reject still deregisters (the leak fix stays fixed).
|
// A post-registration reject still deregisters (the leak fix stays fixed).
|
||||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
ingress.drive(generate_req(
|
ingress.drive(generate_req(
|
||||||
42,
|
42,
|
||||||
SamplingParams {
|
SamplingParams {
|
||||||
@@ -1140,7 +1284,7 @@ mod tests {
|
|||||||
/// sees `Register` then `Deregister`. Regression for RSS growth on bad input.
|
/// sees `Register` then `Deregister`. Regression for RSS growth on bad input.
|
||||||
#[test]
|
#[test]
|
||||||
fn rejected_request_deregisters_from_shard() {
|
fn rejected_request_deregisters_from_shard() {
|
||||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
// top_p = 2.0 is outside (0, 1], so `SamplingParams::normalize` rejects it.
|
// top_p = 2.0 is outside (0, 1], so `SamplingParams::normalize` rejects it.
|
||||||
let bad = SamplingParams {
|
let bad = SamplingParams {
|
||||||
top_p: 2.0,
|
top_p: 2.0,
|
||||||
@@ -1167,7 +1311,7 @@ mod tests {
|
|||||||
/// and kills the scheduler process (`make_ingress` bounds vocab at 1000).
|
/// and kills the scheduler process (`make_ingress` bounds vocab at 1000).
|
||||||
#[test]
|
#[test]
|
||||||
fn out_of_vocab_input_ids_rejected() {
|
fn out_of_vocab_input_ids_rejected() {
|
||||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
let mut req = generate_req(21, SamplingParams::default());
|
let mut req = generate_req(21, SamplingParams::default());
|
||||||
if let RequestKind::Generate(g) = &mut req.kind {
|
if let RequestKind::Generate(g) = &mut req.kind {
|
||||||
g.input_ids = Some(vec![1, 2_000_000_000]);
|
g.input_ids = Some(vec![1, 2_000_000_000]);
|
||||||
@@ -1185,7 +1329,7 @@ mod tests {
|
|||||||
/// Same guard for negative ids and for `token_ids_logprob` entries.
|
/// Same guard for negative ids and for `token_ids_logprob` entries.
|
||||||
#[test]
|
#[test]
|
||||||
fn negative_and_logprob_token_ids_rejected() {
|
fn negative_and_logprob_token_ids_rejected() {
|
||||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
let mut req = generate_req(22, SamplingParams::default());
|
let mut req = generate_req(22, SamplingParams::default());
|
||||||
if let RequestKind::Generate(g) = &mut req.kind {
|
if let RequestKind::Generate(g) = &mut req.kind {
|
||||||
g.input_ids = Some(vec![-1]);
|
g.input_ids = Some(vec![-1]);
|
||||||
@@ -1196,7 +1340,7 @@ mod tests {
|
|||||||
Ok(_) => panic!("negative token id must not be admitted"),
|
Ok(_) => panic!("negative token id must not be admitted"),
|
||||||
}
|
}
|
||||||
|
|
||||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
let mut req = generate_req(23, SamplingParams::default());
|
let mut req = generate_req(23, SamplingParams::default());
|
||||||
if let RequestKind::Generate(g) = &mut req.kind {
|
if let RequestKind::Generate(g) = &mut req.kind {
|
||||||
g.token_ids_logprob = Some(vec![999_999]);
|
g.token_ids_logprob = Some(vec![999_999]);
|
||||||
@@ -1211,7 +1355,7 @@ mod tests {
|
|||||||
/// A valid request is registered and handed onward — never deregistered.
|
/// A valid request is registered and handed onward — never deregistered.
|
||||||
#[test]
|
#[test]
|
||||||
fn admitted_request_keeps_registration() {
|
fn admitted_request_keeps_registration() {
|
||||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
// Empty map → all sampling defaults, passes normalization.
|
// Empty map → all sampling defaults, passes normalization.
|
||||||
ingress.drive(generate_req(9, SamplingParams::default()));
|
ingress.drive(generate_req(9, SamplingParams::default()));
|
||||||
|
|
||||||
@@ -1229,7 +1373,7 @@ mod tests {
|
|||||||
/// path and deregistered, not leaked.
|
/// path and deregistered, not leaked.
|
||||||
#[test]
|
#[test]
|
||||||
fn tokenize_failure_deregisters_via_ingress() {
|
fn tokenize_failure_deregisters_via_ingress() {
|
||||||
let (ingress, detok_rx, _consumer, tm_tx) = make_ingress();
|
let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress();
|
||||||
// The pool marks a failed encode as `Failed(err)` before returning it.
|
// The pool marks a failed encode as `Failed(err)` before returning it.
|
||||||
let mut req = generate_req(11, SamplingParams::default());
|
let mut req = generate_req(11, SamplingParams::default());
|
||||||
let _ = req
|
let _ = req
|
||||||
@@ -1253,7 +1397,7 @@ mod tests {
|
|||||||
fn abort_deregisters_from_shard() {
|
fn abort_deregisters_from_shard() {
|
||||||
// Aborts arrive on their own unbounded lane now, not the request inbox.
|
// Aborts arrive on their own unbounded lane now, not the request inbox.
|
||||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||||
let (ingress, detok_rx, _consumer, tm_tx) = make_ingress_with_abort(abort_rx);
|
let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress_with_abort(abort_rx);
|
||||||
abort_tx.send(AbortSource::Guard("rid-13".into())).unwrap();
|
abort_tx.send(AbortSource::Guard("rid-13".into())).unwrap();
|
||||||
drop(abort_tx);
|
drop(abort_tx);
|
||||||
drop(tm_tx);
|
drop(tm_tx);
|
||||||
@@ -1270,7 +1414,7 @@ mod tests {
|
|||||||
/// rejected; its registration is untouched.
|
/// rejected; its registration is untouched.
|
||||||
#[test]
|
#[test]
|
||||||
fn tokenized_return_pushes_without_deregister() {
|
fn tokenized_return_pushes_without_deregister() {
|
||||||
let (ingress, detok_rx, _consumer, tm_tx) = make_ingress();
|
let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress();
|
||||||
let mut req = generate_req(15, SamplingParams::default());
|
let mut req = generate_req(15, SamplingParams::default());
|
||||||
// Simulate a successful pool return: ids filled, PreSendValidating.
|
// Simulate a successful pool return: ids filled, PreSendValidating.
|
||||||
if let RequestKind::Generate(g) = &mut req.kind {
|
if let RequestKind::Generate(g) = &mut req.kind {
|
||||||
@@ -1293,7 +1437,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn tokenize_pool_gone_deregisters() {
|
fn tokenize_pool_gone_deregisters() {
|
||||||
// `make_ingress` drops the tok receiver, so `tok.send` fails.
|
// `make_ingress` drops the tok receiver, so `tok.send` fails.
|
||||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
// No ids → NeedsTokenize → Tokenizing branch.
|
// No ids → NeedsTokenize → Tokenizing branch.
|
||||||
let mut req = generate_req(21, SamplingParams::default());
|
let mut req = generate_req(21, SamplingParams::default());
|
||||||
if let RequestKind::Generate(g) = &mut req.kind {
|
if let RequestKind::Generate(g) = &mut req.kind {
|
||||||
@@ -1311,4 +1455,157 @@ mod tests {
|
|||||||
);
|
);
|
||||||
assert!(detok_rx.try_recv().is_err(), "no further shard messages");
|
assert!(detok_rx.try_recv().is_err(), "no further shard messages");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Build a generate request carrying an image. The parked entry and the
|
||||||
|
/// `MmEncoded` resume path agree on identity via the rid string.
|
||||||
|
fn mm_generate_req(rid: &str) -> Request {
|
||||||
|
let (tx, _rx) = mpsc::channel(8);
|
||||||
|
Request {
|
||||||
|
rid: rid.to_string().into(),
|
||||||
|
state: RequestState::Received,
|
||||||
|
sink: EgressSink::Local(tx),
|
||||||
|
kind: RequestKind::Generate(Box::new(GenerateRequest {
|
||||||
|
rid: rid.to_string().into(),
|
||||||
|
text: Some("<image> hi".into()),
|
||||||
|
mm: Some(Box::new(crate::message::MmData {
|
||||||
|
image_data: Some(rmpv::Value::from("data:image/jpeg;base64,xxxx")),
|
||||||
|
..Default::default()
|
||||||
|
})),
|
||||||
|
..Default::default()
|
||||||
|
})),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// An abort while the request is parked for MM cancels it: the pending
|
||||||
|
/// entry is removed, the worker's late result is dropped, and its parked
|
||||||
|
/// sidecar entry is purged — no scheduler work runs for a dead client.
|
||||||
|
#[test]
|
||||||
|
fn abort_cancels_parked_mm_request() {
|
||||||
|
let (mut ingress, _detok_rx, consumer, _tm_tx, mm_rx) = make_ingress();
|
||||||
|
ingress.drive(mm_generate_req("mm-gone"));
|
||||||
|
mm_rx.try_recv().expect("parked to mm pool");
|
||||||
|
|
||||||
|
// The worker parks its result, as it always does before MmEncoded.
|
||||||
|
ingress.mm.sidecar.park(
|
||||||
|
"mm-gone".into(),
|
||||||
|
crate::mm::MmSidecarEntry {
|
||||||
|
features: crate::mm::FeatureStore::Inline(vec![]),
|
||||||
|
grids: vec![],
|
||||||
|
hashes: vec![],
|
||||||
|
offsets: vec![],
|
||||||
|
mrope: vec![],
|
||||||
|
mrope_delta: 0,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
ingress.on_abort(AbortSource::Guard("mm-gone".to_string().into()));
|
||||||
|
assert_eq!(consumer.drain(16).headers.len(), 1, "only the AbortReq");
|
||||||
|
|
||||||
|
// The late result must be dropped, not queued, and the sidecar purged.
|
||||||
|
ingress.on_mm_encoded("mm-gone".to_string().into(), vec![5, 6]);
|
||||||
|
assert!(
|
||||||
|
consumer.drain(16).headers.is_empty(),
|
||||||
|
"cancelled, not queued"
|
||||||
|
);
|
||||||
|
assert!(ingress.mm.sidecar.take("mm-gone").is_none(), "entry purged");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A multimodal request parks in `Encoding` (submitted to the mm worker
|
||||||
|
/// pool, not the tokenizer pool, not the ring) until `MmEncoded` resumes
|
||||||
|
/// it → ring.
|
||||||
|
#[test]
|
||||||
|
fn mm_request_parks_then_mm_encoded_pushes_to_ring() {
|
||||||
|
let (mut ingress, _detok_rx, consumer, _tm_tx, mm_rx) = make_ingress();
|
||||||
|
ingress.drive(mm_generate_req("mm-1"));
|
||||||
|
|
||||||
|
// Submitted to the mm pool with the typed work item; nothing on the ring yet.
|
||||||
|
let sub = mm_rx.try_recv().expect("mm pool must receive the request");
|
||||||
|
assert_eq!(sub.rid.as_str(), "mm-1");
|
||||||
|
assert_eq!(sub.work.text.as_deref(), Some("<image> hi"));
|
||||||
|
assert!(sub.work.input_ids.is_none(), "no client input_ids");
|
||||||
|
assert_eq!(
|
||||||
|
sub.work.image_data.as_ref().and_then(|v| v.as_str()),
|
||||||
|
Some("data:image/jpeg;base64,xxxx")
|
||||||
|
);
|
||||||
|
assert!(consumer.drain(16).headers.is_empty(), "parked, not queued");
|
||||||
|
|
||||||
|
// The worker returns the final expanded ids → pushed to the ring.
|
||||||
|
ingress.on_mm_encoded("mm-1".to_string().into(), vec![5, 6, 7, 8]);
|
||||||
|
let batch = consumer.drain(16);
|
||||||
|
assert_eq!(batch.headers.len(), 1);
|
||||||
|
assert_eq!(
|
||||||
|
batch.lengths,
|
||||||
|
vec![4],
|
||||||
|
"expanded ids ride the columnar cell"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A worker failure rejects the parked request (deregister, no ring push).
|
||||||
|
#[test]
|
||||||
|
fn mm_failure_rejects_parked_request() {
|
||||||
|
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
|
ingress.drive(mm_generate_req("mm-2"));
|
||||||
|
assert!(
|
||||||
|
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { .. })),
|
||||||
|
"registered before parking",
|
||||||
|
);
|
||||||
|
|
||||||
|
ingress.on_mm_failed("mm-2".to_string().into(), "bad image".into());
|
||||||
|
assert!(
|
||||||
|
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid })
|
||||||
|
if rid.as_str() == "mm-2"),
|
||||||
|
"mm failure must deregister",
|
||||||
|
);
|
||||||
|
assert!(consumer.drain(16).headers.is_empty(), "nothing queued");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// On a non-multimodal model (`Mm::enabled == false`), image_data is silently
|
||||||
|
/// ignored and the request tokenizes as plain text — the Python
|
||||||
|
/// TokenizerManager behavior when `mm_processor is None`.
|
||||||
|
#[test]
|
||||||
|
fn mm_fields_ignored_when_disabled() {
|
||||||
|
let (tok_tx, tok_rx) = flume::unbounded();
|
||||||
|
let (detok_tx, _detok_rx) = flume::unbounded();
|
||||||
|
let senders = Senders {
|
||||||
|
tm: flume::unbounded().0,
|
||||||
|
abort: flume::unbounded().0,
|
||||||
|
tok: tok_tx,
|
||||||
|
detok: vec![detok_tx],
|
||||||
|
};
|
||||||
|
let (ingress_producer, _consumer) = ingress_ring(16);
|
||||||
|
let (_tm_tx, tm_rx) = flume::unbounded();
|
||||||
|
let (mm_tx, mm_rx) = flume::unbounded();
|
||||||
|
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||||
|
std::mem::forget(abort_tx);
|
||||||
|
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||||
|
std::mem::forget(sd_tx);
|
||||||
|
let mut ingress = Ingress::new(
|
||||||
|
tm_rx,
|
||||||
|
abort_rx,
|
||||||
|
senders,
|
||||||
|
ingress_producer,
|
||||||
|
test_limits(),
|
||||||
|
test_mm(mm_tx, false),
|
||||||
|
sd_rx,
|
||||||
|
);
|
||||||
|
|
||||||
|
ingress.drive(mm_generate_req("mm-3"));
|
||||||
|
assert!(
|
||||||
|
mm_rx.try_recv().is_err(),
|
||||||
|
"mm disabled: nothing submitted to the mm channel",
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
tok_rx.try_recv().is_ok(),
|
||||||
|
"request must fall through to plain tokenization",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A late mm result for a rid that is no longer parked is dropped without
|
||||||
|
/// panicking (e.g. hash-collision overwrite) — regression guard.
|
||||||
|
#[test]
|
||||||
|
fn late_mm_result_is_dropped() {
|
||||||
|
let (mut ingress, _detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||||
|
ingress.on_mm_encoded("ghost".to_string().into(), vec![1]);
|
||||||
|
ingress.on_mm_failed("ghost".to_string().into(), "boom".into());
|
||||||
|
assert!(consumer.drain(16).headers.is_empty());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -88,5 +88,12 @@ def image_bytes(width, height, seed=0):
|
|||||||
return buffer.getvalue()
|
return buffer.getvalue()
|
||||||
|
|
||||||
|
|
||||||
def spec_json(config, image_token_id=IMAGE_TOKEN_ID):
|
def spec_json(config, image_token_id=IMAGE_TOKEN_ID, resample="aten_u8"):
|
||||||
return json.dumps({"family": "qwen_vl", "image_token_id": image_token_id, **config})
|
return json.dumps(
|
||||||
|
{
|
||||||
|
"family": "qwen_vl",
|
||||||
|
"image_token_id": image_token_id,
|
||||||
|
"resample": resample,
|
||||||
|
**config,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|||||||
@@ -0,0 +1,102 @@
|
|||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from tokenizers import Tokenizer, decoders
|
||||||
|
from tokenizers.models import WordLevel
|
||||||
|
from tokenizers.pre_tokenizers import WhitespaceSplit
|
||||||
|
from transformers import (
|
||||||
|
PreTrainedTokenizerFast,
|
||||||
|
Qwen2VLProcessor,
|
||||||
|
Qwen2VLVideoProcessor,
|
||||||
|
)
|
||||||
|
from transformers.models.qwen2_vl.image_processing_qwen2_vl import (
|
||||||
|
Qwen2VLImageProcessor as HfQwenImageProcessor,
|
||||||
|
)
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import maybe_stub_sgl_kernel
|
||||||
|
|
||||||
|
maybe_stub_sgl_kernel()
|
||||||
|
|
||||||
|
from sglang.srt.multimodal.processors.qwen_vl import ( # noqa: E402
|
||||||
|
QwenVLImageProcessor,
|
||||||
|
)
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=0, suite="base-a-test-cpu", disabled="Qwen test fixtures")
|
||||||
|
|
||||||
|
|
||||||
|
def make_processor(config, image_processor_cls=None):
|
||||||
|
"""A ``QwenVLImageProcessor`` over a tiny hand-built tokenizer.
|
||||||
|
``image_processor_cls`` picks the HF backend; they resample differently."""
|
||||||
|
image_processor_cls = image_processor_cls or HfQwenImageProcessor
|
||||||
|
vocab = [
|
||||||
|
"<unk>",
|
||||||
|
"<|vision_start|>",
|
||||||
|
"<|image_pad|>",
|
||||||
|
"<|vision_end|>",
|
||||||
|
"hello",
|
||||||
|
"<|video_pad|>",
|
||||||
|
"<pad>",
|
||||||
|
]
|
||||||
|
backend = Tokenizer(
|
||||||
|
WordLevel(
|
||||||
|
{token: index for index, token in enumerate(vocab)}, unk_token=vocab[0]
|
||||||
|
)
|
||||||
|
)
|
||||||
|
backend.pre_tokenizer, backend.decoder = WhitespaceSplit(), decoders.Fuse()
|
||||||
|
tokenizer = PreTrainedTokenizerFast(
|
||||||
|
tokenizer_object=backend,
|
||||||
|
unk_token=vocab[0],
|
||||||
|
pad_token=vocab[-1],
|
||||||
|
additional_special_tokens=vocab[1:4] + [vocab[5]],
|
||||||
|
)
|
||||||
|
processor = Qwen2VLProcessor(
|
||||||
|
image_processor=image_processor_cls(**config),
|
||||||
|
video_processor=Qwen2VLVideoProcessor(),
|
||||||
|
tokenizer=tokenizer,
|
||||||
|
)
|
||||||
|
hf_config = SimpleNamespace(
|
||||||
|
model_type="qwen2_5_vl",
|
||||||
|
architectures=["Qwen2_5_VLForConditionalGeneration"],
|
||||||
|
vision_start_token_id=1,
|
||||||
|
image_token_id=2,
|
||||||
|
vision_end_token_id=3,
|
||||||
|
video_token_id=5,
|
||||||
|
vision_config=SimpleNamespace(spatial_merge_size=2, tokens_per_second=2),
|
||||||
|
)
|
||||||
|
server_args = SimpleNamespace(
|
||||||
|
# Non-auto: get_resolved_model_impl would choke on a SimpleNamespace.
|
||||||
|
model_impl="sglang",
|
||||||
|
keep_mm_feature_on_device=False,
|
||||||
|
mm_feature_transport="cpu",
|
||||||
|
disable_fast_image_processor=True,
|
||||||
|
skip_tokenizer_init=False,
|
||||||
|
# Read by NativeMmHost._use_feature_shm (single-rank fixture → the
|
||||||
|
# inline zero-copy transport, like the 1-GPU e2e).
|
||||||
|
tp_size=1,
|
||||||
|
dist_init_addr=None,
|
||||||
|
mm_process_config={},
|
||||||
|
mm_io_worker_num=1,
|
||||||
|
mm_processor_worker_num=1,
|
||||||
|
tokenizer_worker_num=1,
|
||||||
|
base_gpu_id=0,
|
||||||
|
)
|
||||||
|
return QwenVLImageProcessor(
|
||||||
|
hf_config, server_args, processor, None, skip_mm_pool=True
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def snapshot(input_ids, output):
|
||||||
|
return {
|
||||||
|
"input_ids": tuple(input_ids),
|
||||||
|
"grids": tuple(
|
||||||
|
tuple(item.image_grid_thw.flatten().tolist()) for item in output.mm_items
|
||||||
|
),
|
||||||
|
"offsets": tuple(item.offsets[0] for item in output.mm_items),
|
||||||
|
"features": np.concatenate(
|
||||||
|
[item.feature.detach().cpu().numpy() for item in output.mm_items]
|
||||||
|
),
|
||||||
|
"mrope": output.mrope_positions.detach().cpu().numpy(),
|
||||||
|
"delta": int(output.mrope_position_delta.item()),
|
||||||
|
"tokens": (output.im_start_id, output.im_token_id, output.im_end_id),
|
||||||
|
}
|
||||||
@@ -0,0 +1,162 @@
|
|||||||
|
"""End-to-end parity at the scheduler-input boundary.
|
||||||
|
|
||||||
|
`test_preprocess.py` pins the `preprocess` binding; this drives the whole native
|
||||||
|
path — the `process_native_mm` driver, then `NativeMmHost.build_native_mm` — and
|
||||||
|
compares every field the scheduler reads against the Python `mm_processor`.
|
||||||
|
Bitwise, for both HF backends: the Rust resize clones PIL's fixed-point bicubic
|
||||||
|
and ATen's uint8 antialias kernel, so whichever one a server is configured with
|
||||||
|
is reproduced exactly.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import base64
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import unittest
|
||||||
|
from pathlib import Path
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||||
|
|
||||||
|
maybe_stub_sgl_kernel()
|
||||||
|
|
||||||
|
from sglang.srt.managers.rust_server import NativeMmHost # noqa: E402
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||||
|
from _fixtures import make_processor, snapshot # noqa: E402
|
||||||
|
from _mm_rust_utils import PROCESSOR_CONFIGS, image_bytes, load_core # noqa: E402
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=40, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
CORE = load_core()
|
||||||
|
DRIVER = getattr(getattr(CORE, "qwen_vl", None), "process_native_mm", None)
|
||||||
|
|
||||||
|
# The fixture tokenizer's vocab (see `_fixtures.make_processor`):
|
||||||
|
# 1 = <|vision_start|>, 2 = <|image_pad|>, 3 = <|vision_end|>, 4 = "hello".
|
||||||
|
PROMPT_PER_IMAGE = [1, 2, 3, 4]
|
||||||
|
|
||||||
|
|
||||||
|
@unittest.skipUnless(DRIVER, "sglang-mm native Qwen driver not built")
|
||||||
|
class TestQwenE2eParity(CustomTestCase):
|
||||||
|
"""The torchvision backend a default server runs."""
|
||||||
|
|
||||||
|
image_processor = "Qwen2VLImageProcessor"
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
import transformers.models.qwen2_vl as qwen2_vl
|
||||||
|
|
||||||
|
self.processor = make_processor(
|
||||||
|
PROCESSOR_CONFIGS["qwen2_5_vl"], getattr(qwen2_vl, self.image_processor)
|
||||||
|
)
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
self.processor.io_executor.shutdown()
|
||||||
|
self.processor.cpu_executor.shutdown()
|
||||||
|
|
||||||
|
# --- the two paths under comparison ---
|
||||||
|
|
||||||
|
def native_spec(self):
|
||||||
|
"""Resolve the spec through the production gate, so a gate that stops
|
||||||
|
recognizing this image processor fails here too."""
|
||||||
|
from sglang.srt.managers.multimodal_processor import import_processors
|
||||||
|
|
||||||
|
import_processors("sglang.srt.multimodal.processors")
|
||||||
|
# Skip __init__: it would build a processor; reuse the fixture's.
|
||||||
|
host = NativeMmHost.__new__(NativeMmHost)
|
||||||
|
host.model_config = SimpleNamespace(hf_config=self.processor.hf_config)
|
||||||
|
host._processor = self.processor._processor
|
||||||
|
host.server_args = self.processor.server_args
|
||||||
|
spec = host.resolve_native_spec()
|
||||||
|
self.assertIsNotNone(spec, f"gate rejected {self.image_processor}")
|
||||||
|
return spec
|
||||||
|
|
||||||
|
def run_native(self, spec, sources):
|
||||||
|
"""The Rust path: the `process_native_mm` driver, then the drain
|
||||||
|
adapter — the same two steps `RustServer.drain` performs."""
|
||||||
|
ids, features, grids, hashes, offsets, mrope, delta = DRIVER(
|
||||||
|
PROMPT_PER_IMAGE * len(sources), sources, spec.rust_json()
|
||||||
|
)
|
||||||
|
# The shape of Rust's MmHandoff, inline transport (test_build_native_mm
|
||||||
|
# pins the shm shape).
|
||||||
|
handoff = SimpleNamespace(
|
||||||
|
features=features,
|
||||||
|
shm_names=None,
|
||||||
|
grids=grids,
|
||||||
|
hashes=hashes,
|
||||||
|
offsets=offsets,
|
||||||
|
mrope=mrope,
|
||||||
|
mrope_delta=delta,
|
||||||
|
)
|
||||||
|
return snapshot(ids, NativeMmHost.build_native_mm(spec, handoff))
|
||||||
|
|
||||||
|
def run_python(self, sources):
|
||||||
|
"""The reference path: the Python `mm_processor` the scheduler would use."""
|
||||||
|
output = asyncio.run(
|
||||||
|
self.processor.process_mm_data_async(
|
||||||
|
image_data=sources,
|
||||||
|
input_text=PROMPT_PER_IMAGE * len(sources),
|
||||||
|
request_obj=SimpleNamespace(
|
||||||
|
video_data=None, audio_data=None, rid="parity"
|
||||||
|
),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return snapshot(output.input_ids, output)
|
||||||
|
|
||||||
|
def assert_parity(self, spec, sources):
|
||||||
|
rust, python = self.run_native(spec, sources), self.run_python(sources)
|
||||||
|
for field in ("input_ids", "grids", "offsets", "delta", "tokens"):
|
||||||
|
with self.subTest(field=field):
|
||||||
|
self.assertEqual(rust[field], python[field])
|
||||||
|
with self.subTest(field="mrope"):
|
||||||
|
np.testing.assert_array_equal(rust["mrope"], python["mrope"])
|
||||||
|
with self.subTest(field="features"):
|
||||||
|
# Bytes, not allclose: the scheduler gets these float32 buffers verbatim.
|
||||||
|
self.assertEqual(rust["features"].tobytes(), python["features"].tobytes())
|
||||||
|
|
||||||
|
# --- inputs ---
|
||||||
|
|
||||||
|
def source_forms(self, directory):
|
||||||
|
"""One image in each accepted transport form, plus a two-image batch."""
|
||||||
|
first, second = image_bytes(96, 80), image_bytes(112, 88, 1)
|
||||||
|
path = Path(directory) / "image.png"
|
||||||
|
path.write_bytes(first)
|
||||||
|
return {
|
||||||
|
"raw_bytes": [first],
|
||||||
|
"data_url": ["data:image/png;base64," + base64.b64encode(first).decode()],
|
||||||
|
"file_uri": [path.as_uri()],
|
||||||
|
"two_image_batch": [first, second],
|
||||||
|
}
|
||||||
|
|
||||||
|
def test_parity_across_source_forms(self):
|
||||||
|
spec = self.native_spec()
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
for form, sources in self.source_forms(directory).items():
|
||||||
|
with self.subTest(form=form):
|
||||||
|
self.assert_parity(spec, sources)
|
||||||
|
|
||||||
|
|
||||||
|
class TestQwenE2eParityPil(TestQwenE2eParity):
|
||||||
|
"""The PIL backend (`--disable-fast-image-processor`): the same parity
|
||||||
|
suite, plus the transport-invariance check — a property of the native path
|
||||||
|
alone, so running it under one backend is enough."""
|
||||||
|
|
||||||
|
image_processor = "Qwen2VLImageProcessorPil"
|
||||||
|
|
||||||
|
def test_source_form_is_transport_only(self):
|
||||||
|
"""Bytes, a data: URL and a file:// path must yield one identical result."""
|
||||||
|
spec = self.native_spec()
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
forms = self.source_forms(directory)
|
||||||
|
features = {
|
||||||
|
form: self.run_native(spec, forms[form])["features"].tobytes()
|
||||||
|
for form in ("raw_bytes", "data_url", "file_uri")
|
||||||
|
}
|
||||||
|
self.assertEqual(len(set(features.values())), 1, sorted(features))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -0,0 +1,105 @@
|
|||||||
|
"""The mm-item hash contract of the native Qwen path.
|
||||||
|
|
||||||
|
The two paths hash different things, deliberately. Native hashes are `content_hash`
|
||||||
|
over the raw encoded source bytes, computed on an MM worker; the Python path hashes
|
||||||
|
the decoded feature tensor. Both feed `set_pad_value`, so the native drain can skip
|
||||||
|
`hash_feature` on the scheduler loop — the point of precomputing them.
|
||||||
|
|
||||||
|
Field-by-field parity of everything else is `test_e2e_parity.py`.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import base64
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import unittest
|
||||||
|
from pathlib import Path
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||||
|
|
||||||
|
maybe_stub_sgl_kernel()
|
||||||
|
|
||||||
|
from sglang.srt.managers.mm_utils import hash_feature # noqa: E402
|
||||||
|
from sglang.srt.managers.rust_server import NativeMmHost # noqa: E402
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||||
|
from _fixtures import make_processor # noqa: E402
|
||||||
|
from _mm_rust_utils import PROCESSOR_CONFIGS, image_bytes, load_core # noqa: E402
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
CORE = load_core()
|
||||||
|
DRIVER = getattr(getattr(CORE, "qwen_vl", None), "process_native_mm", None)
|
||||||
|
|
||||||
|
|
||||||
|
def raw_bytes(source):
|
||||||
|
"""The encoded bytes behind any accepted source form."""
|
||||||
|
if isinstance(source, bytes):
|
||||||
|
return source
|
||||||
|
if source.startswith("data:"):
|
||||||
|
return base64.b64decode(source.split(",", 1)[1])
|
||||||
|
return Path(source.removeprefix("file://")).read_bytes()
|
||||||
|
|
||||||
|
|
||||||
|
@unittest.skipUnless(DRIVER, "sglang-mm native Qwen driver not built")
|
||||||
|
class TestQwenNativeMmHashes(CustomTestCase):
|
||||||
|
def setUp(self):
|
||||||
|
from sglang.srt.managers.multimodal_processor import import_processors
|
||||||
|
|
||||||
|
import_processors("sglang.srt.multimodal.processors")
|
||||||
|
self.processor = make_processor(PROCESSOR_CONFIGS["qwen2_5_vl"])
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
self.processor.io_executor.shutdown()
|
||||||
|
self.processor.cpu_executor.shutdown()
|
||||||
|
|
||||||
|
def native_hashes(self, sources):
|
||||||
|
"""Per-item hashes the Rust driver returns, via the production gate."""
|
||||||
|
host = NativeMmHost.__new__(NativeMmHost)
|
||||||
|
host.model_config = SimpleNamespace(hf_config=self.processor.hf_config)
|
||||||
|
host._processor = self.processor._processor
|
||||||
|
host.server_args = self.processor.server_args
|
||||||
|
spec = host.resolve_native_spec()
|
||||||
|
self.assertIsNotNone(spec, "gate rejected the fixture processor")
|
||||||
|
input_ids = [t for _ in sources for t in (1, 2, 3, 4)]
|
||||||
|
return DRIVER(input_ids, sources, spec.rust_json())[3]
|
||||||
|
|
||||||
|
def test_native_hashes_the_raw_source_bytes(self):
|
||||||
|
"""Same image in any source form hashes identically, because the hash is
|
||||||
|
over the bytes and not over anything the transport changes."""
|
||||||
|
first, second = image_bytes(96, 80), image_bytes(112, 88, 1)
|
||||||
|
data_url = "data:image/png;base64," + base64.b64encode(first).decode()
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
path = Path(directory) / "image.png"
|
||||||
|
path.write_bytes(first)
|
||||||
|
for sources in ([first], [data_url], [path.as_uri()], [first, second]):
|
||||||
|
with self.subTest(n=len(sources), form=type(sources[0]).__name__):
|
||||||
|
hashes = self.native_hashes(sources)
|
||||||
|
self.assertEqual(
|
||||||
|
list(hashes),
|
||||||
|
[CORE.common.content_hash(raw_bytes(s)) for s in sources],
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_python_hashes_the_feature(self):
|
||||||
|
"""The contrast: `set_pad_value` on the Python path derives its hash from
|
||||||
|
the feature tensor, which is why the native path must precompute one."""
|
||||||
|
output = asyncio.run(
|
||||||
|
self.processor.process_mm_data_async(
|
||||||
|
image_data=[image_bytes(96, 80)],
|
||||||
|
input_text=[1, 2, 3, 4],
|
||||||
|
request_obj=SimpleNamespace(
|
||||||
|
video_data=None, audio_data=None, rid="hash"
|
||||||
|
),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
for item in output.mm_items:
|
||||||
|
expected = hash_feature(item.feature)
|
||||||
|
item.set_pad_value()
|
||||||
|
self.assertEqual(item.hash, expected)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -5,14 +5,13 @@ Covers ``QwenVlProcessor::process_item`` and ``smart_resize`` in
|
|||||||
and ``smart_resize_py`` bindings), against the HF Qwen2-VL image processors
|
and ``smart_resize_py`` bindings), against the HF Qwen2-VL image processors
|
||||||
and the Python ``smart_resize``.
|
and the Python ``smart_resize``.
|
||||||
|
|
||||||
Both HF processors are pinned, because they resample differently and only one
|
Both HF processors are pinned, because they resample differently and either can
|
||||||
of them is what a server actually runs. On transformers 5.x
|
be what a server runs. On transformers 5.x ``Qwen2VLImageProcessor`` is the
|
||||||
``Qwen2VLImageProcessor`` is the torchvision path (the ``Fast`` suffix was
|
torchvision path (the ``Fast`` suffix was dropped) and is what
|
||||||
dropped) and is what ``AutoImageProcessor`` hands SGLang by default;
|
``AutoImageProcessor`` hands SGLang by default; ``Qwen2VLImageProcessorPil`` is
|
||||||
``Qwen2VLImageProcessorPil`` is the PIL path, reachable via
|
the PIL path, reachable via ``--disable-fast-image-processor``. The Rust resize
|
||||||
``--disable-fast-image-processor``. The Rust resize is a bit-exact clone of
|
clones both kernels, so each is asserted with zero tolerance under the
|
||||||
PIL's fixed-point kernel, so the PIL processor is asserted exactly and the
|
``resample`` its spec selects.
|
||||||
torchvision one carries the cross-implementation envelope.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import sys
|
import sys
|
||||||
@@ -43,48 +42,39 @@ SIZES = ((640, 480), (1024, 683), (50, 40), (300, 301))
|
|||||||
|
|
||||||
@unittest.skipUnless(QWEN_CORE, "sglang-mm Qwen binding not built")
|
@unittest.skipUnless(QWEN_CORE, "sglang-mm Qwen binding not built")
|
||||||
class TestQwenImagePreprocess(CustomTestCase):
|
class TestQwenImagePreprocess(CustomTestCase):
|
||||||
def _assert_matches(self, processor, max_diff, mean_diff):
|
def _assert_bit_exact(self, processor, resample):
|
||||||
|
"""Zero tolerance, so any drift in any stage — smart_resize geometry, the
|
||||||
|
fixed-point kernel, rescale/normalize, HF patch order — shows up here
|
||||||
|
instead of being absorbed."""
|
||||||
for family, config in PROCESSOR_CONFIGS.items():
|
for family, config in PROCESSOR_CONFIGS.items():
|
||||||
hf = processor(**config)
|
hf = processor(**config)
|
||||||
for index, size in enumerate(SIZES):
|
for index, size in enumerate(SIZES):
|
||||||
with self.subTest(family=family, size=size):
|
with self.subTest(family=family, size=size):
|
||||||
image = make_image(*size, seed=index)
|
image = make_image(*size, seed=index)
|
||||||
actual, grid = QWEN_CORE.preprocess(
|
actual, grid = QWEN_CORE.preprocess(
|
||||||
image_bytes(*size, seed=index), spec_json(config)
|
image_bytes(*size, seed=index),
|
||||||
|
spec_json(config, resample=resample),
|
||||||
)
|
)
|
||||||
expected = hf(images=[image], return_tensors="pt")
|
expected = hf(images=[image], return_tensors="pt")
|
||||||
self.assertEqual(grid, tuple(expected.image_grid_thw[0].tolist()))
|
self.assertEqual(grid, tuple(expected.image_grid_thw[0].tolist()))
|
||||||
diff = np.abs(
|
np.testing.assert_array_equal(
|
||||||
np.asarray(actual).reshape(expected.pixel_values.shape)
|
np.asarray(actual).reshape(expected.pixel_values.shape),
|
||||||
- expected.pixel_values.numpy()
|
expected.pixel_values.numpy(),
|
||||||
)
|
)
|
||||||
# LessEqual, not Less: max_diff=0.0 is a real bound here.
|
|
||||||
self.assertLessEqual(diff.max(), max_diff)
|
|
||||||
self.assertLessEqual(diff.mean(), mean_diff)
|
|
||||||
|
|
||||||
def test_features_match_pil_processor_exactly(self):
|
def test_features_match_pil_processor_exactly(self):
|
||||||
"""Against the PIL processor the native path is bit-exact, so this is
|
|
||||||
asserted with zero tolerance: every stage (smart_resize geometry,
|
|
||||||
the fixed-point bicubic kernel, rescale/normalize, HF patch order) is
|
|
||||||
pinned, and any drift in any of them shows up here rather than being
|
|
||||||
absorbed by a tolerance."""
|
|
||||||
from transformers.models.qwen2_vl.image_processing_pil_qwen2_vl import (
|
from transformers.models.qwen2_vl.image_processing_pil_qwen2_vl import (
|
||||||
Qwen2VLImageProcessorPil,
|
Qwen2VLImageProcessorPil,
|
||||||
)
|
)
|
||||||
|
|
||||||
self._assert_matches(Qwen2VLImageProcessorPil, max_diff=0.0, mean_diff=0.0)
|
self._assert_bit_exact(Qwen2VLImageProcessorPil, "pil")
|
||||||
|
|
||||||
def test_features_match_torchvision_processor_within_envelope(self):
|
def test_features_match_torchvision_processor_exactly(self):
|
||||||
"""The torchvision processor is what a default server runs, so its
|
|
||||||
divergence is bounded separately — it is a different antialiased-bicubic
|
|
||||||
implementation, which the bit-exact PIL assertion above says nothing
|
|
||||||
about. Measured worst case is max 0.030 / mean 6.7e-5 (≈2 u8 levels
|
|
||||||
after normalize with the qwen2_vl std)."""
|
|
||||||
from transformers.models.qwen2_vl.image_processing_qwen2_vl import (
|
from transformers.models.qwen2_vl.image_processing_qwen2_vl import (
|
||||||
Qwen2VLImageProcessor,
|
Qwen2VLImageProcessor,
|
||||||
)
|
)
|
||||||
|
|
||||||
self._assert_matches(Qwen2VLImageProcessor, max_diff=0.035, mean_diff=1e-3)
|
self._assert_bit_exact(Qwen2VLImageProcessor, "aten_u8")
|
||||||
|
|
||||||
def test_smart_resize_matches_python(self):
|
def test_smart_resize_matches_python(self):
|
||||||
from sglang.srt.multimodal.processors.qwen_vl import smart_resize
|
from sglang.srt.multimodal.processors.qwen_vl import smart_resize
|
||||||
|
|||||||
@@ -0,0 +1,166 @@
|
|||||||
|
"""``NativeMmHost.build_native_mm`` (managers/rust_server.py): the drain-time
|
||||||
|
wrapping contracts — tensors are zero-copy views over the Rust-owned buffers, and
|
||||||
|
pad values come from worker-precomputed hashes, since the scheduler loop must
|
||||||
|
never hash features. Synthetic buffers, so this needs no Rust extension."""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import unittest
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||||
|
|
||||||
|
maybe_stub_sgl_kernel()
|
||||||
|
|
||||||
|
from sglang.srt.managers.rust_server import NativeMmHost, NativeMmSpec # noqa: E402
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
|
||||||
|
class TestBuildNativeMm(CustomTestCase):
|
||||||
|
def setUp(self):
|
||||||
|
# feature_dim == 3 * temporal_patch_size * patch_size**2 == 6.
|
||||||
|
self.spec = NativeMmSpec(
|
||||||
|
family="qwen_vl",
|
||||||
|
feature_shm=False,
|
||||||
|
image_token_id=10,
|
||||||
|
patch_size=1,
|
||||||
|
merge_size=1,
|
||||||
|
temporal_patch_size=2,
|
||||||
|
min_pixels=1,
|
||||||
|
max_pixels=1 << 30,
|
||||||
|
image_mean=(0.0, 0.0, 0.0),
|
||||||
|
image_std=(1.0, 1.0, 1.0),
|
||||||
|
resample="aten_u8",
|
||||||
|
vision_start_token_id=11,
|
||||||
|
vision_end_token_id=12,
|
||||||
|
video_token_id=13,
|
||||||
|
)
|
||||||
|
|
||||||
|
GRIDS = [(1, 2, 2), (1, 1, 1)]
|
||||||
|
HASHES = [101, 202]
|
||||||
|
OFFSETS = [(2, 5), (8, 8)]
|
||||||
|
|
||||||
|
def transport(self, features):
|
||||||
|
"""Inline: the features ride the numpy array itself."""
|
||||||
|
return dict(features=features, shm_names=None)
|
||||||
|
|
||||||
|
def build(self):
|
||||||
|
features = np.arange(30, dtype=np.float32)
|
||||||
|
output = NativeMmHost.build_native_mm(
|
||||||
|
self.spec,
|
||||||
|
SimpleNamespace( # the shape of Rust's MmHandoff
|
||||||
|
grids=self.GRIDS,
|
||||||
|
hashes=self.HASHES,
|
||||||
|
offsets=self.OFFSETS,
|
||||||
|
mrope=np.arange(30, dtype=np.int64),
|
||||||
|
mrope_delta=-3,
|
||||||
|
**self.transport(features),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
return output, features
|
||||||
|
|
||||||
|
def test_wraps_and_slices_native_buffers(self):
|
||||||
|
output, features = self.build()
|
||||||
|
self.assertEqual(
|
||||||
|
[tuple(item.feature.shape) for item in output.mm_items], [(4, 6), (1, 6)]
|
||||||
|
)
|
||||||
|
self.assertEqual([item.hash for item in output.mm_items], [101, 202])
|
||||||
|
self.assertEqual(
|
||||||
|
[item.offsets for item in output.mm_items], [[(2, 5)], [(8, 8)]]
|
||||||
|
)
|
||||||
|
self.assertEqual(tuple(output.mrope_positions.shape), (3, 10))
|
||||||
|
self.assertEqual(output.mrope_position_delta.item(), -3)
|
||||||
|
self.assertEqual(
|
||||||
|
(output.im_start_id, output.im_token_id, output.im_end_id), (11, 10, 12)
|
||||||
|
)
|
||||||
|
features[0] = 99
|
||||||
|
self.assertEqual(output.mm_items[0].feature[0, 0].item(), 99)
|
||||||
|
|
||||||
|
def test_optional_pad_values_use_precomputed_hashes(self):
|
||||||
|
from sglang.srt.managers.schedule_batch import _compute_pad_value
|
||||||
|
|
||||||
|
# The whole point of worker-precomputed hashes is that the scheduler
|
||||||
|
# loop never runs hash_feature — make any call a hard failure.
|
||||||
|
with (
|
||||||
|
patch.dict(os.environ, {"SGLANG_MM_PRECOMPUTE_HASH": "1"}),
|
||||||
|
patch(
|
||||||
|
"sglang.srt.managers.mm_utils.hash_feature",
|
||||||
|
side_effect=AssertionError("scheduler loop must not hash features"),
|
||||||
|
),
|
||||||
|
):
|
||||||
|
output, _ = self.build()
|
||||||
|
self.assertEqual(
|
||||||
|
[item.pad_value for item in output.mm_items],
|
||||||
|
[_compute_pad_value(101), _compute_pad_value(202)],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestBuildNativeMmShm(TestBuildNativeMm):
|
||||||
|
"""The shm entry shape (TP>1): features arrive as named POSIX segments, and
|
||||||
|
each item becomes a ``ShmPointerMMData`` stub whose ``materialize()`` yields
|
||||||
|
that item's slice — and unlinks, taking the cleanup duty exactly once."""
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
super().setUp()
|
||||||
|
self._segments = []
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
# Defensive: unlink anything a failing test left behind.
|
||||||
|
for shm in self._segments:
|
||||||
|
try:
|
||||||
|
shm.close()
|
||||||
|
shm.unlink()
|
||||||
|
except FileNotFoundError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _park(self, features):
|
||||||
|
from multiprocessing import shared_memory
|
||||||
|
|
||||||
|
names, row = [], 0
|
||||||
|
for t, h, w in self.GRIDS:
|
||||||
|
n = t * h * w
|
||||||
|
payload = features[row * 6 : (row + n) * 6].tobytes()
|
||||||
|
shm = shared_memory.SharedMemory(create=True, size=len(payload))
|
||||||
|
shm.buf[:] = payload
|
||||||
|
self._segments.append(shm)
|
||||||
|
names.append(shm.name)
|
||||||
|
row += n
|
||||||
|
return names
|
||||||
|
|
||||||
|
def transport(self, features):
|
||||||
|
"""Shm: the worker parked each item's slice in its own segment."""
|
||||||
|
return dict(features=None, shm_names=self._park(features))
|
||||||
|
|
||||||
|
def test_wraps_and_slices_native_buffers(self):
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.managers.mm_utils import ShmPointerMMData
|
||||||
|
|
||||||
|
output, features = self.build()
|
||||||
|
for item in output.mm_items:
|
||||||
|
self.assertIsInstance(item.feature, ShmPointerMMData)
|
||||||
|
# The stub is a zero-copy view over the segment until materialized.
|
||||||
|
self.assertEqual(
|
||||||
|
[tuple(item.feature.shape) for item in output.mm_items], [(4, 6), (1, 6)]
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
[item.feature.precomputed_hash for item in output.mm_items], self.HASHES
|
||||||
|
)
|
||||||
|
tensors = [item.feature.materialize() for item in output.mm_items]
|
||||||
|
expected = torch.from_numpy(features).reshape(-1, 6)
|
||||||
|
self.assertTrue(torch.equal(tensors[0], expected[:4]))
|
||||||
|
self.assertTrue(torch.equal(tensors[1], expected[4:]))
|
||||||
|
# materialize() unlinked: the names must be gone.
|
||||||
|
from multiprocessing import shared_memory
|
||||||
|
|
||||||
|
for item in output.mm_items:
|
||||||
|
with self.assertRaises(FileNotFoundError):
|
||||||
|
shared_memory.SharedMemory(name=item.feature.shm_name)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -0,0 +1,61 @@
|
|||||||
|
"""The native-MM launch gate's family selection (managers/rust_server.py).
|
||||||
|
|
||||||
|
``NATIVE_MM_FAMILIES`` decides which models the Rust pipeline serves natively;
|
||||||
|
for everything else ``native_mm_family_for`` must return ``None``, which
|
||||||
|
``RustServer.launch`` turns into a hard launch error. Pins that non-Qwen
|
||||||
|
multimodal models — Inkling being the in-tree case — keep their Python
|
||||||
|
processor and never match a native family, so growing the registry cannot
|
||||||
|
silently reroute them.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import unittest
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||||
|
|
||||||
|
maybe_stub_sgl_kernel()
|
||||||
|
|
||||||
|
from sglang.srt.managers.multimodal_processor import ( # noqa: E402
|
||||||
|
get_mm_processor_cls,
|
||||||
|
import_processors,
|
||||||
|
)
|
||||||
|
from sglang.srt.managers.rust_server import native_mm_family_for # noqa: E402
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
|
||||||
|
def processor_cls_for(architecture, model_type):
|
||||||
|
"""Through the production selection, as `resolve_native_spec` calls it."""
|
||||||
|
hf_config = SimpleNamespace(architectures=[architecture], model_type=model_type)
|
||||||
|
return get_mm_processor_cls(hf_config, SimpleNamespace(model_impl="sglang"))
|
||||||
|
|
||||||
|
|
||||||
|
class TestNativeMmGate(CustomTestCase):
|
||||||
|
@classmethod
|
||||||
|
def setUpClass(cls):
|
||||||
|
import_processors("sglang.srt.multimodal.processors")
|
||||||
|
|
||||||
|
def test_qwen_vl_resolves_its_family(self):
|
||||||
|
cls = processor_cls_for("Qwen2_5_VLForConditionalGeneration", "qwen2_5_vl")
|
||||||
|
family = native_mm_family_for(cls, "qwen2_5_vl")
|
||||||
|
self.assertEqual(family and family.name, "qwen_vl")
|
||||||
|
|
||||||
|
def test_inkling_keeps_its_python_processor(self):
|
||||||
|
from sglang.srt.multimodal.processors.inkling import InklingMultimodalProcessor
|
||||||
|
|
||||||
|
cls = processor_cls_for("InklingForConditionalGeneration", "inkling_model")
|
||||||
|
self.assertIs(cls, InklingMultimodalProcessor)
|
||||||
|
self.assertIsNone(native_mm_family_for(cls, "inkling_model"))
|
||||||
|
|
||||||
|
def test_family_requires_both_processor_and_model_type(self):
|
||||||
|
qwen = processor_cls_for("Qwen2_5_VLForConditionalGeneration", "qwen2_5_vl")
|
||||||
|
self.assertIsNone(native_mm_family_for(qwen, "inkling_model"))
|
||||||
|
# Identity, not name: an override class must not match (the
|
||||||
|
# SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE contract).
|
||||||
|
impostor = type("QwenVLImageProcessor", (), {})
|
||||||
|
self.assertIsNone(native_mm_family_for(impostor, "qwen2_5_vl"))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -0,0 +1,54 @@
|
|||||||
|
"""``RustServer._partition_cores`` (managers/rust_server.py): the pool cores must
|
||||||
|
be a *bounded* slice of this rank's allowed cores, not the whole remainder —
|
||||||
|
sibling TP ranks share the NUMA node, so an unbounded mask lets MM preprocessing
|
||||||
|
bursts preempt a sibling's CUDA-launch thread (measured: ~20 ms of ViT wall time
|
||||||
|
per image request on TP4). Pure computation, so no Rust extension needed."""
|
||||||
|
|
||||||
|
import unittest
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||||
|
|
||||||
|
maybe_stub_sgl_kernel()
|
||||||
|
|
||||||
|
from sglang.srt.managers.rust_server import RustServer # noqa: E402
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
|
||||||
|
def partition(node_cores, **kwargs):
|
||||||
|
with patch("os.sched_getaffinity", return_value=set(node_cores), create=True):
|
||||||
|
return RustServer._partition_cores(**kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
class TestPartitionCores(CustomTestCase):
|
||||||
|
def test_pool_is_bounded_not_the_node_remainder(self):
|
||||||
|
# A 120-core NUMA node shared with sibling TP ranks: the pools must
|
||||||
|
# NOT get cores 2..119.
|
||||||
|
launch, pool = partition(range(120), mm_workers=8)
|
||||||
|
self.assertEqual(launch, [0, 1])
|
||||||
|
self.assertEqual(pool, list(range(2, 14))) # max(8, 8 + 4) after reserve
|
||||||
|
|
||||||
|
def test_budget_scales_with_mm_workers_with_a_floor(self):
|
||||||
|
_, pool_text = partition(range(120), mm_workers=0)
|
||||||
|
_, pool_mm = partition(range(120), mm_workers=16)
|
||||||
|
self.assertEqual(len(pool_text), 8) # the floor covers the I/O threads
|
||||||
|
self.assertEqual(len(pool_mm), 20)
|
||||||
|
|
||||||
|
def test_small_allowance_degrades_gracefully(self):
|
||||||
|
# Fewer allowed cores than the budget: take what exists after the
|
||||||
|
# launch reserve, never raise.
|
||||||
|
launch, pool = partition(range(6), mm_workers=8)
|
||||||
|
self.assertEqual(launch, [0]) # min(2, 6 // 4) == 1
|
||||||
|
self.assertEqual(pool, list(range(1, 6)))
|
||||||
|
# Below the split threshold: unpinned.
|
||||||
|
self.assertEqual(partition(range(3), mm_workers=8), (None, None))
|
||||||
|
|
||||||
|
def test_launch_and_pool_cores_are_disjoint(self):
|
||||||
|
launch, pool = partition(range(32), mm_workers=8)
|
||||||
|
self.assertFalse(set(launch) & set(pool))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -0,0 +1,129 @@
|
|||||||
|
"""E2E: Rust tokenizer-manager native multimodal path (``SGLANG_RUST_SERVER=1``).
|
||||||
|
|
||||||
|
Covers what the CPU parity units structurally cannot: the sidecar handoff, the
|
||||||
|
drain ordering, Rust-side tokenization of multimodal prompts, and the rejection
|
||||||
|
of inputs outside the native pipeline's scope (there is no Python fallback).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import importlib.util
|
||||||
|
import io
|
||||||
|
import os
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import requests
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
from sglang.srt.utils import kill_process_tree
|
||||||
|
from sglang.test.ci.ci_register import register_cuda_ci
|
||||||
|
from sglang.test.test_utils import (
|
||||||
|
DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
|
||||||
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||||
|
DEFAULT_URL_FOR_TEST,
|
||||||
|
CustomTestCase,
|
||||||
|
popen_launch_server,
|
||||||
|
)
|
||||||
|
|
||||||
|
register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large")
|
||||||
|
|
||||||
|
IMAGE_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/man_ironing_on_back_of_suv.png"
|
||||||
|
VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||||
|
|
||||||
|
|
||||||
|
def chat_prompt(question, image_count=1):
|
||||||
|
return (
|
||||||
|
f"<|im_start|>user\n{VISION_BLOCK * image_count}{question}<|im_end|>\n"
|
||||||
|
"<|im_start|>assistant\n"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def solid_image_data_url(fmt):
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
Image.fromarray(np.full((64, 64, 3), (255, 0, 0), dtype=np.uint8)).save(
|
||||||
|
buffer, format=fmt
|
||||||
|
)
|
||||||
|
encoded = base64.b64encode(buffer.getvalue()).decode()
|
||||||
|
return f"data:image/{fmt.lower()};base64,{encoded}"
|
||||||
|
|
||||||
|
|
||||||
|
@unittest.skipIf(
|
||||||
|
importlib.util.find_spec("sglang.srt.server._core") is None,
|
||||||
|
"sglang-server rust extension not installed (e.g. AMD suite)",
|
||||||
|
)
|
||||||
|
class TestRustServerNativeMm(CustomTestCase):
|
||||||
|
env = {"SGLANG_RUST_SERVER": "1"}
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def setUpClass(cls):
|
||||||
|
cls.process = popen_launch_server(
|
||||||
|
DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
|
||||||
|
DEFAULT_URL_FOR_TEST,
|
||||||
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||||
|
other_args=["--enable-multimodal", "--mem-fraction-static", "0.8"],
|
||||||
|
env={**os.environ, **cls.env},
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def tearDownClass(cls):
|
||||||
|
if hasattr(cls, "process") and cls.process:
|
||||||
|
kill_process_tree(cls.process.pid)
|
||||||
|
|
||||||
|
def generate(self, prompt, image_data):
|
||||||
|
response = requests.post(
|
||||||
|
DEFAULT_URL_FOR_TEST + "/generate",
|
||||||
|
json={
|
||||||
|
"text": prompt,
|
||||||
|
"image_data": image_data,
|
||||||
|
"sampling_params": {"temperature": 0, "max_new_tokens": 48},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
self.assertEqual(response.status_code, 200, response.text)
|
||||||
|
return response.json()["text"].lower()
|
||||||
|
|
||||||
|
def test_single_image_url(self):
|
||||||
|
text = self.generate(
|
||||||
|
chat_prompt("Describe this image in one sentence."), [IMAGE_URL]
|
||||||
|
)
|
||||||
|
keywords = ("iron", "man", "taxi", "cab", "car", "suv", "street")
|
||||||
|
self.assertTrue(any(w in text for w in keywords), text)
|
||||||
|
|
||||||
|
def test_two_images(self):
|
||||||
|
red = solid_image_data_url("PNG")
|
||||||
|
text = self.generate(
|
||||||
|
chat_prompt("What color is the second image?", image_count=2),
|
||||||
|
[IMAGE_URL, red],
|
||||||
|
)
|
||||||
|
self.assertIn("red", text)
|
||||||
|
|
||||||
|
def test_unsupported_format_is_rejected(self):
|
||||||
|
# An undecodable format must be rejected, never silently answered, and
|
||||||
|
# must not crash the server. PCX is the probe because feature unification
|
||||||
|
# in the server binary (dynamo-parsers → openai-harmony) widens
|
||||||
|
# sglang-mm's jpeg/png/webp/gif/bmp set to every image-crate default, so
|
||||||
|
# the probe has to be a format the image crate does not know at all.
|
||||||
|
response = requests.post(
|
||||||
|
DEFAULT_URL_FOR_TEST + "/generate",
|
||||||
|
json={
|
||||||
|
"text": chat_prompt("What color is this image?"),
|
||||||
|
"image_data": [solid_image_data_url("PCX")],
|
||||||
|
"sampling_params": {"max_new_tokens": 8},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
self.assertIn(response.status_code, (400, 500), response.text)
|
||||||
|
|
||||||
|
def test_corrupt_image_is_rejected(self):
|
||||||
|
response = requests.post(
|
||||||
|
DEFAULT_URL_FOR_TEST + "/generate",
|
||||||
|
json={
|
||||||
|
"text": chat_prompt("Describe this image."),
|
||||||
|
"image_data": ["data:image/png;base64,aW52YWxpZA=="],
|
||||||
|
"sampling_params": {"max_new_tokens": 8},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
# Surfaced as Error::Encode (500); rejected without killing the server.
|
||||||
|
self.assertIn(response.status_code, (400, 500), response.text)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main(verbosity=3)
|
||||||
@@ -0,0 +1,168 @@
|
|||||||
|
"""MMMU accuracy gate for the Rust tokenizer manager's native multimodal path.
|
||||||
|
|
||||||
|
``test_rust_native_mm_e2e.py`` checks that the output is *valid*; this checks that
|
||||||
|
native Rust preprocessing yields *equally good* model inputs. A systematic skew
|
||||||
|
(wrong resample filter, channel order, normalization, patch layout) still reads as
|
||||||
|
fluent text and passes a keyword smoke check, but drops MMMU below the gate.
|
||||||
|
|
||||||
|
The rust server has no ``/v1/chat/completions`` route yet, so the eval drives
|
||||||
|
``/generate`` with hand-rendered Qwen chat prompts instead of lmms-eval's OpenAI
|
||||||
|
client (``MMMUMixin``).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import importlib.util
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
import time
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
import requests
|
||||||
|
|
||||||
|
from sglang.srt.utils import kill_process_tree
|
||||||
|
from sglang.test.ci.ci_register import register_cuda_ci
|
||||||
|
from sglang.test.simple_eval_common import MessageList, SamplerBase
|
||||||
|
from sglang.test.simple_eval_mmmu_vlm import MMMUVLMEval
|
||||||
|
from sglang.test.test_utils import (
|
||||||
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||||
|
DEFAULT_URL_FOR_TEST,
|
||||||
|
CustomTestCase,
|
||||||
|
dump_metric,
|
||||||
|
popen_launch_server,
|
||||||
|
)
|
||||||
|
|
||||||
|
register_cuda_ci(est_time=900, stage="base-b", runner_config="1-gpu-large")
|
||||||
|
|
||||||
|
MODEL = "Qwen/Qwen3.5-0.8B"
|
||||||
|
VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||||
|
|
||||||
|
NUM_EXAMPLES = 100
|
||||||
|
# Calibrated 2026-07-24 on H200: the native path scores 0.37 on this fixed subset
|
||||||
|
# at temperature 0 (two runs), matching the Python reference (0.37, same sampler
|
||||||
|
# and samples). The gate leaves headroom for batching nondeterminism.
|
||||||
|
MMMU_ACCURACY_THRESHOLD = 0.30
|
||||||
|
|
||||||
|
|
||||||
|
class QwenGenerateVisionSampler(SamplerBase):
|
||||||
|
"""Drive ``/generate`` with Qwen chat prompts and ``image_data``.
|
||||||
|
|
||||||
|
``MMMUVLMEval`` emits OpenAI-style messages mixing ``text`` and ``image_url``
|
||||||
|
parts. This sampler renders the Qwen chat format by hand — each image part
|
||||||
|
becomes a ``VISION_BLOCK`` at its original position — and ships the images
|
||||||
|
through ``image_data``.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, base_url: str, max_tokens: int = 1024):
|
||||||
|
self.generate_url = base_url + "/generate"
|
||||||
|
self.max_tokens = max_tokens
|
||||||
|
|
||||||
|
def __call__(self, message_list: MessageList) -> str:
|
||||||
|
segments = []
|
||||||
|
images = []
|
||||||
|
for message in message_list:
|
||||||
|
content = message["content"]
|
||||||
|
parts = (
|
||||||
|
[{"type": "text", "text": content}]
|
||||||
|
if isinstance(content, str)
|
||||||
|
else content
|
||||||
|
)
|
||||||
|
for part in parts:
|
||||||
|
if part["type"] == "image_url":
|
||||||
|
images.append(part["image_url"]["url"])
|
||||||
|
segments.append(VISION_BLOCK)
|
||||||
|
else:
|
||||||
|
segments.append(part["text"])
|
||||||
|
prompt = (
|
||||||
|
"<|im_start|>user\n"
|
||||||
|
+ "".join(segments)
|
||||||
|
+ "<|im_end|>\n<|im_start|>assistant\n"
|
||||||
|
)
|
||||||
|
payload = {
|
||||||
|
"text": prompt,
|
||||||
|
"image_data": images,
|
||||||
|
"sampling_params": {
|
||||||
|
"temperature": 0,
|
||||||
|
"max_new_tokens": self.max_tokens,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
# Retry transient failures but fail loudly when they persist: returning ""
|
||||||
|
# would silently degrade the score and blur the gate.
|
||||||
|
for attempt in range(3):
|
||||||
|
try:
|
||||||
|
response = requests.post(self.generate_url, json=payload, timeout=600)
|
||||||
|
response.raise_for_status()
|
||||||
|
return response.json()["text"]
|
||||||
|
except requests.RequestException:
|
||||||
|
if attempt == 2:
|
||||||
|
raise
|
||||||
|
time.sleep(2**attempt)
|
||||||
|
|
||||||
|
|
||||||
|
@unittest.skipIf(
|
||||||
|
importlib.util.find_spec("sglang.srt.server._core") is None,
|
||||||
|
"sglang-server rust extension not installed (e.g. AMD suite)",
|
||||||
|
)
|
||||||
|
class TestRustNativeMmMMMU(CustomTestCase):
|
||||||
|
@classmethod
|
||||||
|
def setUpClass(cls):
|
||||||
|
# Capture the server log so the test can pin that the native MM
|
||||||
|
# pipeline is active.
|
||||||
|
cls.log_dir = tempfile.TemporaryDirectory()
|
||||||
|
cls.server_logs = tuple(
|
||||||
|
open(os.path.join(cls.log_dir.name, name), "w")
|
||||||
|
for name in ("stdout.log", "stderr.log")
|
||||||
|
)
|
||||||
|
cls.process = popen_launch_server(
|
||||||
|
MODEL,
|
||||||
|
DEFAULT_URL_FOR_TEST,
|
||||||
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||||
|
other_args=["--enable-multimodal", "--mem-fraction-static", "0.8"],
|
||||||
|
env={**os.environ, "SGLANG_RUST_SERVER": "1"},
|
||||||
|
return_stdout_stderr=cls.server_logs,
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def tearDownClass(cls):
|
||||||
|
if hasattr(cls, "process") and cls.process:
|
||||||
|
kill_process_tree(cls.process.pid)
|
||||||
|
if hasattr(cls, "server_logs"):
|
||||||
|
for f in cls.server_logs:
|
||||||
|
f.close()
|
||||||
|
if hasattr(cls, "log_dir"):
|
||||||
|
cls.log_dir.cleanup()
|
||||||
|
|
||||||
|
def _read_server_log(self):
|
||||||
|
text = []
|
||||||
|
for f in self.server_logs:
|
||||||
|
with open(f.name) as reader:
|
||||||
|
text.append(reader.read())
|
||||||
|
return "\n".join(text)
|
||||||
|
|
||||||
|
def test_mmmu_accuracy(self):
|
||||||
|
# Guard the path under test: if the model ever drops off
|
||||||
|
# NATIVE_MM_FAMILIES, launch fails and this names why.
|
||||||
|
self.assertIn(
|
||||||
|
"native MM pipeline enabled",
|
||||||
|
self._read_server_log(),
|
||||||
|
"rust server did not enable the native MM pipeline for "
|
||||||
|
f"{MODEL}; this test must exercise the native path",
|
||||||
|
)
|
||||||
|
|
||||||
|
eval_obj = MMMUVLMEval(num_examples=NUM_EXAMPLES, num_threads=32)
|
||||||
|
sampler = QwenGenerateVisionSampler(base_url=DEFAULT_URL_FOR_TEST)
|
||||||
|
result = eval_obj(sampler)
|
||||||
|
print(f"MMMU metrics: {result.metrics}")
|
||||||
|
dump_metric(
|
||||||
|
"mmmu_score",
|
||||||
|
result.score,
|
||||||
|
labels={"model": MODEL, "eval": "mmmu", "api": "generate-rust-native-mm"},
|
||||||
|
)
|
||||||
|
self.assertGreaterEqual(
|
||||||
|
result.score,
|
||||||
|
MMMU_ACCURACY_THRESHOLD,
|
||||||
|
f"Rust native MM path scored {result.score:.4f} on MMMU, below the "
|
||||||
|
f"{MMMU_ACCURACY_THRESHOLD:.2f} gate",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main(verbosity=3)
|
||||||
Reference in New Issue
Block a user