Add Omni MiniCPM-o GGUF whole-session CPU and Radeon paths

vllm_omni/engine/backends/minicpmo_cpp.py

@@ -0,0 +1,496 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Bounded whole-session MiniCPM-o GGUF stage in Omni's graph control path.
+
+The C++ worker owns weights, KV, TTS and vocoder state for one complete request.
+Omni owns artifact verification, shared-memory admission, request identity,
+terminal output acknowledgement and process cancellation. This stage has no
+incremental playback or state handoff contract.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+import hashlib
+import io
+import os
+import re
+import shutil
+import subprocess
+import tempfile
+import threading
+import time
+import uuid
+import wave
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import psutil
+import torch
+from omni_stage_contracts import StageEvent, StageRequest
+from PIL import Image
+from vllm.outputs import CompletionOutput
+
+from vllm_omni.engine.resource_ledger import ResourceUnavailable
+from vllm_omni.engine.stage_client import StageClientBase
+from vllm_omni.outputs import OmniRequestOutput
+
+_ARTIFACTS = (
+    "MiniCPM-o-4_5-Q4_K_M.gguf",
+    "audio/MiniCPM-o-4_5-audio-F16.gguf",
+    "vision/MiniCPM-o-4_5-vision-F16.gguf",
+    "tts/MiniCPM-o-4_5-tts-F16.gguf",
+    "tts/MiniCPM-o-4_5-projector-F16.gguf",
+    "token2wav-gguf/encoder.gguf",
+    "token2wav-gguf/flow_matching.gguf",
+    "token2wav-gguf/flow_extra.gguf",
+    "token2wav-gguf/hifigan2.gguf",
+    "token2wav-gguf/prompt_cache.gguf",
+)
+_WAV_NAME = re.compile(r"wav_(\d+)\.wav")
+_CHUNK_NAME = re.compile(r"chunk_(\d+)")
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as source:
+        for block in iter(lambda: source.read(8 << 20), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def _wav_format(data: bytes, *, sample_rate: int, max_frames: int) -> int:
+    with wave.open(io.BytesIO(data), "rb") as source:
+        fmt = (source.getnchannels(), source.getsampwidth(), source.getframerate())
+        frames = source.getnframes()
+        if fmt != (1, 2, sample_rate) or not 0 < frames <= max_frames:
+            raise ValueError(f"expected bounded mono PCM16 {sample_rate} Hz WAV")
+        if len(source.readframes(frames)) != frames * 2:
+            raise ValueError("truncated input WAV")
+    return frames
+
+
+def _jpeg_size(data: bytes) -> tuple[int, int]:
+    with Image.open(io.BytesIO(data)) as image:
+        if image.format != "JPEG":
+            raise ValueError("MiniCPM-o image input must be JPEG")
+        size = image.size
+        image.verify()
+    return size
+
+
+def _tree_bytes(path: Path) -> int:
+    total = 0
+    for item in path.rglob("*"):
+        try:
+            if item.is_file():
+                total += item.stat().st_size
+        except FileNotFoundError:
+            pass  # The worker may atomically replace an output while we sample.
+    return total
+
+
+def _read_result(work_dir: Path, max_audio_s: float, max_wav_bytes: int,
+                 max_text_bytes: int) -> tuple[str, bytes, int]:
+    round_dir = work_dir / "tools" / "omni" / "output" / "round_000"
+    wav_dir = round_dir / "tts_wav"
+    if not (wav_dir / "generation_done.flag").is_file():
+        raise RuntimeError("MiniCPM-o audio generation is incomplete")
+    chunks = sorted(
+        wav_dir.glob("wav_*.wav"),
+        key=lambda path: int(_WAV_NAME.fullmatch(path.name).group(1)),
+    )
+    if not chunks or len(chunks) > 256 or [item.name for item in chunks] != [
+        f"wav_{index}.wav" for index in range(len(chunks))
+    ]:
+        raise RuntimeError("MiniCPM-o WAV chunk sequence has gaps or is unbounded")
+    pcm_parts: list[bytes] = []
+    frames = 0
+    for path in chunks:
+        with wave.open(str(path), "rb") as source:
+            fmt = (source.getnchannels(), source.getsampwidth(), source.getframerate())
+            count = source.getnframes()
+            if fmt != (1, 2, 24000) or not 0 < count <= 24000 * 10:
+                raise RuntimeError(f"MiniCPM-o invalid WAV chunk: {path.name}")
+            pcm = source.readframes(count)
+        if len(pcm) != count * 2:
+            raise RuntimeError(f"MiniCPM-o truncated WAV chunk: {path.name}")
+        frames += count
+        if frames > max_audio_s * 24000 or frames * 2 > max_wav_bytes:
+            raise ResourceUnavailable("MiniCPM-o speech exceeds admitted output bound")
+        pcm_parts.append(pcm)
+    pcm = b"".join(pcm_parts)
+    if not np.any(np.frombuffer(pcm, dtype="<i2")):
+        raise RuntimeError("MiniCPM-o generated silent speech")
+    text_paths = sorted(
+        (round_dir / "llm_debug").glob("chunk_*/llm_text.txt"),
+        key=lambda path: int(_CHUNK_NAME.fullmatch(path.parent.name).group(1)),
+    )
+    text = "".join(path.read_text(encoding="utf-8") for path in text_paths).strip()
+    if not text or len(text.encode("utf-8")) > max_text_bytes:
+        raise ResourceUnavailable("MiniCPM-o text is empty or exceeds admitted bound")
+    return text, pcm, frames
+
+
+class MiniCPMOCppStageClient(StageClientBase):
+    def __init__(self, metadata, config: dict, ledger, reservation) -> None:
+        for name, value in vars(metadata).items():
+            setattr(self, name, value)
+        self.stage_type = "graph"
+        self._ledger, self._reservation = ledger, reservation
+        self._generation = uuid.uuid4().hex
+        self._closed = False
+        self._active: str | None = None
+        self._task: asyncio.Task | None = None
+        self._output: OmniRequestOutput | None = None
+        self._proc: subprocess.Popen | None = None
+        self._lock = threading.Lock()
+        self._epoch = 0
+        self._config = dict(config)
+        self._placement = str(config["placement"])
+        self._memory_pool = str(config.get("memory_pool", "host_ram"))
+        self._context_tokens = int(config.get("context_tokens", 2048))
+        self._max_new_tokens = int(config.get("max_new_tokens", 128))
+        self._max_input_bytes = int(config.get("max_input_bytes", 2 << 20))
+        self._max_wav_bytes = int(config.get("max_wav_bytes", 2 << 20))
+        self._max_text_bytes = int(config.get("max_text_bytes", 16 << 10))
+        self._max_work_bytes = int(config.get("max_work_bytes", 64 << 20))
+        self._max_input_audio_s = float(config.get("max_input_audio_s", 10))
+        self._max_output_audio_s = float(config.get("max_output_audio_s", 20))
+        self._request_timeout_s = float(config.get("request_timeout_s", 300))
+        self._t2w_wait_seconds = int(config.get("t2w_wait_seconds", 180))
+        self._retain_workdir = bool(config.get("retain_workdir", False))
+        try:
+            if self._placement not in ("cpu", "radeon-hybrid"):
+                raise ValueError("MiniCPM-o placement must be cpu or radeon-hybrid")
+            if self._placement == "radeon-hybrid" and (
+                config.get("ggml_vk_visible_devices") != "1"
+                or config.get("expected_gpu_name") != "AMD Radeon(TM) 890M Graphics"
+            ):
+                raise ValueError("Radeon stage requires the pinned Vulkan device name and selector")
+            if not (
+                0 < self._max_new_tokens < self._context_tokens
+                and 0 < self._max_input_audio_s <= 30
+                and 0 < self._max_output_audio_s <= 120
+                and 0 < self._max_input_bytes <= 16 << 20
+                and 0 < self._max_wav_bytes <= self._max_work_bytes
+                and 0 < self._max_text_bytes <= self._max_work_bytes
+                and 0 < self._t2w_wait_seconds < self._request_timeout_s <= 3600
+            ):
+                raise ValueError("invalid MiniCPM-o context, media, time or output bound")
+            if self._memory_pool not in reservation.demands:
+                raise ResourceUnavailable("MiniCPM-o shared-RAM pool absent from reservation")
+            demand = int(reservation.demands[self._memory_pool])
+            if psutil.virtual_memory().available < demand:
+                raise ResourceUnavailable("available host/shared RAM below MiniCPM-o reservation")
+            self._model_dir = Path(config["model_dir"]).resolve(strict=True)
+            self._binary = Path(config["cli_bin"]).resolve(strict=True)
+            self._reference_wav = Path(config["reference_wav"]).resolve(strict=True)
+            expected = dict(config["artifact_sha256"])
+            if set(expected) != set(_ARTIFACTS):
+                raise ValueError("MiniCPM-o requires the exact ten-artifact manifest")
+            self._artifacts = {name: self._model_dir / name for name in _ARTIFACTS}
+            for name, path in self._artifacts.items():
+                if not path.is_file() or _sha256(path) != str(expected[name]).lower():
+                    raise ValueError(f"MiniCPM-o GGUF differs from declared hash: {name}")
+            for path, digest in (
+                (self._binary, str(config["cli_sha256"]).lower()),
+                (self._reference_wav, str(config["reference_sha256"]).lower()),
+            ):
+                if _sha256(path) != digest:
+                    raise ValueError(f"MiniCPM-o executable/reference differs: {path.name}")
+            _wav_format(self._reference_wav.read_bytes(), sample_rate=16000,
+                        max_frames=16000 * 10)
+            overhead = int(config["memory_overhead_bytes"])
+            total_weights = sum(path.stat().st_size for path in self._artifacts.values())
+            if overhead <= 0 or total_weights + overhead + self._max_work_bytes > demand:
+                raise ResourceUnavailable("MiniCPM-o weights, state/workspace, I/O and headroom exceed reservation")
+            help_result = subprocess.run(
+                [str(self._binary), "--help"], stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT, timeout=10, check=True,
+                creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
+            )
+            if not all(flag in help_result.stdout for flag in (
+                b"--max-new-tokens", b"--t2w-device", b"--t2w-wait-seconds"
+            )):
+                raise ValueError("MiniCPM-o CLI lacks bounded-generation or completion support")
+            self._work_root = Path(config["work_root"]).resolve()
+            self._log_root = Path(config["log_root"]).resolve()
+            self._work_root.mkdir(parents=True, exist_ok=True)
+            self._log_root.mkdir(parents=True, exist_ok=True)
+            self.execution_plan = {
+                "backend": "external.minicpmo.gguf.v1",
+                "stage_id": self.stage_id,
+                "worker_generation": self._generation,
+                "model_revision": str(config["model_revision"]),
+                "artifact_sha256": {key: str(value).lower() for key, value in expected.items()},
+                "cli_sha256": str(config["cli_sha256"]).lower(),
+                "requested_placement": self._placement,
+                "ggml_vk_visible_devices": config.get("ggml_vk_visible_devices"),
+                "shared_ram_reservation": dict(reservation.demands),
+                "weights_bytes": total_weights,
+                "memory_overhead_bytes": overhead,
+                "context_tokens": self._context_tokens,
+                "max_new_tokens": self._max_new_tokens,
+                "request_capacity": 1,
+                "stateful_session": "C++-owned within one request; discarded after completion",
+                "placement_evidence": "pending first request",
+            }
+        except BaseException:
+            self.shutdown()
+            raise
+
+    async def add_request_async(self, request_id: str, prompt: Any, params: Any = None) -> None:
+        self.check_health()
+        if self._active is not None:
+            raise ResourceUnavailable("MiniCPM-o stage has an unacknowledged request; capacity is one")
+        if not isinstance(prompt, dict):
+            raise ValueError("MiniCPM-o request must contain audio_wav and image_jpeg bytes")
+        audio, image = prompt.get("audio_wav"), prompt.get("image_jpeg")
+        if not isinstance(audio, (bytes, bytearray)) or not isinstance(image, (bytes, bytearray)):
+            raise ValueError("MiniCPM-o media must use bounded byte buffers")
+        audio, image = bytes(audio), bytes(image)
+        if not 0 < len(audio) <= self._max_input_bytes or not 0 < len(image) <= self._max_input_bytes:
+            raise ResourceUnavailable("MiniCPM-o input exceeds admitted transfer bound")
+        _wav_format(audio, sample_rate=16000,
+                    max_frames=int(self._max_input_audio_s * 16000))
+        if _jpeg_size(image) != (448, 448):
+            raise ValueError("MiniCPM-o v1 stage accepts only a 448x448 JPEG")
+        self._epoch += 1
+        request = StageRequest(request_id, self.stage_id, self._epoch, self._generation)
+        self._active = request_id
+        self._task = asyncio.create_task(
+            self._run(request, audio, image), name=f"minicpmo-{request_id}"
+        )
+
+    async def _run(self, request: StageRequest, audio: bytes, image: bytes) -> None:
+        try:
+            result = await asyncio.to_thread(self._execute, request, audio, image)
+            samples = np.frombuffer(result["pcm"], dtype="<i2").astype(np.float32) / 32768.0
+            completion = CompletionOutput(0, result["text"], [], None, None, finish_reason="stop")
+            completion.multimodal_output = {"audio": torch.from_numpy(samples), "sr": 24000}
+            event = StageEvent(
+                request.request_id, self.stage_id, request.epoch, 1, "audio",
+                self._generation, terminal=True,
+            )
+            output = OmniRequestOutput(
+                request_id=request.request_id,
+                prompt="[local audio+image]",
+                stage_id=self.stage_id,
+                final_output_type="audio",
+                outputs=[completion],
+                _custom_output={
+                    "stage_event": dataclasses.asdict(event),
+                    "audio_metadata": {
+                        "sample_rate_hz": 24000,
+                        "channels": 1,
+                        "pcm_frames": result["frames"],
+                        "duration_s": result["frames"] / 24000,
+                        "pcm_sha256": hashlib.sha256(result["pcm"]).hexdigest(),
+                    },
+                    "input_sha256": {
+                        "audio": hashlib.sha256(audio).hexdigest(),
+                        "image": hashlib.sha256(image).hexdigest(),
+                    },
+                },
+                metrics={
+                    "minicpmo_whole_request_wall_s": result["wall_s"],
+                    "peak_sampled_rss_bytes": result["peak_rss"],
+                    "worker_log": str(result["log_path"]),
+                },
+            )
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            output = OmniRequestOutput.from_error(request.request_id, str(exc))
+            output.stage_id = self.stage_id
+        if not self._closed and self._epoch == request.epoch:
+            loop = asyncio.get_running_loop()
+            output._stage_release = lambda: loop.call_soon_threadsafe(
+                self.acknowledge, request.request_id, request.epoch, request.worker_generation
+            )
+            self._output = output
+
+    def _execute(self, request: StageRequest, audio: bytes, image: bytes) -> dict:
+        work_dir = Path(tempfile.mkdtemp(prefix=f"omni-minicpmo-{request.epoch}-", dir=self._work_root))
+        if not work_dir.resolve().is_relative_to(self._work_root):
+            raise RuntimeError("MiniCPM-o work directory escaped its declared root")
+        log_path = self._log_root / f"minicpmo-{request.epoch}-{uuid.uuid4().hex}.log"
+        started = time.perf_counter()
+        peak_rss = 0
+        try:
+            shutil.copyfile(self._reference_wav, work_dir / "fixture_0000.wav")
+            (work_dir / "fixture_0001.wav").write_bytes(audio)
+            (work_dir / "fixture_0001.jpg").write_bytes(image)
+            command = [
+                str(self._binary), "-m", str(self._artifacts[_ARTIFACTS[0]]),
+                "-ngl", "0" if self._placement == "cpu" else "99",
+                "-c", str(self._context_tokens),
+                "--max-new-tokens", str(self._max_new_tokens),
+                "--omni", "--t2w-device", "cpu",
+                "--t2w-wait-seconds", str(self._t2w_wait_seconds),
+                "--ref-audio", str(self._reference_wav),
+                "--test", str(work_dir / "fixture_"), "2",
+            ]
+            env = os.environ.copy()
+            env["GGML_VK_VISIBLE_DEVICES"] = (
+                "" if self._placement == "cpu" else "1"
+            )
+            with log_path.open("wb") as log:
+                with self._lock:
+                    if self._closed or self._epoch != request.epoch:
+                        raise RuntimeError("MiniCPM-o request was cancelled before launch")
+                    self._proc = subprocess.Popen(
+                        command, cwd=work_dir, stdout=log, stderr=subprocess.STDOUT,
+                        env=env,
+                        creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
+                    )
+                    proc = self._proc
+                process = psutil.Process(proc.pid)
+                deadline = time.monotonic() + self._request_timeout_s
+                while proc.poll() is None:
+                    if time.monotonic() >= deadline:
+                        self._kill_process(proc)
+                        raise TimeoutError("MiniCPM-o complete request timed out")
+                    if log_path.stat().st_size + _tree_bytes(work_dir) > self._max_work_bytes:
+                        self._kill_process(proc)
+                        raise ResourceUnavailable("MiniCPM-o worker exceeded admitted file/I/O bound")
+                    try:
+                        peak_rss = max(peak_rss, process.memory_info().rss)
+                    except psutil.NoSuchProcess:
+                        pass
+                    if peak_rss > self._reservation.demands[self._memory_pool]:
+                        self._kill_process(proc)
+                        raise ResourceUnavailable("MiniCPM-o worker RSS exceeded shared-RAM reservation")
+                    try:
+                        proc.wait(timeout=0.1)
+                    except subprocess.TimeoutExpired:
+                        pass
+                if proc.returncode != 0:
+                    raise RuntimeError(f"MiniCPM-o CLI exited {proc.returncode}; see {log_path}")
+            log_text = log_path.read_text(encoding="utf-8", errors="replace")
+            if self._placement == "cpu":
+                required = (
+                    "GPU layers: 0", "vision_ctx: vision using CPU backend",
+                    "Token2Wav device: cpu",
+                )
+            else:
+                required = (
+                    "GPU layers: 99",
+                    "using device Vulkan0 (AMD Radeon(TM) 890M Graphics)",
+                    "offloaded 37/37 layers to GPU",
+                    "vision_ctx: vision using Vulkan0 backend",
+                    "TTS model: loading with n_gpu_layers=0",
+                    "Token2Wav: non-CUDA backend, vocoder using CPU",
+                )
+            if any(item not in log_text for item in required):
+                raise RuntimeError("REFUSE_DEVICE_PLACEMENT: MiniCPM-o did not use declared CPU/Radeon route")
+            if not all(item in log_text for item in (
+                "stream_prefill(index=1): processing user audio:",
+                "stream_prefill: prefilled 1 vision chunks",
+                "Audio generation completed.",
+            )):
+                raise RuntimeError("MiniCPM-o complete input/output path not observed")
+            text, pcm, frames = _read_result(
+                work_dir, self._max_output_audio_s,
+                self._max_wav_bytes, self._max_text_bytes,
+            )
+            self.execution_plan["placement_evidence"] = {
+                "observed": self._placement,
+                "worker_pid": proc.pid,
+                "peak_sampled_rss_bytes": peak_rss,
+                "worker_log": str(log_path),
+            }
+            return {
+                "text": text, "pcm": pcm, "frames": frames,
+                "peak_rss": peak_rss, "wall_s": time.perf_counter() - started,
+                "log_path": log_path,
+            }
+        finally:
+            with self._lock:
+                proc = self._proc
+            if proc is not None:
+                self._kill_process(proc)
+                with self._lock:
+                    if self._proc is proc:
+                        self._proc = None
+            if not self._retain_workdir:
+                if not work_dir.resolve().is_relative_to(self._work_root):
+                    raise RuntimeError("refusing cleanup outside MiniCPM-o work root")
+                shutil.rmtree(work_dir)
+
+    @staticmethod
+    def _kill_process(proc: subprocess.Popen) -> bool:
+        try:
+            if proc.poll() is None:
+                proc.terminate()
+                try:
+                    proc.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    proc.kill()
+                    proc.wait(timeout=5)
+        except OSError:
+            return False
+        return proc.poll() is not None
+
+    def get_graph_output_nowait(self) -> OmniRequestOutput | None:
+        output, self._output = self._output, None
+        return output
+
+    def acknowledge(self, request_id: str, epoch: int, generation: str) -> None:
+        if (
+            self._active == request_id and epoch == self._epoch
+            and generation == self._generation
+            and (self._task is None or self._task.done())
+        ):
+            self._active = None
+            self._task = None
+
+    async def abort_requests_async(self, request_ids: list[str]) -> None:
+        if self._active not in request_ids:
+            return
+        self._epoch += 1
+        self._output = None
+        with self._lock:
+            self._closed = True
+            proc = self._proc
+        drained = await asyncio.to_thread(self._kill_process, proc) if proc else True
+        if self._task is not None and not self._task.done():
+            done, _ = await asyncio.wait({self._task}, timeout=5)
+            if done:
+                self._ledger.release(self._reservation, drained=drained)
+            else:
+                self._ledger.release(self._reservation, drained=False)
+                self._task.add_done_callback(
+                    lambda task: self._ledger.release(
+                        self._reservation, drained=drained and not task.cancelled()
+                    )
+                )
+        self._active = None
+
+    def check_health(self) -> None:
+        if self._closed:
+            from vllm.v1.engine.exceptions import EngineDeadError
+
+            raise EngineDeadError()
+
+    async def collective_rpc_async(self, method, timeout=None, args=(), kwargs=None):
+        raise NotImplementedError(f"MiniCPM-o backend does not implement collective RPC {method}")
+
+    def shutdown(self) -> None:
+        with self._lock:
+            self._closed = True
+            proc = self._proc
+        drained = self._kill_process(proc) if proc else True
+        self._output = None
+        task = self._task
+        if task is not None and not task.done():
+            self._ledger.release(self._reservation, drained=False)
+            task.add_done_callback(
+                lambda done: self._ledger.release(
+                    self._reservation, drained=drained and not done.cancelled()
+                )
+            )
+        else:
+            self._ledger.release(self._reservation, drained=drained)