[Refactor] New EPD (#30398)
Co-authored-by: Yuang Chen <1131578721@qq.com> Co-authored-by: Yuang Chen <cya539102@antgroup.com> Co-authored-by: ZhengWG <zwg0606@gmail.com>
This commit is contained in:
siyu
co-authored by
Yuang Chen
Yuang Chen
ZhengWG
Yuang Chen
Yuang Chen
ZhengWG
parent
6a12583679
commit
8a123cbd0e
+1
-2
@@ -16,8 +16,7 @@
|
||||
/python/sglang/srt/constrained @hnyls2002 @DarkSharpness @JustinTong0323
|
||||
/python/sglang/srt/disaggregation @ByronHsu @hnyls2002 @ShangmingCai @HaiShaw @Duyi-Wang @sogalin
|
||||
/python/sglang/srt/disaggregation/ascend @ping1jing2 @iforgetmyname
|
||||
/python/sglang/srt/disaggregation/encode_receiver.py @ShangmingCai @liusy58 @ZhengWG @gty111
|
||||
/python/sglang/srt/disaggregation/encode_server.py @ShangmingCai @liusy58 @ZhengWG @gty111
|
||||
/python/sglang/srt/disaggregation/encoder @ShangmingCai @liusy58 @ZhengWG @gty111
|
||||
/python/sglang/srt/disaggregation/mori @Duyi-Wang @kkHuang-amd @HaiShaw @Lzy17 @billishyahao
|
||||
/python/sglang/srt/distributed @yizhang2077 @merrymercy @ch-wan
|
||||
/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py @ShangmingCai @stmatengss
|
||||
|
||||
@@ -18,13 +18,13 @@ def run_server(server_args):
|
||||
if server_args.encoder_only:
|
||||
# For encoder disaggregation
|
||||
if server_args.smg_grpc_mode or server_args.grpc_mode:
|
||||
from sglang.srt.disaggregation.encode_grpc_server import (
|
||||
from sglang.srt.disaggregation.encoder.grpc_server import (
|
||||
serve_grpc_encoder,
|
||||
)
|
||||
|
||||
asyncio.run(serve_grpc_encoder(server_args))
|
||||
else:
|
||||
from sglang.srt.disaggregation.encode_server import launch_server
|
||||
from sglang.srt.disaggregation.encoder.http_server import launch_server
|
||||
|
||||
launch_server(server_args)
|
||||
elif server_args.smg_grpc_mode:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+15
-22
@@ -21,11 +21,7 @@ from grpc_health.v1 import health_pb2, health_pb2_grpc
|
||||
from grpc_reflection.v1alpha import reflection
|
||||
from smg_grpc_proto import sglang_encoder_pb2, sglang_encoder_pb2_grpc
|
||||
|
||||
from sglang.srt.disaggregation.encode_server import (
|
||||
MMEncoder,
|
||||
handle_scheduler_receive_url_request,
|
||||
launch_encoder,
|
||||
)
|
||||
from sglang.srt.disaggregation.encoder.server import MMEncoder, launch_encoder
|
||||
from sglang.srt.managers.io_struct import async_sock_send, wrap_as_pickle
|
||||
from sglang.srt.managers.schedule_batch import Modality
|
||||
from sglang.srt.runtime_context import get_disagg
|
||||
@@ -92,6 +88,7 @@ class SGLangEncoderServer(SGLangEncoderServicer):
|
||||
try:
|
||||
request_dict = {
|
||||
"mm_items": list(request.mm_items),
|
||||
"modality": Modality.IMAGE.name,
|
||||
"req_id": request.req_id,
|
||||
"num_parts": request.num_parts,
|
||||
"part_idx": request.part_idx,
|
||||
@@ -99,21 +96,17 @@ class SGLangEncoderServer(SGLangEncoderServicer):
|
||||
for socket in self.send_sockets:
|
||||
await async_sock_send(socket, wrap_as_pickle(request_dict))
|
||||
|
||||
# gRPC encode is image-only; encoder.encode() requires modality
|
||||
# gRPC encode is image-only; the request follows the configured
|
||||
# cache and transfer backend.
|
||||
(
|
||||
nbytes,
|
||||
embedding_len,
|
||||
embedding_dim,
|
||||
error_msg,
|
||||
error_code,
|
||||
) = await self.encoder.encode(
|
||||
mm_items=list(request.mm_items),
|
||||
modality=Modality.IMAGE,
|
||||
req_id=request.req_id,
|
||||
num_parts=request.num_parts,
|
||||
part_idx=request.part_idx,
|
||||
)
|
||||
) = await self.encoder.encode_request(request_dict, Modality.IMAGE)
|
||||
if error_msg is not None:
|
||||
await self.encoder.release_request(request.req_id)
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(error_msg)
|
||||
return sglang_encoder_pb2.EncodeResponse()
|
||||
@@ -140,7 +133,7 @@ class SGLangEncoderServer(SGLangEncoderServicer):
|
||||
)
|
||||
)
|
||||
await asyncio.gather(*tasks)
|
||||
self.encoder.embedding_to_send.pop(request.req_id, None)
|
||||
await self.encoder.release_request(request.req_id)
|
||||
return sglang_encoder_pb2.EncodeResponse()
|
||||
elif get_disagg().encoder_transfer_backend == "zmq_to_tokenizer":
|
||||
embedding_port = (
|
||||
@@ -151,7 +144,7 @@ class SGLangEncoderServer(SGLangEncoderServicer):
|
||||
prefill_host=request.prefill_host,
|
||||
embedding_port=embedding_port,
|
||||
)
|
||||
self.encoder.embedding_to_send.pop(request.req_id, None)
|
||||
await self.encoder.release_request(request.req_id)
|
||||
return sglang_encoder_pb2.EncodeResponse()
|
||||
|
||||
return sglang_encoder_pb2.EncodeResponse()
|
||||
@@ -159,6 +152,7 @@ class SGLangEncoderServer(SGLangEncoderServicer):
|
||||
except Exception as e:
|
||||
logger.error(f"Encode error: {e}")
|
||||
traceback.print_exc()
|
||||
await self.encoder.release_request(request.req_id)
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(str(e))
|
||||
return sglang_encoder_pb2.EncodeResponse()
|
||||
@@ -176,12 +170,13 @@ class SGLangEncoderServer(SGLangEncoderServicer):
|
||||
request.buffer_address if request.buffer_address else None
|
||||
),
|
||||
)
|
||||
self.encoder.embedding_to_send.pop(request.req_id, None)
|
||||
await self.encoder.release_request(request.req_id)
|
||||
return sglang_encoder_pb2.SendResponse()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Send error: {e}")
|
||||
traceback.print_exc()
|
||||
await self.encoder.release_request(request.req_id)
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(str(e))
|
||||
return sglang_encoder_pb2.SendResponse()
|
||||
@@ -190,12 +185,10 @@ class SGLangEncoderServer(SGLangEncoderServicer):
|
||||
self, request: sglang_encoder_pb2.SchedulerReceiveUrlRequest, context
|
||||
) -> sglang_encoder_pb2.SchedulerReceiveUrlResponse:
|
||||
try:
|
||||
await handle_scheduler_receive_url_request(
|
||||
{
|
||||
"req_id": request.req_id,
|
||||
"receive_count": request.receive_count,
|
||||
"receive_url": request.receive_url,
|
||||
}
|
||||
await self.encoder.register_embedding_destinations(
|
||||
request.req_id,
|
||||
request.receive_count,
|
||||
[request.receive_url],
|
||||
)
|
||||
return sglang_encoder_pb2.SchedulerReceiveUrlResponse()
|
||||
|
||||
@@ -0,0 +1,653 @@
|
||||
"""HTTP API layer for the EPD encoder server.
|
||||
|
||||
This module is designed to be replaceable by a Rust implementation.
|
||||
It contains the FastAPI application, HTTP route handlers, HTTP lifecycle, and
|
||||
response conversion. Backend scheduling and process management are provided by
|
||||
the protocol-neutral :mod:`runtime` module.
|
||||
|
||||
GPU tensor operations remain in :mod:`server.MMEncoder`.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from http import HTTPStatus
|
||||
from typing import Annotated, List, Optional
|
||||
|
||||
import requests as http_requests
|
||||
import uvicorn
|
||||
import zmq
|
||||
from fastapi import Body, FastAPI
|
||||
from fastapi.responses import ORJSONResponse, Response
|
||||
|
||||
import sglang.srt.disaggregation.encoder.server as server_module
|
||||
from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
|
||||
from sglang.srt.disaggregation.encoder.runtime import (
|
||||
DPDispatcher,
|
||||
EncoderRuntime,
|
||||
EncoderScheduler,
|
||||
execute_encode_pipeline,
|
||||
launch_dp_runtime,
|
||||
launch_local_runtime,
|
||||
)
|
||||
from sglang.srt.disaggregation.encoder.server import (
|
||||
EncoderProfiler,
|
||||
MMEncoder,
|
||||
MMError,
|
||||
)
|
||||
from sglang.srt.managers.io_struct import (
|
||||
ProfileReq,
|
||||
ProfileReqType,
|
||||
sock_send,
|
||||
wrap_as_pickle,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import Modality
|
||||
from sglang.srt.runtime_context import (
|
||||
get_disagg,
|
||||
get_observability,
|
||||
get_parallel,
|
||||
get_serving,
|
||||
publish,
|
||||
)
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils import (
|
||||
add_prometheus_middleware,
|
||||
configure_logger,
|
||||
)
|
||||
from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
HEALTH_CHECK_TIMEOUT = 30
|
||||
|
||||
# Minimal 32x32 black PNG for health check dummy encode
|
||||
MINIMUM_PNG_PICTURE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAACXBIWXMAAA7EAAAOxAGVKw4bAAAAbUlEQVRYhe3VsQ2AMAxE0Y/lIgNQULD/OqyCMgCihCKSG4yRuKuiNH6JLsoEbMACOGBcua9HOR7Y6w6swBwMy0qLTpkeI77qdEBpBFAHBBDAGH8WrwJKI4AAegUCfAKgEgpQDvh3CR3oQCuav58qlAw73kKCSgAAAABJRU5ErkJggg=="
|
||||
|
||||
# Minimal WAV: 16kHz mono 16-bit PCM, 160 samples (0.01s) of silence
|
||||
MINIMUM_WAV_SILENCE_BASE64 = "UklGRmQBAABXQVZFZm10IBAAAAABAAEAgD4AAAB9AAACABAAZGF0YUABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=="
|
||||
|
||||
encoder: Optional[MMEncoder] = None
|
||||
send_sockets: List[zmq.Socket] = []
|
||||
encoder_scheduler: Optional[EncoderScheduler] = None
|
||||
local_runtime: Optional[EncoderRuntime] = None
|
||||
|
||||
# DP mode (--dp-size > 1): the protocol-neutral runtime owns worker processes
|
||||
# and ZMQ; HTTP only keeps the dispatcher handle used by route handlers.
|
||||
dp_dispatcher: Optional["DPDispatcher"] = None
|
||||
|
||||
|
||||
def is_health_check_request(rid: Optional[str]) -> bool:
|
||||
return isinstance(rid, str) and rid.startswith(HEALTH_CHECK_RID_PREFIX)
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def _lifespan(app: FastAPI):
|
||||
if dp_dispatcher is not None:
|
||||
dp_dispatcher.start()
|
||||
yield
|
||||
return
|
||||
if local_runtime is not None:
|
||||
local_runtime.start()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if local_runtime is not None:
|
||||
await local_runtime.stop()
|
||||
|
||||
|
||||
app = FastAPI(lifespan=_lifespan)
|
||||
|
||||
|
||||
def _register_encoder_url_with_bootstrap(server_args: ServerArgs):
|
||||
"""Asynchronously register this encoder with each bootstrap URL.
|
||||
|
||||
Spawns a daemon thread that retries each URL independently with bounded
|
||||
backoff. The encoder's own startup is not blocked: if some bootstrap
|
||||
server is slow or unreachable, only the background worker waits.
|
||||
|
||||
Inspired by ``_ensure_prefill_info`` in disaggregation/decode.py: each
|
||||
target keeps its own retry count and is retried at a fixed interval
|
||||
instead of serialising sleeps in a single thread.
|
||||
"""
|
||||
|
||||
host = server_args.host
|
||||
if not host or host in ("0.0.0.0", "::"):
|
||||
host = get_local_ip_auto(server_args.host)
|
||||
scheme = "https" if server_args.ssl_certfile else "http"
|
||||
encoder_url = NetworkAddress(host, server_args.port).to_url(scheme)
|
||||
payload = {"url": encoder_url}
|
||||
bootstrap_urls = list(server_args.encoder_register_urls)
|
||||
if not bootstrap_urls:
|
||||
return
|
||||
|
||||
max_retries = 30
|
||||
retry_interval = 5.0
|
||||
request_timeout = 5.0
|
||||
|
||||
def _try_register_once(bootstrap_url: str) -> bool:
|
||||
try:
|
||||
resp = http_requests.post(
|
||||
f"{bootstrap_url}/register_encoder_url",
|
||||
json=payload,
|
||||
timeout=request_timeout,
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
logger.info(
|
||||
f"Registered encoder URL '{encoder_url}' with bootstrap "
|
||||
f"at {bootstrap_url}"
|
||||
)
|
||||
return True
|
||||
logger.warning(
|
||||
f"Bootstrap {bootstrap_url} returned {resp.status_code}: {resp.text}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Register attempt to {bootstrap_url} failed: {e}")
|
||||
return False
|
||||
|
||||
def _worker():
|
||||
pending = list(bootstrap_urls)
|
||||
retry_count = {url: 0 for url in pending}
|
||||
while pending:
|
||||
still_pending = []
|
||||
for bootstrap_url in pending:
|
||||
if _try_register_once(bootstrap_url):
|
||||
continue
|
||||
retry_count[bootstrap_url] += 1
|
||||
if retry_count[bootstrap_url] >= max_retries:
|
||||
logger.error(
|
||||
f"Giving up on bootstrap {bootstrap_url} after "
|
||||
f"{max_retries} attempts. Encoder discovery via this "
|
||||
f"bootstrap will be incomplete."
|
||||
)
|
||||
continue
|
||||
still_pending.append(bootstrap_url)
|
||||
pending = still_pending
|
||||
if pending:
|
||||
time.sleep(retry_interval)
|
||||
|
||||
threading.Thread(
|
||||
target=_worker, daemon=True, name="encoder-bootstrap-register"
|
||||
).start()
|
||||
|
||||
|
||||
def _unregister_encoder_url_from_bootstrap(server_args: ServerArgs):
|
||||
host = server_args.host
|
||||
if not host or host in ("0.0.0.0", "::"):
|
||||
host = get_local_ip_auto(server_args.host)
|
||||
scheme = "https" if server_args.ssl_certfile else "http"
|
||||
encoder_url = NetworkAddress(host, server_args.port).to_url(scheme)
|
||||
payload = {"url": encoder_url}
|
||||
|
||||
for bootstrap_url in server_args.encoder_register_urls:
|
||||
try:
|
||||
resp = http_requests.delete(
|
||||
f"{bootstrap_url}/unregister_encoder_url",
|
||||
json=payload,
|
||||
timeout=2.0,
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
logger.info(
|
||||
f"Unregistered encoder URL '{encoder_url}' from "
|
||||
f"bootstrap at {bootstrap_url}"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Bootstrap {bootstrap_url} returned "
|
||||
f"{resp.status_code} on unregister: {resp.text}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Unregister from {bootstrap_url} failed: {e}")
|
||||
|
||||
|
||||
def launch_server(server_args: ServerArgs):
|
||||
global dp_dispatcher, encoder, encoder_scheduler, local_runtime, send_sockets
|
||||
|
||||
configure_logger(server_args, prefix=" encode_server")
|
||||
# Publish before the launch path reads configuration; each encoder built
|
||||
# below re-projects the same object in its process.
|
||||
publish(server_args, role="encoder")
|
||||
if get_parallel().dp_size > 1:
|
||||
dp_dispatcher = launch_dp_runtime(server_args)
|
||||
# runtime initializes multiprocess metrics before spawning;
|
||||
# HTTP only exposes their endpoint.
|
||||
if get_observability().enable_metrics:
|
||||
add_prometheus_middleware(app)
|
||||
else:
|
||||
local_runtime = launch_local_runtime(server_args)
|
||||
# Compatibility aliases for the existing HTTP request path. Runtime is
|
||||
# now the sole constructor and lifecycle owner of these objects.
|
||||
encoder = local_runtime.encoder
|
||||
encoder_scheduler = local_runtime.scheduler
|
||||
send_sockets = local_runtime.send_sockets
|
||||
if get_observability().enable_metrics:
|
||||
add_prometheus_middleware(app)
|
||||
|
||||
# Register this encoder's URL with prefill server(s) if configured.
|
||||
if get_disagg().encoder_register_urls:
|
||||
import atexit
|
||||
|
||||
_register_encoder_url_with_bootstrap(server_args)
|
||||
atexit.register(_unregister_encoder_url_from_bootstrap, server_args)
|
||||
|
||||
uvicorn.run(app, host=get_serving().host, port=get_serving().port)
|
||||
|
||||
|
||||
def _summarise_dp_broadcast(results: List[dict]) -> Response:
|
||||
# Treat missing/None content as failure so a stuck rank doesn't hide
|
||||
# behind the others' "ok". Status = the most severe per-rank error code
|
||||
# (5xx beats 4xx) rather than a blanket 400, so a worker's 500/503/504
|
||||
# isn't misreported as a client error.
|
||||
msgs: List[str] = []
|
||||
error_codes: List[int] = []
|
||||
for r in results:
|
||||
content = r.get("content")
|
||||
if isinstance(content, dict):
|
||||
msgs.append(content.get("msg", ""))
|
||||
if not content.get("ok"):
|
||||
# Worker ran but reported a logical failure; no transport code,
|
||||
# so treat as a bad request (matches the non-DP profile path).
|
||||
error_codes.append(int(r.get("_error_code") or HTTPStatus.BAD_REQUEST))
|
||||
else:
|
||||
msgs.append(r.get("_error", "unknown error"))
|
||||
error_codes.append(
|
||||
int(r.get("_error_code") or HTTPStatus.INTERNAL_SERVER_ERROR)
|
||||
)
|
||||
status_code = 200 if not error_codes else max(error_codes)
|
||||
return Response(
|
||||
content="\n".join(msgs) + "\n",
|
||||
status_code=status_code,
|
||||
)
|
||||
|
||||
|
||||
@app.post("/encode")
|
||||
async def handle_encode_request(request: dict):
|
||||
req_id = request["req_id"]
|
||||
start_time = time.monotonic()
|
||||
time_stats_json = request.pop("time_stats_json", None)
|
||||
if dp_dispatcher is not None:
|
||||
if time_stats_json:
|
||||
request = dict(request)
|
||||
request["time_stats_json"] = time_stats_json
|
||||
try:
|
||||
result = await dp_dispatcher.dispatch(request)
|
||||
except MMError as e:
|
||||
# Surface MMError.code (503 when all workers dead) instead of
|
||||
# FastAPI's default 500.
|
||||
logger.error(f"DP dispatch refused req_id={req_id}: {e}")
|
||||
return ORJSONResponse(
|
||||
status_code=int(e.code),
|
||||
content={"status": "error", "message": str(e), "req_id": req_id},
|
||||
)
|
||||
if result.get("_error"):
|
||||
error_type = result.get("_error_type", "")
|
||||
# `or` (not `dict.get(key, default)`) so explicit None falls back too.
|
||||
status_code = result.get("_error_code") or (
|
||||
HTTPStatus.BAD_REQUEST
|
||||
if error_type == "ValueError"
|
||||
else HTTPStatus.INTERNAL_SERVER_ERROR
|
||||
)
|
||||
logger.error(f"DP worker error for req_id={req_id}: {result['_error']}")
|
||||
return ORJSONResponse(
|
||||
status_code=status_code,
|
||||
content={
|
||||
"status": "error",
|
||||
"message": result["_error"],
|
||||
"req_id": req_id,
|
||||
},
|
||||
)
|
||||
elapsed = time.monotonic() - start_time
|
||||
logger.info(
|
||||
f"[{req_id}] /encode completed in {elapsed:.3f}s, "
|
||||
f"modality={request.get('modality', 'image')}"
|
||||
)
|
||||
content = result.get("content")
|
||||
return ORJSONResponse(content=content)
|
||||
|
||||
try:
|
||||
if time_stats_json:
|
||||
request["time_stats_json"] = time_stats_json
|
||||
content = await execute_encode_pipeline(
|
||||
encoder,
|
||||
encoder_scheduler,
|
||||
request,
|
||||
send_sockets=send_sockets,
|
||||
)
|
||||
elapsed = time.monotonic() - start_time
|
||||
logger.info(
|
||||
f"[{req_id}] /encode completed in {elapsed:.3f}s, "
|
||||
f"modality={request.get('modality', 'image')}"
|
||||
)
|
||||
return ORJSONResponse(content=content)
|
||||
except asyncio.TimeoutError:
|
||||
return ORJSONResponse(
|
||||
status_code=HTTPStatus.GATEWAY_TIMEOUT,
|
||||
content={
|
||||
"status": "error",
|
||||
"message": "encoder batch timed out",
|
||||
"req_id": req_id,
|
||||
},
|
||||
)
|
||||
except MMError as e:
|
||||
return ORJSONResponse(
|
||||
status_code=int(e.code),
|
||||
content={"status": "error", "message": str(e), "req_id": req_id},
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
logger.error(f"Unexpected error in encoder logic for {req_id}: {error_msg}")
|
||||
return ORJSONResponse(
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
||||
content={
|
||||
"status": "error",
|
||||
"message": error_msg,
|
||||
"req_id": req_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@app.post("/send")
|
||||
async def handle_send_request(request: dict):
|
||||
"""Mooncake-only: drive the RDMA push of a staged embedding. The zmq
|
||||
backends deliver embeddings inline during /encode and never call /send."""
|
||||
req_id = request["req_id"]
|
||||
receive_count = request.get("receive_count")
|
||||
if dp_dispatcher is not None:
|
||||
try:
|
||||
result = await dp_dispatcher.dispatch_send(request)
|
||||
except MMError as e:
|
||||
logger.error(f"DP dispatch_send refused req_id={req_id}: {e}")
|
||||
return Response(
|
||||
content=f"Encoder DP worker send error: {e}",
|
||||
status_code=int(e.code),
|
||||
)
|
||||
if result.get("_error"):
|
||||
status_code = result.get("_error_code") or int(
|
||||
HTTPStatus.INTERNAL_SERVER_ERROR
|
||||
)
|
||||
logger.error(
|
||||
f"DP worker send error for req_id={req_id}: {result['_error']}"
|
||||
)
|
||||
return Response(
|
||||
content=f"Encoder DP worker send error: {result['_error']}",
|
||||
status_code=status_code,
|
||||
)
|
||||
return ORJSONResponse(content=result.get("content"))
|
||||
sent = await encoder.send(
|
||||
req_id=req_id,
|
||||
prefill_host=request["prefill_host"],
|
||||
embedding_port=request["embedding_port"],
|
||||
session_id=request["session_id"],
|
||||
buffer_address=request["buffer_address"],
|
||||
)
|
||||
if not sent:
|
||||
# No transfer happened: fail fast rather than 200 + a phantom count.
|
||||
return ORJSONResponse(
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
||||
content={
|
||||
"status": "error",
|
||||
"message": f"no staged embedding for req_id={req_id} (already released)",
|
||||
"req_id": req_id,
|
||||
},
|
||||
)
|
||||
# Sibling ranks share this embedding, so free it only once all have sent.
|
||||
# No count means a pre-refcount decoder: leave it to the sweep, as when
|
||||
# some rank never sends at all.
|
||||
if receive_count:
|
||||
await server_module.meta_registry.note_send_done(req_id, receive_count)
|
||||
return ORJSONResponse(content=None)
|
||||
|
||||
|
||||
@app.post("/scheduler_receive_meta_data")
|
||||
async def handle_scheduler_receive_meta_data(request: dict):
|
||||
"""Decoder pull endpoint for the per-part encode metadata. Blocks until the
|
||||
encode publishes its sizes, so a pull that beats the encode simply waits."""
|
||||
req_id = request["req_id"]
|
||||
if dp_dispatcher is not None:
|
||||
try:
|
||||
result = await dp_dispatcher.dispatch_wait_metadata(request)
|
||||
except MMError as e:
|
||||
return ORJSONResponse(
|
||||
status_code=int(e.code),
|
||||
content={"status": "error", "message": str(e), "req_id": req_id},
|
||||
)
|
||||
if result.get("_error"):
|
||||
return ORJSONResponse(
|
||||
status_code=result.get("_error_code")
|
||||
or int(HTTPStatus.INTERNAL_SERVER_ERROR),
|
||||
content={
|
||||
"status": "error",
|
||||
"message": result["_error"],
|
||||
"req_id": req_id,
|
||||
},
|
||||
)
|
||||
meta = result.get("content")
|
||||
else:
|
||||
try:
|
||||
meta = await server_module.meta_registry.wait(req_id)
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(f"[{req_id}] /scheduler_receive_meta_data timed out")
|
||||
return ORJSONResponse(
|
||||
status_code=HTTPStatus.GATEWAY_TIMEOUT,
|
||||
content={
|
||||
"status": "error",
|
||||
"message": "encode metadata not ready",
|
||||
"req_id": req_id,
|
||||
},
|
||||
)
|
||||
if meta is None or meta.get("error") is not None:
|
||||
message = meta["error"] if meta else "encode metadata missing"
|
||||
return ORJSONResponse(
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
||||
content={"status": "error", "message": message, "req_id": req_id},
|
||||
)
|
||||
return ORJSONResponse(
|
||||
content={
|
||||
"req_id": req_id,
|
||||
"part_idx": request["part_idx"],
|
||||
"embedding_size": meta["embedding_size"],
|
||||
"embedding_len": meta["embedding_len"],
|
||||
"embedding_dim": meta["embedding_dim"],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@app.post("/scheduler_receive_url")
|
||||
async def handle_scheduler_receive_url_request(request: dict):
|
||||
if dp_dispatcher is not None:
|
||||
try:
|
||||
result = await dp_dispatcher.dispatch_register_destinations(request)
|
||||
except MMError as e:
|
||||
return ORJSONResponse(
|
||||
status_code=int(e.code),
|
||||
content={
|
||||
"status": "error",
|
||||
"message": str(e),
|
||||
"req_id": request["req_id"],
|
||||
},
|
||||
)
|
||||
if result.get("_error"):
|
||||
return ORJSONResponse(
|
||||
status_code=result.get("_error_code")
|
||||
or int(HTTPStatus.INTERNAL_SERVER_ERROR),
|
||||
content={
|
||||
"status": "error",
|
||||
"message": result["_error"],
|
||||
"req_id": request["req_id"],
|
||||
},
|
||||
)
|
||||
return ORJSONResponse(content=None)
|
||||
if encoder is None:
|
||||
return ORJSONResponse(
|
||||
status_code=HTTPStatus.SERVICE_UNAVAILABLE,
|
||||
content={
|
||||
"status": "error",
|
||||
"message": "encoder not ready",
|
||||
"req_id": request["req_id"],
|
||||
},
|
||||
)
|
||||
try:
|
||||
await encoder.register_embedding_destinations(
|
||||
request["req_id"],
|
||||
request["receive_count"],
|
||||
[request["receive_url"]],
|
||||
)
|
||||
except MMError as e:
|
||||
return ORJSONResponse(
|
||||
status_code=int(e.code),
|
||||
content={
|
||||
"status": "error",
|
||||
"message": str(e),
|
||||
"req_id": request["req_id"],
|
||||
},
|
||||
)
|
||||
return ORJSONResponse(content=None)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
@app.get("/health_generate")
|
||||
async def health_generate():
|
||||
"""
|
||||
Health check endpoint for the encoder server.
|
||||
Performs a dummy encode to verify the encoder is functional.
|
||||
Returns 200 if the encoder is healthy, 503 otherwise.
|
||||
"""
|
||||
if dp_dispatcher is not None:
|
||||
# Strict: any dead (exited) rank fails health → orchestrator restarts.
|
||||
if not dp_dispatcher.all_ranks_alive:
|
||||
return Response(status_code=503)
|
||||
# Process-liveness (proc.sentinel) can't see a worker that's alive but
|
||||
# wedged (hung GPU / NCCL deadlock / stalled ZMQ). Probe every rank with
|
||||
# a tiny dummy encode; each worker runs it only when idle and otherwise
|
||||
# reports healthy at once, keeping the probe off the GPU under load.
|
||||
try:
|
||||
results = await dp_dispatcher.broadcast(
|
||||
{"_dp_type": "health_encode"},
|
||||
timeout=HEALTH_CHECK_TIMEOUT,
|
||||
)
|
||||
except MMError:
|
||||
return Response(status_code=503)
|
||||
if any(r.get("_error") for r in results):
|
||||
return Response(status_code=503)
|
||||
return Response(status_code=200)
|
||||
if encoder is None:
|
||||
return Response(status_code=503)
|
||||
|
||||
# Pick the first available modality for the dummy encode
|
||||
if encoder.supports_modality(Modality.IMAGE):
|
||||
mm_items = [f"data:image/png;base64,{MINIMUM_PNG_PICTURE_BASE64}"]
|
||||
modality = Modality.IMAGE
|
||||
elif encoder.supports_modality(Modality.AUDIO):
|
||||
mm_items = [f"data:audio/wav;base64,{MINIMUM_WAV_SILENCE_BASE64}"]
|
||||
modality = Modality.AUDIO
|
||||
else:
|
||||
# No processor available, fall back to liveness check only
|
||||
return Response(status_code=200)
|
||||
|
||||
try:
|
||||
# uuid keeps rids unique across workers; a bare time.time() can collide.
|
||||
req_id = f"{HEALTH_CHECK_RID_PREFIX}_{uuid.uuid4().hex}"
|
||||
|
||||
dummy_request = {
|
||||
"mm_items": mm_items,
|
||||
"modality": modality.name,
|
||||
"req_id": req_id,
|
||||
"num_parts": 1,
|
||||
"part_idx": 0,
|
||||
}
|
||||
|
||||
# A health encode participates in the same TP collectives as a real
|
||||
# request. Serialize its broadcast and rank-0 forward with every other
|
||||
# collective dispatch, then recheck whether traffic made the probe
|
||||
# unnecessary while it waited for the lock.
|
||||
async with encoder.encode_dispatch_lock:
|
||||
if encoder.has_pending_embeddings():
|
||||
return Response(status_code=200)
|
||||
for socket in send_sockets:
|
||||
sock_send(socket, wrap_as_pickle(dummy_request))
|
||||
|
||||
_, _, _, error_msg, _ = await asyncio.wait_for(
|
||||
encoder.encode(
|
||||
mm_items=mm_items,
|
||||
modality=modality,
|
||||
req_id=req_id,
|
||||
num_parts=1,
|
||||
part_idx=0,
|
||||
),
|
||||
timeout=HEALTH_CHECK_TIMEOUT,
|
||||
)
|
||||
|
||||
# Clean up stored embedding
|
||||
await encoder.release_request(req_id)
|
||||
|
||||
if error_msg:
|
||||
logger.error(f"Encoder health check failed: {error_msg}")
|
||||
return Response(status_code=503)
|
||||
|
||||
return Response(status_code=200)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(f"Encoder health check timed out after {HEALTH_CHECK_TIMEOUT}s")
|
||||
return Response(status_code=503)
|
||||
except Exception as e:
|
||||
logger.error(f"Encoder health check failed: {e}")
|
||||
return Response(status_code=503)
|
||||
|
||||
|
||||
@app.api_route("/start_profile", methods=["GET", "POST"])
|
||||
async def start_profile_async(obj: Annotated[Optional[ProfileReq], Body()] = None):
|
||||
if dp_dispatcher is not None:
|
||||
if obj is not None:
|
||||
obj.req_type = ProfileReqType.START_PROFILE
|
||||
try:
|
||||
results = await dp_dispatcher.broadcast(
|
||||
{"_dp_type": "start_profile", "profile_req": obj}
|
||||
)
|
||||
except MMError as e:
|
||||
return Response(content=f"{e}\n", status_code=int(e.code))
|
||||
return _summarise_dp_broadcast(results)
|
||||
if encoder is None:
|
||||
return Response(content="encoder not ready\n", status_code=503)
|
||||
req = obj or ProfileReq()
|
||||
req.req_type = ProfileReqType.START_PROFILE
|
||||
for socket in send_sockets:
|
||||
sock_send(socket, req)
|
||||
if encoder.profiler is None:
|
||||
encoder.profiler = EncoderProfiler(encoder.rank)
|
||||
ok, msg = encoder.profiler.start(req)
|
||||
if ok:
|
||||
detail = (
|
||||
f"Start profiling. output_dir={encoder.profiler.output_dir} "
|
||||
f"profile_id={encoder.profiler.profile_id}\n"
|
||||
)
|
||||
return Response(content=detail, status_code=200)
|
||||
return Response(
|
||||
content=(msg or "Start profiling failed.\n"), status_code=HTTPStatus.BAD_REQUEST
|
||||
)
|
||||
|
||||
|
||||
@app.api_route("/stop_profile", methods=["GET", "POST"])
|
||||
async def stop_profile_async():
|
||||
if dp_dispatcher is not None:
|
||||
try:
|
||||
results = await dp_dispatcher.broadcast({"_dp_type": "stop_profile"})
|
||||
except MMError as e:
|
||||
return Response(content=f"{e}\n", status_code=int(e.code))
|
||||
return _summarise_dp_broadcast(results)
|
||||
if encoder is None:
|
||||
return Response(content="encoder not ready\n", status_code=503)
|
||||
if encoder.profiler is None:
|
||||
return Response(
|
||||
content="profiling not initialized\n", status_code=HTTPStatus.BAD_REQUEST
|
||||
)
|
||||
req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE)
|
||||
for socket in send_sockets:
|
||||
sock_send(socket, req)
|
||||
ok, msg = encoder.profiler.stop()
|
||||
if ok:
|
||||
return Response(content="Stop profiling.\n", status_code=200)
|
||||
return Response(
|
||||
content=(msg or "Stop profiling failed.\n"), status_code=HTTPStatus.BAD_REQUEST
|
||||
)
|
||||
@@ -0,0 +1,827 @@
|
||||
"""CPU-bound multimodal preprocessing for the EPD encoder.
|
||||
|
||||
This module is designed to be replaceable by a Rust implementation.
|
||||
It handles all CPU-bound work: media I/O (image/video/audio loading),
|
||||
HF processor calls, config validation, and related helper computations.
|
||||
GPU tensor operations remain in :mod:`server.MMEncoder`.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.managers.schedule_batch import Modality
|
||||
from sglang.srt.multimodal.cache import parse_content_hash, snapshot_media
|
||||
from sglang.srt.multimodal.encoder_preprocessing import (
|
||||
EncoderMediaProcessorConfig,
|
||||
EncoderPreprocessOutput,
|
||||
invoke_encoder_preprocessor,
|
||||
)
|
||||
from sglang.srt.multimodal.processors.qwen_vl import preprocess_video
|
||||
from sglang.srt.runtime_context import (
|
||||
get_device,
|
||||
get_mm,
|
||||
get_model,
|
||||
get_parallel,
|
||||
get_serving,
|
||||
)
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils import (
|
||||
CLIENT_MEDIA_EXCEPTIONS,
|
||||
load_audio,
|
||||
load_image,
|
||||
load_video,
|
||||
)
|
||||
from sglang.srt.utils.hf_transformers_utils import resolve_image_processor_backend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_mm_grid_attrs = {
|
||||
# Kimi K2.5/K3 HF processors use grid_thws (see base_processor.ATTR_NAME_TO_MODALITY).
|
||||
Modality.IMAGE: ("image_grid_thw", "image_grid_hws", "grid_thws"),
|
||||
Modality.VIDEO: ("video_grid_thw",),
|
||||
Modality.AUDIO: ("audio_feature_lens_raw",),
|
||||
}
|
||||
|
||||
|
||||
def _convert(data):
|
||||
if isinstance(data, torch.Tensor):
|
||||
return data
|
||||
elif isinstance(data, np.ndarray):
|
||||
return torch.tensor(data)
|
||||
elif isinstance(data, list) and isinstance(data[0], np.ndarray):
|
||||
return torch.tensor(np.array(data))
|
||||
elif isinstance(data, list) and isinstance(data[0], (int, float)):
|
||||
return torch.tensor(data)
|
||||
else:
|
||||
return data
|
||||
|
||||
|
||||
def _get_original_image_size(image):
|
||||
"""Return an image's original (width, height) before encoder preprocessing."""
|
||||
if isinstance(image, dict):
|
||||
image = image.get("image")
|
||||
if isinstance(image, torch.Tensor):
|
||||
if image.ndim < 2:
|
||||
raise ValueError(f"Invalid image tensor shape: {tuple(image.shape)}")
|
||||
return [int(image.shape[-1]), int(image.shape[-2])]
|
||||
if hasattr(image, "size"):
|
||||
width, height = image.size
|
||||
return [int(width), int(height)]
|
||||
raise TypeError(f"Cannot determine original image size from {type(image)}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class EncoderPreprocessResult:
|
||||
mm_inputs: dict
|
||||
grid_thw: Union[torch.Tensor, List]
|
||||
token_counts: List[int]
|
||||
|
||||
|
||||
class EncoderPreprocessor:
|
||||
"""CPU-bound multimodal preprocessing pipeline.
|
||||
|
||||
Takes raw media URLs / base64 data and produces HF processor output dicts
|
||||
(CPU tensors). The GPU model is never touched here — only the HF
|
||||
image/video/audio processors are invoked.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
server_args : ServerArgs
|
||||
Server configuration (model path, processor flags, etc.).
|
||||
model_config : ModelConfig
|
||||
Model configuration (hf_config, hidden_size, etc.).
|
||||
model_preprocessor : callable, optional
|
||||
Optional model-specific preprocessor (``model.preprocess_mm_for_encoder``).
|
||||
When provided, overrides the default HF processor path for the given
|
||||
modality.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_args: ServerArgs,
|
||||
model_config: ModelConfig,
|
||||
encoder_media_processor_config: EncoderMediaProcessorConfig,
|
||||
model_preprocessor: Optional[Callable] = None,
|
||||
):
|
||||
self.server_args = server_args
|
||||
self.model_config = model_config
|
||||
self._model_preprocessor = model_preprocessor
|
||||
self.encoder_media_processor_config = encoder_media_processor_config
|
||||
self.model_type = getattr(
|
||||
model_config.hf_config, "model_type", "unknown"
|
||||
).lower()
|
||||
|
||||
self.device = get_device().device
|
||||
|
||||
use_image_processor_gpu = envs.SGLANG_ENCODER_IMAGE_PROCESSOR_USE_GPU.get()
|
||||
self.use_image_processor_gpu = (
|
||||
use_image_processor_gpu
|
||||
and resolve_image_processor_backend(server_args) != "pil"
|
||||
)
|
||||
|
||||
self._load_mm_processor(server_args)
|
||||
self._supported_modalities = frozenset(
|
||||
modality
|
||||
for modality, processor in (
|
||||
(Modality.IMAGE, self.image_processor),
|
||||
(Modality.VIDEO, self.video_processor),
|
||||
(Modality.AUDIO, self.audio_processor),
|
||||
)
|
||||
if processor is not None or self._model_preprocessor is not None
|
||||
)
|
||||
self._build_vision_config(get_mm().mm_process_config)
|
||||
self.model_audio_sr = self._resolve_audio_sr()
|
||||
logger.info(f"Resolved model audio sample rate: {self.model_audio_sr} Hz")
|
||||
|
||||
self.preproc_executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=envs.SGLANG_ENCODER_PREPROC_WORKERS.get()
|
||||
)
|
||||
self.io_executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=int(os.environ.get("SGLANG_ENCODER_MM_LOAD_WORKERS", 4))
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# HF Processor Loading
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_mm_processor(self, server_args: ServerArgs):
|
||||
from transformers import AutoImageProcessor, AutoVideoProcessor
|
||||
|
||||
image_processor_backend = resolve_image_processor_backend(server_args)
|
||||
image_processor_kwargs = (
|
||||
{}
|
||||
if image_processor_backend == "auto"
|
||||
else {"backend": image_processor_backend}
|
||||
)
|
||||
try:
|
||||
self.image_processor = AutoImageProcessor.from_pretrained(
|
||||
get_serving().tokenizer_path or get_model().model_path,
|
||||
trust_remote_code=server_args.trust_remote_code,
|
||||
revision=server_args.revision,
|
||||
**image_processor_kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load image processor: {e}")
|
||||
self.image_processor = None
|
||||
|
||||
try:
|
||||
self.video_processor = AutoVideoProcessor.from_pretrained(
|
||||
get_serving().tokenizer_path or get_model().model_path,
|
||||
trust_remote_code=server_args.trust_remote_code,
|
||||
revision=server_args.revision,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load video processor: {e}")
|
||||
self.video_processor = None
|
||||
|
||||
try:
|
||||
_audio_proc = AutoProcessor.from_pretrained(
|
||||
get_serving().tokenizer_path or get_model().model_path,
|
||||
trust_remote_code=server_args.trust_remote_code,
|
||||
revision=server_args.revision,
|
||||
)
|
||||
if not hasattr(_audio_proc, "feature_extractor"):
|
||||
logger.warning(
|
||||
"Loaded AutoProcessor has no feature_extractor attribute, "
|
||||
"audio processing will be unavailable."
|
||||
)
|
||||
self.audio_processor = None
|
||||
else:
|
||||
self.audio_processor = _audio_proc
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load audio processor: {e}")
|
||||
self.audio_processor = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Config Validation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_vision_config(self, mm_process_config):
|
||||
self.vision_config = (
|
||||
mm_process_config.get("vision_config", {})
|
||||
if mm_process_config is not None
|
||||
else {}
|
||||
)
|
||||
for modality_str in ["image", "video", "audio"]:
|
||||
if not self.vision_config.get(modality_str, None):
|
||||
self.vision_config[modality_str] = {}
|
||||
if self.use_image_processor_gpu:
|
||||
self.vision_config[modality_str]["device"] = self.device
|
||||
|
||||
if modality_str == "video":
|
||||
video_defaults = {"fps": 2.0, "max_frames": 768, "min_frames": 4}
|
||||
for k, v in video_defaults.items():
|
||||
self.vision_config["video"].setdefault(k, v)
|
||||
|
||||
if modality_str == "audio":
|
||||
if "return_attention_mask" not in self.vision_config["audio"]:
|
||||
self.vision_config["audio"]["return_attention_mask"] = True
|
||||
if "padding" not in self.vision_config["audio"]:
|
||||
if self.model_type == "qwen2_audio":
|
||||
self.vision_config["audio"]["padding"] = "max_length"
|
||||
else:
|
||||
self.vision_config["audio"]["padding"] = True
|
||||
if "truncation" not in self.vision_config["audio"]:
|
||||
if (
|
||||
hasattr(self, "audio_processor")
|
||||
and self.audio_processor is not None
|
||||
):
|
||||
if self.audio_processor.__class__.__name__ in {
|
||||
"Gemma3nProcessor",
|
||||
"GlmAsrProcessor",
|
||||
"Qwen2AudioProcessor",
|
||||
"Qwen3OmniMoeProcessor",
|
||||
}:
|
||||
self.vision_config["audio"]["truncation"] = False
|
||||
|
||||
def _resolve_audio_sr(self) -> int:
|
||||
def _read(obj, attr):
|
||||
if obj is None:
|
||||
return None
|
||||
if isinstance(obj, dict):
|
||||
return obj.get(attr)
|
||||
return getattr(obj, attr, None)
|
||||
|
||||
audio_cfg = self.vision_config.get("audio", {})
|
||||
sr = audio_cfg.get("audio_sampling_rate")
|
||||
if sr:
|
||||
return int(sr)
|
||||
|
||||
hf_cfg = self.model_config.hf_config
|
||||
thinker_cfg = _read(hf_cfg, "thinker_config")
|
||||
pc = _read(thinker_cfg, "processor_config") or _read(hf_cfg, "processor_config")
|
||||
sr = _read(pc, "audio_sampling_rate")
|
||||
if sr:
|
||||
return int(sr)
|
||||
ac = _read(thinker_cfg, "audio_config") or _read(hf_cfg, "audio_config")
|
||||
for attr in ("sampling_rate", "sample_rate"):
|
||||
sr = _read(ac, attr)
|
||||
if sr:
|
||||
return int(sr)
|
||||
|
||||
sr = audio_cfg.get("sampling_rate")
|
||||
if sr:
|
||||
return int(sr)
|
||||
logger.warning(
|
||||
"No audio sampling rate found in mm_config or hf_config; "
|
||||
"falling back to 16000 Hz. If the model expects a different SR "
|
||||
"(e.g. MiMo-V2 defaults to 24000), audio will be warped."
|
||||
)
|
||||
return 16000
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Media I/O
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_single_item(
|
||||
self,
|
||||
data,
|
||||
modality: Modality,
|
||||
frame_count_limit=None,
|
||||
discard_alpha_channel=True,
|
||||
):
|
||||
from sglang.srt.disaggregation.encoder.server import BadRequestError, MMError
|
||||
|
||||
media_metadata = {}
|
||||
content_hash = None
|
||||
if isinstance(data, dict):
|
||||
if "url" not in data:
|
||||
return data
|
||||
media_metadata = {key: value for key, value in data.items() if key != "url"}
|
||||
content_hash = parse_content_hash(data.get("content_hash"))
|
||||
data = data["url"]
|
||||
try:
|
||||
if modality == Modality.IMAGE:
|
||||
if content_hash is not None:
|
||||
snapshot = snapshot_media(data)
|
||||
if snapshot.content_digest != content_hash:
|
||||
raise BadRequestError(
|
||||
"Encoder media content hash mismatch: "
|
||||
f"expected {content_hash}, got {snapshot.content_digest}"
|
||||
)
|
||||
data = snapshot.data
|
||||
gpu_image_decode = (
|
||||
self.encoder_media_processor_config.image_decode_mode
|
||||
if self.use_image_processor_gpu
|
||||
else False
|
||||
)
|
||||
img, _ = load_image(data, gpu_image_decode)
|
||||
if (
|
||||
discard_alpha_channel
|
||||
and not isinstance(img, torch.Tensor)
|
||||
and img.mode != "RGB"
|
||||
):
|
||||
img = img.convert("RGB")
|
||||
if (
|
||||
media_metadata
|
||||
and self.encoder_media_processor_config.preserve_media_metadata
|
||||
):
|
||||
return {
|
||||
"type": "image",
|
||||
"image": img,
|
||||
**media_metadata,
|
||||
}
|
||||
return img
|
||||
elif modality == Modality.VIDEO:
|
||||
return load_video(data, frame_count_limit)
|
||||
elif modality == Modality.AUDIO:
|
||||
return load_audio(data, self.model_audio_sr)
|
||||
|
||||
except MMError:
|
||||
raise
|
||||
except CLIENT_MEDIA_EXCEPTIONS as e:
|
||||
# Not ValueError: the DP envelope classifies by `.code`, which only
|
||||
# MMError carries.
|
||||
raise BadRequestError(f"Error while loading data {data}: {e}") from e
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Error while loading data {data}: {e}")
|
||||
|
||||
def _submit_data_loading_tasks(self, items, modalities):
|
||||
futures = []
|
||||
task_info = []
|
||||
|
||||
for data, modality in zip(items, modalities):
|
||||
if modality is not None:
|
||||
futures.append(
|
||||
self.io_executor.submit(
|
||||
self._load_single_item,
|
||||
data,
|
||||
modality,
|
||||
)
|
||||
)
|
||||
task_info.append((modality, data))
|
||||
return futures, task_info
|
||||
|
||||
async def _flatten_and_load_data_by_modality(self, mm_items, modality):
|
||||
if not isinstance(mm_items, (list, tuple)):
|
||||
futures, _ = self._submit_data_loading_tasks([mm_items], [modality])
|
||||
return await asyncio.wrap_future(futures[0])
|
||||
|
||||
if len(mm_items) > 0 and isinstance(mm_items[0], (list, tuple)):
|
||||
flat_data = []
|
||||
flat_indices = []
|
||||
for group_idx, item_group in enumerate(mm_items):
|
||||
for item in item_group:
|
||||
flat_data.append(item)
|
||||
flat_indices.append(group_idx)
|
||||
|
||||
futures, _ = self._submit_data_loading_tasks(
|
||||
flat_data, [modality] * len(flat_data)
|
||||
)
|
||||
|
||||
async_futures = [asyncio.wrap_future(f) for f in futures]
|
||||
results = await asyncio.gather(*async_futures)
|
||||
|
||||
nested_results = [[] for _ in range(len(mm_items))]
|
||||
for idx, result in zip(flat_indices, results):
|
||||
nested_results[idx].append(result)
|
||||
|
||||
return nested_results
|
||||
|
||||
else:
|
||||
futures, _ = self._submit_data_loading_tasks(
|
||||
mm_items, [modality] * len(mm_items)
|
||||
)
|
||||
async_futures = [asyncio.wrap_future(f) for f in futures]
|
||||
return await asyncio.gather(*async_futures)
|
||||
|
||||
async def _flatten_and_load_images(self, mm_items):
|
||||
return await self._flatten_and_load_data_by_modality(mm_items, Modality.IMAGE)
|
||||
|
||||
async def _flatten_and_load_videos(self, mm_items):
|
||||
if not isinstance(mm_items, (list, tuple)):
|
||||
mm_items = [mm_items]
|
||||
|
||||
futures, _ = self._submit_data_loading_tasks(
|
||||
mm_items, [Modality.VIDEO] * len(mm_items)
|
||||
)
|
||||
async_futures = [asyncio.wrap_future(f) for f in futures]
|
||||
video_items = await asyncio.gather(*async_futures)
|
||||
|
||||
video_processor_kwargs = {}
|
||||
if "qwen" in self.model_type:
|
||||
video_processed = [
|
||||
await preprocess_video(
|
||||
video, video_config=self.vision_config.get("video", {})
|
||||
)
|
||||
for video in video_items
|
||||
]
|
||||
videos, video_metadata = map(list, zip(*video_processed))
|
||||
video_processor_kwargs["do_sample_frames"] = False
|
||||
if video_metadata:
|
||||
video_processor_kwargs["video_metadata"] = video_metadata
|
||||
return videos, video_processor_kwargs
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Video processing is not supported for {self.model_type} model."
|
||||
)
|
||||
|
||||
async def _flatten_and_load_audios(self, mm_items):
|
||||
return await self._flatten_and_load_data_by_modality(mm_items, Modality.AUDIO)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# HF Processor Calls
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def process_mm_items(
|
||||
self, mm_items, modality: Modality
|
||||
) -> EncoderPreprocessResult:
|
||||
"""Process multimodal items through the HF processor pipeline.
|
||||
|
||||
Returns the ``mm_inputs`` dict produced by the HF image/video/audio
|
||||
processor, its normalized grid metadata, and one output token count per
|
||||
grid entry. Does not look up ``get_feature_fn``; that stays in
|
||||
:class:`MMEncoder`.
|
||||
"""
|
||||
if modality == Modality.IMAGE:
|
||||
mm_inputs = await self._process_image_items(
|
||||
mm_items, self._model_preprocessor
|
||||
)
|
||||
elif modality == Modality.VIDEO:
|
||||
mm_inputs = await self._process_video_items(
|
||||
mm_items, self._model_preprocessor
|
||||
)
|
||||
elif modality == Modality.AUDIO:
|
||||
mm_inputs = await self._process_audio_items(
|
||||
mm_items, self._model_preprocessor
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported modality: {modality}")
|
||||
grid_thw = self._get_mm_grid_dim(mm_inputs, modality)
|
||||
token_counts = [self.get_num_tokens(grid, modality) for grid in grid_thw]
|
||||
return EncoderPreprocessResult(
|
||||
mm_inputs=mm_inputs,
|
||||
grid_thw=grid_thw,
|
||||
token_counts=token_counts,
|
||||
)
|
||||
|
||||
def supports_modality(self, modality: Modality) -> bool:
|
||||
return modality in self._supported_modalities
|
||||
|
||||
async def process_batch_mm_items(
|
||||
self, requests: List[dict], modality: Modality
|
||||
) -> tuple[EncoderPreprocessResult, List[int]]:
|
||||
"""Flatten requests, run the processor once, and return batch layout."""
|
||||
flat_items, items_per_req = self._flatten_batch_requests(requests, modality)
|
||||
result = await self.process_mm_items(flat_items, modality)
|
||||
return result, items_per_req
|
||||
|
||||
def _flatten_batch_requests(
|
||||
self, requests: List[dict], modality: Modality
|
||||
) -> tuple[List, List[int]]:
|
||||
# items_per_req counts grid entries (post-expansion) so per-request
|
||||
# slicing of grid_dim/final_slices stays aligned for processors that
|
||||
# expand one leaf into multiple grids (e.g. Kimi-VL/K2.5/K3 dict-of-images).
|
||||
flat_items = []
|
||||
items_per_req = []
|
||||
for req in requests:
|
||||
leaves = self._flatten_nested_items(req["mm_items"])
|
||||
flat_items.extend(leaves)
|
||||
items_per_req.append(sum(self._grid_count_per_leaf(leaves, modality)))
|
||||
return flat_items, items_per_req
|
||||
|
||||
async def _process_image_items(self, mm_items, model_preprocessor):
|
||||
if not (self.image_processor or model_preprocessor):
|
||||
raise ValueError("No image processor available")
|
||||
images = await self._flatten_and_load_images(mm_items)
|
||||
if self.model_type in ["kimi_k25", "kimi_k3", "kimi_vl"]:
|
||||
images = self._normalize_kimi_encoder_images(images)
|
||||
original_image_sizes = [_get_original_image_size(item) for item in images]
|
||||
if model_preprocessor:
|
||||
processor_output = invoke_encoder_preprocessor(
|
||||
model_preprocessor,
|
||||
images,
|
||||
Modality.IMAGE,
|
||||
self.vision_config,
|
||||
image_processor=self.image_processor,
|
||||
use_gpu_preprocessing=self.use_image_processor_gpu,
|
||||
)
|
||||
if (
|
||||
isinstance(processor_output, EncoderPreprocessOutput)
|
||||
and processor_output.materialize_local_items is not None
|
||||
):
|
||||
parallel = get_parallel()
|
||||
await asyncio.get_running_loop().run_in_executor(
|
||||
self.preproc_executor,
|
||||
processor_output.materialize_for_rank,
|
||||
parallel.attn_tp_rank,
|
||||
parallel.attn_tp_size,
|
||||
)
|
||||
return processor_output
|
||||
image_config = self.vision_config.get("image", {})
|
||||
processor_input = await asyncio.get_running_loop().run_in_executor(
|
||||
self.preproc_executor,
|
||||
functools.partial(self.image_processor, images=images, **image_config),
|
||||
)
|
||||
if self.model_type == "kimi_k3":
|
||||
processor_input["original_image_sizes"] = original_image_sizes
|
||||
return processor_input
|
||||
|
||||
async def _process_video_items(self, mm_items, model_preprocessor):
|
||||
if model_preprocessor:
|
||||
return model_preprocessor(mm_items, Modality.VIDEO, self.vision_config)
|
||||
if not self.video_processor:
|
||||
raise ValueError("No video processor available")
|
||||
|
||||
videos, video_processor_kwargs = await self._flatten_and_load_videos(mm_items)
|
||||
processor_input = await asyncio.get_running_loop().run_in_executor(
|
||||
self.preproc_executor,
|
||||
functools.partial(
|
||||
self.video_processor, videos=videos, **video_processor_kwargs
|
||||
),
|
||||
)
|
||||
|
||||
if (
|
||||
self.model_type
|
||||
in [
|
||||
"qwen3_vl",
|
||||
"qwen3_vl_moe",
|
||||
"qwen3_5",
|
||||
"qwen3_5_moe",
|
||||
"intern_s2_preview",
|
||||
]
|
||||
and video_processor_kwargs.get("video_metadata", None) is not None
|
||||
):
|
||||
video_metadata = video_processor_kwargs["video_metadata"]
|
||||
try:
|
||||
merge_size = (
|
||||
self.model_config.hf_config.vision_config.spatial_merge_size
|
||||
)
|
||||
except (AttributeError, KeyError):
|
||||
merge_size = 2
|
||||
video_timestamps = []
|
||||
for metadata in video_metadata:
|
||||
video_fps = metadata.get("fps", None) or 24
|
||||
frames_indices = metadata.get("frames_indices", None)
|
||||
timestamps = self._calculate_timestamps(
|
||||
frames_indices, video_fps, merge_size
|
||||
)
|
||||
video_timestamps.append(timestamps)
|
||||
processor_input["video_timestamps"] = video_timestamps
|
||||
elif (
|
||||
self.model_type in ["qwen2_5_vl", "qwen2_5_omni", "qwen3_omni_moe"]
|
||||
and processor_input.get("video_grid_thw", None) is not None
|
||||
):
|
||||
video_grid_thw = processor_input["video_grid_thw"]
|
||||
try:
|
||||
temporal_patch_size = self.video_processor.temporal_patch_size
|
||||
except AttributeError:
|
||||
temporal_patch_size = 2
|
||||
fps_list = [
|
||||
self.vision_config.get("video", {}).get("fps", None) or 2
|
||||
] * len(video_grid_thw)
|
||||
second_per_grid_ts = [(temporal_patch_size / fps) for fps in fps_list]
|
||||
second_per_grid_ts_tensor = torch.tensor(
|
||||
second_per_grid_ts, dtype=torch.float32
|
||||
)
|
||||
processor_input["second_per_grid_ts"] = second_per_grid_ts_tensor
|
||||
|
||||
return processor_input
|
||||
|
||||
async def _process_audio_items(self, mm_items, model_preprocessor):
|
||||
audios = await self._flatten_and_load_audios(mm_items)
|
||||
|
||||
if model_preprocessor:
|
||||
return model_preprocessor(audios, Modality.AUDIO, self.vision_config)
|
||||
|
||||
if not self.audio_processor:
|
||||
raise ValueError("No audio processor available")
|
||||
|
||||
audio_config = self.vision_config.get("audio", {})
|
||||
processor_input = await asyncio.get_running_loop().run_in_executor(
|
||||
self.preproc_executor,
|
||||
functools.partial(
|
||||
self.audio_processor.feature_extractor, audios, **audio_config
|
||||
),
|
||||
)
|
||||
processor_input["feature_attention_mask"] = processor_input.pop(
|
||||
"attention_mask"
|
||||
)
|
||||
input_lengths = torch.tensor(
|
||||
processor_input["feature_attention_mask"].sum(-1), dtype=torch.long
|
||||
)
|
||||
processor_input["audio_feature_lens_raw"] = input_lengths
|
||||
output_lengths = self._get_feat_extract_output_lengths(input_lengths)
|
||||
processor_input["audio_feature_lens"] = output_lengths
|
||||
return processor_input
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Audio Feature Length Computation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _get_feat_extract_output_lengths(self, feature_lens):
|
||||
if self.model_type in ["qwen2_audio", "qwen2_5_omni"]:
|
||||
input_length = (feature_lens - 1) // 2 + 1
|
||||
return (input_length - 2) // 2 + 1
|
||||
elif self.model_type in ["qwen3_asr", "qwen3_omni_moe"]:
|
||||
input_lengths_leave = feature_lens % 100
|
||||
feat_lengths = (input_lengths_leave - 1) // 2 + 1
|
||||
output_lengths = (
|
||||
((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (feature_lens // 100) * 13
|
||||
)
|
||||
return output_lengths
|
||||
elif self.model_type == "mimo_v2":
|
||||
return feature_lens
|
||||
else:
|
||||
logger.warning(
|
||||
f"Fallback to original HF audio sample logic for {self.model_type}"
|
||||
)
|
||||
input_length = (feature_lens - 1) // 2 + 1
|
||||
return (input_length - 2) // 2 + 1
|
||||
|
||||
def _get_mm_grid_dim(self, mm_inputs: dict, modality: Modality):
|
||||
# Kimi K2.5/K3 vision processors only emit `grid_thws`; prefer it over generic keys
|
||||
# so we never pick a mis-typed or stale `image_grid_hws` field from kwargs.
|
||||
attrs = _mm_grid_attrs[modality]
|
||||
model_type = (self.model_type or "").lower()
|
||||
if modality == Modality.IMAGE:
|
||||
# Kimi K2.5/K3 emit grid_thws, while Kimi-VL emits image_grid_hws.
|
||||
# Other model types keep the generic attr order above.
|
||||
if model_type in ("kimi_k25", "kimi_k3"):
|
||||
attrs = ("grid_thws", "image_grid_thw", "image_grid_hws")
|
||||
elif model_type == "kimi_vl":
|
||||
attrs = ("image_grid_hws", "image_grid_thw", "grid_thws")
|
||||
|
||||
for attr in attrs:
|
||||
if attr in mm_inputs and mm_inputs[attr] is not None:
|
||||
return _convert(mm_inputs[attr])
|
||||
raise ValueError(
|
||||
f"Grid dim ({_mm_grid_attrs[modality]}) not found in {mm_inputs}"
|
||||
)
|
||||
|
||||
def get_num_patches(
|
||||
self, grid: Union[torch.Tensor, List[int]], modality: Modality
|
||||
) -> int:
|
||||
"""Calculate number of raw patches (before merge/sampling). Used for pixel_values slicing."""
|
||||
if modality == Modality.AUDIO:
|
||||
return int(grid.item())
|
||||
if self.model_type == "kimi_vl" and modality == Modality.IMAGE:
|
||||
h, w = self._kimi_hw_from_patch_grid(grid)
|
||||
return h * w
|
||||
return int(grid[0] * grid[1] * grid[2])
|
||||
|
||||
@staticmethod
|
||||
def _kimi_hw_from_patch_grid(
|
||||
grid: Union[torch.Tensor, np.ndarray, List[int], Tuple[int, ...]],
|
||||
) -> Tuple[int, int]:
|
||||
"""Extract (height, width) from Kimi 2D or 3D patch-grid metadata."""
|
||||
if isinstance(grid, torch.Tensor):
|
||||
values = grid.flatten().tolist()
|
||||
elif isinstance(grid, np.ndarray):
|
||||
values = grid.reshape(-1).tolist()
|
||||
else:
|
||||
values = np.asarray(grid).reshape(-1).tolist()
|
||||
|
||||
if len(values) not in (2, 3):
|
||||
raise ValueError(
|
||||
f"Invalid Kimi image grid metadata: {values}; "
|
||||
"expected [h, w] or [t, h, w]"
|
||||
)
|
||||
return int(values[-2]), int(values[-1])
|
||||
|
||||
def _kimi_tokens_from_patch_grid(self, grid: Union[torch.Tensor, List[int]]) -> int:
|
||||
"""Calculate Kimi image tokens from either 2D or 3D patch metadata."""
|
||||
h, w = self._kimi_hw_from_patch_grid(grid)
|
||||
merge_h, merge_w = self.model_config.hf_config.vision_config.merge_kernel_size
|
||||
return (h * w) // (merge_h * merge_w)
|
||||
|
||||
def get_num_tokens(
|
||||
self, grid: Union[torch.Tensor, List[int]], modality: Modality
|
||||
) -> int:
|
||||
"""Compatibility helper for callers that still provide patch grids."""
|
||||
if modality == Modality.AUDIO:
|
||||
input_length = self.get_num_patches(grid, modality)
|
||||
return self._get_feat_extract_output_lengths(input_length)
|
||||
else:
|
||||
if (
|
||||
self.model_type in ["kimi_k25", "kimi_k3", "kimi_vl"]
|
||||
and modality == Modality.IMAGE
|
||||
):
|
||||
return self._kimi_tokens_from_patch_grid(grid)
|
||||
merge_size = getattr(self.image_processor, "merge_size", 2)
|
||||
return self.get_num_patches(grid, modality) // (merge_size**2)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Video Timestamp Computation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _calculate_timestamps(self, indices, video_fps: float, merge_size: int = 2):
|
||||
if not isinstance(indices, list):
|
||||
indices = indices.tolist()
|
||||
if len(indices) % merge_size != 0:
|
||||
indices.extend(
|
||||
indices[-1] for _ in range(merge_size - len(indices) % merge_size)
|
||||
)
|
||||
timestamps = [idx / video_fps for idx in indices]
|
||||
timestamps = [
|
||||
(timestamps[i] + timestamps[i + merge_size - 1]) / 2
|
||||
for i in range(0, len(timestamps), merge_size)
|
||||
]
|
||||
return timestamps
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Kimi Normalization
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _normalize_kimi_encoder_images(self, images):
|
||||
"""Normalize Kimi image inputs for the image processor call."""
|
||||
from PIL import Image as PILImage
|
||||
|
||||
def wrap_one(img):
|
||||
if isinstance(img, dict) and img.get("type") in ("image", "video_chunk"):
|
||||
return [img]
|
||||
if isinstance(img, PILImage.Image):
|
||||
return [{"type": "image", "image": img}]
|
||||
return [img]
|
||||
|
||||
if not images:
|
||||
return images
|
||||
|
||||
# Disagg may supply nested lists from grouped routing.
|
||||
images = self._flatten_nested_items(images)
|
||||
|
||||
if self.model_type == "kimi_vl":
|
||||
normalized = []
|
||||
for img in images:
|
||||
if (
|
||||
isinstance(img, dict)
|
||||
and img.get("type") == "image"
|
||||
and "image" in img
|
||||
):
|
||||
inner = img["image"]
|
||||
if isinstance(inner, (list, tuple)):
|
||||
normalized.extend(self._flatten_nested_items(inner))
|
||||
else:
|
||||
normalized.append(inner)
|
||||
else:
|
||||
normalized.append(img)
|
||||
return normalized
|
||||
|
||||
# Kimi-K2.5/K3 vision processors expect media dicts.
|
||||
normalized = []
|
||||
for img in images:
|
||||
wrapped = wrap_one(img)
|
||||
for media in wrapped:
|
||||
if (
|
||||
isinstance(media, dict)
|
||||
and media.get("type") == "image"
|
||||
and isinstance(media.get("image"), (list, tuple))
|
||||
):
|
||||
for inner in self._flatten_nested_items(media["image"]):
|
||||
normalized.append({**media, "image": inner})
|
||||
else:
|
||||
normalized.append(media)
|
||||
|
||||
return normalized
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Utility Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _flatten_nested_items(items):
|
||||
if not isinstance(items, (list, tuple)):
|
||||
return [items]
|
||||
|
||||
flat = []
|
||||
for item in items:
|
||||
if isinstance(item, (list, tuple)):
|
||||
flat.extend(EncoderPreprocessor._flatten_nested_items(item))
|
||||
else:
|
||||
flat.append(item)
|
||||
return flat
|
||||
|
||||
def _grid_count_per_leaf(self, leaves: List, modality: Modality) -> List[int]:
|
||||
"""Number of grid entries each leaf produces under the model's processor.
|
||||
|
||||
Most processors map 1 leaf -> 1 grid. Kimi-VL/K2.5/K3 image processors expand
|
||||
a leaf shaped {"type": "image", "image": [pil1, pil2, ...]} into N grids.
|
||||
"""
|
||||
if (
|
||||
self.model_type not in ("kimi_k25", "kimi_k3", "kimi_vl")
|
||||
or modality != Modality.IMAGE
|
||||
):
|
||||
return [1] * len(leaves)
|
||||
|
||||
def count(leaf):
|
||||
if (
|
||||
isinstance(leaf, dict)
|
||||
and leaf.get("type") == "image"
|
||||
and isinstance(leaf.get("image"), (list, tuple))
|
||||
):
|
||||
return len(self._flatten_nested_items(leaf["image"]))
|
||||
return 1
|
||||
|
||||
return [count(leaf) for leaf in leaves]
|
||||
+628
-592
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -317,7 +317,6 @@ class GenerateReqInput:
|
||||
# For EPD-disaggregated inference
|
||||
need_wait_for_mm_inputs: Optional[bool] = None
|
||||
num_items_assigned: Optional[Dict[Modality, List[int]]] = None
|
||||
mm_data_mooncake: Optional[List[Any]] = None
|
||||
# Snapshot of encoder URLs at the time tokenizer-side computed
|
||||
# ``num_items_assigned``.
|
||||
encoder_urls: Optional[List[str]] = None
|
||||
@@ -1024,10 +1023,6 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
|
||||
|
||||
need_wait_for_mm_inputs: Optional[bool] = None
|
||||
num_items_assigned: Optional[Dict[Modality, List[int]]] = None
|
||||
# Pickled Optional[List[{"url": MultimodalDataInputItem, "modality": Modality}]]
|
||||
# from MMReceiverBase._extract_url_data. "url" is ImageData.url,
|
||||
# dict["url"] when present, or the original raw multimodal item.
|
||||
mm_data_mooncake: Optional[PickleWrapper] = None
|
||||
# Encoder URL snapshot frozen at tokenizer-side dispatch time so that
|
||||
# encoder_idx assignments stay consistent in the scheduler subprocess.
|
||||
# Internal IPC only.
|
||||
@@ -1044,11 +1039,9 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
|
||||
cache_salt: Optional[str] = None
|
||||
|
||||
def wrap_pickle_fields(self):
|
||||
self.mm_data_mooncake = wrap_as_pickle(self.mm_data_mooncake)
|
||||
self.time_stats = wrap_as_pickle(self.time_stats)
|
||||
|
||||
def unwrap_pickle_fields(self):
|
||||
self.mm_data_mooncake = unwrap_from_pickle(self.mm_data_mooncake)
|
||||
self.time_stats = unwrap_from_pickle(self.time_stats)
|
||||
|
||||
|
||||
|
||||
@@ -709,6 +709,7 @@ def general_mm_embed_routine(
|
||||
if (
|
||||
isinstance(precomputed_embeddings, torch.Tensor)
|
||||
and precomputed_embeddings.is_cuda
|
||||
and not mm_item.keep_device_embedding
|
||||
):
|
||||
mm_item.precomputed_embeddings = (
|
||||
precomputed_embeddings.to(
|
||||
|
||||
@@ -343,6 +343,8 @@ class MultimodalDataItem(msgspec.Struct, kw_only=True, dict=True, array_like=Tru
|
||||
# the precomputed embeddings, passed as final encoder embeddings
|
||||
# One and only one of the feature and precomputed_embeddings will be empty
|
||||
precomputed_embeddings: Optional[MultimodalDataValue] = None
|
||||
# Keep precomputed_embeddings on GPU after use (EPD pool/GPU receive path)
|
||||
keep_device_embedding: bool = False
|
||||
|
||||
# Processor-owned tensors/arrays/scalars/transports. msgspec rejects a
|
||||
# precise union with multiple custom types, but accepts Ext-decoded values
|
||||
|
||||
@@ -82,7 +82,7 @@ from sglang.srt.disaggregation.decode import (
|
||||
from sglang.srt.disaggregation.decode_kvcache_offload_manager import (
|
||||
DecodeKVCacheOffloadManager,
|
||||
)
|
||||
from sglang.srt.disaggregation.encode_receiver import create_mm_receiver
|
||||
from sglang.srt.disaggregation.encoder.receiver import create_mm_receiver
|
||||
from sglang.srt.disaggregation.prefill import (
|
||||
PrefillBootstrapQueue,
|
||||
SchedulerDisaggregationPrefillMixin,
|
||||
@@ -4539,6 +4539,10 @@ class Scheduler(
|
||||
self._pending_chunked_abort_req = chunked_req
|
||||
|
||||
# todo hisparse, release resources for abort requests in hisparse coordinator
|
||||
# Abort requests still waiting for encoder embeddings (EPD language-only)
|
||||
if self.mm_receiver is not None:
|
||||
self.mm_receiver.abort_waiting_requests(recv_req)
|
||||
|
||||
# Delete requests in the waiting queue
|
||||
to_del = []
|
||||
for i, req in enumerate(self.waiting_queue):
|
||||
|
||||
@@ -47,7 +47,7 @@ from fastapi import BackgroundTasks
|
||||
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
|
||||
from sglang.srt.disaggregation.encode_receiver import create_mm_receiver
|
||||
from sglang.srt.disaggregation.encoder.receiver import create_mm_receiver
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry
|
||||
@@ -671,7 +671,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
||||
# Encoder Disaggregation
|
||||
self.encoder_bootstrap_server = None
|
||||
if self.server_args.language_only:
|
||||
from sglang.srt.disaggregation.encode_receiver import (
|
||||
from sglang.srt.disaggregation.encoder.receiver import (
|
||||
EncoderBootstrapServer,
|
||||
)
|
||||
|
||||
@@ -1409,7 +1409,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
||||
need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs,
|
||||
num_items_assigned=obj.num_items_assigned,
|
||||
multi_item_delimiter_indices=obj.multi_item_delimiter_indices,
|
||||
mm_data_mooncake=obj.mm_data_mooncake,
|
||||
encoder_urls=obj.encoder_urls,
|
||||
)
|
||||
elif isinstance(obj, EmbeddingReqInput):
|
||||
|
||||
@@ -1021,7 +1021,7 @@ class MiMoProcessor:
|
||||
"num_video_tokens": num_media_tokens_per_grid,
|
||||
"segment_audio_token_len": segment_audio_token_len,
|
||||
"segment_audio": segment_audio,
|
||||
# Used by encode_server to trim audio_encoder output.
|
||||
# Used by encoder.server to trim audio_encoder output.
|
||||
"audio_start_token_idx": audio_start_token_idx,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -206,7 +206,6 @@ class RequestLogger:
|
||||
"image_data",
|
||||
"audio_data",
|
||||
"video_data",
|
||||
"mm_data_mooncake",
|
||||
"lora_path",
|
||||
"sampling_params",
|
||||
}
|
||||
@@ -220,7 +219,6 @@ class RequestLogger:
|
||||
"image_data",
|
||||
"audio_data",
|
||||
"video_data",
|
||||
"mm_data_mooncake",
|
||||
"lora_path",
|
||||
}
|
||||
out_skip_names = {"text", "output_ids", "embedding"}
|
||||
|
||||
@@ -10,7 +10,7 @@ import zmq
|
||||
from prometheus_client.parser import text_string_to_metric_families
|
||||
from prometheus_client.samples import Sample
|
||||
|
||||
from sglang.srt.disaggregation.encode_server import MINIMUM_PNG_PICTURE_BASE64
|
||||
from sglang.srt.disaggregation.encoder.http_server import MINIMUM_PNG_PICTURE_BASE64
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.srt.utils.network import get_zmq_socket_on_host
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
@@ -4,7 +4,7 @@ import unittest
|
||||
from array import array
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.disaggregation.encode_receiver import MMReceiverBase
|
||||
from sglang.srt.disaggregation.encoder.receiver import MMReceiverBase
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.sampling.sampling_params import SamplingParams
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
@@ -1,30 +1,53 @@
|
||||
import asyncio
|
||||
import pickle
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.disaggregation.encode_receiver import EmbeddingData
|
||||
from sglang.srt.disaggregation.encode_server import MMEncoder, _get_mm_grid_dim
|
||||
from sglang.srt.disaggregation.encoder.preprocessor import EncoderPreprocessor
|
||||
from sglang.srt.disaggregation.encoder.receiver import EmbeddingData
|
||||
from sglang.srt.disaggregation.encoder.runtime import execute_encode_pipeline
|
||||
from sglang.srt.disaggregation.encoder.server import (
|
||||
EncoderDelivery,
|
||||
InternalError,
|
||||
MMEncoder,
|
||||
MooncakeDelivery,
|
||||
ReqState,
|
||||
SendDestination,
|
||||
ZmqDelivery,
|
||||
meta_registry,
|
||||
rid_to_cond,
|
||||
rid_to_receive_count,
|
||||
rid_to_receive_endpoint,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import Modality
|
||||
from sglang.srt.utils.common import safe_pickle_loads
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestKimiVLEPDGrid(unittest.TestCase):
|
||||
class TestEncoderPreprocessorKimiGrid(CustomTestCase):
|
||||
@staticmethod
|
||||
def _make_encoder(model_type="kimi_vl"):
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.model_type = model_type
|
||||
encoder.model_config = SimpleNamespace(
|
||||
def _make_preprocessor(model_type="kimi_vl"):
|
||||
preprocessor = EncoderPreprocessor.__new__(EncoderPreprocessor)
|
||||
preprocessor.model_type = model_type
|
||||
preprocessor.model_config = SimpleNamespace(
|
||||
hf_config=SimpleNamespace(
|
||||
vision_config=SimpleNamespace(merge_kernel_size=(2, 2))
|
||||
)
|
||||
)
|
||||
return encoder
|
||||
preprocessor.image_processor = SimpleNamespace(merge_size=2)
|
||||
preprocessor._model_preprocessor = None
|
||||
return preprocessor
|
||||
|
||||
@staticmethod
|
||||
def _make_encoder():
|
||||
return MMEncoder.__new__(MMEncoder)
|
||||
|
||||
def test_kimi_vl_prefers_and_normalizes_hw_grid(self):
|
||||
mm_inputs = {
|
||||
@@ -33,7 +56,7 @@ class TestKimiVLEPDGrid(unittest.TestCase):
|
||||
"grid_thws": torch.tensor([[1, 10, 15]]),
|
||||
}
|
||||
|
||||
grid = _get_mm_grid_dim(mm_inputs, Modality.IMAGE, "kimi_vl")
|
||||
grid = self._make_preprocessor()._get_mm_grid_dim(mm_inputs, Modality.IMAGE)
|
||||
|
||||
self.assertIsInstance(grid, torch.Tensor)
|
||||
torch.testing.assert_close(grid, torch.tensor([[40, 60]]))
|
||||
@@ -44,49 +67,51 @@ class TestKimiVLEPDGrid(unittest.TestCase):
|
||||
"grid_thws": np.array([[1, 10, 15]], dtype=np.int64),
|
||||
}
|
||||
|
||||
grid = _get_mm_grid_dim(mm_inputs, Modality.IMAGE, "kimi_k25")
|
||||
grid = self._make_preprocessor("kimi_k25")._get_mm_grid_dim(
|
||||
mm_inputs, Modality.IMAGE
|
||||
)
|
||||
|
||||
torch.testing.assert_close(grid, torch.tensor([[1, 10, 15]]))
|
||||
|
||||
def test_kimi_vl_2d_grid_counting_and_slicing(self):
|
||||
preprocessor = self._make_preprocessor()
|
||||
encoder = self._make_encoder()
|
||||
grids = torch.tensor([[40, 60], [20, 40]])
|
||||
embedding = torch.arange(800 * 2).reshape(800, 2)
|
||||
|
||||
self.assertEqual(
|
||||
encoder.get_num_patches(grids[0], Modality.IMAGE),
|
||||
preprocessor.get_num_patches(grids[0], Modality.IMAGE),
|
||||
2400,
|
||||
)
|
||||
self.assertEqual(
|
||||
encoder.get_num_tokens(grids[0], Modality.IMAGE),
|
||||
preprocessor.get_num_tokens(grids[0], Modality.IMAGE),
|
||||
600,
|
||||
)
|
||||
|
||||
slices = encoder.slice_embedding(embedding, grids, Modality.IMAGE)
|
||||
slices = encoder.slice_embedding(embedding, [600, 200])
|
||||
|
||||
self.assertEqual([item.shape for item in slices], [(600, 2), (200, 2)])
|
||||
torch.testing.assert_close(slices[0], embedding[:600])
|
||||
torch.testing.assert_close(slices[1], embedding[600:])
|
||||
|
||||
def test_kimi_3d_grid_remains_supported(self):
|
||||
encoder = self._make_encoder()
|
||||
preprocessor = self._make_preprocessor()
|
||||
grid = torch.tensor([1, 40, 60])
|
||||
|
||||
self.assertEqual(encoder.get_num_patches(grid, Modality.IMAGE), 2400)
|
||||
self.assertEqual(encoder.get_num_tokens(grid, Modality.IMAGE), 600)
|
||||
self.assertEqual(preprocessor.get_num_patches(grid, Modality.IMAGE), 2400)
|
||||
self.assertEqual(preprocessor.get_num_tokens(grid, Modality.IMAGE), 600)
|
||||
|
||||
def test_kimi_k25_3d_patch_counting_is_unchanged(self):
|
||||
encoder = self._make_encoder("kimi_k25")
|
||||
preprocessor = self._make_preprocessor("kimi_k25")
|
||||
grid = torch.tensor([2, 12, 16])
|
||||
|
||||
self.assertEqual(encoder.get_num_patches(grid, Modality.IMAGE), 384)
|
||||
self.assertEqual(encoder.get_num_tokens(grid, Modality.IMAGE), 48)
|
||||
self.assertEqual(preprocessor.get_num_patches(grid, Modality.IMAGE), 384)
|
||||
self.assertEqual(preprocessor.get_num_tokens(grid, Modality.IMAGE), 48)
|
||||
|
||||
def test_grid_metadata_is_safe_to_deserialize(self):
|
||||
grid = _get_mm_grid_dim(
|
||||
grid = self._make_preprocessor()._get_mm_grid_dim(
|
||||
{"image_grid_hws": np.array([[40, 60]], dtype=np.int64)},
|
||||
Modality.IMAGE,
|
||||
"kimi_vl",
|
||||
)
|
||||
embedding_data = EmbeddingData(
|
||||
req_id="test-request",
|
||||
@@ -104,5 +129,411 @@ class TestKimiVLEPDGrid(unittest.TestCase):
|
||||
torch.testing.assert_close(restored.grid_dim, torch.tensor([[40, 60]]))
|
||||
|
||||
|
||||
class TestEncoderDelivery(CustomTestCase):
|
||||
def test_contract_has_two_direct_implementations(self):
|
||||
self.assertEqual(EncoderDelivery.__abstractmethods__, {"send", "release"})
|
||||
self.assertEqual(
|
||||
set(EncoderDelivery.__subclasses__()),
|
||||
{
|
||||
MooncakeDelivery,
|
||||
ZmqDelivery,
|
||||
},
|
||||
)
|
||||
|
||||
def test_zmq_delivery_cleanup_is_configurable(self):
|
||||
async def run():
|
||||
req_id = "test-zmq-delivery-cleanup"
|
||||
rid_to_receive_endpoint[req_id] = {"127.0.0.1:1"}
|
||||
rid_to_receive_count[req_id] = 1
|
||||
rid_to_cond[req_id] = asyncio.Condition()
|
||||
state = ReqState(req_id)
|
||||
encoder = SimpleNamespace()
|
||||
|
||||
await ZmqDelivery(encoder, cleanup_receive_state=False).release(state)
|
||||
self.assertIn(req_id, rid_to_receive_endpoint)
|
||||
self.assertIn(req_id, rid_to_receive_count)
|
||||
self.assertIn(req_id, rid_to_cond)
|
||||
|
||||
await ZmqDelivery(encoder, cleanup_receive_state=True).release(state)
|
||||
self.assertNotIn(req_id, rid_to_receive_endpoint)
|
||||
self.assertNotIn(req_id, rid_to_receive_count)
|
||||
self.assertNotIn(req_id, rid_to_cond)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_preprocess_metadata_precedes_embedding(self):
|
||||
async def run():
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.rank = 0
|
||||
encoder.req_states = {}
|
||||
encoder._embedding_dims = {Modality.IMAGE: 8}
|
||||
encoder._embedding_dtype = torch.float16
|
||||
encoder._element_size = 2
|
||||
first_state = encoder._acquire_encode_ref("req-0")
|
||||
second_state = encoder._acquire_encode_ref("req-1")
|
||||
ctx = SimpleNamespace(
|
||||
req_id="req-0",
|
||||
modality=Modality.IMAGE,
|
||||
items_per_req=[1, 2],
|
||||
preprocess_result=SimpleNamespace(
|
||||
token_counts=[2, 3, 4],
|
||||
grid_thw=[[1, 2, 3], [1, 4, 5], [1, 6, 7]],
|
||||
),
|
||||
)
|
||||
requests = [
|
||||
{"req_id": "req-0", "num_parts": 2, "part_idx": 0},
|
||||
{"req_id": "req-1", "num_parts": 2, "part_idx": 1},
|
||||
]
|
||||
|
||||
publish = AsyncMock()
|
||||
with patch.object(meta_registry, "publish", publish):
|
||||
await encoder._publish_preprocess_metadata(ctx, requests)
|
||||
|
||||
self.assertIs(encoder.req_states["req-0"], first_state)
|
||||
self.assertIs(encoder.req_states["req-1"], second_state)
|
||||
self.assertEqual(first_state.embedding_data.shape, [2, 8])
|
||||
self.assertEqual(second_state.embedding_data.shape, [7, 8])
|
||||
self.assertEqual(first_state.embedding_data.grid_dim, [[1, 2, 3]])
|
||||
self.assertEqual(
|
||||
second_state.embedding_data.grid_dim,
|
||||
[[1, 4, 5], [1, 6, 7]],
|
||||
)
|
||||
self.assertEqual(first_state.embedding_data.dtype, torch.float16)
|
||||
self.assertEqual(second_state.embedding_data.dtype, torch.float16)
|
||||
self.assertFalse(first_state.embedding_ready.is_set())
|
||||
self.assertFalse(second_state.embedding_ready.is_set())
|
||||
self.assertEqual(
|
||||
publish.await_args_list,
|
||||
[
|
||||
unittest.mock.call("req-0", 32, 2, 8),
|
||||
unittest.mock.call("req-1", 112, 7, 8),
|
||||
],
|
||||
)
|
||||
await encoder._release_encode_ref(first_state)
|
||||
await encoder._release_encode_ref(second_state)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_mooncake_embedding_is_ready_only_after_cuda_sync(self):
|
||||
class FakeCudaEmbedding:
|
||||
shape = (2, 4)
|
||||
dtype = torch.float16
|
||||
nbytes = 16
|
||||
is_cuda = True
|
||||
device = "cuda:0"
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self
|
||||
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.rank = 0
|
||||
events = []
|
||||
encoder._stage_embedding = lambda mm_data: events.append("ready")
|
||||
ctx = SimpleNamespace(
|
||||
req_id="req",
|
||||
modality=Modality.IMAGE,
|
||||
items_per_req=[1],
|
||||
preprocess_result=SimpleNamespace(
|
||||
token_counts=[2],
|
||||
grid_thw=[[1, 2, 3]],
|
||||
),
|
||||
aux_data={},
|
||||
use_global_cache=True,
|
||||
)
|
||||
stream = SimpleNamespace(synchronize=lambda: events.append("sync"))
|
||||
|
||||
with patch.object(torch.cuda, "current_stream", return_value=stream):
|
||||
encoder._stage_embeddings(
|
||||
ctx,
|
||||
[{"req_id": "req", "num_parts": 1, "part_idx": 0}],
|
||||
FakeCudaEmbedding(),
|
||||
keep_on_gpu=True,
|
||||
)
|
||||
|
||||
self.assertEqual(events, ["sync", "ready"])
|
||||
|
||||
def test_stage_embedding_does_not_resurrect_missing_state(self):
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.req_states = {}
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
InternalError, "No request state exists while encoding request: req"
|
||||
):
|
||||
encoder._stage_embedding(
|
||||
EmbeddingData(
|
||||
"req",
|
||||
1,
|
||||
0,
|
||||
None,
|
||||
Modality.IMAGE,
|
||||
embedding=torch.ones((1, 1)),
|
||||
)
|
||||
)
|
||||
|
||||
self.assertNotIn("req", encoder.req_states)
|
||||
|
||||
def test_stage_embedding_requires_active_encode(self):
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
state = ReqState("req")
|
||||
encoder.req_states = {"req": state}
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
InternalError, "Request state has no active encode work: req"
|
||||
):
|
||||
encoder._stage_embedding(
|
||||
EmbeddingData(
|
||||
"req",
|
||||
1,
|
||||
0,
|
||||
None,
|
||||
Modality.IMAGE,
|
||||
embedding=torch.ones((1, 1)),
|
||||
)
|
||||
)
|
||||
|
||||
self.assertIs(encoder.req_states["req"], state)
|
||||
self.assertIsNone(state.embedding_data)
|
||||
self.assertFalse(state.embedding_ready.is_set())
|
||||
|
||||
def test_release_during_encode_is_deferred_without_resurrecting_state(self):
|
||||
async def run():
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.rank = 0
|
||||
encoder.req_states = {}
|
||||
encoder.delivery = SimpleNamespace(release=AsyncMock())
|
||||
|
||||
state = encoder._acquire_encode_ref("req")
|
||||
state.embedding_data = EmbeddingData(
|
||||
"req",
|
||||
1,
|
||||
0,
|
||||
None,
|
||||
Modality.IMAGE,
|
||||
embedding_shape=[1, 1],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
discard = AsyncMock()
|
||||
with patch.object(meta_registry, "discard", discard):
|
||||
await encoder.release_request("req")
|
||||
self.assertTrue(state.release_requested)
|
||||
self.assertIn("req", encoder.req_states)
|
||||
encoder.delivery.release.assert_not_awaited()
|
||||
|
||||
embedding = torch.ones((1, 1))
|
||||
encoder._stage_embedding(
|
||||
EmbeddingData(
|
||||
"req",
|
||||
1,
|
||||
0,
|
||||
None,
|
||||
Modality.IMAGE,
|
||||
embedding=embedding,
|
||||
)
|
||||
)
|
||||
await encoder._release_encode_ref(state)
|
||||
|
||||
encoder.delivery.release.assert_awaited_once_with(state)
|
||||
discard.assert_awaited_once_with("req")
|
||||
self.assertIsNone(state.embedding_data)
|
||||
self.assertNotIn("req", encoder.req_states)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_error_metadata_survives_buffer_release_for_waiter(self):
|
||||
async def run():
|
||||
req_id = "test-error-metadata-waiter"
|
||||
await meta_registry.discard(req_id)
|
||||
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.req_states = {}
|
||||
encoder.delivery = SimpleNamespace(release=AsyncMock())
|
||||
state = ReqState(
|
||||
req_id,
|
||||
EmbeddingData(
|
||||
req_id,
|
||||
1,
|
||||
0,
|
||||
None,
|
||||
Modality.IMAGE,
|
||||
error_msg="encode failed",
|
||||
),
|
||||
)
|
||||
state.embedding_ready.set()
|
||||
encoder.req_states[req_id] = state
|
||||
|
||||
waiter = asyncio.create_task(meta_registry.wait(req_id))
|
||||
await asyncio.sleep(0)
|
||||
try:
|
||||
await meta_registry.publish(req_id, 0, 0, 0, error="encode failed")
|
||||
await encoder.release_request(req_id, preserve_metadata=True)
|
||||
meta = await asyncio.wait_for(waiter, timeout=1)
|
||||
self.assertEqual(meta, {"error": "encode failed"})
|
||||
self.assertNotIn(req_id, encoder.req_states)
|
||||
finally:
|
||||
if not waiter.done():
|
||||
waiter.cancel()
|
||||
await meta_registry.discard(req_id)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_zmq_pipeline_sends_only_after_encode_completes(self):
|
||||
async def run():
|
||||
events = []
|
||||
finish_encode = asyncio.Event()
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.transfer_backend = "zmq_to_tokenizer"
|
||||
|
||||
async def encode(**kwargs):
|
||||
events.append("metadata_published")
|
||||
await finish_encode.wait()
|
||||
events.append("encode_completed")
|
||||
return 16, 2, 4, None, None
|
||||
|
||||
async def send(**kwargs):
|
||||
events.append("send")
|
||||
return True
|
||||
|
||||
async def release_request(req_id, **kwargs):
|
||||
events.append("release")
|
||||
|
||||
encoder.encode = AsyncMock(side_effect=encode)
|
||||
encoder.send = AsyncMock(side_effect=send)
|
||||
encoder.release_request = AsyncMock(side_effect=release_request)
|
||||
|
||||
publish = AsyncMock()
|
||||
request = {
|
||||
"req_id": "req",
|
||||
"mm_items": ["item"],
|
||||
"modality": "image",
|
||||
"num_parts": 1,
|
||||
"part_idx": 0,
|
||||
"prefill_host": "127.0.0.1",
|
||||
"embedding_port": 1234,
|
||||
}
|
||||
with patch.object(meta_registry, "publish", publish):
|
||||
task = asyncio.create_task(
|
||||
execute_encode_pipeline(encoder, None, request)
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
self.assertEqual(events, ["metadata_published"])
|
||||
encoder.send.assert_not_awaited()
|
||||
|
||||
finish_encode.set()
|
||||
self.assertIsNone(await task)
|
||||
|
||||
self.assertEqual(
|
||||
events,
|
||||
["metadata_published", "encode_completed", "send", "release"],
|
||||
)
|
||||
publish.assert_awaited_once_with("req", 16, 2, 4)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_send_waits_for_embedding_published_by_encode(self):
|
||||
async def run():
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.rank = 0
|
||||
encoder.req_states = {}
|
||||
state = encoder._acquire_encode_ref("req")
|
||||
state.embedding_data = EmbeddingData(
|
||||
"req",
|
||||
1,
|
||||
0,
|
||||
None,
|
||||
Modality.IMAGE,
|
||||
embedding_shape=[1, 1],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
delivered = []
|
||||
|
||||
async def send(current_state, destination):
|
||||
delivered.append(await encoder._wait_for_embedding(current_state))
|
||||
|
||||
encoder.delivery = SimpleNamespace(
|
||||
send=AsyncMock(side_effect=send),
|
||||
release=AsyncMock(),
|
||||
)
|
||||
send_task = asyncio.create_task(
|
||||
encoder.send_to_destination(state, SendDestination("127.0.0.1:1"))
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
self.assertFalse(send_task.done())
|
||||
|
||||
embedding = torch.ones((1, 1))
|
||||
encoder._stage_embedding(
|
||||
EmbeddingData(
|
||||
"req",
|
||||
1,
|
||||
0,
|
||||
None,
|
||||
Modality.IMAGE,
|
||||
embedding=embedding,
|
||||
)
|
||||
)
|
||||
await send_task
|
||||
await encoder._release_encode_ref(state)
|
||||
|
||||
with patch.object(meta_registry, "discard", AsyncMock()):
|
||||
await encoder.release_request("req")
|
||||
|
||||
self.assertEqual(len(delivered), 1)
|
||||
self.assertIs(delivered[0].embedding, embedding)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_release_waits_for_send_then_clears_embedding(self):
|
||||
async def run():
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.req_states = {}
|
||||
|
||||
send_started = asyncio.Event()
|
||||
finish_send = asyncio.Event()
|
||||
|
||||
async def send(state, destination):
|
||||
send_started.set()
|
||||
await finish_send.wait()
|
||||
|
||||
embedding_seen_by_release = []
|
||||
|
||||
async def release(state):
|
||||
embedding_seen_by_release.append(state.embedding_data.embedding)
|
||||
|
||||
encoder.delivery = SimpleNamespace(
|
||||
send=AsyncMock(side_effect=send),
|
||||
release=AsyncMock(side_effect=release),
|
||||
)
|
||||
embedding = torch.ones((1, 1))
|
||||
state = ReqState(
|
||||
"req",
|
||||
EmbeddingData("req", 1, 0, None, Modality.IMAGE, embedding=embedding),
|
||||
)
|
||||
state.embedding_ready.set()
|
||||
encoder.req_states[state.req_id] = state
|
||||
|
||||
send_task = asyncio.create_task(
|
||||
encoder.send_to_destination(state, SendDestination("127.0.0.1:1"))
|
||||
)
|
||||
await send_started.wait()
|
||||
release_task = asyncio.create_task(encoder.release_request("req"))
|
||||
await asyncio.sleep(0)
|
||||
|
||||
encoder.delivery.release.assert_not_awaited()
|
||||
self.assertIs(state.embedding_data.embedding, embedding)
|
||||
|
||||
finish_send.set()
|
||||
await send_task
|
||||
await release_task
|
||||
|
||||
encoder.delivery.release.assert_awaited_once_with(state)
|
||||
self.assertEqual(len(embedding_seen_by_release), 1)
|
||||
self.assertIs(embedding_seen_by_release[0], embedding)
|
||||
self.assertIsNone(state.embedding_data)
|
||||
self.assertNotIn("req", encoder.req_states)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -3,7 +3,8 @@ import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.disaggregation import encode_server
|
||||
from sglang.srt.disaggregation.encoder import http_server
|
||||
from sglang.srt.managers.schedule_batch import Modality
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||
@@ -17,18 +18,27 @@ class _FakeEncoder:
|
||||
self.encode_dispatch_lock = asyncio.Lock()
|
||||
self.encode_calls = []
|
||||
|
||||
def has_pending_embeddings(self):
|
||||
return bool(self.embedding_to_send)
|
||||
|
||||
def supports_modality(self, modality):
|
||||
return modality == Modality.IMAGE
|
||||
|
||||
async def encode(self, **kwargs):
|
||||
self.encode_calls.append(kwargs)
|
||||
return 1, 1, 1, None, None
|
||||
|
||||
async def release_request(self, _req_id):
|
||||
return None
|
||||
|
||||
|
||||
def _install_tp_encoder(monkeypatch, encoder):
|
||||
broadcasts = []
|
||||
monkeypatch.setattr(encode_server, "dp_dispatcher", None)
|
||||
monkeypatch.setattr(encode_server, "encoder", encoder)
|
||||
monkeypatch.setattr(encode_server, "send_sockets", [object()])
|
||||
monkeypatch.setattr(http_server, "dp_dispatcher", None)
|
||||
monkeypatch.setattr(http_server, "encoder", encoder)
|
||||
monkeypatch.setattr(http_server, "send_sockets", [object()])
|
||||
monkeypatch.setattr(
|
||||
encode_server,
|
||||
http_server,
|
||||
"sock_send",
|
||||
lambda socket, payload: broadcasts.append((socket, payload)),
|
||||
)
|
||||
@@ -41,7 +51,7 @@ def test_health_encode_waits_for_collective_dispatch_lock(monkeypatch):
|
||||
broadcasts = _install_tp_encoder(monkeypatch, encoder)
|
||||
await encoder.encode_dispatch_lock.acquire()
|
||||
|
||||
task = asyncio.create_task(encode_server.health_generate())
|
||||
task = asyncio.create_task(http_server.health_generate())
|
||||
await asyncio.sleep(0)
|
||||
assert broadcasts == []
|
||||
assert encoder.encode_calls == []
|
||||
@@ -61,7 +71,7 @@ def test_health_encode_rechecks_busy_state_after_waiting(monkeypatch):
|
||||
broadcasts = _install_tp_encoder(monkeypatch, encoder)
|
||||
await encoder.encode_dispatch_lock.acquire()
|
||||
|
||||
task = asyncio.create_task(encode_server.health_generate())
|
||||
task = asyncio.create_task(http_server.health_generate())
|
||||
await asyncio.sleep(0)
|
||||
encoder.embedding_to_send["real-request"] = object()
|
||||
encoder.encode_dispatch_lock.release()
|
||||
|
||||
@@ -3,7 +3,7 @@ import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.disaggregation.encode_server import (
|
||||
from sglang.srt.disaggregation.encoder.runtime import (
|
||||
EncoderScheduler,
|
||||
PendingRequest,
|
||||
_resolve_encoder_batch_policy,
|
||||
|
||||
@@ -15,14 +15,18 @@ import zmq.asyncio
|
||||
from fastapi import HTTPException
|
||||
from PIL import Image
|
||||
|
||||
from sglang.srt.disaggregation.encode_receiver import (
|
||||
from sglang.srt.disaggregation.encoder.preprocessor import (
|
||||
EncoderPreprocessor,
|
||||
EncoderPreprocessResult,
|
||||
)
|
||||
from sglang.srt.disaggregation.encoder.receiver import (
|
||||
EmbeddingData,
|
||||
MMReceiverHTTP,
|
||||
MultiModalEmbeddingData,
|
||||
_encoder_media_item,
|
||||
_select_mm_processor_prompt,
|
||||
)
|
||||
from sglang.srt.disaggregation.encode_server import MMEncoder, _get_mm_grid_dim
|
||||
from sglang.srt.disaggregation.encoder.server import MMEncoder
|
||||
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
|
||||
from sglang.srt.managers.tokenizer_manager import (
|
||||
_reject_missing_dispatched_encoder_embedding,
|
||||
@@ -220,16 +224,19 @@ def test_epd_allows_local_processing_when_request_was_not_dispatched():
|
||||
def _encoder(model_type="kimi_k3"):
|
||||
encoder = MMEncoder.__new__(MMEncoder)
|
||||
encoder.model_type = model_type
|
||||
encoder.model_config = SimpleNamespace(
|
||||
preprocessor = EncoderPreprocessor.__new__(EncoderPreprocessor)
|
||||
preprocessor.model_type = model_type
|
||||
preprocessor.model_config = SimpleNamespace(
|
||||
hf_config=SimpleNamespace(
|
||||
vision_config=SimpleNamespace(merge_kernel_size=(2, 2))
|
||||
)
|
||||
)
|
||||
encoder.encoder_media_processor_config = (
|
||||
preprocessor.encoder_media_processor_config = (
|
||||
KimiK3ForConditionalGeneration.encoder_media_processor_config
|
||||
if model_type == "kimi_k3"
|
||||
else EncoderMediaProcessorConfig()
|
||||
)
|
||||
encoder.preprocessor = preprocessor
|
||||
return encoder
|
||||
|
||||
|
||||
@@ -237,11 +244,11 @@ def test_kimi_k3_encoder_normalizes_pillow_images_to_media_dicts():
|
||||
image = Image.new("RGB", (2, 2))
|
||||
encoder = _encoder()
|
||||
|
||||
assert encoder._grid_count_per_leaf(
|
||||
assert encoder.preprocessor._grid_count_per_leaf(
|
||||
[image, {"type": "image", "image": [image, image]}], Modality.IMAGE
|
||||
) == [1, 2]
|
||||
|
||||
normalized = encoder._normalize_kimi_encoder_images(
|
||||
normalized = encoder.preprocessor._normalize_kimi_encoder_images(
|
||||
[image, {"type": "image", "image": [image, image]}]
|
||||
)
|
||||
assert len(normalized) == 3
|
||||
@@ -258,14 +265,15 @@ def test_kimi_k3_encoder_passes_media_dicts_to_image_processor():
|
||||
return {"pixel_values": torch.ones(1, 3), "grid_thws": [[1, 1, 1]]}
|
||||
|
||||
encoder = _encoder()
|
||||
encoder.image_processor = image_processor
|
||||
encoder.vision_config = {"image": {"return_tensors": "pt"}}
|
||||
encoder._flatten_and_load_images = AsyncMock(return_value=[image])
|
||||
encoder.preproc_executor = ThreadPoolExecutor(max_workers=1)
|
||||
preprocessor = encoder.preprocessor
|
||||
preprocessor.image_processor = image_processor
|
||||
preprocessor.vision_config = {"image": {"return_tensors": "pt"}}
|
||||
preprocessor._flatten_and_load_images = AsyncMock(return_value=[image])
|
||||
preprocessor.preproc_executor = ThreadPoolExecutor(max_workers=1)
|
||||
try:
|
||||
output = asyncio.run(encoder._process_image_items([image], None))
|
||||
output = asyncio.run(preprocessor._process_image_items([image], None))
|
||||
finally:
|
||||
encoder.preproc_executor.shutdown()
|
||||
preprocessor.preproc_executor.shutdown()
|
||||
|
||||
assert "pixel_values" in output
|
||||
assert output["original_image_sizes"] == [[3, 2]]
|
||||
@@ -354,27 +362,28 @@ def test_kimi_k3_epd_model_preprocessor_receives_image_processor():
|
||||
return prepare_kimi_k3_encoder_inputs(mm_data, image_processor)
|
||||
|
||||
encoder = _encoder()
|
||||
encoder.image_processor = image_processor
|
||||
encoder.use_image_processor_gpu = False
|
||||
encoder.vision_config = {"image": {"return_tensors": "pt"}}
|
||||
encoder._flatten_and_load_images = AsyncMock(return_value=[image])
|
||||
encoder.preproc_executor = ThreadPoolExecutor(max_workers=1)
|
||||
preprocessor = encoder.preprocessor
|
||||
preprocessor.image_processor = image_processor
|
||||
preprocessor.use_image_processor_gpu = False
|
||||
preprocessor.vision_config = {"image": {"return_tensors": "pt"}}
|
||||
preprocessor._flatten_and_load_images = AsyncMock(return_value=[image])
|
||||
preprocessor.preproc_executor = ThreadPoolExecutor(max_workers=1)
|
||||
try:
|
||||
with patch(
|
||||
"sglang.srt.disaggregation.encode_server.get_parallel",
|
||||
"sglang.srt.disaggregation.encoder.preprocessor.get_parallel",
|
||||
return_value=SimpleNamespace(attn_tp_rank=0, attn_tp_size=1),
|
||||
):
|
||||
output = asyncio.run(
|
||||
encoder._process_image_items([image], model_preprocessor)
|
||||
preprocessor._process_image_items([image], model_preprocessor)
|
||||
)
|
||||
finally:
|
||||
encoder.preproc_executor.shutdown()
|
||||
preprocessor.preproc_executor.shutdown()
|
||||
|
||||
assert len(calls) == 1
|
||||
assert calls[0][0][0] == {"type": "image", "image": image}
|
||||
assert calls[0][1:] == (
|
||||
Modality.IMAGE,
|
||||
encoder.vision_config,
|
||||
preprocessor.vision_config,
|
||||
image_processor,
|
||||
False,
|
||||
)
|
||||
@@ -481,13 +490,13 @@ def test_kimi_k3_epd_selects_matching_jpeg_decode_mode(
|
||||
):
|
||||
expected = torch.zeros((3, 2, 3), dtype=torch.uint8)
|
||||
encoder = _encoder()
|
||||
encoder.use_image_processor_gpu = use_image_processor_gpu
|
||||
encoder.preprocessor.use_image_processor_gpu = use_image_processor_gpu
|
||||
|
||||
with patch(
|
||||
"sglang.srt.disaggregation.encode_server.load_image",
|
||||
"sglang.srt.disaggregation.encoder.preprocessor.load_image",
|
||||
return_value=(expected, None),
|
||||
) as load:
|
||||
output = encoder._load_single_item(b"jpeg", Modality.IMAGE)
|
||||
output = encoder.preprocessor._load_single_item(b"jpeg", Modality.IMAGE)
|
||||
|
||||
assert output is expected
|
||||
load.assert_called_once_with(b"jpeg", expected_decode_mode)
|
||||
@@ -498,13 +507,13 @@ def test_kimi_k3_epd_verifies_content_hash_before_decode():
|
||||
digest = snapshot_media(payload).content_digest
|
||||
expected = torch.zeros((3, 2, 3), dtype=torch.uint8)
|
||||
encoder = _encoder()
|
||||
encoder.use_image_processor_gpu = False
|
||||
encoder.preprocessor.use_image_processor_gpu = False
|
||||
|
||||
with patch(
|
||||
"sglang.srt.disaggregation.encode_server.load_image",
|
||||
"sglang.srt.disaggregation.encoder.preprocessor.load_image",
|
||||
return_value=(expected, None),
|
||||
) as load:
|
||||
output = encoder._load_single_item(
|
||||
output = encoder.preprocessor._load_single_item(
|
||||
{"url": payload, "content_hash": digest}, Modality.IMAGE
|
||||
)
|
||||
|
||||
@@ -589,8 +598,9 @@ def test_kimi_k3_encoder_prefers_grid_thws_and_uses_temporal_pool_length():
|
||||
stale_grid = torch.tensor([[1, 2, 2]])
|
||||
mm_inputs = {"grid_thws": grid_thws, "image_grid_thw": stale_grid}
|
||||
|
||||
assert _get_mm_grid_dim(mm_inputs, Modality.IMAGE, "kimi_k3") is grid_thws
|
||||
assert _encoder().get_num_tokens(grid_thws[0], Modality.IMAGE) == 24
|
||||
preprocessor = _encoder().preprocessor
|
||||
assert preprocessor._get_mm_grid_dim(mm_inputs, Modality.IMAGE) is grid_thws
|
||||
assert preprocessor.get_num_tokens(grid_thws[0], Modality.IMAGE) == 24
|
||||
|
||||
|
||||
def test_kimi_k3_encoder_splits_cross_request_batch_into_single_grid_items():
|
||||
@@ -606,12 +616,14 @@ def test_kimi_k3_encoder_splits_cross_request_batch_into_single_grid_items():
|
||||
|
||||
output = encoder._encode_missing(
|
||||
feature,
|
||||
{"pixel_values": feature, "grid_thws": grid_thws},
|
||||
EncoderPreprocessResult(
|
||||
mm_inputs={"pixel_values": feature, "grid_thws": grid_thws},
|
||||
grid_thw=grid_thws,
|
||||
token_counts=[1, 2, 2],
|
||||
),
|
||||
indices=[2, 0, 1],
|
||||
modality=Modality.IMAGE,
|
||||
get_feature_fn=get_feature_fn,
|
||||
grid_thw=grid_thws,
|
||||
keep_on_gpu=True,
|
||||
)
|
||||
|
||||
items = captured["items"]
|
||||
@@ -643,7 +655,7 @@ def test_encoder_preprocessed_items_follow_dp_owner_selection_order():
|
||||
{"pixel_values": [item.feature for item in items], "grid_thws": grid_thws},
|
||||
mm_items=items,
|
||||
)
|
||||
embeddings = torch.arange(4, dtype=torch.float32).reshape(4, 1)
|
||||
embeddings = torch.arange(3, dtype=torch.float32).reshape(3, 1)
|
||||
captured = {}
|
||||
|
||||
def get_feature_fn(selected_items):
|
||||
@@ -652,17 +664,19 @@ def test_encoder_preprocessed_items_follow_dp_owner_selection_order():
|
||||
|
||||
output = encoder._encode_missing(
|
||||
mm_inputs["pixel_values"],
|
||||
mm_inputs,
|
||||
EncoderPreprocessResult(
|
||||
mm_inputs=mm_inputs,
|
||||
grid_thw=grid_thws,
|
||||
token_counts=[1, 2, 2],
|
||||
),
|
||||
indices=[2, 0],
|
||||
modality=Modality.IMAGE,
|
||||
get_feature_fn=get_feature_fn,
|
||||
grid_thw=grid_thws,
|
||||
keep_on_gpu=True,
|
||||
)
|
||||
|
||||
assert captured["items"] == [items[2], items[0]]
|
||||
assert [part.shape[0] for part in output] == [2, 1]
|
||||
torch.testing.assert_close(torch.cat(output), embeddings[:3])
|
||||
torch.testing.assert_close(torch.cat(output), embeddings)
|
||||
|
||||
|
||||
def test_encoder_preprocessed_items_hash_individually():
|
||||
@@ -767,6 +781,8 @@ def test_epd_encoder_reuses_scheduler_zmq_peer():
|
||||
)
|
||||
with config_override as server_args:
|
||||
encoder.server_args = server_args
|
||||
encoder.transfer_backend = "zmq_to_scheduler"
|
||||
encoder.use_mooncake = False
|
||||
encoder.send_timeout = 3
|
||||
encoder.context = context
|
||||
encoder.scheduler_send_sockets = {}
|
||||
@@ -841,6 +857,8 @@ def test_epd_encoder_pipelines_zero_copy_sends_per_peer():
|
||||
)
|
||||
with config_override as server_args:
|
||||
encoder.server_args = server_args
|
||||
encoder.transfer_backend = "zmq_to_scheduler"
|
||||
encoder.use_mooncake = False
|
||||
encoder.send_timeout = 1
|
||||
encoder.context = FakeContext(socket)
|
||||
encoder.scheduler_send_sockets = {}
|
||||
|
||||
@@ -108,7 +108,7 @@ _CONFIGURED_SIZE_CALL_SITES = {
|
||||
"the consumer count is configured fan-out arithmetic (tp_size // "
|
||||
"dp_size), which is what the record answered before"
|
||||
),
|
||||
("srt/disaggregation/encode_server.py", "configured_tp_size"): (
|
||||
("srt/disaggregation/encoder/runtime.py", "configured_tp_size"): (
|
||||
"the encode server's launch entry sizes its workers before it has "
|
||||
"spawned any of them"
|
||||
),
|
||||
|
||||
@@ -77,8 +77,8 @@ _KNOWN_ENTRIES = frozenset(
|
||||
"run_data_parallel_controller_process",
|
||||
),
|
||||
("srt/ray/scheduler_actor.py", "__init__"),
|
||||
("srt/disaggregation/encode_server.py", "__init__"),
|
||||
("srt/disaggregation/encode_server.py", "launch_server"),
|
||||
("srt/disaggregation/encoder/server.py", "__init__"),
|
||||
("srt/disaggregation/encoder/http_server.py", "launch_server"),
|
||||
("srt/managers/tokenizer_manager.py", "__init__"),
|
||||
("srt/entrypoints/engine.py", "_launch_subprocesses"),
|
||||
(
|
||||
|
||||
Reference in New Issue
Block a user