254 lines
9.7 KiB
Python
254 lines
9.7 KiB
Python
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
|
|
|
register_cuda_ci(est_time=150, suite="stage-b-test-large-1-gpu")
|
|
register_amd_ci(est_time=270, suite="stage-b-test-small-1-gpu-amd")
|
|
"""
|
|
Usage:
|
|
python3 -m unittest test_vision_chunked_prefill.TestVisionChunkedPrefill.test_chunked_prefill
|
|
"""
|
|
|
|
import io
|
|
import logging
|
|
import os
|
|
import time
|
|
import unittest
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Union
|
|
|
|
import numpy as np
|
|
import pybase64
|
|
import requests
|
|
from PIL import Image
|
|
|
|
from sglang.srt.utils import kill_process_tree
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
DEFAULT_URL_FOR_TEST,
|
|
CustomTestCase,
|
|
calculate_rouge_l,
|
|
popen_launch_server,
|
|
)
|
|
|
|
# Configure logging to help diagnose CI timeouts
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s - %(levelname)s - %(message)s",
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TestVisionChunkedPrefill(CustomTestCase):
|
|
|
|
def prepare_video_messages(self, video_path, max_frames_num=8):
|
|
# We import decord here to avoid a strange Segmentation fault (core dumped) issue.
|
|
# The following import order will cause Segmentation fault.
|
|
# import decord
|
|
# from transformers import AutoTokenizer
|
|
from decord import VideoReader, cpu
|
|
|
|
vr = VideoReader(video_path, ctx=cpu(0))
|
|
total_frame_num = len(vr)
|
|
uniform_sampled_frames = np.linspace(
|
|
0, total_frame_num - 1, max_frames_num, dtype=int
|
|
)
|
|
frame_idx = uniform_sampled_frames.tolist()
|
|
frames = vr.get_batch(frame_idx).asnumpy()
|
|
|
|
base64_frames = []
|
|
for frame in frames:
|
|
pil_img = Image.fromarray(frame)
|
|
buff = io.BytesIO()
|
|
pil_img.save(buff, format="JPEG")
|
|
base64_str = pybase64.b64encode(buff.getvalue()).decode("utf-8")
|
|
base64_frames.append(base64_str)
|
|
|
|
messages = [{"role": "user", "content": []}]
|
|
frame_format = {
|
|
"type": "image_url",
|
|
"image_url": {"url": "data:image/jpeg;base64,{}"},
|
|
"modalities": "video",
|
|
}
|
|
|
|
for base64_frame in base64_frames:
|
|
frame_format["image_url"]["url"] = "data:image/jpeg;base64,{}".format(
|
|
base64_frame
|
|
)
|
|
messages[0]["content"].append(frame_format.copy())
|
|
|
|
prompt = {"type": "text", "text": "Please describe the video briefly."}
|
|
messages[0]["content"].append(prompt)
|
|
|
|
return messages
|
|
|
|
def get_prompt_from_messages(self, messages):
|
|
text = (
|
|
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
|
"<|im_start|>user\n"
|
|
)
|
|
image_data = []
|
|
for content in messages[0]["content"]:
|
|
if content["type"] == "image_url":
|
|
text += "<image>\n"
|
|
image_data.append(content["image_url"]["url"])
|
|
text += "Please describe the video briefly.<|im_end|>\n<|im_start|>assistant\n"
|
|
return text, image_data
|
|
|
|
def generate(self, text, image_data):
|
|
num_images = len(image_data) if image_data else 0
|
|
logger.info(f"Starting generate request with {num_images} images")
|
|
start_time = time.time()
|
|
response = requests.post(
|
|
self.base_url + "/generate",
|
|
json={
|
|
"text": text,
|
|
"image_data": image_data,
|
|
"sampling_params": {
|
|
"temperature": 0,
|
|
"max_new_tokens": 32,
|
|
"no_stop_trim": True,
|
|
"skip_special_tokens": False,
|
|
},
|
|
"modalities": ["multi-images"],
|
|
},
|
|
timeout=120, # Add timeout to prevent hanging indefinitely
|
|
).json()
|
|
elapsed = time.time() - start_time
|
|
logger.info(f"Generate request completed in {elapsed:.2f}s")
|
|
return response["text"]
|
|
|
|
def generate_for_video(self, batch, num_frame) -> Union[str, list[str]]:
|
|
logger.info(
|
|
f"generate_for_video called with batch={batch}, num_frame={num_frame}"
|
|
)
|
|
|
|
# prepare the video input about Steven introducing ipod nano
|
|
url = "https://raw.githubusercontent.com/evolvinglmms-lab/sglang/dev/onevision_local/assets/jobs.mp4"
|
|
cache_dir = os.path.expanduser("~/.cache")
|
|
file_path = os.path.join(cache_dir, "jobs.mp4")
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|
if not os.path.exists(file_path):
|
|
logger.info(f"Downloading video from {url}")
|
|
start_time = time.time()
|
|
response = requests.get(url, timeout=60)
|
|
response.raise_for_status()
|
|
with open(file_path, "wb") as f:
|
|
f.write(response.content)
|
|
elapsed = time.time() - start_time
|
|
logger.info(
|
|
f"Video downloaded in {elapsed:.2f}s, size={len(response.content)} bytes"
|
|
)
|
|
else:
|
|
logger.info(f"Using cached video at {file_path}")
|
|
|
|
if not batch:
|
|
assert isinstance(num_frame, int)
|
|
logger.info(f"Processing single video with {num_frame} frames")
|
|
messages = self.prepare_video_messages(file_path, max_frames_num=num_frame)
|
|
text, image_data = self.get_prompt_from_messages(messages)
|
|
return self.generate(text, image_data)
|
|
else:
|
|
assert isinstance(num_frame, list)
|
|
logger.info(f"Processing batch of videos with frame counts: {num_frame}")
|
|
func_args = []
|
|
for max_frames_num in num_frame:
|
|
messages = self.prepare_video_messages(
|
|
file_path,
|
|
max_frames_num=max_frames_num,
|
|
)
|
|
text, image_data = self.get_prompt_from_messages(messages)
|
|
func_args.append((text, image_data))
|
|
|
|
logger.info(f"Starting batch generation with {len(func_args)} requests")
|
|
with ThreadPoolExecutor(max_workers=10) as executor:
|
|
responses = list(executor.map(lambda p: self.generate(*p), func_args))
|
|
logger.info(f"Batch generation completed")
|
|
|
|
return responses
|
|
|
|
def launch_server(self, chunked_prefill_size) -> int:
|
|
# launch server
|
|
model = "lmms-lab/llava-onevision-qwen2-7b-ov"
|
|
# model = "meta-llama/Llama-3.2-11B-Vision-Instruct"
|
|
self.base_url = DEFAULT_URL_FOR_TEST
|
|
process = popen_launch_server(
|
|
model,
|
|
self.base_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=[
|
|
"--chunked-prefill-size",
|
|
f"{chunked_prefill_size}",
|
|
],
|
|
)
|
|
return process.pid
|
|
|
|
def _test_chunked_prefill(self, batches, num_frames):
|
|
logger.info("=" * 60)
|
|
logger.info("Starting chunked prefill test")
|
|
logger.info("=" * 60)
|
|
|
|
# Chunked
|
|
logger.info("Phase 1: Testing with chunked_prefill_size=1024")
|
|
chunked_server_pid = self.launch_server(chunked_prefill_size=1024)
|
|
logger.info(f"Chunked server started with pid={chunked_server_pid}")
|
|
try:
|
|
outputs_chunked = []
|
|
for i, (batch, num_frame) in enumerate(zip(batches, num_frames)):
|
|
logger.info(f"Chunked test iteration {i+1}/{len(batches)}")
|
|
output_chunked = self.generate_for_video(
|
|
batch=batch, num_frame=num_frame
|
|
)
|
|
outputs_chunked += [output_chunked]
|
|
logger.info(f"Chunked test iteration {i+1} completed")
|
|
finally:
|
|
logger.info(f"Killing chunked server pid={chunked_server_pid}")
|
|
kill_process_tree(chunked_server_pid)
|
|
logger.info("Chunked server killed")
|
|
time.sleep(4)
|
|
|
|
# None-chunked
|
|
logger.info("Phase 2: Testing with chunked_prefill_size=-1 (no chunking)")
|
|
try:
|
|
no_chunked_server_pid = self.launch_server(chunked_prefill_size=-1)
|
|
logger.info(f"Non-chunked server started with pid={no_chunked_server_pid}")
|
|
outputs_no_chunked = []
|
|
for i, (batch, num_frame) in enumerate(zip(batches, num_frames)):
|
|
logger.info(f"Non-chunked test iteration {i+1}/{len(batches)}")
|
|
output_no_chunked = self.generate_for_video(
|
|
batch=batch, num_frame=num_frame
|
|
)
|
|
outputs_no_chunked += [output_no_chunked]
|
|
logger.info(f"Non-chunked test iteration {i+1} completed")
|
|
|
|
finally:
|
|
logger.info(f"Killing non-chunked server pid={no_chunked_server_pid}")
|
|
kill_process_tree(no_chunked_server_pid)
|
|
logger.info("Non-chunked server killed")
|
|
time.sleep(4)
|
|
|
|
for output_chunked, output_no_chunked in zip(
|
|
outputs_chunked, outputs_no_chunked
|
|
):
|
|
print("output with chunked prefill:")
|
|
print(output_chunked)
|
|
print("output without chunked prefill:")
|
|
print(output_no_chunked)
|
|
self.assertEqual(len(output_chunked), len(output_no_chunked))
|
|
rouge_scores = calculate_rouge_l(output_chunked, output_no_chunked)
|
|
avg_score = sum(rouge_scores) / len(rouge_scores)
|
|
print(f"ROUGE-L scores: {rouge_scores}")
|
|
print(f"Average ROUGE-L score: {avg_score:.4f}")
|
|
# Allow for occasional divergence in one item while maintaining overall output quality
|
|
self.assertGreater(
|
|
avg_score,
|
|
0.90,
|
|
f"Average ROUGE-L score too low: {avg_score:.4f}. "
|
|
f"Individual scores: {rouge_scores}",
|
|
)
|
|
|
|
def test_chunked_prefill(self):
|
|
self._test_chunked_prefill(batches=[False, True], num_frames=[1, [2, 6, 8, 10]])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|