from __future__ import annotations import os import shutil import subprocess from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path import pytest import torch from tests.conftest import OmniServer, OmniServerParams def pytest_addoption(parser): group = parser.getgroup("accuracy-e2e") group.addoption("--gebench-root", action="store", default=None, help="Local GEBench dataset root") group.addoption("--gedit-root", action="store", default=None, help="Local GEdit-Bench dataset root") group.addoption( "--gebench-model", action="store", default="Qwen/Qwen-Image-2512", help="Generate model for GEBench smoke" ) group.addoption( "--gedit-model", action="store", default="Qwen/Qwen-Image-Edit", help="Generate model for GEdit-Bench smoke" ) group.addoption( "--accuracy-judge-model", action="store", default="QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ", help="Judge model path", ) group.addoption("--accuracy-gpu", action="store", default="0", help="Single GPU id used sequentially") group.addoption("--gebench-port", action="store", type=int, default=8093, help="Generate port for GEBench") group.addoption("--gedit-port", action="store", type=int, default=8093, help="Generate port for GEdit-Bench") group.addoption( "--gebench-samples-per-type", action="store", type=int, default=10, help="Balanced sample count per GEBench type", ) group.addoption( "--gedit-samples-per-group", action="store", type=int, default=20, help="Balanced sample count per GEdit task group", ) group.addoption("--accuracy-workers", action="store", type=int, default=1, help="Worker count for accuracy benches") def _hf_cache_root() -> Path: return Path(os.environ.get("HF_HOME", "/root/.cache/huggingface")) def _dataset_cache_dirs(dataset_id: str) -> list[Path]: cache_root = _hf_cache_root() / "hub" / f"datasets--{dataset_id.replace('/', '--')}" / "snapshots" if not cache_root.exists(): return [] return sorted( (path for path in cache_root.iterdir() if path.is_dir()), key=lambda path: path.stat().st_mtime, reverse=True ) def _ensure_dataset_snapshot(dataset_id: str) -> Path: candidates = _dataset_cache_dirs(dataset_id) if candidates: return candidates[0] subprocess.run( ["huggingface-cli", "download", "--repo-type", "dataset", dataset_id], check=True, ) candidates = _dataset_cache_dirs(dataset_id) if not candidates: raise FileNotFoundError( f"Dataset {dataset_id} was downloaded but no snapshot was found under {_hf_cache_root()}" ) return candidates[0] def _resolve_dataset_root(request: pytest.FixtureRequest, option_name: str, dataset_id: str) -> Path: value = request.config.getoption(option_name) if value: path = Path(value) if not path.exists(): pytest.skip(f"Dataset path does not exist: {path}") return path return _ensure_dataset_snapshot(dataset_id) @dataclass class AccuracyServerConfig: generate_params: OmniServerParams judge_params: OmniServerParams run_level: str model_prefix: str @contextmanager def generate_server(self): params = self.generate_params model = self.model_prefix + params.model server_args = params.server_args or [] if params.use_omni: server_args = ["--stage-init-timeout", "120", *server_args] with OmniServer( model, server_args, port=params.port, env_dict=params.env_dict, use_omni=params.use_omni, ) as server: yield server @contextmanager def judge_server(self): params = self.judge_params model = self.model_prefix + params.model server_args = params.server_args or [] with OmniServer( model, server_args, port=params.port, env_dict=params.env_dict, use_omni=params.use_omni, ) as server: yield server @pytest.fixture(scope="session") def gebench_dataset_root(request: pytest.FixtureRequest) -> Path: return _resolve_dataset_root(request, "gebench_root", "stepfun-ai/GEBench") @pytest.fixture(scope="session") def gedit_dataset_root(request: pytest.FixtureRequest) -> Path: return _resolve_dataset_root(request, "gedit_root", "stepfun-ai/GEdit-Bench") @pytest.fixture(scope="session") def accuracy_workers(request: pytest.FixtureRequest) -> int: return int(request.config.getoption("accuracy_workers")) @pytest.fixture(scope="session") def gebench_samples_per_type(request: pytest.FixtureRequest) -> int: return int(request.config.getoption("gebench_samples_per_type")) @pytest.fixture(scope="session") def gedit_samples_per_group(request: pytest.FixtureRequest) -> int: return int(request.config.getoption("gedit_samples_per_group")) @pytest.fixture(scope="session") def accuracy_artifact_root() -> Path: root = Path(__file__).resolve().parent / "artifacts" root.mkdir(parents=True, exist_ok=True) return root def reset_artifact_dir(path: Path) -> Path: if path.exists(): shutil.rmtree(path) path.mkdir(parents=True, exist_ok=True) return path def infer_model_label(model: str) -> str: label = Path(model.rstrip("/\\")).name or "model" return "".join(char if char.isalnum() or char in {"-", "_"} else "_" for char in label) def _build_accuracy_server_config( *, generate_model: str, judge_model: str, shared_gpu: str, port: int, run_level: str, model_prefix: str, ) -> AccuracyServerConfig: if torch.cuda.device_count() < 1: pytest.skip("Need at least 1 CUDA GPU for accuracy benchmark smoke tests.") if not generate_model: pytest.skip("No generate model configured for accuracy benchmark test.") generate_server_args = ["--num-gpus", "1"] judge_server_args = [ "--max-model-len", "32768", "--gpu-memory-utilization", "0.8", ] judge_env = {"CUDA_VISIBLE_DEVICES": shared_gpu} return AccuracyServerConfig( generate_params=OmniServerParams( model=generate_model, port=port, server_args=generate_server_args, env_dict={"CUDA_VISIBLE_DEVICES": shared_gpu}, use_omni=True, ), judge_params=OmniServerParams( model=judge_model, port=port, server_args=judge_server_args, env_dict=judge_env, use_omni=False, ), run_level=run_level, model_prefix=model_prefix, ) @pytest.fixture def gebench_accuracy_servers( request: pytest.FixtureRequest, run_level: str, model_prefix: str, ) -> AccuracyServerConfig: return _build_accuracy_server_config( generate_model=request.config.getoption("gebench_model"), judge_model=request.config.getoption("accuracy_judge_model"), shared_gpu=str(request.config.getoption("accuracy_gpu")), port=int(request.config.getoption("gebench_port")), run_level=run_level, model_prefix=model_prefix, ) @pytest.fixture def gedit_accuracy_servers( request: pytest.FixtureRequest, run_level: str, model_prefix: str, ) -> AccuracyServerConfig: return _build_accuracy_server_config( generate_model=request.config.getoption("gedit_model"), judge_model=request.config.getoption("accuracy_judge_model"), shared_gpu=str(request.config.getoption("accuracy_gpu")), port=int(request.config.getoption("gedit_port")), run_level=run_level, model_prefix=model_prefix, )