chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Apply fixture-kit mode to MMMUVLMMixin (#15615)

Browse Source
This commit is contained in:
lif
2025-12-28 17:22:05 +08:00
committed by GitHub
parent 8fab48952e
commit 5969be2f06
6 changed files with 236 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
python/sglang/srt/environ.py
View File

@@ -6,6 +6,29 @@ from enum import IntEnum
from typing import Any
@contextmanager
def temp_set_env(**env_vars: dict[str, Any]):
"""Temporarily set non-sglang environment variables, e.g. OPENAI_API_KEY"""
for key in env_vars:
if key.startswith("SGLANG_") or key.startswith("SGL_"):
raise ValueError("temp_set_env should not be used for sglang env vars")
backup = {key: os.environ.get(key) for key in env_vars}
try:
for key, value in env_vars.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = str(value)
yield
finally:
for key, value in backup.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
class EnvField:
_allow_set_name = True

137
python/sglang/test/kits/mmmu_vlm_kit.py
View File

@@ -2,8 +2,10 @@ import glob
import json
import os
import subprocess
import tempfile
from types import SimpleNamespace
from sglang.srt.environ import temp_set_env
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -16,8 +18,116 @@ from sglang.test.test_utils import (
DEFAULT_MEM_FRACTION_STATIC = 0.8
class MMMUVLMTestBase(CustomTestCase):
# TODO: split the MMMUVLMTestBase into a fixture and a mixin
class MMMUMixin:
"""Mixin for MMMU evaluation.
Use with MMMUServerBase for single-model tests:
class TestMyModel(MMMUMixin, MMMUServerBase):
model = "my/model"
accuracy = 0.4
"""
accuracy: float
mmmu_args: list[str] = []
# For OpenAI API settings
api_key = "sk-123456"
def run_mmmu_eval(
self: CustomTestCase,
model_version: str,
output_path: str,
):
"""
Evaluate a VLM on the MMMU validation set with lmms-eval.
Only `model_version` (checkpoint) and `chat_template` vary;
We are focusing only on the validation set due to resource constraints.
"""
# -------- fixed settings --------
model = "openai_compatible"
tp = 1
tasks = "mmmu_val"
batch_size = 64
log_suffix = "openai_compatible"
os.makedirs(output_path, exist_ok=True)
# -------- compose --model_args --------
model_args = f'model_version="{model_version}",' f"tp={tp}"
# -------- build command list --------
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),
*self.mmmu_args,
]
# Set OpenAI API key and base URL environment variables.
# Needed for lmms-eval to work.
with temp_set_env(
OPENAI_API_KEY=self.api_key,
OPENAI_API_BASE=f"{self.base_url}/v1",
):
subprocess.run(
cmd,
check=True,
timeout=3600,
)
def test_mmmu(self: CustomTestCase):
"""Run MMMU evaluation test."""
with tempfile.TemporaryDirectory() as output_path:
# Run evaluation
self.run_mmmu_eval(self.model, output_path)
# Get the result file
# Search recursively for JSON result files (lmms-eval v0.4.1+ creates subdirectories)
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:
raise FileNotFoundError(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"Result: {result}")
# Process the result
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
print(f"Model {self.model} achieved accuracy: {mmmu_accuracy:.4f}")
# Assert performance meets expected threshold
self.assertGreaterEqual(
mmmu_accuracy,
self.accuracy,
f"Model {self.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({self.accuracy:.4f})",
)
class MMMUMultiModelTestBase(CustomTestCase):
"""Base class for multi-model MMMU tests.
This class is for tests that need to evaluate multiple models,
starting and stopping a server for each model within the test method.
For single-model tests, use MMMUMixin with MMMUServerBase instead.
"""
parsed_args = None # Class variable to store args
other_args = []
mmmu_args = []
@@ -34,10 +144,27 @@ class MMMUVLMTestBase(CustomTestCase):
mem_fraction_static=DEFAULT_MEM_FRACTION_STATIC
)
# Save original environment variables for restoration in tearDownClass
cls._original_openai_api_key = os.environ.get("OPENAI_API_KEY")
cls._original_openai_api_base = os.environ.get("OPENAI_API_BASE")
# Set OpenAI API key and base URL environment variables. Needed for lmm-evals to work.
os.environ["OPENAI_API_KEY"] = cls.api_key
os.environ["OPENAI_API_BASE"] = f"{cls.base_url}/v1"
@classmethod
def tearDownClass(cls):
# Restore original environment variables
if cls._original_openai_api_key is not None:
os.environ["OPENAI_API_KEY"] = cls._original_openai_api_key
elif "OPENAI_API_KEY" in os.environ:
del os.environ["OPENAI_API_KEY"]
if cls._original_openai_api_base is not None:
os.environ["OPENAI_API_BASE"] = cls._original_openai_api_base
elif "OPENAI_API_BASE" in os.environ:
del os.environ["OPENAI_API_BASE"]
def run_mmmu_eval(
self,
model_version: str,
@@ -46,7 +173,7 @@ class MMMUVLMTestBase(CustomTestCase):
env: dict | None = None,
):
"""
Evaluate a VLM on the MMMU validation set with lmms‑eval.
Evaluate a VLM on the MMMU validation set with lmms-eval.
Only `model_version` (checkpoint) and `chat_template` vary;
We are focusing only on the validation set due to resource constraints.
"""
@@ -232,3 +359,7 @@ class MMMUVLMTestBase(CustomTestCase):
print(f"Error reading {tag.lower()} file: {e}")
return "\n".join(output_lines)
# Backward compatibility alias
MMMUVLMTestBase = MMMUMultiModelTestBase

68
python/sglang/test/server_fixtures/mmmu_fixture.py Normal file
View File

@@ -0,0 +1,68 @@
import logging
import os
import time
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
logger = logging.getLogger(__name__)
# Set default mem_fraction_static to 0.8
DEFAULT_MEM_FRACTION_STATIC = 0.8
class MMMUServerBase(CustomTestCase):
"""Server fixture for MMMU VLM tests.
This fixture handles server lifecycle for single-model MMMU tests.
For multi-model tests that need to start/stop servers within test methods,
use MMMUMultiModelTestBase instead.
"""
model = None
base_url = DEFAULT_URL_FOR_TEST
timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
other_args: list[str] = []
mem_fraction_static: float = DEFAULT_MEM_FRACTION_STATIC
@classmethod
def setUpClass(cls):
assert cls.model is not None, "Please set cls.model in subclass"
# Prepare environment variables
process_env = os.environ.copy()
process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1"
# Build server args with MMMU-specific settings
server_args = [
"--trust-remote-code",
"--cuda-graph-max-bs",
"64",
"--enable-multimodal",
"--mem-fraction-static",
str(cls.mem_fraction_static),
*cls.other_args,
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=cls.timeout,
api_key=cls.api_key,
other_args=server_args,
env=process_env,
)
@classmethod
def tearDownClass(cls):
if cls.process is not None and cls.process.poll() is None:
try:
kill_process_tree(cls.process.pid)
except Exception as e:
logger.error(f"Error killing process: {e}")
time.sleep(2)