[EPD][VLM] support video/audio input (#17824)
Co-authored-by: siyu <liusy58@linux.alibaba.com>
This commit is contained in:
Zheng Wengang
@@ -1,10 +1,13 @@
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import grpc
|
||||
import openai
|
||||
import zmq
|
||||
from grpc_health.v1 import health_pb2, health_pb2_grpc
|
||||
|
||||
@@ -21,10 +24,595 @@ from sglang.test.test_utils import (
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
from sglang.test.vlm_utils import (
|
||||
AUDIO_TRUMP_SPEECH_URL,
|
||||
IMAGE_MAN_IRONING_URL,
|
||||
IMAGE_SGL_LOGO_URL,
|
||||
VIDEO_JOBS_URL,
|
||||
)
|
||||
|
||||
# Omni model for local testing; override via env var EPD_OMNI_MODEL
|
||||
DEFAULT_OMNI_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
|
||||
|
||||
|
||||
register_cuda_ci(est_time=150, suite="stage-c-test-4-gpu-h100")
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
is_in_ci(),
|
||||
"Omni model EPD test with image, video, and audio modalities, running locally only",
|
||||
)
|
||||
class TestEPDDisaggregationOmni(PDDisaggregationServerBase):
|
||||
"""
|
||||
EPD disaggregation test for omni models (e.g. Qwen3-Omni). Covers image, video,
|
||||
and audio when server_type=http (encoder_transfer_backend: mooncake/zmq_to_scheduler/zmq_to_tokenizer).
|
||||
When server_type=grpc, only image is tested (gRPC encode is image-only).
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = os.environ.get("EPD_OMNI_MODEL", DEFAULT_OMNI_MODEL)
|
||||
cls.server_type = os.environ.get("EPD_ENCODE_SERVER_TYPE", "http")
|
||||
assert cls.server_type in (
|
||||
"grpc",
|
||||
"http",
|
||||
), f"Invalid EPD_ENCODE_SERVER_TYPE: {cls.server_type}"
|
||||
cls.encoder_transfer_backend = os.environ.get(
|
||||
"EPD_ENCODER_TRANSFER_BACKEND", "zmq_to_scheduler"
|
||||
)
|
||||
assert cls.encoder_transfer_backend in (
|
||||
"mooncake",
|
||||
"zmq_to_scheduler",
|
||||
"zmq_to_tokenizer",
|
||||
), f"Invalid EPD_ENCODER_TRANSFER_BACKEND: {cls.encoder_transfer_backend}"
|
||||
cls.enable_global_cache = (
|
||||
os.environ.get("MOONCAKE_MASTER") is not None
|
||||
or os.environ.get("MOONCAKE_CLIENT") is not None
|
||||
)
|
||||
if cls.server_type == "grpc":
|
||||
cls.encode_port = f"{int(cls.lb_port) + 305}"
|
||||
cls.encode_url = f"grpc://{cls.base_host}:{cls.encode_port}"
|
||||
else:
|
||||
cls.encode_port = f"{int(cls.lb_port) + 300}"
|
||||
cls.encode_url = f"http://{cls.base_host}:{cls.encode_port}"
|
||||
|
||||
cls.image_man_ironing = IMAGE_MAN_IRONING_URL
|
||||
cls.image_sgl_logo = IMAGE_SGL_LOGO_URL
|
||||
cls.video_jobs = VIDEO_JOBS_URL
|
||||
cls.audio_trump = AUDIO_TRUMP_SPEECH_URL
|
||||
|
||||
print(
|
||||
f"Setting up EPD Omni: model={cls.model}, encode={cls.encode_port}, "
|
||||
f"prefill={cls.prefill_port}, decode={cls.decode_port}, "
|
||||
f"server_type={cls.server_type}, backend={cls.encoder_transfer_backend}, "
|
||||
f"global_cache={cls.enable_global_cache}"
|
||||
)
|
||||
print(f"Data URLs: image={cls.image_man_ironing}, audio={cls.audio_trump}")
|
||||
|
||||
cls.start_encode()
|
||||
prefill_thread = threading.Thread(target=cls.start_prefill)
|
||||
decode_thread = threading.Thread(target=cls.start_decode)
|
||||
prefill_thread.start()
|
||||
decode_thread.start()
|
||||
prefill_thread.join()
|
||||
decode_thread.join()
|
||||
|
||||
if cls.server_type == "grpc":
|
||||
cls._wait_grpc_ready(cls.base_host, cls.encode_port, cls.process_encode)
|
||||
else:
|
||||
cls.wait_server_ready(
|
||||
cls.encode_url + "/health", process=cls.process_encode
|
||||
)
|
||||
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
|
||||
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
cls.api_key = "sk-123456"
|
||||
os.environ["OPENAI_API_KEY"] = cls.api_key
|
||||
os.environ["OPENAI_API_BASE"] = f"{cls.lb_url}/v1"
|
||||
|
||||
@classmethod
|
||||
def start_encode(cls):
|
||||
if cls.server_type == "grpc":
|
||||
cls.encode_stdout = io.StringIO()
|
||||
cls.encode_stderr = io.StringIO()
|
||||
cls.process_encode = subprocess.Popen(
|
||||
[
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang.launch_server",
|
||||
"--model-path",
|
||||
cls.model,
|
||||
"--host",
|
||||
cls.base_host,
|
||||
"--port",
|
||||
cls.encode_port,
|
||||
"--trust-remote-code",
|
||||
"--encoder-only",
|
||||
"--grpc-mode",
|
||||
"--encoder-transfer-backend",
|
||||
"zmq_to_scheduler",
|
||||
"--tp",
|
||||
"1",
|
||||
]
|
||||
)
|
||||
else:
|
||||
encode_args = [
|
||||
"--trust-remote-code",
|
||||
"--encoder-only",
|
||||
"--encoder-transfer-backend",
|
||||
cls.encoder_transfer_backend,
|
||||
"--tp",
|
||||
"1",
|
||||
"--port",
|
||||
cls.encode_port,
|
||||
]
|
||||
if cls.enable_global_cache:
|
||||
encode_args.append("--enable-mm-global-cache")
|
||||
cls.encode_stdout = io.StringIO()
|
||||
cls.encode_stderr = io.StringIO()
|
||||
cls.process_encode = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=cls.encode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=encode_args,
|
||||
return_stdout_stderr=(cls.encode_stdout, cls.encode_stderr),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--language-only",
|
||||
"--encoder-urls",
|
||||
cls.encode_url,
|
||||
"--encoder-transfer-backend",
|
||||
(
|
||||
"zmq_to_scheduler"
|
||||
if cls.server_type == "grpc"
|
||||
else cls.encoder_transfer_backend
|
||||
),
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"1",
|
||||
"--port",
|
||||
cls.prefill_port,
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
prefill_env = os.environ.copy()
|
||||
if cls.server_type == "grpc":
|
||||
prefill_env["SGLANG_ENCODER_MM_RECEIVER_MODE"] = "grpc"
|
||||
cls.process_prefill = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
env=prefill_env,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"2",
|
||||
"--port",
|
||||
cls.decode_port,
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
for process in [
|
||||
cls.process_lb,
|
||||
cls.process_decode,
|
||||
cls.process_prefill,
|
||||
cls.process_encode,
|
||||
]:
|
||||
if process:
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"Error killing process: {e}")
|
||||
|
||||
@staticmethod
|
||||
def _wait_grpc_ready(
|
||||
host, port, process, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
|
||||
):
|
||||
deadline = time.time() + timeout
|
||||
channel = grpc.insecure_channel(f"{host}:{port}")
|
||||
stub = health_pb2_grpc.HealthStub(channel)
|
||||
try:
|
||||
while time.time() < deadline:
|
||||
if process.poll() is not None:
|
||||
raise RuntimeError(
|
||||
f"gRPC encoder exited with code {process.returncode}"
|
||||
)
|
||||
try:
|
||||
response = stub.Check(
|
||||
health_pb2.HealthCheckRequest(service=""), timeout=2
|
||||
)
|
||||
if response.status == health_pb2.HealthCheckResponse.SERVING:
|
||||
return
|
||||
except grpc.RpcError:
|
||||
pass
|
||||
time.sleep(1)
|
||||
finally:
|
||||
channel.close()
|
||||
raise RuntimeError(f"gRPC encoder not ready at {host}:{port} within {timeout}s")
|
||||
|
||||
# ---- helpers ----
|
||||
|
||||
def _client(self):
|
||||
return openai.Client(api_key=self.api_key, base_url=f"{self.lb_url}/v1")
|
||||
|
||||
def _skip_if_grpc(self, msg="gRPC encode is image-only"):
|
||||
"""Skip this test when encode server is gRPC (image-only)."""
|
||||
if self.server_type == "grpc":
|
||||
self.skipTest(msg)
|
||||
|
||||
def _parse_cache_log(self):
|
||||
"""Parse encode server logs and return list of (local_hits, global_hits, misses)
|
||||
tuples from '=== Multi-Level Cache Check ===' lines."""
|
||||
log = self.encode_stdout.getvalue() + self.encode_stderr.getvalue()
|
||||
pattern = re.compile(
|
||||
r"Multi-Level Cache Check.*?"
|
||||
r"Local Hits:\s*(\d+).*?"
|
||||
r"Global Hits:\s*(\d+).*?"
|
||||
r"Misses.*?:\s*(\d+)"
|
||||
)
|
||||
return [(int(m[1]), int(m[2]), int(m[3])) for m in pattern.finditer(log)]
|
||||
|
||||
# ---- image ----
|
||||
def test_image(self):
|
||||
client = self._client()
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": self.image_man_ironing},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Describe this image in a sentence.",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=256,
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
print(f"[Omni EPD] Image response:\n{text}")
|
||||
self.assertIsNotNone(text)
|
||||
self.assertGreater(len(text), 0)
|
||||
|
||||
text_lower = text.lower()
|
||||
self.assertTrue(
|
||||
any(w in text_lower for w in ("man", "person", "driver")),
|
||||
f"Image response should mention a person: {text}",
|
||||
)
|
||||
self.assertTrue(
|
||||
any(w in text_lower for w in ("iron", "cloth", "hang", "holding")),
|
||||
f"Image response should mention ironing/clothes: {text}",
|
||||
)
|
||||
|
||||
def test_image_cache_hit(self):
|
||||
"""Send the same image twice; the second request should hit the global-mm-cache."""
|
||||
self._skip_if_grpc("gRPC encode is image-only; cache test uses HTTP path")
|
||||
if not self.enable_global_cache:
|
||||
self.skipTest("global-mm-cache not enabled (MOONCAKE_MASTER not set)")
|
||||
client = self._client()
|
||||
baseline = len(self._parse_cache_log())
|
||||
for i in range(2):
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": self.image_sgl_logo},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What is shown in this image?",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=128,
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
print(f"[Omni EPD] Image cache-hit round {i}: {text}")
|
||||
self.assertIsNotNone(text)
|
||||
self.assertGreater(len(text), 0)
|
||||
time.sleep(1)
|
||||
|
||||
entries = self._parse_cache_log()[baseline:]
|
||||
print(f"[Omni EPD] Image cache log entries: {entries}")
|
||||
self.assertGreaterEqual(
|
||||
len(entries), 2, "Expected at least 2 cache-check log entries"
|
||||
)
|
||||
local_hits, global_hits, _ = entries[-1]
|
||||
self.assertGreater(
|
||||
local_hits + global_hits,
|
||||
0,
|
||||
f"Second image request should have cache hits, got: {entries[-1]}",
|
||||
)
|
||||
|
||||
# ---- video ----
|
||||
def test_video(self):
|
||||
self._skip_if_grpc()
|
||||
client = self._client()
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe the video."},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": self.video_jobs},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
max_tokens=8192,
|
||||
stream=False,
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
print(f"[Omni EPD] Video response:\n{text}")
|
||||
self.assertIsNotNone(text)
|
||||
self.assertGreater(len(text), 0)
|
||||
|
||||
text_lower = text.lower()
|
||||
self.assertTrue(
|
||||
any(
|
||||
w in text_lower
|
||||
for w in ("ipod", "device", "microphone", "smartphone", "phone")
|
||||
),
|
||||
f"Video response should mention a device: {text}",
|
||||
)
|
||||
self.assertTrue(
|
||||
any(
|
||||
w in text_lower
|
||||
for w in (
|
||||
"man",
|
||||
"person",
|
||||
"individual",
|
||||
"speaker",
|
||||
"presenter",
|
||||
"steve",
|
||||
"hand",
|
||||
"hands",
|
||||
)
|
||||
),
|
||||
f"Video response should mention a person: {text}",
|
||||
)
|
||||
self.assertTrue(
|
||||
any(
|
||||
w in text_lower
|
||||
for w in (
|
||||
"present",
|
||||
"presenting",
|
||||
"examine",
|
||||
"examining",
|
||||
"display",
|
||||
"displaying",
|
||||
"hold",
|
||||
"holding",
|
||||
"gestur",
|
||||
"speak",
|
||||
"speaking",
|
||||
)
|
||||
),
|
||||
f"Video response should mention an action: {text}",
|
||||
)
|
||||
|
||||
def test_video_cache_hit(self):
|
||||
"""Send the same video twice; the second request should hit the global-mm-cache."""
|
||||
self._skip_if_grpc()
|
||||
if not self.enable_global_cache:
|
||||
self.skipTest("global-mm-cache not enabled (MOONCAKE_MASTER not set)")
|
||||
client = self._client()
|
||||
baseline = len(self._parse_cache_log())
|
||||
for i in range(2):
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe the video."},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": self.video_jobs},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
max_tokens=256,
|
||||
stream=False,
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
print(f"[Omni EPD] Video cache-hit round {i}: {text}")
|
||||
self.assertIsNotNone(text)
|
||||
self.assertGreater(len(text), 0)
|
||||
time.sleep(1)
|
||||
|
||||
entries = self._parse_cache_log()[baseline:]
|
||||
print(f"[Omni EPD] Video cache log entries: {entries}")
|
||||
self.assertGreaterEqual(
|
||||
len(entries), 2, "Expected at least 2 cache-check log entries"
|
||||
)
|
||||
local_hits, global_hits, _ = entries[-1]
|
||||
self.assertGreater(
|
||||
local_hits + global_hits,
|
||||
0,
|
||||
f"Second video request should have cache hits, got: {entries[-1]}",
|
||||
)
|
||||
|
||||
# ---- audio ----
|
||||
|
||||
def test_audio(self):
|
||||
self._skip_if_grpc()
|
||||
client = self._client()
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": self.audio_trump},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Listen to this audio and write down the audio transcription in English.",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=256,
|
||||
stream=False,
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
print(f"[Omni EPD] Audio response:\n{text}")
|
||||
self.assertIsNotNone(text)
|
||||
self.assertGreater(len(text), 0)
|
||||
|
||||
text_lower = text.lower()
|
||||
for keyword in ("thank you", "leader"):
|
||||
self.assertIn(
|
||||
keyword,
|
||||
text_lower,
|
||||
f"Audio response should contain '{keyword}': {text}",
|
||||
)
|
||||
|
||||
def test_audio_cache_hit(self):
|
||||
"""Send the same audio twice; the second request should hit the global-mm-cache."""
|
||||
self._skip_if_grpc()
|
||||
if not self.enable_global_cache:
|
||||
self.skipTest("global-mm-cache not enabled (MOONCAKE_MASTER not set)")
|
||||
client = self._client()
|
||||
baseline = len(self._parse_cache_log())
|
||||
for i in range(2):
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": self.audio_trump},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What is this audio about?",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=128,
|
||||
stream=False,
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
print(f"[Omni EPD] Audio cache-hit round {i}: {text}")
|
||||
self.assertIsNotNone(text)
|
||||
self.assertGreater(len(text), 0)
|
||||
time.sleep(1)
|
||||
|
||||
entries = self._parse_cache_log()[baseline:]
|
||||
print(f"[Omni EPD] Audio cache log entries: {entries}")
|
||||
self.assertGreaterEqual(
|
||||
len(entries), 2, "Expected at least 2 cache-check log entries"
|
||||
)
|
||||
local_hits, global_hits, _ = entries[-1]
|
||||
self.assertGreater(
|
||||
local_hits + global_hits,
|
||||
0,
|
||||
f"Second audio request should have cache hits, got: {entries[-1]}",
|
||||
)
|
||||
|
||||
# ---- mixed modality ----
|
||||
|
||||
def test_mixed_image_audio_video(self):
|
||||
"""Image + audio + video in one request to test multi-modal routing."""
|
||||
self._skip_if_grpc()
|
||||
client = self._client()
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": self.image_man_ironing},
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": self.audio_trump},
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": self.video_jobs},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"I have an image, an audio clip, and a video, which are not related at all. "
|
||||
"Please: 1. Describe the image in a sentence, "
|
||||
"2. Summarize the audio content briefly, "
|
||||
"3. Describe what happens in the video."
|
||||
),
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=512,
|
||||
stream=False,
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
print(f"[Omni EPD] Mixed image+audio+video response:\n{text}")
|
||||
self.assertIsNotNone(text)
|
||||
self.assertGreater(len(text), 0)
|
||||
|
||||
text_lower = text.lower()
|
||||
self.assertTrue(
|
||||
any(w in text_lower for w in ("man", "person", "iron", "cloth")),
|
||||
f"Mixed response should describe the image: {text}",
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "Skipping in CI to reduce multi-GPU runtime")
|
||||
class TestEPDDisaggregationOneEncoder(PDDisaggregationServerBase):
|
||||
"""Test EPD disaggregation with single encode server"""
|
||||
|
||||
Reference in New Issue
Block a user