Files
vllm-omni/vllm_omni/engine/backends/minicpmo_cpp.py
T

497 lines
23 KiB
Python

# 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)