benchscope 1.0.0__py3-none-any.whl

benchscope/datasets.py

@@ -0,0 +1,200 @@
+"""数据集管理：ShareGPT 自动下载（modelscope）、自定义数据集上传与本地路径。"""
+from __future__ import annotations
+
+import json
+import logging
+import shutil
+import urllib.parse
+from pathlib import Path
+from typing import Optional
+
+import requests
+
+from benchscope.constants import SHAREGPT_DATASET_ID
+
+log = logging.getLogger("benchscope.datasets")
+
+# 期望的 ShareGPT 文件（下载源）与转换后的 jsonl 文件名
+SHAREGPT_SOURCE_CANDIDATES = [
+    "ShareGPT_V3_unfiltered_cleaned_split.json",
+    "ShareGPT_V3_unfiltered_cleaned_split.jsonl",
+    "sharegpt_v3_unfiltered_cleaned_split.jsonl",
+    "sharegpt.jsonl",
+]
+SHAREGPT_JSONL_NAME = "ShareGPT_V3_unfiltered_cleaned_split.jsonl"
+
+MODELSCOPE_API = "https://modelscope.cn/api/v1/datasets/{namespace}/{name}/repo"
+
+
+def _split_dataset_id(dataset_id: str) -> tuple[str, str]:
+    parts = dataset_id.split("/", 1)
+    if len(parts) != 2:
+        raise ValueError(f"无效的 modelscope 数据集 id: {dataset_id}")
+    return parts[0], parts[1]
+
+
+def list_dataset_files(dataset_id: str, revision: str = "master", timeout: int = 30) -> list[dict]:
+    """通过 modelscope HTTP API 列出数据集文件（/repo/tree）。"""
+    namespace, name = _split_dataset_id(dataset_id)
+    url = MODELSCOPE_API.format(namespace=namespace, name=name) + "/tree"
+    params = {"Revision": revision, "Root": "", "Recursive": "true"}
+    resp = requests.get(url, params=params, timeout=timeout)
+    resp.raise_for_status()
+    data = resp.json()
+    files = (data.get("Data") or {}).get("Files") or []
+    out = []
+    for f in files:
+        if isinstance(f, dict) and f.get("Type") == "blob":
+            out.append(f)
+    return out
+
+
+def convert_json_array_to_jsonl(src: Path, dst: Path) -> Path:
+    """流式把 JSON 数组文件转换为 jsonl（每行一个对象），内存占用低。
+
+    逐字符状态机跟踪对象边界与字符串（正确处理字符串内的 {} 与转义）；
+    每个对象归一化为单行 json 输出（vllm/sglang 按行 json.loads 解析）。
+    """
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dst.with_suffix(dst.suffix + ".part")
+    with open(src, "r", encoding="utf-8", errors="replace") as fin, open(tmp, "w", encoding="utf-8") as fout:
+        obj_chars: list[str] = []
+        depth = 0
+        in_string = False
+        escaped = False
+        for chunk in iter(lambda: fin.read(65536), ""):
+            for ch in chunk:
+                if in_string:
+                    if escaped:
+                        escaped = False
+                    elif ch == "\\":
+                        escaped = True
+                    elif ch == '"':
+                        in_string = False
+                else:
+                    if ch == '"':
+                        in_string = True
+                    elif ch == "{":
+                        depth += 1
+                        if depth == 1:
+                            obj_chars = ["{"]
+                            continue
+                    elif ch == "}":
+                        depth -= 1
+                        if depth == 0:
+                            try:
+                                obj = json.loads("".join(obj_chars) + "}")
+                                fout.write(json.dumps(obj, ensure_ascii=False) + "\n")
+                            except Exception:
+                                pass  # 跳过损坏对象
+                            obj_chars = []
+                            continue
+                if depth > 0:
+                    obj_chars.append(ch)
+    tmp.replace(dst)
+    return dst
+
+
+def download_file(
+    dataset_id: str, file_path: str, dest: Path, revision: str = "master", timeout: int = 60
+) -> Path:
+    """下载 modelscope 数据集中的单个文件。"""
+    namespace, name = _split_dataset_id(dataset_id)
+    url = MODELSCOPE_API.format(namespace=namespace, name=name)
+    params = {"Revision": revision, "FilePath": file_path}
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_suffix(dest.suffix + ".part")
+    with requests.get(url, params=params, stream=True, timeout=timeout) as resp:
+        resp.raise_for_status()
+        with open(tmp, "wb") as f:
+            for chunk in resp.iter_content(chunk_size=1 << 20):
+                f.write(chunk)
+    tmp.replace(dest)
+    return dest
+
+
+def ensure_sharegpt(datasets_dir: Path, force: bool = False) -> Path:
+    """确保 ShareGPT jsonl 已下载并转换缓存，返回 jsonl 路径。"""
+    cache_dir = datasets_dir / "sharegpt"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    jsonl_path = cache_dir / SHAREGPT_JSONL_NAME
+
+    # 已缓存的 jsonl
+    if not force and jsonl_path.exists() and jsonl_path.stat().st_size > 0:
+        return jsonl_path
+
+    # 已下载但未转换的源文件
+    source = None
+    if not force:
+        for cand in SHAREGPT_SOURCE_CANDIDATES:
+            p = cache_dir / cand
+            if p.exists() and p.stat().st_size > 0:
+                source = p
+                break
+        if source and jsonl_path.exists() and jsonl_path.stat().st_size > 0:
+            return jsonl_path
+
+    # 尝试 modelscope SDK（若安装）
+    try:
+        from modelscope import snapshot_download  # type: ignore
+
+        root = Path(snapshot_download(SHAREGPT_DATASET_ID, cache_dir=str(cache_dir)))
+        if source is None:
+            for cand in SHAREGPT_SOURCE_CANDIDATES:
+                p = root / cand
+                if p.exists():
+                    source = p
+                    break
+        if source is None:
+            jsonls = sorted(root.rglob("*.json*"))
+            if jsonls:
+                source = jsonls[0]
+    except ImportError:
+        pass
+    except Exception as e:
+        log.warning("modelscope SDK 下载失败，回退 HTTP 下载: %s", e)
+
+    # HTTP API 方式
+    if source is None:
+        files = list_dataset_files(SHAREGPT_DATASET_ID)
+        cands = []
+        for f in files:
+            path = f.get("Path") or ""
+            if path.endswith((".json", ".jsonl")):
+                cands.append((path, f.get("Size") or 0))
+        if not cands:
+            raise FileNotFoundError(f"数据集 {SHAREGPT_DATASET_ID} 中未找到数据文件: {files[:5]}")
+        # 优先文件名含 sharegpt，其次最大文件
+        cands.sort(key=lambda x: (0 if "sharegpt" in x[0].lower() else 1, -x[1]))
+        candidate, cand_size = cands[0]
+        source = cache_dir / Path(candidate).name
+        if not source.exists() or source.stat().st_size == 0 or force:
+            log.info("正在从 modelscope 下载 %s（约 %.0f MB）...", candidate, cand_size / 1e6)
+            download_file(SHAREGPT_DATASET_ID, candidate, source, timeout=300)
+
+    # 转换为 jsonl
+    if jsonl_path.exists() and jsonl_path.stat().st_size > 0 and not force:
+        return jsonl_path
+    if source.suffix.lower() == ".jsonl":
+        if source != jsonl_path:
+            import shutil as _shutil
+
+            _shutil.copyfile(source, jsonl_path)
+    else:
+        log.info("正在转换 %s 为 jsonl ...", source.name)
+        convert_json_array_to_jsonl(source, jsonl_path)
+    return jsonl_path
+
+
+def is_valid_jsonl(path: Path, sample: int = 2000) -> bool:
+    """粗略校验 jsonl 文件格式。"""
+    try:
+        with open(path, "r", encoding="utf-8", errors="replace") as f:
+            for _ in range(sample):
+                line = f.readline()
+                if not line:
+                    break
+                json.loads(line)
+        return True
+    except Exception:
+        return False

benchscope/gpu.py

@@ -0,0 +1,39 @@
+"""GPU 信息自动获取（nvidia-smi），失败时回退手动配置。"""
+from __future__ import annotations
+
+import shutil
+import subprocess
+from typing import Optional
+
+
+def detect_gpu() -> Optional[dict]:
+    """尝试自动获取 GPU 型号与数量。成功返回 {"name":..., "count":N}，失败返回 None。"""
+    nvidia_smi = shutil.which("nvidia-smi")
+    if not nvidia_smi:
+        return None
+    try:
+        result = subprocess.run(
+            [nvidia_smi, "--query-gpu=name,count", "--format=csv,noheader,nounits"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode != 0:
+            return None
+        lines = [ln.strip() for ln in result.stdout.splitlines() if ln.strip()]
+        if not lines:
+            return None
+        names = [ln.split(",")[0].strip() for ln in lines]
+        counts = []
+        for ln in lines:
+            parts = ln.split(",")
+            if len(parts) > 1:
+                try:
+                    counts.append(int(parts[1].strip()))
+                except ValueError:
+                    pass
+        name = names[0] if names else ""
+        count = sum(counts) if counts else len(lines)
+        return {"name": name, "count": count}
+    except Exception:
+        return None

benchscope/parser.py

@@ -0,0 +1,93 @@
+"""bench 输出解析：同时解析 mean 与 P99 两套指标。
+
+兼容 vLLM `vllm bench serve` 与 SGLang `sglang.bench_serving` 两种输出格式。
+"""
+from __future__ import annotations
+
+import re
+from typing import Optional
+
+# 指标键（统一命名）
+METRIC_KEYS = [
+    "output",
+    "peakoutput",
+    "total",
+    "ttft",
+    "itl",
+    "tpot",
+]
+
+# 每个指标在输出中的可能标签（vllm / sglang 写法不同）
+_PATTERNS: list[tuple[str, str, re.Pattern]] = []
+
+
+def _add(kind: str, metric: str, regex: str):
+    _PATTERNS.append((kind, metric, re.compile(regex, re.IGNORECASE)))
+
+
+# ---- vLLM bench serve ----
+_add("mean", "output", r"Output token throughput \(tok/s\):\s+([\d.]+)")
+_add("mean", "peakoutput", r"Peak output token throughput \(tok/s\):\s+([\d.]+)")
+_add("mean", "total", r"Total token throughput \(tok/s\):\s+([\d.]+)")
+_add("mean", "ttft", r"Mean TTFT \(ms\):\s+([\d.]+)")
+_add("p99", "ttft", r"P99 TTFT \(ms\):\s+([\d.]+)")
+_add("mean", "tpot", r"Mean TPOT \(ms\):\s+([\d.]+)")
+_add("p99", "tpot", r"P99 TPOT \(ms\):\s+([\d.]+)")
+_add("mean", "itl", r"Mean ITL \(ms\):\s+([\d.]+)")
+_add("p99", "itl", r"P99 ITL \(ms\):\s+([\d.]+)")
+_add("concurrency", "concurrency", r"Maximum request concurrency:\s+(\d+)")
+
+# ---- SGLang bench_serving ----
+_add("mean", "output", r"Output token throughput \(tok/s\):\s+([\d.]+)")
+_add("mean", "total", r"Total token throughput \(tok/s\):\s+([\d.]+)")
+_add("mean", "ttft", r"Time to first token \(TTFT\) mean \(ms\):\s+([\d.]+)")
+_add("p99", "ttft", r"Time to first token \(TTFT\) p99 \(ms\):\s+([\d.]+)")
+_add("mean", "tpot", r"Time per output token \(TPOT\) mean \(ms\):\s+([\d.]+)")
+_add("p99", "tpot", r"Time per output token \(TPOT\) p99 \(ms\):\s+([\d.]+)")
+_add("mean", "itl", r"Inter-token latency \(ITL\) mean \(ms\):\s+([\d.]+)")
+_add("p99", "itl", r"Inter-token latency \(ITL\) p99 \(ms\):\s+([\d.]+)")
+
+# 请求吞吐（req/s），sglang/vllm 均输出
+_add("mean", "req_per_s", r"Request throughput \(req/s\):\s+([\d.]+)")
+
+
+def parse_metrics(output: str) -> dict:
+    """解析 bench 输出，返回统一指标字典。
+
+    返回结构:
+    {
+      "concurrency": int,
+      "output": float, "peakoutput": float, "total": float,
+      "ttft_mean": float, "tpot_mean": float, "itl_mean": float,
+      "ttft_p99": float, "tpot_p99": float, "itl_p99": float,
+      "req_per_s": float, "raw": str(完整原始输出)
+    }
+    """
+    metrics: dict = {}
+    for kind, metric, pattern in _PATTERNS:
+        m = pattern.search(output)
+        if m:
+            key = metric if kind == "concurrency" else f"{metric}_{kind}"
+            try:
+                value = float(m.group(1)) if "." in m.group(1) else int(m.group(1))
+            except ValueError:
+                value = float(m.group(1))
+            metrics[key] = value
+
+    # 兼容只有 mean 没有 p99 的情况：p99 缺省回退 mean
+    for metric in ("ttft", "tpot", "itl"):
+        mean_key = f"{metric}_mean"
+        p99_key = f"{metric}_p99"
+        if mean_key in metrics and p99_key not in metrics:
+            metrics[p99_key] = metrics[mean_key]
+
+    # 单用户 QPS = 1000 / tpot_mean（README 定义）
+    if "tpot_mean" in metrics and metrics["tpot_mean"] > 0:
+        metrics["single_user"] = round(1000.0 / metrics["tpot_mean"], 2)
+    metrics["raw"] = output
+    return metrics
+
+
+def has_result_block(output: str) -> bool:
+    """判断输出中是否包含结果块（可能成功或失败）。"""
+    return "Serving Benchmark Result" in output or "Benchmark Result" in output

benchscope/server/__init__.py

@@ -0,0 +1,3 @@
+from benchscope.server.app import create_app
+
+__all__ = ["create_app"]

benchscope/server/api_config.py

@@ -0,0 +1,97 @@
+"""配置 / 模型 / GPU / 状态 相关 API。"""
+from __future__ import annotations
+
+import logging
+
+import requests
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel
+
+from benchscope.constants import STATUS_READY
+from benchscope.gpu import detect_gpu
+from benchscope.server.state import state
+
+log = logging.getLogger("benchscope.api_config")
+
+router = APIRouter(prefix="/api/config", tags=["config"])
+
+
+@router.get("")
+def get_config():
+    return state.config.snapshot()
+
+
+class ConfigPatch(BaseModel):
+    api: dict | None = None
+    gpu: dict | None = None
+    logs_dir: str | None = None
+    datasets_dir: str | None = None
+    tpot_threshold_ms: float | None = None
+    request_rate: str | None = None
+    bench_commands: dict | None = None
+    framework: str | None = None
+
+
+@router.post("")
+def update_config(patch: ConfigPatch):
+    data = patch.model_dump(exclude_none=True)
+    state.config.update(data)
+    return state.config.snapshot()
+
+
+@router.get("/status")
+def get_status():
+    snap = state.monitor.check_once(broadcast=False)
+    snap["web"] = STATUS_READY
+    return snap
+
+
+@router.get("/models")
+def get_models():
+    """返回推理服务当前模型列表（优先用状态缓存，必要时实时探测）。"""
+    snap = state.monitor.check_once(broadcast=False)
+    if not snap["models"] and snap["error"]:
+        raise HTTPException(status_code=502, detail=f"推理服务不可达: {snap['error']}")
+    return {"models": snap["models"], "inference": snap["inference"], "error": snap["error"]}
+
+
+class ConnTest(BaseModel):
+    base_url: str
+    endpoint: str = "/v1/chat/completions"
+    api_key: str = ""
+    extra_headers: dict = {}
+
+
+@router.post("/test-connection")
+def test_connection(req: ConnTest):
+    base = req.base_url.rstrip("/")
+    headers = {}
+    if req.api_key:
+        headers["Authorization"] = f"Bearer {req.api_key}"
+    headers.update(req.extra_headers or {})
+    try:
+        resp = requests.get(f"{base}/v1/models", headers=headers, timeout=6)
+        resp.raise_for_status()
+        models = [m.get("id") for m in resp.json().get("data", []) if m.get("id")]
+        return {"ok": True, "models": models}
+    except Exception as e:
+        return {"ok": False, "error": str(e)[:300]}
+
+
+@router.get("/gpu")
+def get_gpu():
+    auto = detect_gpu()
+    cfg = state.config.get("gpu", {"auto": True, "name": "", "count": 8})
+    return {"auto_detected": auto, "config": cfg}
+
+
+@router.get("/params/{framework}")
+def get_params(framework: str):
+    """返回指定框架的可配置参数定义（前端表单渲染用）。"""
+    if framework == "sglang":
+        from benchscope.benches.sglang_bench import CURATED_PARAMS
+    elif framework == "vllm":
+        from benchscope.benches.vllm_bench import CURATED_PARAMS
+    else:
+        raise HTTPException(status_code=404, detail="未知框架")
+    return {"framework": framework, "params": [p.__dict__ for p in CURATED_PARAMS]}