Add native Windows CPU Qwen TTS Omni stage and profile evidence

benchmarks/edge_harness/experiments/probe_omni_qwen_tts_cpu.py

@@ -0,0 +1,216 @@
+#!/usr/bin/env python3
+"""Profile the isolated official Qwen3-TTS CPU wrapper through Omni."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import hashlib
+import json
+import math
+import time
+import wave
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+
+
+def _rank(values: list[float], fraction: float) -> float:
+    return sorted(values)[math.ceil(len(values) * fraction) - 1]
+
+
+def _save_wav(path: Path, audio) -> None:
+    pcm = np.rint(audio.numpy() * 32768).astype("<i2").tobytes()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with wave.open(str(path), "wb") as wav:
+        wav.setnchannels(1)
+        wav.setsampwidth(2)
+        wav.setframerate(24000)
+        wav.writeframes(pcm)
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    for name in (
+        "python_bin", "python_sha256", "model_dir", "overlay_dir", "talker_sha256",
+        "tokenizer_sha256", "server_log", "output_report", "output_wav",
+    ):
+        parser.add_argument("--" + name.replace("_", "-"), required=True)
+    parser.add_argument("--warmups", type=int, default=1)
+    parser.add_argument("--repeats", type=int, default=20)
+    parser.add_argument("--reserve-gib", type=int, default=10)
+    parser.add_argument("--abort-check", action="store_true")
+    args = parser.parse_args()
+    if args.warmups < 0 or args.repeats <= 0 or not 0 < args.reserve_gib <= 16:
+        parser.error("invalid profile or budget")
+
+    from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy
+    from vllm_omni.engine.stage_runtime import StageRuntime
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_CPU_WHOLE_PIPELINE
+
+    backend = {
+        "name": "external.qwen_tts.cpu.v1",
+        "python_bin": args.python_bin,
+        "python_sha256": args.python_sha256,
+        "model_dir": args.model_dir,
+        "overlay_dir": args.overlay_dir,
+        "talker_sha256": args.talker_sha256,
+        "tokenizer_sha256": args.tokenizer_sha256,
+        "log_file": args.server_log,
+        "memory_overhead_bytes": 6 << 30,
+        "expected_torch": "2.13.0+cu130",
+        "max_text_bytes": 4096,
+        "max_wav_bytes": 4 << 20,
+        "max_audio_s": 10,
+        "request_timeout_s": 180,
+        "start_timeout_s": 180,
+    }
+    deploy = DeployConfig(
+        async_chunk=False,
+        stages=[StageDeployConfig(
+            stage_id=0,
+            backend=backend,
+            resource_budget={
+                "capacities": {"host_ram": 16 << 30},
+                "demands": {"host_ram": args.reserve_gib << 30},
+            },
+        )],
+    )
+    pipeline = QWEN3_TTS_CPU_WHOLE_PIPELINE
+    configs = [stage.to_omegaconf() for stage in merge_pipeline_deploy(pipeline, deploy)]
+    runtime = StageRuntime(configs, "local-qwen-tts-cpu", "", stage_init_timeout=180, async_chunk=False)
+    report = {
+        "scope": "complete resident native Windows CPU Qwen3-TTS text-to-WAV through Omni",
+        "warmup_count": args.warmups,
+        "measured_count": args.repeats,
+        "memory_budget": deploy.stages[0].resource_budget,
+    }
+    sample_task = None
+    stop_sampling = asyncio.Event()
+    try:
+        started = time.perf_counter()
+        runtime.initialize()
+        report["startup_s"] = time.perf_counter() - started
+        pool = runtime.stage_pools[0]
+        report["execution_plan"] = pool.stage_client.execution_plan
+        state = SimpleNamespace(sampling_params_list=[None])
+        import psutil
+
+        memory_samples = []
+
+        async def sample_memory() -> None:
+            while not stop_sampling.is_set():
+                try:
+                    root = psutil.Process(pool.stage_client._proc.pid)
+                    members = [root, *root.children(recursive=True)]
+                    memories = [process.memory_full_info() for process in members]
+                    memory_samples.append({
+                        "unix": time.time(),
+                        "rss_bytes": sum(int(memory.rss) for memory in memories),
+                        "private_bytes": sum(
+                            int(getattr(memory, "private", memory.rss)) for memory in memories
+                        ),
+                        "host_available_bytes": psutil.virtual_memory().available,
+                        "pids": [process.pid for process in members],
+                    })
+                except psutil.Error:
+                    break
+                await asyncio.sleep(0.25)
+
+        sample_task = asyncio.create_task(sample_memory())
+
+        async def request_one(text: str, *, save: bool = False) -> dict:
+            request_id = f"cpu-tts-{time.monotonic_ns()}"
+            started = time.perf_counter()
+            await pool.submit_initial(request_id, state, {"text": text, "voice": "Ryan", "seed": 42})
+            deadline = time.monotonic() + 180
+            while True:
+                output = pool.poll_graph_output(0)
+                if output is not None:
+                    break
+                if time.monotonic() >= deadline:
+                    raise TimeoutError("Omni Qwen CPU TTS stage request timed out")
+                await asyncio.sleep(0.005)
+            wall_s = time.perf_counter() - started
+            try:
+                if output.error:
+                    raise RuntimeError(output.error)
+                payload = output.multimodal_output
+                audio, sample_rate = payload["audio"], payload["sr"]
+                metadata = output.custom_output["audio_metadata"]
+                assert sample_rate == 24000 and audio.ndim == 1
+                assert len(audio) == metadata["pcm_frames"] and len(audio) > 0
+                assert float(audio.square().mean().sqrt()) > 1e-5
+                if save:
+                    _save_wav(Path(args.output_wav), audio)
+                pcm = np.rint(audio.numpy() * 32768).astype("<i2").tobytes()
+                assert hashlib.sha256(pcm).hexdigest() == metadata["pcm_sha256"]
+                return {
+                    "text": text,
+                    "wall_s": wall_s,
+                    "frames": len(audio),
+                    "duration_s": metadata["duration_s"],
+                    "pcm_sha256": metadata["pcm_sha256"],
+                    "stage_event": output.custom_output["stage_event"],
+                }
+            finally:
+                output.release_stage_buffers()
+                await asyncio.sleep(0)
+
+        first = "Hello from the local computer."
+        second = "The blue car is parked beside the library."
+        report["checks"] = [await request_one(first), await request_one(second)]
+        expected = (
+            "c8a49b5c73efd642a1eb04be973d6aefdc8ed9cfd18b34ff67690b8464906bed",
+            "4edc7a62f749b1fde5ea39bedbc15a35f291ab1db8d24d7b7eb3df7270885e16",
+        )
+        report["standalone_pcm_parity"] = [
+            row["pcm_sha256"] == digest for row, digest in zip(report["checks"], expected)
+        ]
+        if report["standalone_pcm_parity"] != [True, True]:
+            raise RuntimeError("CPU Omni PCM differs from the same checkpoint's standalone output")
+        report["warmups"] = [await request_one(first) for _ in range(args.warmups)]
+        report["measured"] = [await request_one(first, save=i == 0) for i in range(args.repeats)]
+        walls = [row["wall_s"] for row in report["measured"]]
+        report["nearest_rank_p50_wall_s"] = _rank(walls, 0.5)
+        report["nearest_rank_p95_wall_s"] = _rank(walls, 0.95)
+        report["all_same_pcm_sha256"] = len({row["pcm_sha256"] for row in report["measured"]}) == 1
+        report["memory_samples"] = memory_samples
+        report["sampled_max_rss_bytes"] = max((row["rss_bytes"] for row in memory_samples), default=None)
+        report["sampled_max_private_bytes"] = max(
+            (row["private_bytes"] for row in memory_samples), default=None
+        )
+        if args.abort_check:
+            request_id = f"cpu-tts-abort-{time.monotonic_ns()}"
+            await pool.submit_initial(request_id, state, {"text": second, "voice": "Ryan", "seed": 42})
+            await asyncio.sleep(0.01)
+            await pool.abort_requests([request_id])
+            await asyncio.sleep(0.1)
+            report["abort_check"] = {
+                "stale_output": pool.poll_graph_output(0) is not None,
+                "worker_exited": pool.stage_client._proc.poll() is not None,
+                "ledger_after_abort": runtime.resource_ledger.snapshot(),
+            }
+            if report["abort_check"]["stale_output"] or not report["abort_check"]["worker_exited"]:
+                raise RuntimeError("CPU TTS cancellation left stale output or live worker")
+        report["status"] = "passed"
+    except BaseException as exc:
+        report["status"] = "failed"
+        report["error"] = f"{type(exc).__name__}: {exc}"
+        raise
+    finally:
+        stop_sampling.set()
+        if sample_task is not None:
+            await sample_task
+        runtime.shutdown()
+        if runtime.resource_ledger is not None:
+            report["ledger_after_shutdown"] = runtime.resource_ledger.snapshot()
+        destination = Path(args.output_report)
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        destination.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
+        print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

benchmarks/edge_harness/experiments/probe_omni_qwen_tts_cpu_entrypoint.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""Verify public AsyncOmni WAV delivery from the isolated CPU Qwen3-TTS worker."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import hashlib
+import json
+import tempfile
+import time
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    for name in (
+        "python_bin", "python_sha256", "model_dir", "overlay_dir", "talker_sha256",
+        "tokenizer_sha256", "server_log", "output_report",
+    ):
+        parser.add_argument("--" + name.replace("_", "-"), required=True)
+    args = parser.parse_args()
+
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_CPU_WHOLE_PIPELINE
+
+    backend = {
+        "name": "external.qwen_tts.cpu.v1",
+        "python_bin": args.python_bin,
+        "python_sha256": args.python_sha256,
+        "model_dir": args.model_dir,
+        "overlay_dir": args.overlay_dir,
+        "talker_sha256": args.talker_sha256,
+        "tokenizer_sha256": args.tokenizer_sha256,
+        "log_file": args.server_log,
+        "memory_overhead_bytes": 6 << 30,
+        "expected_torch": "2.13.0+cu130",
+        "max_text_bytes": 4096,
+        "max_wav_bytes": 4 << 20,
+        "max_audio_s": 10,
+        "request_timeout_s": 180,
+        "start_timeout_s": 180,
+    }
+    report = {"entrypoint": "AsyncOmni.generate", "pipeline": QWEN3_TTS_CPU_WHOLE_PIPELINE.model_type}
+    engine = None
+    try:
+        with tempfile.TemporaryDirectory(prefix="omni-qwen-cpu-tts-") as temporary:
+            deployment = Path(temporary) / "deploy.yaml"
+            deployment.write_text(yaml.safe_dump({
+                "pipeline": QWEN3_TTS_CPU_WHOLE_PIPELINE.model_type,
+                "async_chunk": False,
+                "stages": [{
+                    "stage_id": 0,
+                    "backend": backend,
+                    "resource_budget": {
+                        "capacities": {"host_ram": 16 << 30},
+                        "demands": {"host_ram": 10 << 30},
+                    },
+                }],
+            }), encoding="utf-8")
+            started = time.perf_counter()
+            engine = AsyncOmni(
+                model=str(Path(args.model_dir).resolve(strict=True)),
+                deploy_config=str(deployment), stage_init_timeout=180, init_timeout=240,
+            )
+            report["startup_s"] = time.perf_counter() - started
+            outputs = []
+            started = time.perf_counter()
+            async for output in engine.generate(
+                {"text": "Hello from the local computer.", "voice": "Ryan", "seed": 42},
+                request_id="public-cpu-tts-1",
+            ):
+                if output.error:
+                    raise RuntimeError(output.error)
+                audio = output.multimodal_output["audio"]
+                pcm = np.rint(audio.numpy() * 32768).astype("<i2").tobytes()
+                outputs.append({
+                    "request_id": output.request_id,
+                    "sample_rate": output.multimodal_output["sr"],
+                    "frames": len(audio),
+                    "pcm_sha256": hashlib.sha256(pcm).hexdigest(),
+                    "stage_event": output.custom_output.get("stage_event"),
+                })
+            report["request_wall_s"] = time.perf_counter() - started
+            report["outputs"] = outputs
+            assert len(outputs) == 1 and outputs[0]["request_id"] == "public-cpu-tts-1"
+            assert outputs[0]["sample_rate"] == 24000 and outputs[0]["frames"] == 109440
+            assert outputs[0]["pcm_sha256"] == "c8a49b5c73efd642a1eb04be973d6aefdc8ed9cfd18b34ff67690b8464906bed"
+            report["status"] = "passed"
+    except BaseException as exc:
+        report["status"] = "failed"
+        report["error"] = f"{type(exc).__name__}: {exc}"
+        raise
+    finally:
+        if engine is not None:
+            engine.shutdown()
+        destination = Path(args.output_report)
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        destination.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
+        print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    asyncio.run(main())