[Refactor][WIP] Add missing test for modelrunner (G1/N) (#5333)

Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

tools/baselines/omni_hash_baseline.py

@@ -0,0 +1,276 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Record / verify a deterministic text+audio output hash baseline for an Omni model.
+
+This is the cheap "output is bit-identical" gate for the behaviour-preserving
+worker/model-runner changes: record a baseline on the bench box
+BEFORE any code moves, then re-run with ``--check`` after each refactor.
+
+Determinism: **batch size 1** (one prompt per ``generate`` call) + **fixed seed**.
+Thinker and code2wav run greedy (temperature 0); the talker's speech sampler needs
+temperature to emit valid codec tokens, so it runs at its natural sampling with a
+fixed seed (reproducible via seed + batch-1, not pure-greedy). 3 fixed prompts, run
+with async output OFF and ON (6 entries total). Hashes are only meaningful WITHIN a single recorded
+environment (driver/cuda/torch/gpu), so the baseline JSON records the env and
+``--check`` refuses to compare across a changed env unless ``--force-env``.
+
+Baseline JSON schema (results keyed by BOTH prompt_id and async_mode so the
+matrix is self-identifying)::
+
+    {
+      "model": "...",
+      "commit": "<git sha>",
+      "env": {"driver": "<nvidia driver>", "cuda": "<cuda build>", "torch": "...", "gpu": "..."},
+      "results": [
+        {"prompt_id": 0, "async_mode": false, "text_sha": "...", "wav_sha": "..."},
+        ...  # exactly len(FIXED_PROMPTS) * 2 entries
+      ]
+    }
+
+The file content is JSON, but committed baselines use a NON-``.json`` extension
+(``.baseline``) because the repo ``.gitignore`` excludes ``*.json``; ``--check``
+loads by content, so the extension does not matter to the script.
+
+Usage::
+
+    # record (run twice, confirm the two runs are identical before committing)
+    python tools/baselines/omni_hash_baseline.py <model> --yaml <deploy.yaml> \
+        --out tools/baselines/<model>_hash.baseline
+    # verify after a refactor (same bench box)
+    python tools/baselines/omni_hash_baseline.py <model> --yaml <deploy.yaml> \
+        --check tools/baselines/<model>_hash.baseline
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import subprocess
+import sys
+
+import torch
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.omni import Omni
+
+SEED = 42
+DEFAULT_SYSTEM = (
+    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, "
+    "capable of perceiving auditory and visual inputs, as well as generating "
+    "text and speech."
+)
+
+# 3 fixed prompts — do NOT edit once a baseline is recorded (prompt_id is positional).
+FIXED_PROMPTS = [
+    "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.",
+    "Describe how ocean tides work. Answer in 15 words.",
+    "List three benefits of unit tests. Answer in 15 words.",
+]
+ASYNC_MODES = (False, True)
+
+
+def _text_prompt(question: str) -> dict:
+    return {
+        "prompt": (
+            f"<|im_start|>system\n{DEFAULT_SYSTEM}<|im_end|>\n"
+            f"<|im_start|>user\n{question}<|im_end|>\n"
+            f"<|im_start|>assistant\n"
+        )
+    }
+
+
+def _sha(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def _git_commit() -> str:
+    try:
+        return subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()
+    except Exception:
+        return "unknown"
+
+
+def _driver_version() -> str:
+    """The actual NVIDIA driver version (NOT torch.version.cuda, which is the build)."""
+    try:
+        out = subprocess.check_output(["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], text=True)
+        return out.strip().splitlines()[0].strip()
+    except Exception:
+        return "unknown"
+
+
+def _env_fingerprint() -> dict:
+    gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu"
+    return {
+        "driver": _driver_version(),
+        "cuda": str(torch.version.cuda),
+        "torch": torch.__version__,
+        "gpu": gpu,
+    }
+
+
+def _sampling_params() -> list[SamplingParams]:
+    # thinker + code2wav: greedy. talker: seeded (temperature>0) — its speech sampler
+    # needs temperature to emit valid codec tokens; the fixed seed keeps it reproducible.
+    thinker = SamplingParams(
+        temperature=0.0, top_p=1.0, top_k=-1, max_tokens=128, seed=SEED, detokenize=True, repetition_penalty=1.1
+    )
+    talker = SamplingParams(
+        temperature=0.9, top_p=0.8, top_k=40, max_tokens=256, seed=SEED, detokenize=True, repetition_penalty=1.05
+    )
+    code2wav = SamplingParams(
+        temperature=0.0, top_p=1.0, top_k=-1, max_tokens=256, seed=SEED, detokenize=True, repetition_penalty=1.1
+    )
+    return [thinker, talker, code2wav]
+
+
+def _audio_bytes(audio) -> bytes:
+    """Deterministic bytes for one request's audio, handling chunked (list) output.
+
+    async-chunk mode can emit a list of audio chunks instead of a single tensor;
+    concatenate them along the last axis in emit order.
+    """
+    if isinstance(audio, (list, tuple)):
+        chunks = [a for a in audio if a is not None]
+        if not chunks:
+            return b""
+        audio = torch.cat([torch.as_tensor(a).detach().to(torch.float32).cpu() for a in chunks], dim=-1)
+    if audio is None:
+        return b""
+    return torch.as_tensor(audio).detach().to(torch.float32).cpu().numpy().tobytes()
+
+
+def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]:
+    """One deterministic pass, batch size 1 per prompt; returns per-prompt hashes."""
+    # ``async_chunk`` is the repo's documented async-output toggle for streaming
+    # (see tests/e2e/accuracy/qwen3_omni). Confirm the exact knob at bench time.
+    omni = Omni(
+        model,
+        stage_configs_path=yaml_path,
+        stage_init_timeout=1800,
+        init_timeout=3600,
+        async_chunk=async_mode,
+    )
+    results: list[dict] = []
+    try:
+        for prompt_id, question in enumerate(FIXED_PROMPTS):
+            text: str | None = None
+            wav: bytes = b""
+            # Batch size 1: submit a single prompt so scheduling is deterministic.
+            for step in omni.generate([_text_prompt(question)], _sampling_params(), py_generator=False):
+                out = step.request_output
+                if step.final_output_type == "text":
+                    text = out.outputs[0].text  # exact text, no strip -> truly bit-identical
+                elif step.final_output_type == "audio":
+                    wav = _audio_bytes(out.outputs[0].multimodal_output["audio"])
+            if text is None:
+                raise RuntimeError(f"prompt_id={prompt_id} async={async_mode}: produced no text output")
+            if not wav:
+                raise RuntimeError(f"prompt_id={prompt_id} async={async_mode}: produced no/empty audio output")
+            results.append(
+                {
+                    "prompt_id": prompt_id,
+                    "async_mode": async_mode,
+                    "text_sha": _sha(text.encode("utf-8")),
+                    "wav_sha": _sha(wav),
+                }
+            )
+    finally:
+        omni.close()
+    return results
+
+
+def _record(model: str, yaml_path: str) -> dict:
+    results: list[dict] = []
+    for async_mode in ASYNC_MODES:
+        results.extend(_run_once(model, yaml_path, async_mode=async_mode))
+    return {"model": model, "commit": _git_commit(), "env": _env_fingerprint(), "results": results}
+
+
+def _key(record: dict):
+    return (record["prompt_id"], record["async_mode"])
+
+
+def _expected_keys() -> set:
+    return {(pid, mode) for pid in range(len(FIXED_PROMPTS)) for mode in ASYNC_MODES}
+
+
+def _check(model: str, yaml_path: str, baseline_path: str, *, force_env: bool) -> int:
+    with open(baseline_path) as fh:
+        baseline = json.load(fh)
+
+    now = _record(model, yaml_path)
+    if not force_env and baseline.get("env") != now["env"]:
+        print(
+            f"ENV MISMATCH: baseline={baseline.get('env')} now={now['env']} — hashes are "
+            "env-specific. Re-record on this box or pass --force-env.",
+            file=sys.stderr,
+        )
+        return 2
+
+    problems: list[str] = []
+
+    def _index_unique(records: list[dict], label: str) -> dict:
+        by_key: dict = {}
+        for record in records:
+            key = _key(record)
+            if key in by_key:
+                problems.append(f"{label}: duplicate entry for (prompt_id={key[0]}, async={key[1]})")
+            by_key[key] = record
+        return by_key
+
+    base_by_key = _index_unique(baseline["results"], "baseline")
+    now_by_key = _index_unique(now["results"], "current-run")
+    expected = _expected_keys()
+
+    # Guard against duplicate-collapse: the record count must match the key count.
+    if len(baseline["results"]) != len(expected):
+        problems.append(f"baseline has {len(baseline['results'])} records, expected exactly {len(expected)}")
+    if set(now_by_key) != expected:
+        problems.append(f"current-run key set {sorted(set(now_by_key))} != expected {sorted(expected)}")
+    if set(base_by_key) != set(now_by_key):
+        missing = sorted(set(base_by_key) - set(now_by_key))
+        extra = sorted(set(now_by_key) - set(base_by_key))
+        problems.append(f"baseline/run key-set mismatch: missing-from-run={missing} extra-in-run={extra}")
+    for key in sorted(set(base_by_key) & set(now_by_key)):
+        b, r = base_by_key[key], now_by_key[key]
+        if (b["text_sha"], b["wav_sha"]) != (r["text_sha"], r["wav_sha"]):
+            problems.append(
+                f"HASH MISMATCH (prompt_id={key[0]}, async={key[1]}): "
+                f"text {b['text_sha'][:8]}!={r['text_sha'][:8]} / wav {b['wav_sha'][:8]}!={r['wav_sha'][:8]}"
+            )
+
+    if problems:
+        for problem in problems:
+            print(problem, file=sys.stderr)
+        return 1
+    print(f"OK: all {len(expected)} (prompt_id, async_mode) entries match the baseline.")
+    return 0
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("model")
+    p.add_argument("--yaml", required=True, help="Stage deploy config path")
+    p.add_argument("--out", help="Write a new baseline JSON to this path")
+    p.add_argument("--check", metavar="BASELINE_JSON", help="Verify current output against this baseline")
+    p.add_argument("--force-env", action="store_true", help="Compare hashes even if the recorded env differs")
+    args = p.parse_args()
+
+    if args.check:
+        return _check(args.model, args.yaml, args.check, force_env=args.force_env)
+
+    baseline = _record(args.model, args.yaml)
+    payload = json.dumps(baseline, indent=2)
+    if args.out:
+        with open(args.out, "w") as fh:
+            fh.write(payload)
+        print(f"Wrote baseline: {args.out}")
+    else:
+        print(payload)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tools/baselines/qwen3_omni_hash.baseline

@@ -0,0 +1,48 @@
+{
+  "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+  "commit": "24eb4ab89cc2ad3f8cdd83f57b66ce7786a0d65c",
+  "env": {
+    "driver": "570.133.20",
+    "cuda": "13.0",
+    "torch": "2.11.0+cu130",
+    "gpu": "NVIDIA L20X"
+  },
+  "results": [
+    {
+      "prompt_id": 0,
+      "async_mode": false,
+      "text_sha": "e4bf6ed906c0dada879874a9cc09f9ba3a4399149262a831a01fd53e6a8fd043",
+      "wav_sha": "20a203a149005f948afd485bbea780fa06c109f3689ebb6b3110f51b3237bd35"
+    },
+    {
+      "prompt_id": 1,
+      "async_mode": false,
+      "text_sha": "2e8a6b141a9747bd68aadc2931f278b89712ad7ad7ae12b25b39e193b241bf68",
+      "wav_sha": "1839a759bb5d8671b6352b2e3110dab5edfa718e38ff073d379176201f1eac8c"
+    },
+    {
+      "prompt_id": 2,
+      "async_mode": false,
+      "text_sha": "e0aa5a513a5162024f35afee3830aeb67f2e378381726adb9d3a66b984aa652d",
+      "wav_sha": "74338df3e8357bc9d0c05bf141b63c45251b12807de34495097edb07fc7b19f8"
+    },
+    {
+      "prompt_id": 0,
+      "async_mode": true,
+      "text_sha": "e4bf6ed906c0dada879874a9cc09f9ba3a4399149262a831a01fd53e6a8fd043",
+      "wav_sha": "e09f862f54b553cc7c64e1631b821a768167768812b3f2cd3d0c69d0da991ee7"
+    },
+    {
+      "prompt_id": 1,
+      "async_mode": true,
+      "text_sha": "2e8a6b141a9747bd68aadc2931f278b89712ad7ad7ae12b25b39e193b241bf68",
+      "wav_sha": "9115d481f62a4b771c4c3d8b89d5c5a01149f002096289ab1da3f034fa388a2b"
+    },
+    {
+      "prompt_id": 2,
+      "async_mode": true,
+      "text_sha": "e0aa5a513a5162024f35afee3830aeb67f2e378381726adb9d3a66b984aa652d",
+      "wav_sha": "6a5985c9d4ca1596518c291228cd29239e76fcaf310c1cdb5921276b7efce801"
+    }
+  ]
+}