sgl-project--sglang/python/sglang/test/kits/mmmu_vlm_kit.py

import glob
import json
import os
import shutil
import subprocess
import tempfile
from pathlib import Path
from types import SimpleNamespace

from sglang.srt.utils import kill_process_tree
from sglang.srt.utils.common import temp_set_env
from sglang.test.test_utils import (
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
    CustomTestCase,
    dump_metric,
    popen_launch_server,
)

# Set default mem_fraction_static to 0.8
DEFAULT_MEM_FRACTION_STATIC = 0.8


def _is_mmmu_parquet_corruption(error_output: str) -> bool:
    """Check if error is due to MMMU parquet file corruption."""
    return (
        "ArrowInvalid" in error_output
        and "Parquet magic bytes not found" in error_output
        and ("MMMU" in error_output or "lmms-lab--MMMU" in error_output)
    )


def _cleanup_mmmu_dataset_cache():
    """Clean up corrupted MMMU dataset cache to allow fresh download."""
    # Priority 1: Check CI convention path /hf_home first (used in Docker containers)
    ci_hf_home = Path("/hf_home/hub/datasets--lmms-lab--MMMU")
    if ci_hf_home.exists():
        mmmu_cache_path = ci_hf_home
    else:
        # Priority 2: Use HF_HOME env var or default user cache
        hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
        mmmu_cache_path = Path(hf_home) / "hub" / "datasets--lmms-lab--MMMU"

    if mmmu_cache_path.exists():
        print(f"Detected corrupted MMMU parquet cache. Cleaning up: {mmmu_cache_path}")
        try:
            shutil.rmtree(mmmu_cache_path)
            print(f"Successfully removed corrupted cache: {mmmu_cache_path}")
            return True
        except OSError as e:
            print(f"Warning: Failed to remove cache {mmmu_cache_path}: {e}")
            return False
    else:
        print(f"MMMU cache not found at {mmmu_cache_path}, skipping cleanup")
        return False


def _run_lmms_eval_with_retry(cmd: list[str], timeout: int = 3600) -> None:
    """Run lmms_eval command with automatic retry on MMMU parquet corruption."""
    try:
        result = subprocess.run(
            cmd,
            check=True,
            timeout=timeout,
            capture_output=True,
            text=True,
        )
        # Check for errors in output even if exit code is 0
        # lmms_eval sometimes returns 0 even when errors occur
        combined_output = result.stdout + result.stderr
        if _is_mmmu_parquet_corruption(combined_output):
            print(
                "Detected MMMU parquet corruption error in output. Attempting recovery..."
            )
            if _cleanup_mmmu_dataset_cache():
                print("Retrying lmms_eval with fresh download...")
                with temp_set_env(
                    HF_HUB_OFFLINE="0",
                    HF_DATASETS_DOWNLOAD_MODE="force_redownload",
                ):
                    retry_result = subprocess.run(
                        cmd, check=True, timeout=timeout, capture_output=True, text=True
                    )
                    # Print retry output
                    if retry_result.stdout:
                        print(retry_result.stdout, end="")
                    if retry_result.stderr:
                        print(retry_result.stderr, end="")
            else:
                print(
                    f"Failed to cleanup corrupted MMMU cache. Output from lmms_eval:\nStdout:\n{result.stdout}\nStderr:\n{result.stderr}"
                )
                raise RuntimeError("Failed to cleanup corrupted MMMU cache")
        else:
            # Print captured output to maintain visibility of successful runs
            if result.stdout:
                print(result.stdout, end="")
            if result.stderr:
                print(result.stderr, end="")
    except subprocess.CalledProcessError as e:
        error_output = e.stderr + e.stdout
        if _is_mmmu_parquet_corruption(error_output):
            print("Detected MMMU parquet corruption error. Attempting recovery...")
            if _cleanup_mmmu_dataset_cache():
                print("Retrying lmms_eval with fresh download...")
                with temp_set_env(
                    HF_HUB_OFFLINE="0",
                    HF_DATASETS_DOWNLOAD_MODE="force_redownload",
                ):
                    retry_result = subprocess.run(
                        cmd, check=True, timeout=timeout, capture_output=True, text=True
                    )
                    # Print retry output
                    if retry_result.stdout:
                        print(retry_result.stdout, end="")
                    if retry_result.stderr:
                        print(retry_result.stderr, end="")
            else:
                print(
                    f"Failed to cleanup corrupted MMMU cache. Error from lmms_eval:\nStdout:\n{e.stdout}\nStderr:\n{e.stderr}"
                )
                raise
        else:
            print(
                f"lmms_eval failed with an unhandled error.\nStdout:\n{e.stdout}\nStderr:\n{e.stderr}"
            )
            raise


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",
        ):
            _run_lmms_eval_with_retry(cmd)

    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)

            result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)

            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}")

            dump_metric(
                "mmmu_score",
                mmmu_accuracy,
                labels={"model": self.model, "eval": "mmmu", "api": "lmms-eval"},
            )

            # 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 = []

    @classmethod
    def setUpClass(cls):
        # Removed argument parsing from here
        cls.base_url = DEFAULT_URL_FOR_TEST
        cls.api_key = "sk-123456"
        cls.time_out = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH

        if cls.parsed_args is None:
            cls.parsed_args = SimpleNamespace(
                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,
        output_path: str,
        *,
        env: dict | None = None,
    ):
        """
        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,
        ]

        _run_lmms_eval_with_retry(cmd)

    def _run_vlm_mmmu_test(
        self,
        model,
        output_path,
        test_name="",
        custom_env=None,
        log_level="info",
        capture_output=False,
    ):
        """
        Common method to run VLM MMMU benchmark test.

        Args:
            model: Model to test
            output_path: Path for output logs
            test_name: Optional test name for logging
            custom_env: Optional custom environment variables
            log_level: Log level for server (default: "info")
            capture_output: Whether to capture server stdout/stderr
        """
        print(f"\nTesting model: {model.model}{test_name}")

        process = None
        mmmu_accuracy = 0  # Initialize to handle potential exceptions
        server_output = ""

        try:
            # Prepare environment variables
            process_env = os.environ.copy()
            if custom_env:
                process_env.update(custom_env)
            # if test vlm with cuda_ipc feature, open this env_var
            process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1"

            # Prepare stdout/stderr redirection if needed
            stdout_file = None
            stderr_file = None
            if capture_output:
                stdout_file = open("/tmp/server_stdout.log", "w")
                stderr_file = open("/tmp/server_stderr.log", "w")

            # Launch server for testing
            process = popen_launch_server(
                model.model,
                base_url=self.base_url,
                timeout=self.time_out,
                api_key=self.api_key,
                other_args=[
                    "--trust-remote-code",
                    "--cuda-graph-max-bs-decode",
                    "64",
                    "--enable-multimodal",
                    "--mem-fraction-static",
                    str(self.parsed_args.mem_fraction_static),  # Use class variable
                    "--log-level",
                    log_level,
                    *self.other_args,
                ],
                env=process_env,
                return_stdout_stderr=(
                    (stdout_file, stderr_file) if capture_output else None
                ),
            )

            # Run evaluation
            self.run_mmmu_eval(model.model, output_path)

            result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)

            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{test_name}\n: {result}")

            # Process the result
            mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
            print(
                f"Model {model.model} achieved accuracy{test_name}: {mmmu_accuracy:.4f}"
            )

            dump_metric(
                "mmmu_score",
                mmmu_accuracy,
                labels={"model": model.model, "eval": "mmmu", "api": "lmms-eval"},
            )

            # Capture server output if requested
            if capture_output and process:
                server_output = self._read_output_from_files()

            # Assert performance meets expected threshold
            self.assertGreaterEqual(
                mmmu_accuracy,
                model.mmmu_accuracy,
                f"Model {model.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({model.mmmu_accuracy:.4f}){test_name}",
            )

            return server_output

        except Exception as e:
            print(f"Error testing {model.model}{test_name}: {e}")
            self.fail(f"Test failed for {model.model}{test_name}: {e}")

        finally:
            # Ensure process cleanup happens regardless of success/failure
            if process is not None and process.poll() is None:
                print(f"Cleaning up process {process.pid}")
                try:
                    kill_process_tree(process.pid)
                except Exception as e:
                    print(f"Error killing process: {e}")

            # clean up temporary files
            if capture_output:
                if stdout_file:
                    stdout_file.close()
                if stderr_file:
                    stderr_file.close()
                for filename in ["/tmp/server_stdout.log", "/tmp/server_stderr.log"]:
                    try:
                        if os.path.exists(filename):
                            os.remove(filename)
                    except Exception as e:
                        print(f"Error removing {filename}: {e}")

    def _read_output_from_files(self):
        output_lines = []

        log_files = [
            ("/tmp/server_stdout.log", "[STDOUT]"),
            ("/tmp/server_stderr.log", "[STDERR]"),
        ]
        for filename, tag in log_files:
            try:
                if os.path.exists(filename):
                    with open(filename, "r") as f:
                        for line in f:
                            output_lines.append(f"{tag} {line.rstrip()}")
            except Exception as e:
                print(f"Error reading {tag.lower()} file: {e}")

        return "\n".join(output_lines)


# Backward compatibility alias
MMMUVLMTestBase = MMMUMultiModelTestBase