1361 lines
45 KiB
Python
1361 lines
45 KiB
Python
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
|
|
|
|
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_cuda_ci
|
|
from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
|
|
from sglang.test.server_fixtures.disaggregation_fixture import (
|
|
PDDisaggregationServerBase,
|
|
)
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
|
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
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"""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super().setUpClass()
|
|
cls.model = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
|
|
cls.encode_port = f"{int(cls.lb_port) + 300}"
|
|
cls.encode_url = f"http://{cls.base_host}:{cls.encode_port}"
|
|
|
|
print(
|
|
f"Setting up EPD (one encoder): encode={cls.encode_port}, "
|
|
f"prefill={cls.prefill_port}, decode={cls.decode_port}"
|
|
)
|
|
|
|
# Start servers in order: encode -> prefill/decode
|
|
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()
|
|
|
|
# Wait for all servers to be ready
|
|
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()
|
|
|
|
# Set OpenAI API key and base URL environment variables. Needed for lmms-eval to work.
|
|
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):
|
|
"""Start encode server for multimodal processing"""
|
|
encode_args = [
|
|
"--trust-remote-code",
|
|
"--encoder-only",
|
|
"--encoder-transfer-backend",
|
|
"zmq_to_scheduler",
|
|
"--tp",
|
|
"1",
|
|
"--port",
|
|
cls.encode_port,
|
|
"--enable-prefix-mm-cache",
|
|
]
|
|
cls.process_encode = popen_launch_server(
|
|
cls.model,
|
|
base_url=cls.encode_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=encode_args,
|
|
)
|
|
|
|
@classmethod
|
|
def start_prefill(cls):
|
|
"""Start prefill server with language model only"""
|
|
prefill_args = [
|
|
"--trust-remote-code",
|
|
"--language-only",
|
|
"--encoder-urls",
|
|
cls.encode_url,
|
|
"--encoder-transfer-backend",
|
|
"zmq_to_scheduler",
|
|
"--disaggregation-mode",
|
|
"prefill",
|
|
"--tp",
|
|
"1",
|
|
"--base-gpu-id",
|
|
"1",
|
|
"--port",
|
|
cls.prefill_port,
|
|
]
|
|
prefill_args += cls.transfer_backend + cls.rdma_devices
|
|
cls.process_prefill = popen_launch_server(
|
|
cls.model,
|
|
base_url=cls.prefill_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=prefill_args,
|
|
)
|
|
|
|
@classmethod
|
|
def start_decode(cls):
|
|
"""Start decode server"""
|
|
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):
|
|
"""Clean up all processes"""
|
|
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}")
|
|
|
|
def run_mmmu_eval(self, model_version: str, output_path: str, limit: str = "50"):
|
|
"""
|
|
Evaluate a VLM on the MMMU validation set with lmms-eval.
|
|
Reference: test_vlm_models.py
|
|
|
|
Args:
|
|
model_version: Model version/checkpoint to evaluate
|
|
output_path: Path to save evaluation results
|
|
limit: Number of samples to evaluate (default: "50" for CI time constraints)
|
|
"""
|
|
model = "openai_compatible"
|
|
tp = 1
|
|
tasks = "mmmu_val"
|
|
batch_size = 32
|
|
log_suffix = "openai_compatible"
|
|
os.makedirs(output_path, exist_ok=True)
|
|
|
|
model_args = f'model_version="{model_version}",tp={tp}'
|
|
|
|
cmd = [
|
|
"python3",
|
|
"-m",
|
|
"lmms_eval",
|
|
"--model",
|
|
model,
|
|
"--model_args",
|
|
model_args,
|
|
"--tasks",
|
|
tasks,
|
|
"--batch_size",
|
|
str(batch_size),
|
|
"--log_samples",
|
|
"--log_samples_suffix",
|
|
log_suffix,
|
|
"--output_path",
|
|
str(output_path),
|
|
"--limit",
|
|
limit,
|
|
]
|
|
|
|
_run_lmms_eval_with_retry(cmd, timeout=3600)
|
|
|
|
def test_mmmu(self):
|
|
"""Test MMMU evaluation with EPD disaggregation"""
|
|
import glob
|
|
import json
|
|
|
|
output_path = "./logs/epd_one_encoder_mmmu"
|
|
self.run_mmmu_eval(self.model, output_path)
|
|
|
|
# Get the result file
|
|
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
|
|
if not result_files:
|
|
result_files = glob.glob(f"{output_path}/*.json")
|
|
|
|
if not result_files:
|
|
self.fail(f"No JSON result files found in {output_path}")
|
|
|
|
result_file_path = result_files[0]
|
|
with open(result_file_path, "r") as f:
|
|
result = json.load(f)
|
|
print(f"MMMU result: {result}")
|
|
|
|
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
|
|
print(f"MMMU accuracy: {mmmu_accuracy:.4f}")
|
|
|
|
# for qwen2.5-vl-3b-instruct, the accuracy is 0.40
|
|
self.assertGreater(mmmu_accuracy, 0.40)
|
|
|
|
|
|
class TestEPDDisaggregationMultiEncoders(PDDisaggregationServerBase):
|
|
"""
|
|
Test EPD disaggregation with multiple encode servers for load balancing.
|
|
Both encode servers run on GPU 0 (different ports) for testing load distribution.
|
|
"""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super().setUpClass()
|
|
cls.model = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
|
|
cls.encode_port1 = f"{int(cls.lb_port) + 300}"
|
|
cls.encode_port2 = f"{int(cls.lb_port) + 301}"
|
|
cls.encode_url1 = f"http://{cls.base_host}:{cls.encode_port1}"
|
|
cls.encode_url2 = f"http://{cls.base_host}:{cls.encode_port2}"
|
|
|
|
print(
|
|
f"Setting up EPD (multiple encoders): encode1={cls.encode_port1}, "
|
|
f"encode2={cls.encode_port2}, prefill={cls.prefill_port}, decode={cls.decode_port}"
|
|
)
|
|
|
|
# Start two encode servers on GPU 0/1
|
|
encode1_thread = threading.Thread(
|
|
target=cls.start_encode_server, args=(cls.encode_port1, 0)
|
|
)
|
|
encode2_thread = threading.Thread(
|
|
target=cls.start_encode_server, args=(cls.encode_port2, 1)
|
|
)
|
|
encode1_thread.start()
|
|
encode2_thread.start()
|
|
encode1_thread.join()
|
|
encode2_thread.join()
|
|
|
|
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()
|
|
|
|
cls.wait_server_ready(cls.encode_url1 + "/health", process=cls.process_encode1)
|
|
cls.wait_server_ready(cls.encode_url2 + "/health", process=cls.process_encode2)
|
|
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()
|
|
|
|
# Set OpenAI API key and base URL environment variables. Needed for lmms-eval to work.
|
|
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_server(cls, port, gpu_id):
|
|
"""Start an encode server on specific port and GPU"""
|
|
encode_args = [
|
|
"--trust-remote-code",
|
|
"--encoder-only",
|
|
"--encoder-transfer-backend",
|
|
"zmq_to_scheduler",
|
|
"--tp",
|
|
"1",
|
|
"--port",
|
|
port,
|
|
"--enable-prefix-mm-cache",
|
|
]
|
|
# Only set base-gpu-id if not using GPU 0
|
|
if gpu_id != 0:
|
|
encode_args.extend(["--base-gpu-id", str(gpu_id)])
|
|
|
|
process = popen_launch_server(
|
|
cls.model,
|
|
base_url=f"http://{cls.base_host}:{port}",
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=encode_args,
|
|
)
|
|
if port == cls.encode_port1:
|
|
cls.process_encode1 = process
|
|
else:
|
|
cls.process_encode2 = process
|
|
|
|
@classmethod
|
|
def start_prefill(cls):
|
|
"""Start prefill server with multiple encode URLs"""
|
|
prefill_args = [
|
|
"--trust-remote-code",
|
|
"--language-only",
|
|
"--encoder-urls",
|
|
cls.encode_url1,
|
|
cls.encode_url2,
|
|
"--encoder-transfer-backend",
|
|
"zmq_to_scheduler",
|
|
"--disaggregation-mode",
|
|
"prefill",
|
|
"--tp",
|
|
"1",
|
|
"--base-gpu-id",
|
|
"2",
|
|
"--port",
|
|
cls.prefill_port,
|
|
]
|
|
prefill_args += cls.transfer_backend + cls.rdma_devices
|
|
cls.process_prefill = popen_launch_server(
|
|
cls.model,
|
|
base_url=cls.prefill_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=prefill_args,
|
|
)
|
|
|
|
@classmethod
|
|
def start_decode(cls):
|
|
"""Start decode server"""
|
|
decode_args = [
|
|
"--trust-remote-code",
|
|
"--disaggregation-mode",
|
|
"decode",
|
|
"--tp",
|
|
"1",
|
|
"--base-gpu-id",
|
|
"3",
|
|
"--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):
|
|
"""Clean up all processes"""
|
|
for process in [
|
|
cls.process_lb,
|
|
cls.process_decode,
|
|
cls.process_prefill,
|
|
cls.process_encode1,
|
|
cls.process_encode2,
|
|
]:
|
|
if process:
|
|
try:
|
|
kill_process_tree(process.pid)
|
|
except Exception as e:
|
|
print(f"Error killing process: {e}")
|
|
|
|
def run_mmmu_eval(self, model_version: str, output_path: str, limit: str = "50"):
|
|
"""
|
|
Evaluate a VLM on the MMMU validation set with lmms-eval.
|
|
Reference: test_vlm_models.py
|
|
|
|
Args:
|
|
model_version: Model version/checkpoint to evaluate
|
|
output_path: Path to save evaluation results
|
|
limit: Number of samples to evaluate (default: "50" for CI time constraints)
|
|
"""
|
|
model = "openai_compatible"
|
|
tp = 1
|
|
tasks = "mmmu_val"
|
|
batch_size = 32
|
|
log_suffix = "openai_compatible"
|
|
os.makedirs(output_path, exist_ok=True)
|
|
|
|
model_args = f'model_version="{model_version}",tp={tp}'
|
|
|
|
cmd = [
|
|
"python3",
|
|
"-m",
|
|
"lmms_eval",
|
|
"--model",
|
|
model,
|
|
"--model_args",
|
|
model_args,
|
|
"--tasks",
|
|
tasks,
|
|
"--batch_size",
|
|
str(batch_size),
|
|
"--log_samples",
|
|
"--log_samples_suffix",
|
|
log_suffix,
|
|
"--output_path",
|
|
str(output_path),
|
|
"--limit",
|
|
limit,
|
|
]
|
|
|
|
_run_lmms_eval_with_retry(cmd, timeout=3600)
|
|
|
|
def test_mmmu(self):
|
|
"""Test MMMU evaluation with EPD disaggregation (multiple encoders)"""
|
|
import glob
|
|
import json
|
|
|
|
output_path = "./logs/epd_multi_encoder_mmmu"
|
|
self.run_mmmu_eval(self.model, output_path)
|
|
|
|
# Get the result file
|
|
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
|
|
if not result_files:
|
|
result_files = glob.glob(f"{output_path}/*.json")
|
|
|
|
if not result_files:
|
|
self.fail(f"No JSON result files found in {output_path}")
|
|
|
|
result_file_path = result_files[0]
|
|
with open(result_file_path, "r") as f:
|
|
result = json.load(f)
|
|
print(f"MMMU result (multi encoder): {result}")
|
|
|
|
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
|
|
print(f"MMMU accuracy (multi encoder): {mmmu_accuracy:.4f}")
|
|
# for qwen2.5-vl-3b-instruct, the accuracy is 0.40
|
|
self.assertGreater(mmmu_accuracy, 0.40)
|
|
|
|
|
|
@unittest.skipIf(is_in_ci(), "Skipping in CI to reduce multi-GPU runtime")
|
|
class TestEPDDisaggregationGrpcEncoderMMMU(PDDisaggregationServerBase):
|
|
"""Test MMMU evaluation with gRPC encoder in EPD mode."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super().setUpClass()
|
|
cls.model = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
|
|
cls.encode_port = f"{int(cls.lb_port) + 304}"
|
|
cls.encode_url = f"grpc://{cls.base_host}:{cls.encode_port}"
|
|
|
|
print(
|
|
f"Setting up gRPC EPD (one encoder): encode={cls.encode_port}, "
|
|
f"prefill={cls.prefill_port}, decode={cls.decode_port}"
|
|
)
|
|
|
|
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()
|
|
|
|
cls.wait_grpc_ready(cls.base_host, cls.encode_port, cls.process_encode)
|
|
cls.wait_server_ready(cls.prefill_url + "/health")
|
|
cls.wait_server_ready(cls.decode_url + "/health")
|
|
|
|
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):
|
|
encode_command = [
|
|
"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",
|
|
"--base-gpu-id",
|
|
"0",
|
|
"--enable-prefix-mm-cache",
|
|
]
|
|
cls.process_encode = subprocess.Popen(encode_command)
|
|
|
|
@classmethod
|
|
def start_prefill(cls):
|
|
prefill_args = [
|
|
"--trust-remote-code",
|
|
"--language-only",
|
|
"--encoder-urls",
|
|
cls.encode_url,
|
|
"--encoder-transfer-backend",
|
|
"zmq_to_scheduler",
|
|
"--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()
|
|
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,
|
|
)
|
|
|
|
@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 server 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 server not ready at {host}:{port} within {timeout}s"
|
|
)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
os.environ.pop("SGLANG_ENCODER_MM_RECEIVER_MODE", None)
|
|
os.environ.pop("OPENAI_API_KEY", None)
|
|
os.environ.pop("OPENAI_API_BASE", None)
|
|
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}")
|
|
|
|
def run_mmmu_eval(self, model_version: str, output_path: str, limit: str = "50"):
|
|
model = "openai_compatible"
|
|
tp = 1
|
|
tasks = "mmmu_val"
|
|
batch_size = 32
|
|
log_suffix = "openai_compatible"
|
|
os.makedirs(output_path, exist_ok=True)
|
|
|
|
model_args = f'model_version="{model_version}",tp={tp}'
|
|
|
|
cmd = [
|
|
"python3",
|
|
"-m",
|
|
"lmms_eval",
|
|
"--model",
|
|
model,
|
|
"--model_args",
|
|
model_args,
|
|
"--tasks",
|
|
tasks,
|
|
"--batch_size",
|
|
str(batch_size),
|
|
"--log_samples",
|
|
"--log_samples_suffix",
|
|
log_suffix,
|
|
"--output_path",
|
|
str(output_path),
|
|
"--limit",
|
|
limit,
|
|
]
|
|
|
|
_run_lmms_eval_with_retry(cmd, timeout=3600)
|
|
|
|
def test_mmmu(self):
|
|
import glob
|
|
import json
|
|
|
|
output_path = "./logs/epd_grpc_encoder_mmmu"
|
|
self.run_mmmu_eval(self.model, output_path)
|
|
|
|
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
|
|
if not result_files:
|
|
result_files = glob.glob(f"{output_path}/*.json")
|
|
|
|
if not result_files:
|
|
self.fail(f"No JSON result files found in {output_path}")
|
|
|
|
result_file_path = result_files[0]
|
|
with open(result_file_path, "r") as f:
|
|
result = json.load(f)
|
|
print(f"MMMU result (grpc encoder): {result}")
|
|
|
|
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
|
|
print(f"MMMU accuracy (grpc encoder): {mmmu_accuracy:.4f}")
|
|
# for qwen2.5-vl-3b-instruct, the accuracy is 0.40
|
|
self.assertGreater(mmmu_accuracy, 0.40)
|
|
|
|
|
|
@unittest.skipIf(is_in_ci(), "Skipping in CI to reduce multi-GPU runtime")
|
|
class TestEPDDisaggregationGrpcEncoderOnly(PDDisaggregationServerBase):
|
|
"""Test gRPC encoder server integration with zmq_to_scheduler transfers."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super().setUpClass()
|
|
os.environ["SGLANG_ENCODER_MM_RECEIVER_MODE"] = "grpc"
|
|
cls.model = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
|
|
cls.encode_port = f"{int(cls.lb_port) + 302}"
|
|
|
|
print(f"Setting up gRPC EPD encoder: encode={cls.encode_port}")
|
|
|
|
cls.start_encode()
|
|
cls.wait_grpc_ready(cls.base_host, cls.encode_port, cls.process_encode)
|
|
|
|
@classmethod
|
|
def start_encode(cls):
|
|
encode_command = [
|
|
"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",
|
|
"--base-gpu-id",
|
|
"0",
|
|
"--enable-prefix-mm-cache",
|
|
]
|
|
cls.process_encode = subprocess.Popen(encode_command)
|
|
|
|
@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 server 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 server not ready at {host}:{port} within {timeout}s"
|
|
)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
os.environ.pop("SGLANG_ENCODER_MM_RECEIVER_MODE", None)
|
|
if cls.process_encode:
|
|
try:
|
|
kill_process_tree(cls.process_encode.pid)
|
|
except Exception as e:
|
|
print(f"Error killing process: {e}")
|
|
super().tearDownClass()
|
|
|
|
def test_grpc_encoder_zmq_to_scheduler(self):
|
|
from smg_grpc_proto import sglang_encoder_pb2, sglang_encoder_pb2_grpc
|
|
|
|
context = zmq.Context()
|
|
recv_port, recv_socket = get_zmq_socket_on_host(
|
|
context, zmq.PULL, host=self.base_host
|
|
)
|
|
channel = grpc.insecure_channel(f"{self.base_host}:{self.encode_port}")
|
|
stub = sglang_encoder_pb2_grpc.SglangEncoderStub(channel)
|
|
req_id = f"grpc-epd-{int(time.time() * 1000)}"
|
|
image_path = os.path.abspath("examples/assets/example_image.png")
|
|
|
|
try:
|
|
stub.SchedulerReceiveUrl(
|
|
sglang_encoder_pb2.SchedulerReceiveUrlRequest(
|
|
req_id=req_id,
|
|
receive_url=f"{self.base_host}:{recv_port}",
|
|
receive_count=1,
|
|
),
|
|
timeout=60,
|
|
)
|
|
stub.Encode(
|
|
sglang_encoder_pb2.EncodeRequest(
|
|
mm_items=[image_path],
|
|
req_id=req_id,
|
|
num_parts=1,
|
|
part_idx=0,
|
|
),
|
|
timeout=300,
|
|
)
|
|
|
|
poller = zmq.Poller()
|
|
poller.register(recv_socket, zmq.POLLIN)
|
|
socks = dict(poller.poll(60000))
|
|
self.assertIn(
|
|
recv_socket,
|
|
socks,
|
|
"No embedding payload received from gRPC encoder server",
|
|
)
|
|
parts = recv_socket.recv_multipart()
|
|
self.assertTrue(parts, "Empty embedding payload from gRPC encoder server")
|
|
finally:
|
|
recv_socket.close()
|
|
context.term()
|
|
channel.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|