diff --git a/examples/offline_inference/mage_vl/end2end.py b/examples/offline_inference/mage_vl/end2end.py new file mode 100644 index 000000000..4a79a9f99 --- /dev/null +++ b/examples/offline_inference/mage_vl/end2end.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Run Mage-VL offline end-to-end inference through the upstream Mage scripts. + +Mage-VL currently ships its custom Mage-ViT and codec processors as +Transformers remote code in the Microsoft Mage repository. This example keeps +the vLLM-Omni entrypoint stable while delegating the checkpoint-specific forward +path to that upstream implementation. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from importlib.util import find_spec +from pathlib import Path + +DEFAULT_MODEL = "microsoft/Mage-VL" +DEFAULT_IMAGE_QUESTION = "Describe this image in detail." +DEFAULT_VIDEO_QUESTION = "Describe this video." + + +def _default_mage_repo() -> Path | None: + value = os.getenv("MAGE_REPO_DIR") + if value: + return Path(value) + sibling = Path(__file__).resolve().parents[3].parent / "Mage" + if sibling.exists(): + return sibling + return None + + +def _resolve_mage_repo(value: str | None) -> Path: + repo = Path(value) if value else _default_mage_repo() + if repo is None: + raise ValueError("Set --mage-repo or MAGE_REPO_DIR to a checkout of https://github.com/microsoft/Mage.") + repo = repo.expanduser().resolve() + if not (repo / "mage_vl" / "inference_base.py").is_file(): + raise FileNotFoundError(f"Missing Mage-VL scripts under {repo}. Expected mage_vl/inference_base.py.") + return repo + + +def _default_media_path(mage_repo: Path, task: str) -> Path: + if task == "image": + return mage_repo / "mage_vl" / "assets" / "examples" / "dog.jpg" + return mage_repo / "mage_vl" / "assets" / "examples" / "soccer-broadcast.mp4" + + +def _resolve_media_path(path: str | None, mage_repo: Path, task: str) -> Path: + resolved = Path(path).expanduser().resolve() if path else _default_media_path(mage_repo, task) + if not resolved.is_file(): + raise FileNotFoundError(f"Input media file not found: {resolved}") + return resolved + + +def build_command(args: argparse.Namespace, mage_repo: Path, media_path: Path) -> list[str]: + if args.task == "streaming": + return [ + sys.executable, + str(mage_repo / "mage_vl" / "inference_streaming.py"), + "--video", + str(media_path), + "--checkpoint", + args.model, + "--video_backend", + args.video_backend, + "--num_frames", + str(args.num_frames), + "--cur_fps", + str(args.cur_fps), + "--segment_sec", + str(args.segment_sec), + "--max_new_tokens", + str(args.max_new_tokens), + "--max_segments", + str(args.max_segments), + "--gate_threshold", + str(args.gate_threshold), + "--device", + args.device, + "--attn_impl", + args.attn_impl, + ] + + cmd = [ + sys.executable, + str(mage_repo / "mage_vl" / "inference_base.py"), + "--mode", + "offline", + "--question", + args.question or (DEFAULT_IMAGE_QUESTION if args.task == "image" else DEFAULT_VIDEO_QUESTION), + "--model", + args.model, + "--max-new-tokens", + str(args.max_new_tokens), + ] + if args.task == "image": + cmd.extend(["--image", str(media_path)]) + return cmd + + cmd.extend( + [ + "--video", + str(media_path), + "--video-backend", + args.video_backend, + "--num-frames", + str(args.num_frames), + "--max-pixels", + str(args.max_pixels), + ] + ) + if args.video_backend == "codec": + cmd.extend(["--codec-engine", args.codec_engine]) + return cmd + + +def check_runtime_dependencies(task: str) -> None: + required = ["torch"] + if task == "streaming": + required.append("mamba_ssm") + missing = [name for name in required if find_spec(name) is None] + if not missing: + return + + install_hint = ( + "Install Mage-VL runtime dependencies in this environment first:\n" + " pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n" + " grep -v '^mamba-ssm' /path/to/Mage/mage_vl/requirements.txt > /tmp/mage_vl_requirements_no_mamba.txt\n" + " pip install -r /tmp/mage_vl_requirements_no_mamba.txt\n" + " pip install 'causal-conv1d>=1.4.0' --no-build-isolation\n" + " pip install 'mamba-ssm>=2.2' --no-build-isolation" + ) + raise RuntimeError(f"Missing Python packages: {', '.join(missing)}.\n{install_hint}") + + +def run_command(cmd: list[str], cwd: Path, timeout: int) -> subprocess.CompletedProcess[str]: + return subprocess.run( + cmd, + cwd=str(cwd), + input="y\n", + text=True, + capture_output=True, + timeout=timeout, + check=False, + ) + + +def write_outputs( + *, + args: argparse.Namespace, + mage_repo: Path, + media_path: Path, + cmd: list[str], + result: subprocess.CompletedProcess[str] | None, + elapsed_s: float | None, +) -> Path: + output_dir = Path(args.output_dir).expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + summary = { + "task": args.task, + "model": args.model, + "mage_repo": str(mage_repo), + "media_path": str(media_path), + "command": cmd, + "dry_run": args.dry_run, + "elapsed_s": elapsed_s, + "returncode": None if result is None else result.returncode, + } + if result is not None: + stdout = result.stdout.strip() + stderr = result.stderr.strip() + summary["stdout"] = stdout + summary["stderr"] = stderr + (output_dir / "stdout.txt").write_text(result.stdout, encoding="utf-8") + (output_dir / "stderr.txt").write_text(result.stderr, encoding="utf-8") + summary_path = output_dir / "summary.json" + summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + return summary_path + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Mage-VL offline end-to-end inference example.") + parser.add_argument( + "--mage-repo", + default=None, + help="Path to a Microsoft Mage checkout. Defaults to MAGE_REPO_DIR.", + ) + parser.add_argument("--model", default=DEFAULT_MODEL, help="HF model id or local checkpoint path.") + parser.add_argument("--task", choices=("image", "video", "streaming"), default="image") + parser.add_argument("--image-path", default=None, help="Image input. Defaults to the Mage sample image.") + parser.add_argument("--video-path", default=None, help="Video input. Defaults to the Mage sample video.") + parser.add_argument("--question", default=None, help="Question for image/video offline QA.") + parser.add_argument("--video-backend", choices=("frames", "codec"), default="frames") + parser.add_argument("--codec-engine", choices=("traditional", "neural"), default="traditional") + parser.add_argument("--num-frames", type=int, default=32) + parser.add_argument("--max-pixels", type=int, default=150000) + parser.add_argument("--max-new-tokens", type=int, default=256) + parser.add_argument("--cur-fps", type=float, default=2.0) + parser.add_argument("--segment-sec", type=float, default=8.0) + parser.add_argument("--max-segments", type=int, default=0) + parser.add_argument("--gate-threshold", type=float, default=0.5) + parser.add_argument("--device", default="cuda") + parser.add_argument("--attn-impl", choices=("flash_attention_2", "sdpa", "eager"), default="flash_attention_2") + parser.add_argument("--timeout", type=int, default=1800) + parser.add_argument("--output-dir", default="outputs/mage_vl") + parser.add_argument("--dry-run", action="store_true", help="Build the command and summary without executing it.") + args = parser.parse_args() + if args.num_frames <= 0: + parser.error("--num-frames must be positive") + if args.segment_sec <= 0: + parser.error("--segment-sec must be positive") + if args.task == "image" and args.video_path: + parser.error("--video-path cannot be used with --task image") + if args.task in {"video", "streaming"} and args.image_path: + parser.error("--image-path cannot be used with video tasks") + return args + + +def main() -> None: + args = parse_args() + mage_repo = _resolve_mage_repo(args.mage_repo) + media_arg = args.image_path if args.task == "image" else args.video_path + media_path = _resolve_media_path(media_arg, mage_repo, args.task) + cmd = build_command(args, mage_repo, media_path) + + result = None + elapsed_s = None + if not args.dry_run: + check_runtime_dependencies(args.task) + start = time.perf_counter() + result = run_command(cmd, cwd=mage_repo, timeout=args.timeout) + elapsed_s = time.perf_counter() - start + + summary_path = write_outputs( + args=args, + mage_repo=mage_repo, + media_path=media_path, + cmd=cmd, + result=result, + elapsed_s=elapsed_s, + ) + print(f"summary saved to {summary_path}") + if result is not None: + if result.stdout: + print(result.stdout, end="") + if result.returncode != 0: + if result.stderr: + print(result.stderr, file=sys.stderr, end="") + raise SystemExit(result.returncode) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/mage_vl/duplex_client.py b/examples/online_serving/mage_vl/duplex_client.py new file mode 100644 index 000000000..283b7038c --- /dev/null +++ b/examples/online_serving/mage_vl/duplex_client.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import json +import time +from pathlib import Path + +import websockets + + +async def run(args: argparse.Namespace) -> None: + headers = {"Authorization": f"Bearer {args.auth_token}"} if args.auth_token else None + started = time.perf_counter() + async with websockets.connect( + f"{args.url}?session_id={args.session_id}", + additional_headers=headers, + max_size=args.max_message_bytes, + open_timeout=args.timeout, + ) as websocket: + created = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout)) + print(json.dumps(created, ensure_ascii=False)) + await websocket.send( + json.dumps( + { + "type": "input.append", + "modality": "text", + "data": args.prompt, + } + ) + ) + encoded = base64.b64encode(Path(args.video).read_bytes()).decode("ascii") + await websocket.send( + json.dumps( + { + "type": "input.append", + "modality": "video", + "data": { + "kind": "video/mp4", + "segment_id": args.segment_id, + "pts_ms": args.pts_ms, + "duration_ms": args.duration_ms, + "data": {"video_base64": encoded}, + }, + } + ) + ) + while True: + event = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout)) + print(json.dumps(event, ensure_ascii=False)) + if event.get("type") in {"response.done", "error"}: + break + await websocket.send(json.dumps({"type": "close"})) + print(f"elapsed_s={time.perf_counter() - started:.3f}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Send one Mage-VL video window over the duplex WebSocket.") + parser.add_argument("--url", default="ws://127.0.0.1:8090/v1/mage-vl/duplex") + parser.add_argument("--video", required=True) + parser.add_argument("--prompt", default="Describe this video segment in detail.") + parser.add_argument("--session-id", default="mage-vl-example") + parser.add_argument("--segment-id", default="segment-0") + parser.add_argument("--pts-ms", type=int, default=0) + parser.add_argument("--duration-ms", type=int, default=4000) + parser.add_argument("--auth-token") + parser.add_argument("--timeout", type=float, default=120.0) + parser.add_argument("--max-message-bytes", type=int, default=16 * 1024 * 1024) + args = parser.parse_args() + asyncio.run(run(args)) + + +if __name__ == "__main__": + main() diff --git a/recipes/Microsoft/Mage-VL.md b/recipes/Microsoft/Mage-VL.md new file mode 100644 index 000000000..6963adec1 --- /dev/null +++ b/recipes/Microsoft/Mage-VL.md @@ -0,0 +1,114 @@ +# Mage-VL + +> Codec-native image/video understanding plus proactive full-duplex streaming. + +## Summary + +- Vendor: Microsoft +- Model: [`microsoft/Mage-VL`](https://huggingface.co/microsoft/Mage-VL) +- Task: Image understanding, video understanding, codec-native video understanding, + and event-gated streaming commentary +- Mode: Full-duplex WebSocket server and session adapter in + `vllm_omni/experimental/fullduplex/mage_vl` +- Maintainer: Community + +## Status + +This experimental integration provides the Mage-VL adapter contract: bounded causal +visual windows, codec-window input, rolling gate evaluation, explicit query handling, +per-session state isolation, and stale-output cancellation through the shared +full-duplex runtime. + +The production transport exposes `WebSocket /v1/mage-vl/duplex`, with bearer-token +authentication, bounded concurrent sessions and message sizes, idle timeouts, +disconnect cancellation, per-session state cleanup, and serialized access to the +shared Transformers checkpoint. Native vLLM model execution remains follow-up work. + +The default `frames` backend uses decoded frames and has no codec preprocessing +requirement. The optional `codec` backend requires `ffmpeg` and `ffprobe`, plus the +`cv-preinfer` executable provided by `codec-video-prep>=0.2.5`, on `PATH`. Server +startup fails with an actionable error when `cv-preinfer` is unavailable; use +`--video-backend frames` when codec preprocessing is not installed. + +## Serving + +```bash +python -m vllm_omni.experimental.fullduplex.mage_vl.serving.server \ + --model microsoft/Mage-VL --host 0.0.0.0 --port 8090 \ + --auth-token "$MAGE_VL_API_KEY" --max-sessions 32 \ + --video-backend frames --num-frames 8 --target-fps 1 \ + --gate-threshold 0.5 --attn-impl sdpa + +python examples/online_serving/mage_vl/duplex_client.py \ + --url ws://127.0.0.1:8090/v1/mage-vl/duplex \ + --auth-token "$MAGE_VL_API_KEY" \ + --video /path/to/segment.mp4 \ + --prompt "Describe this video segment in detail." +``` + +Each input window is a JSON `input.append` event whose `data.data.video_base64` +contains an MP4 segment. A text `input.append` before the video requests an explicit +answer; without text, the StreamMind gate decides whether to respond proactively. +The server emits `session.created`, `response.created`, one or more +`response.delta`, and `response.done` events. Clients can send `response.cancel` for +barge-in and `close` for graceful shutdown. + +The adapter can still be wired to another backend through `--adapter module:factory`, +including: + +- a Transformers `trust_remote_code=True` backend for offline/local testing; +- the model author's SGLang-compatible serving path; +- a future native vLLM `mage_vl` backend. + +## Adapter Shape + +```python +from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig +from vllm_omni.experimental.fullduplex.mage_vl import MageVLDuplexAdapter, MageVLDuplexRuntime + + +adapter = MageVLDuplexAdapter(gate=run_mage_gate, generate=run_mage_decoder) +session = DuplexSession( + "stream-1", + DuplexSessionConfig( + input_modalities=("codec_window", "video", "image", "text"), + output_modalities=("text",), + proactive=True, + ), +) +runtime = MageVLDuplexRuntime(session, adapter) +await runtime.run(input_events, emit_event) +``` + +Input events use the shared full-duplex protocol. For decoded video/image, pass the frame +payload directly. For codec-native streaming, pass `modality="codec_window"` and include +fields such as `segment_id`, `pts_ms`, `duration_ms`, `codec`, and `metadata`. + +```python +{ + "type": "input.append", + "modality": "codec_window", + "data": { + "kind": "h264", + "segment_id": "seg-42", + "pts_ms": 42000, + "duration_ms": 1000, + "codec": { + "motion_vectors": ..., + "residual_energy": ..., + }, + "metadata": {"gop": "P"}, + }, +} +``` + +The gate callable receives the current rolling window and returns either a +`MageVLGateDecision`, a mapping with `should_respond` / `text` / `event_id`, a bool, or +a text string. Repeated `event_id` values are de-duplicated at the session adapter. + +## Verification + +```bash +pytest tests/e2e/features/fullduplex/test_mage_vl_adapter.py \ + tests/examples/offline_inference/test_mage_vl.py +``` diff --git a/tests/e2e/features/fullduplex/test_mage_vl_adapter.py b/tests/e2e/features/fullduplex/test_mage_vl_adapter.py new file mode 100644 index 000000000..4c23536b4 --- /dev/null +++ b/tests/e2e/features/fullduplex/test_mage_vl_adapter.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import pytest + +from vllm_omni.experimental.fullduplex.core import protocol as ev +from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig +from vllm_omni.experimental.fullduplex.mage_vl import ( + MageVLCodecWindow, + MageVLDuplexAdapter, + MageVLDuplexRuntime, + MageVLGateDecision, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +async def _feed(events): + for event in events: + yield event + + +def _collector(): + out: list[dict] = [] + + async def emit(event: dict) -> None: + out.append(event) + + return out, emit + + +def test_mage_vl_capabilities_match_backend_contract(): + capabilities = MageVLDuplexAdapter().capabilities() + + assert capabilities.input_modalities == frozenset({"text", "video"}) + assert capabilities.output_modalities == frozenset({"text"}) + assert capabilities.proactive + + +@pytest.mark.asyncio +async def test_mage_vl_proactive_gate_speaks_once_for_event(): + async def gate(session, windows): + del session + if windows[-1].segment_id == "goal": + return MageVLGateDecision(True, text="Goal scored.", event_id="goal-1", score=0.91) + return MageVLGateDecision(False) + + adapter = MageVLDuplexAdapter(gate=gate, window_size=1) + session = DuplexSession( + "mage", + DuplexSessionConfig( + input_modalities=("video", "text"), + output_modalities=("text",), + proactive=True, + ), + ) + rt = MageVLDuplexRuntime(session, adapter) + out, emit = _collector() + + await rt.run( + _feed( + [ + {"type": ev.INPUT_APPEND, "modality": "video", "data": {"segment_id": "warmup", "frames": ["f0"]}}, + {"type": ev.INPUT_APPEND, "modality": "video", "data": {"segment_id": "goal", "frames": ["f1"]}}, + {"type": ev.INPUT_APPEND, "modality": "video", "data": {"segment_id": "goal", "frames": ["f2"]}}, + {"type": ev.CLOSE}, + ] + ), + emit, + ) + + deltas = [event["data"] for event in out if event["type"] == ev.RESPONSE_DELTA] + assert deltas == ["Goal scored."] + + +@pytest.mark.asyncio +async def test_mage_vl_query_bypasses_silent_gate_and_preserves_codec_metadata(): + seen: dict[str, object] = {} + + async def gate(session, windows): + del session, windows + return {"should_respond": False} + + async def generate(session, windows, query, gate_decision): + del session, gate_decision + seen["query"] = query + seen["window"] = windows[-1] + yield "The player is moving downfield." + + adapter = MageVLDuplexAdapter(gate=gate, generate=generate, window_size=1) + session = DuplexSession( + "mage", + DuplexSessionConfig( + input_modalities=("video", "codec_window", "text"), + output_modalities=("text",), + proactive=True, + ), + ) + rt = MageVLDuplexRuntime(session, adapter) + out, emit = _collector() + + await rt.run( + _feed( + [ + {"type": ev.INPUT_APPEND, "modality": "text", "data": "What changed?"}, + { + "type": ev.INPUT_APPEND, + "modality": "video", + "data": { + "kind": "h264", + "segment_id": "seg-7", + "pts_ms": 7000, + "duration_ms": 1000, + "codec": {"motion_vectors": "mv", "residual_energy": "res"}, + "metadata": {"gop": "P"}, + }, + }, + {"type": ev.CLOSE}, + ] + ), + emit, + ) + + window = seen["window"] + assert seen["query"] == "What changed?" + assert isinstance(window, MageVLCodecWindow) + assert window.kind == "h264" + assert window.segment_id == "seg-7" + assert window.pts_ms == 7000 + assert window.metadata == {"gop": "P"} + assert [event["data"] for event in out if event["type"] == ev.RESPONSE_DELTA] == ["The player is moving downfield."] + + +@pytest.mark.asyncio +async def test_mage_vl_session_state_is_isolated(): + seen: dict[str, tuple[MageVLCodecWindow, ...]] = {} + + async def generate(session, windows, query, gate_decision): + del query, gate_decision + seen[session.session_id] = tuple(windows) + yield session.session_id + + adapter = MageVLDuplexAdapter(generate=generate, window_size=1) + session_a = DuplexSession( + "mage-a", + DuplexSessionConfig(input_modalities=("video", "text"), output_modalities=("text",), proactive=True), + ) + session_b = DuplexSession( + "mage-b", + DuplexSessionConfig(input_modalities=("video", "text"), output_modalities=("text",), proactive=True), + ) + + await adapter.on_input(session_a, "video", {"segment_id": "a", "frames": ["a0"]}) + await adapter.on_input(session_b, "video", {"segment_id": "b", "frames": ["b0"]}) + await adapter.on_input(session_a, "text", "describe a") + await adapter.on_input(session_b, "text", "describe b") + assert adapter.should_respond(session_a) + assert adapter.should_respond(session_b) + + async for _ in adapter.respond(session_a): + pass + async for _ in adapter.respond(session_b): + pass + assert [window.segment_id for window in seen["mage-a"]] == ["a"] + assert [window.segment_id for window in seen["mage-b"]] == ["b"] diff --git a/tests/e2e/features/fullduplex/test_mage_vl_serving.py b/tests/e2e/features/fullduplex/test_mage_vl_serving.py new file mode 100644 index 000000000..a565a136e --- /dev/null +++ b/tests/e2e/features/fullduplex/test_mage_vl_serving.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import asyncio +import time + +import pytest +from fastapi.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +from vllm_omni.experimental.fullduplex.mage_vl import MageVLDuplexAdapter +from vllm_omni.experimental.fullduplex.mage_vl.serving import backend as backend_module +from vllm_omni.experimental.fullduplex.mage_vl.serving.server import ( + MageVLServingConfig, + create_app, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_mage_vl_codec_backend_requires_preprocessor(monkeypatch): + monkeypatch.setattr(backend_module.shutil, "which", lambda _name: None) + + with pytest.raises(RuntimeError, match="cv-preinfer"): + backend_module.MageVLTransformersBackend(video_backend="codec") + + +def test_mage_vl_frames_backend_does_not_require_preprocessor(monkeypatch): + def unexpected_lookup(_name): + raise AssertionError("frames backend must not inspect codec dependencies") + + monkeypatch.setattr(backend_module.shutil, "which", unexpected_lookup) + backend = backend_module.MageVLTransformersBackend(video_backend="frames") + + assert backend.video_backend == "frames" + + +def _factory(): + async def gate(_session, windows): + return { + "should_respond": windows[-1].segment_id == "goal", + "event_id": "goal-1", + "score": 0.9, + } + + async def generate(_session, windows, query, _decision): + yield f"{query or 'proactive'}:{windows[-1].segment_id}" + + return MageVLDuplexAdapter(gate=gate, generate=generate, window_size=1) + + +def test_mage_vl_health_ready_and_models(): + with TestClient(create_app(_factory)) as client: + assert client.get("/health").json() == {"status": "ok"} + assert client.get("/ready").json() == {"status": "ready", "active_sessions": 0} + assert client.get("/v1/models").json()["data"][0]["id"] == "microsoft/Mage-VL" + + +def test_mage_vl_websocket_runs_gate_and_generation(): + with TestClient(create_app(_factory)) as client: + with client.websocket_connect("/v1/mage-vl/duplex?session_id=test-session") as websocket: + created = websocket.receive_json() + assert created["type"] == "session.created" + assert created["session_id"] == "test-session" + websocket.send_json( + { + "type": "input.append", + "modality": "video", + "data": {"segment_id": "goal", "frames": ["frame-1"]}, + } + ) + assert websocket.receive_json()["type"] == "response.created" + delta = websocket.receive_json() + assert delta == { + "type": "response.delta", + "response_index": 1, + "modality": "text", + "data": "proactive:goal", + } + assert websocket.receive_json()["type"] == "response.done" + websocket.send_json({"type": "close"}) + + +def test_mage_vl_websocket_requires_bearer_token(): + app = create_app(_factory, MageVLServingConfig(auth_token="secret")) + with TestClient(app) as client: + with pytest.raises(WebSocketDisconnect) as denied: + with client.websocket_connect("/v1/mage-vl/duplex"): + pass + assert denied.value.code == 4401 + + with client.websocket_connect( + "/v1/mage-vl/duplex", + headers={"Authorization": "Bearer secret"}, + ) as websocket: + assert websocket.receive_json()["type"] == "session.created" + websocket.send_json({"type": "close"}) + + +def test_mage_vl_rejects_duplicate_session_and_releases_disconnect(): + app = create_app(_factory, MageVLServingConfig(max_sessions=1)) + with TestClient(app) as client: + with client.websocket_connect("/v1/mage-vl/duplex?session_id=one") as first: + assert first.receive_json()["session_id"] == "one" + with pytest.raises(WebSocketDisconnect) as denied: + with client.websocket_connect("/v1/mage-vl/duplex?session_id=two") as second: + second.receive_json() + assert denied.value.code == 4429 + first.send_json({"type": "close"}) + + with client.websocket_connect("/v1/mage-vl/duplex?session_id=two") as second: + assert second.receive_json()["session_id"] == "two" + second.send_json({"type": "close"}) + + +def test_mage_vl_rejects_malformed_event(): + with TestClient(create_app(_factory)) as client: + with client.websocket_connect("/v1/mage-vl/duplex") as websocket: + websocket.receive_json() + websocket.send_text("[]") + event = websocket.receive_json() + assert event["type"] == "error" + assert "JSON object" in event["message"] + + +def test_mage_vl_cancel_is_not_blocked_by_slow_gate(): + async def slow_gate(_session, _windows): + await asyncio.sleep(2) + return {"should_respond": True, "event_id": "late"} + + def factory(): + return MageVLDuplexAdapter(gate=slow_gate, window_size=1) + + with TestClient(create_app(factory)) as client: + with client.websocket_connect("/v1/mage-vl/duplex") as websocket: + websocket.receive_json() + websocket.send_json( + { + "type": "input.append", + "modality": "video", + "data": {"segment_id": "slow", "frames": ["frame"]}, + } + ) + started = time.monotonic() + websocket.send_json({"type": "response.cancel"}) + event = websocket.receive_json() + assert time.monotonic() - started < 0.5 + assert event["type"] == "response.cancelled" + websocket.send_json({"type": "close"}) diff --git a/tests/examples/offline_inference/test_mage_vl.py b/tests/examples/offline_inference/test_mage_vl.py new file mode 100644 index 000000000..53bd141c3 --- /dev/null +++ b/tests/examples/offline_inference/test_mage_vl.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.example] + +REPO_ROOT = Path(__file__).resolve().parents[3] +EXAMPLE_SCRIPT = REPO_ROOT / "examples" / "offline_inference" / "mage_vl" / "end2end.py" + + +def _fake_mage_repo(tmp_path: Path) -> Path: + repo = tmp_path / "Mage" + mage_vl = repo / "mage_vl" + mage_vl.mkdir(parents=True) + (mage_vl / "inference_base.py").write_text("print('base')\n", encoding="utf-8") + (mage_vl / "inference_streaming.py").write_text("print('streaming')\n", encoding="utf-8") + return repo + + +def test_mage_vl_offline_example_builds_codec_command(tmp_path: Path) -> None: + mage_repo = _fake_mage_repo(tmp_path) + video_path = tmp_path / "input.mp4" + video_path.write_bytes(b"fake") + output_dir = tmp_path / "outputs" + + subprocess.run( + [ + sys.executable, + str(EXAMPLE_SCRIPT), + "--mage-repo", + str(mage_repo), + "--task", + "video", + "--video-path", + str(video_path), + "--video-backend", + "codec", + "--codec-engine", + "traditional", + "--output-dir", + str(output_dir), + "--dry-run", + ], + check=True, + ) + + with open(output_dir / "summary.json", encoding="utf-8") as f: + summary = json.load(f) + + assert summary["task"] == "video" + assert summary["dry_run"] is True + assert "--video-backend" in summary["command"] + assert "codec" in summary["command"] + assert "--codec-engine" in summary["command"] + assert "traditional" in summary["command"] + + +def test_mage_vl_offline_example_builds_streaming_command(tmp_path: Path) -> None: + mage_repo = _fake_mage_repo(tmp_path) + video_path = tmp_path / "input.mp4" + video_path.write_bytes(b"fake") + output_dir = tmp_path / "outputs" + + subprocess.run( + [ + sys.executable, + str(EXAMPLE_SCRIPT), + "--mage-repo", + str(mage_repo), + "--task", + "streaming", + "--video-path", + str(video_path), + "--max-segments", + "1", + "--output-dir", + str(output_dir), + "--dry-run", + ], + check=True, + ) + + with open(output_dir / "summary.json", encoding="utf-8") as f: + summary = json.load(f) + + command = summary["command"] + assert summary["task"] == "streaming" + assert str(mage_repo / "mage_vl" / "inference_streaming.py") in command + assert "--checkpoint" in command + assert "--max_segments" in command diff --git a/vllm_omni/experimental/fullduplex/core/adapter.py b/vllm_omni/experimental/fullduplex/core/adapter.py index 3fe400d49..c447a0b34 100644 --- a/vllm_omni/experimental/fullduplex/core/adapter.py +++ b/vllm_omni/experimental/fullduplex/core/adapter.py @@ -41,3 +41,10 @@ class DuplexAdapter(ABC): async def on_barge_in(self, session: DuplexSession) -> None: ... async def on_playback_ack(self, session: DuplexSession, cursor: int) -> None: ... + + async def on_close(self, session: DuplexSession) -> None: + del session + + async def flush(self, session: DuplexSession) -> None: + """Wait for deferred input processing before an explicit close.""" + del session diff --git a/vllm_omni/experimental/fullduplex/core/runtime.py b/vllm_omni/experimental/fullduplex/core/runtime.py index 86ca011ef..dc6094770 100644 --- a/vllm_omni/experimental/fullduplex/core/runtime.py +++ b/vllm_omni/experimental/fullduplex/core/runtime.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -8,7 +8,7 @@ import contextlib from collections.abc import AsyncIterator, Awaitable, Callable from vllm_omni.experimental.fullduplex.core import protocol as ev -from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter +from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter, DuplexCapability from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexState Emit = Callable[[dict], Awaitable[None]] @@ -19,45 +19,97 @@ class DuplexRuntime: self.session = session self.adapter = adapter self._capabilities = adapter.capabilities() + self._response_task: asyncio.Task[None] | None = None + self._response_lock = asyncio.Lock() + self._closed = False + + @property + def capabilities(self) -> DuplexCapability: + return self._capabilities async def run(self, inputs: AsyncIterator[dict], emit: Emit) -> None: - task: asyncio.Task | None = None - async for event in inputs: - etype = event.get("type") - if etype == ev.INPUT_APPEND: - modality = event.get("modality", "") - if modality not in self._capabilities.input_modalities: - await emit(ev.error(f"unsupported input modality: {modality}")) - continue - await self.adapter.on_input(self.session, modality, event.get("data")) - self.session.state = DuplexState.LISTENING - if self.session.config.proactive and self.adapter.should_respond(self.session): - task = await self._start_response(task, emit) - elif etype in (ev.INPUT_COMMIT, ev.RESPONSE_CREATE): - if self.adapter.should_respond(self.session): - task = await self._start_response(task, emit) - elif etype == ev.RESPONSE_CANCEL: - await self._barge_in() - await self._cancel(task) - task = None - await emit(ev.cancelled(self.session.response_index)) - elif etype == ev.PLAYBACK_ACK: - await self.adapter.on_playback_ack(self.session, int(event.get("cursor", 0))) - elif etype == ev.CLOSE: - break - if task is not None and not task.done(): - await task - self.session.state = DuplexState.CLOSED + graceful_close = False + try: + async for event in inputs: + etype = event.get("type") + if etype == ev.INPUT_APPEND: + modality = event.get("modality", "") + if modality not in self._capabilities.input_modalities: + await emit(ev.error(f"unsupported input modality: {modality}")) + continue + await self.adapter.on_input(self.session, modality, event.get("data")) + self.session.state = DuplexState.LISTENING + if self.session.config.proactive and self.adapter.should_respond(self.session): + await self._start_response(emit) + elif self.session.config.proactive: + await self._on_input_waiting(modality, emit) + elif etype in (ev.INPUT_COMMIT, ev.RESPONSE_CREATE): + if self.adapter.should_respond(self.session): + await self._start_response(emit) + elif etype == ev.RESPONSE_CANCEL: + await self.cancel_response() + await emit(ev.cancelled(self.session.response_index)) + elif etype == ev.PLAYBACK_ACK: + await self.adapter.on_playback_ack(self.session, int(event.get("cursor", 0))) + elif etype == ev.CLOSE: + graceful_close = True + await self._stop_background() + await self._on_graceful_close(emit) + break + finally: + await self._stop_background() + task = self._response_task + if graceful_close and task is not None and not task.done(): + await task + else: + await self.cancel_response() + await self._close_adapter() + self.session.state = DuplexState.CLOSED + self._closed = True - async def _start_response(self, prev: asyncio.Task | None, emit: Emit) -> asyncio.Task: + async def _start_response(self, emit: Emit) -> None: # A new response supersedes any in-flight one. Mark the old response stale and # cancel it (barge-in) instead of awaiting its full stream — otherwise the input # loop would block until the old response finishes and could not service a # cancel / barge-in event in time, breaking the full-duplex contract. - if prev is not None and not prev.done(): + async with self._response_lock: + prev = self._response_task + if prev is not None and not prev.done(): + await self._barge_in() + await self._cancel(prev) + self._response_task = asyncio.create_task(self._respond(emit)) + + async def cancel_response(self) -> None: + if self._closed: + return + async with self._response_lock: await self._barge_in() - await self._cancel(prev) - return asyncio.create_task(self._respond(emit)) + await self._cancel(self._response_task) + self._response_task = None + + async def close(self) -> None: + if self._closed: + return + await self._stop_background() + await self.cancel_response() + await self._close_adapter() + self.session.state = DuplexState.CLOSED + self._closed = True + + async def _on_input_waiting(self, modality: str, emit: Emit) -> None: + del modality, emit + + async def _stop_background(self) -> None: ... + + async def _on_graceful_close(self, emit: Emit) -> None: + del emit + + async def _close_adapter(self) -> None: + # Keep runtimes compatible with existing duck-typed adapters that + # predate the optional DuplexAdapter.on_close lifecycle hook. + on_close = getattr(self.adapter, "on_close", None) + if on_close is not None: + await on_close(self.session) @staticmethod async def _cancel(task: asyncio.Task | None) -> None: @@ -80,6 +132,8 @@ class DuplexRuntime: if not self.session.is_stale(epoch): await emit(ev.done(response_index)) self.session.end_response() + if self._response_task is asyncio.current_task(): + self._response_task = None async def _barge_in(self) -> None: self.session.barge_in() diff --git a/vllm_omni/experimental/fullduplex/mage_vl/__init__.py b/vllm_omni/experimental/fullduplex/mage_vl/__init__.py new file mode 100644 index 000000000..f805d544d --- /dev/null +++ b/vllm_omni/experimental/fullduplex/mage_vl/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from vllm_omni.experimental.fullduplex.mage_vl.adapter import ( + MageVLCodecWindow, + MageVLDuplexAdapter, + MageVLDuplexRuntime, + MageVLGateDecision, +) + +__all__ = [ + "MageVLCodecWindow", + "MageVLDuplexAdapter", + "MageVLDuplexRuntime", + "MageVLGateDecision", +] diff --git a/vllm_omni/experimental/fullduplex/mage_vl/adapter.py b/vllm_omni/experimental/fullduplex/mage_vl/adapter.py new file mode 100644 index 000000000..5648dc7eb --- /dev/null +++ b/vllm_omni/experimental/fullduplex/mage_vl/adapter.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import asyncio +import contextlib +import inspect +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter, DuplexCapability, OutputChunk +from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime, Emit +from vllm_omni.experimental.fullduplex.core.session import DuplexSession + + +@dataclass(frozen=True) +class MageVLCodecWindow: + """One causal Mage-VL visual window. + + A window may carry decoded frames, codec-native side information, or both. + The adapter deliberately keeps these fields opaque so traditional H.264/HEVC + metadata and neural codec tensors can share the same full-duplex session API. + """ + + data: Any + kind: str = "frames" + segment_id: str | None = None + pts_ms: int | None = None + duration_ms: int | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class MageVLGateDecision: + should_respond: bool + text: str = "" + event_id: str | None = None + score: float | None = None + reason: str = "" + + +GateResult = MageVLGateDecision | Mapping[str, Any] | bool | str | None +GateFn = Callable[[DuplexSession, Sequence[MageVLCodecWindow]], Awaitable[GateResult] | GateResult] +GenerateResult = AsyncIterator[str | OutputChunk] | Sequence[str | OutputChunk] | str | OutputChunk | Awaitable[Any] +GenerateFn = Callable[ + [DuplexSession, Sequence[MageVLCodecWindow], str | None, MageVLGateDecision | None], + GenerateResult, +] + + +class MageVLDuplexAdapter(DuplexAdapter): + """Full-duplex session adapter for Mage-VL's proactive streaming mode. + + This adapter owns the serving/session contract around Mage-VL: bounded causal + visual memory, rolling-window cognition-gate evaluation, event de-duplication, + explicit user queries, and stale-output cleanup through the shared runtime. + Model execution is injected through ``gate`` and ``generate`` callables so the + same adapter can front the remote-code Transformers path, vLLM once native + ``mage_vl`` support lands, or an OpenAI-compatible backend. + """ + + def __init__( + self, + *, + gate: GateFn | None = None, + generate: GenerateFn | None = None, + window_size: int = 4, + max_windows: int = 64, + output_modality: str = "text", + ) -> None: + if window_size <= 0: + raise ValueError("window_size must be positive") + if max_windows <= 0: + raise ValueError("max_windows must be positive") + if window_size > max_windows: + raise ValueError("window_size cannot exceed max_windows") + self._gate = gate + self._generate = generate + self._window_size = window_size + self._max_windows = max_windows + self._output_modality = output_modality + # Adapter instances are shared by the serving runtime. Keep mutable + # cognition state per session so concurrent sessions cannot leak visual + # windows, queries, or event ids into one another. + self._session_state: dict[str, _MageVLSessionState] = {} + + def capabilities(self) -> DuplexCapability: + return DuplexCapability( + # The production Transformers backend currently consumes encoded + # video windows. Do not advertise modalities it cannot materialize. + input_modalities=frozenset({"text", "video"}), + output_modalities=frozenset({self._output_modality}), + proactive=True, + ) + + async def on_input(self, session: DuplexSession, modality: str, data: Any) -> None: + state = self._state_for(session) + if modality == "text": + text = str(data or "").strip() + if text: + state.pending_query = text + return + + window = _coerce_window(modality, data) + state.windows.append(window) + if len(state.windows) > self._max_windows: + del state.windows[: len(state.windows) - self._max_windows] + self._schedule_gate(session) + + def should_respond(self, session: DuplexSession) -> bool: + state = self._state_for(session) + if state.pending_query and state.windows: + return True + return bool(state.pending_gate and state.pending_gate.should_respond) + + async def respond(self, session: DuplexSession) -> AsyncIterator[OutputChunk]: + state = self._state_for(session) + windows = tuple(state.windows[-self._window_size :]) + query = state.pending_query + gate = state.pending_gate + state.pending_query = None + state.pending_gate = None + + if self._generate is None: + if gate and gate.text: + yield OutputChunk(self._output_modality, gate.text) + return + + async for chunk in _iter_generate_result(self._generate(session, windows, query, gate)): + yield _as_output_chunk(chunk, self._output_modality) + + async def on_barge_in(self, session: DuplexSession) -> None: + state = self._state_for(session) + state.pending_query = None + state.pending_gate = None + for task in tuple(state.gate_tasks): + task.cancel() + while not state.gate_ready.empty(): + state.gate_ready.get_nowait() + + async def on_close(self, session: DuplexSession) -> None: + state = self._session_state.get(session.session_id) + if state is not None and state.gate_tasks: + await asyncio.gather(*tuple(state.gate_tasks), return_exceptions=True) + self._session_state.pop(session.session_id, None) + + async def flush(self, session: DuplexSession) -> None: + state = self._session_state.get(session.session_id) + if state is not None and state.gate_tasks: + await asyncio.gather(*tuple(state.gate_tasks), return_exceptions=True) + + async def wait_for_gate(self, session: DuplexSession) -> None: + await self._state_for(session).gate_ready.get() + + async def _evaluate_gate(self, session: DuplexSession) -> None: + state = self._state_for(session) + if self._gate is None or len(state.windows) < self._window_size: + return + decision = _coerce_gate_decision( + await _maybe_await(self._gate(session, tuple(state.windows[-self._window_size :]))) + ) + if not decision.should_respond: + return + if decision.event_id and decision.event_id in state.seen_events: + return + if decision.event_id: + state.seen_events.add(decision.event_id) + state.pending_gate = decision + state.gate_ready.put_nowait(None) + + def _schedule_gate(self, session: DuplexSession) -> None: + state = self._state_for(session) + if self._gate is None or len(state.windows) < self._window_size: + return + task = asyncio.create_task(self._evaluate_gate(session)) + state.gate_tasks.add(task) + task.add_done_callback(state.gate_tasks.discard) + + def _state_for(self, session: DuplexSession) -> _MageVLSessionState: + return self._session_state.setdefault(session.session_id, _MageVLSessionState()) + + +@dataclass +class _MageVLSessionState: + windows: list[MageVLCodecWindow] = field(default_factory=list) + pending_query: str | None = None + pending_gate: MageVLGateDecision | None = None + seen_events: set[str] = field(default_factory=set) + gate_tasks: set[asyncio.Task[None]] = field(default_factory=set) + gate_ready: asyncio.Queue[None] = field(default_factory=asyncio.Queue) + + +class MageVLDuplexRuntime(DuplexRuntime): + """Shared turn runtime with a model-owned deferred-gate watcher.""" + + adapter: MageVLDuplexAdapter + + def __init__(self, session: DuplexSession, adapter: MageVLDuplexAdapter) -> None: + super().__init__(session, adapter) + self._gate_watch_task: asyncio.Task[None] | None = None + + async def _on_input_waiting(self, modality: str, emit: Emit) -> None: + if modality == "text" or self._gate_watch_task is not None: + return + self._gate_watch_task = asyncio.create_task(self._watch_gate(emit)) + + async def _watch_gate(self, emit: Emit) -> None: + while True: + await self.adapter.wait_for_gate(self.session) + if self.adapter.should_respond(self.session): + await self._start_response(emit) + # Let respond() consume the accepted gate before checking the + # next notification, so stacked decisions cannot supersede it. + await asyncio.sleep(0) + + async def _stop_background(self) -> None: + task = self._gate_watch_task + self._gate_watch_task = None + if task is not None and not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + async def _on_graceful_close(self, emit: Emit) -> None: + await self.adapter.flush(self.session) + if self.session.config.proactive and self.adapter.should_respond(self.session): + await self._start_response(emit) + + +def _coerce_window(modality: str, data: Any) -> MageVLCodecWindow: + if isinstance(data, MageVLCodecWindow): + return data + if isinstance(data, Mapping): + metadata = data.get("metadata") or {} + if not isinstance(metadata, Mapping): + metadata = {"metadata": metadata} + payload = data.get("data", data.get("frames", data.get("frame", data.get("codec")))) + return MageVLCodecWindow( + data=payload, + kind=str(data.get("kind", modality)), + segment_id=_optional_str(data.get("segment_id")), + pts_ms=_optional_int(data.get("pts_ms")), + duration_ms=_optional_int(data.get("duration_ms")), + metadata=metadata, + ) + return MageVLCodecWindow(data=data, kind=modality) + + +def _coerce_gate_decision(result: GateResult) -> MageVLGateDecision: + if isinstance(result, MageVLGateDecision): + return result + if isinstance(result, Mapping): + return MageVLGateDecision( + should_respond=bool(result.get("should_respond", result.get("respond", result.get("open", False)))), + text=str(result.get("text", "")), + event_id=_optional_str(result.get("event_id")), + score=_optional_float(result.get("score")), + reason=str(result.get("reason", "")), + ) + if isinstance(result, str): + text = result.strip() + return MageVLGateDecision(should_respond=bool(text), text=text) + return MageVLGateDecision(should_respond=bool(result)) + + +async def _iter_generate_result(result: GenerateResult) -> AsyncIterator[str | OutputChunk]: + resolved: Any = await _maybe_await(result) + if isinstance(resolved, str | OutputChunk): + yield resolved + return + if hasattr(resolved, "__aiter__"): + async for item in resolved: + yield item + return + for item in resolved: + yield item + + +async def _maybe_await(value: Any) -> Any: + if inspect.isawaitable(value): + return await value + return value + + +def _as_output_chunk(chunk: str | OutputChunk, default_modality: str) -> OutputChunk: + if isinstance(chunk, OutputChunk): + return chunk + return OutputChunk(default_modality, chunk) + + +def _optional_str(value: Any) -> str | None: + return None if value is None else str(value) + + +def _optional_int(value: Any) -> int | None: + if value is None: + return None + return int(value) + + +def _optional_float(value: Any) -> float | None: + if value is None: + return None + return float(value) diff --git a/vllm_omni/experimental/fullduplex/mage_vl/serving/__init__.py b/vllm_omni/experimental/fullduplex/mage_vl/serving/__init__.py new file mode 100644 index 000000000..0642234cd --- /dev/null +++ b/vllm_omni/experimental/fullduplex/mage_vl/serving/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from vllm_omni.experimental.fullduplex.mage_vl.serving.backend import MageVLTransformersBackend +from vllm_omni.experimental.fullduplex.mage_vl.serving.server import ( + MageVLServingConfig, + create_app, + load_adapter_factory, +) + +__all__ = ["MageVLServingConfig", "MageVLTransformersBackend", "create_app", "load_adapter_factory"] diff --git a/vllm_omni/experimental/fullduplex/mage_vl/serving/backend.py b/vllm_omni/experimental/fullduplex/mage_vl/serving/backend.py new file mode 100644 index 000000000..062edc63e --- /dev/null +++ b/vllm_omni/experimental/fullduplex/mage_vl/serving/backend.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Transformers backend for the Mage-VL production WebSocket transport.""" + +from __future__ import annotations + +import asyncio +import binascii +import contextlib +import shutil +import tempfile +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +import pybase64 as base64 + +from vllm_omni.experimental.fullduplex.core.session import DuplexSession +from vllm_omni.experimental.fullduplex.mage_vl import ( + MageVLCodecWindow, + MageVLDuplexAdapter, + MageVLGateDecision, +) + +_STREAMING_PROMPT = "Please describe the video content in detail based on the provided information." + + +class MageVLTransformersBackend: + """One shared checkpoint with serialized GPU access and isolated adapters.""" + + def __init__( + self, + checkpoint: str = "microsoft/Mage-VL", + *, + device: str = "cuda", + attn_impl: str = "sdpa", + video_backend: str = "frames", + num_frames: int = 8, + target_fps: float = 1.0, + max_new_tokens: int = 128, + gate_threshold: float = 0.5, + window_size: int = 4, + max_windows: int = 64, + ) -> None: + if video_backend not in {"frames", "codec"}: + raise ValueError("video_backend must be frames or codec") + if video_backend == "codec" and shutil.which("cv-preinfer") is None: + raise RuntimeError( + "video_backend='codec' requires cv-preinfer from codec-video-prep>=0.2.5; " + "install it and ensure cv-preinfer is on PATH, or use video_backend='frames'" + ) + if num_frames <= 0 or target_fps <= 0 or max_new_tokens <= 0: + raise ValueError("num_frames, target_fps, and max_new_tokens must be positive") + self.checkpoint = checkpoint + self.device = device + self.attn_impl = attn_impl + self.video_backend = video_backend + self.num_frames = num_frames + self.target_fps = target_fps + self.max_new_tokens = max_new_tokens + self.gate_threshold = gate_threshold + self.window_size = window_size + self.max_windows = max_windows + self.model: Any = None + self.processor: Any = None + self._inference_lock = asyncio.Lock() + + def load(self) -> None: + import torch + from transformers import AutoModelForCausalLM, AutoProcessor + + self.processor = AutoProcessor.from_pretrained(self.checkpoint, trust_remote_code=True) + self.model = ( + AutoModelForCausalLM.from_pretrained( + self.checkpoint, + trust_remote_code=True, + dtype=torch.bfloat16, + attn_implementation=self.attn_impl, + ) + .to(self.device) + .eval() + ) + if not hasattr(self.model, "streammind_gate_forward_segments"): + raise AttributeError("Mage-VL checkpoint does not expose streammind_gate_forward_segments()") + + def adapter_factory(self) -> MageVLDuplexAdapter: + if self.model is None or self.processor is None: + raise RuntimeError("MageVLTransformersBackend.load() must be called before serving") + return MageVLDuplexAdapter( + gate=self.gate, + generate=self.generate, + window_size=self.window_size, + max_windows=self.max_windows, + ) + + async def gate( + self, + _session: DuplexSession, + windows: Sequence[MageVLCodecWindow], + ) -> MageVLGateDecision: + score = await self._run_inference(self._gate_sync, windows) + latest = windows[-1] + return MageVLGateDecision( + should_respond=score >= self.gate_threshold, + event_id=latest.segment_id, + score=score, + reason="streammind_gate", + ) + + async def generate( + self, + _session: DuplexSession, + windows: Sequence[MageVLCodecWindow], + query: str | None, + _decision: MageVLGateDecision | None, + ) -> str: + return await self._run_inference(self._generate_sync, windows[-1], query or _STREAMING_PROMPT) + + async def _run_inference(self, function, *args): + """Serialize GPU calls and drain worker threads before cancellation.""" + async with self._inference_lock: + worker = asyncio.create_task(asyncio.to_thread(function, *args)) + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + # Cancelling an asyncio wrapper cannot stop a CUDA worker thread. + # Retain the lock until it really exits so a replacement request + # cannot race the same shared model. + with contextlib.suppress(Exception, asyncio.CancelledError): + await worker + raise + + def _gate_sync(self, windows: Sequence[MageVLCodecWindow]) -> float: + import torch + + prepared: list[tuple[dict[str, Any], Path | None]] = [] + try: + prepared = [self._prepare(window, _STREAMING_PROMPT) for window in windows] + visual = [self._gate_inputs(item[0]) for item in prepared] + with torch.inference_mode(): + logits = self.model.streammind_gate_forward_segments(visual)[0] + lengths = [int(segment["image_grid_thw"][:, 0].sum()) for segment in visual] + boundary = sum(lengths) - 1 + return float(torch.softmax(logits[boundary].float(), dim=-1)[1].item()) + finally: + for _, temporary_path in prepared: + _remove_temporary(temporary_path) + + def _generate_sync(self, window: MageVLCodecWindow, query: str) -> str: + import torch + + inputs, temporary_path = self._prepare(window, query) + try: + with torch.inference_mode(): + output = self.model.generate( + **inputs, + max_new_tokens=self.max_new_tokens, + do_sample=False, + ) + new_tokens = output[0, inputs["input_ids"].shape[1] :] + return self.processor.tokenizer.decode(new_tokens, skip_special_tokens=True).strip() + finally: + _remove_temporary(temporary_path) + + def _prepare(self, window: MageVLCodecWindow, query: str) -> tuple[dict[str, Any], Path | None]: + video_path, temporary_path = _materialize_video(window.data) + messages = [ + { + "role": "user", + "content": [{"type": "video"}, {"type": "text", "text": query}], + } + ] + prompt = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + kwargs: dict[str, Any] = { + "text": [prompt], + "videos": [str(video_path)], + "video_backend": self.video_backend, + "return_tensors": "pt", + "padding": False, + } + if self.video_backend == "codec": + kwargs.update(codec_config={"patch": 16, "max_pixels": 150000}, max_pixels=150000) + else: + kwargs.update(num_frames=self.num_frames, target_fps=self.target_fps) + try: + inputs = self.processor(**kwargs) + inputs = { + key: ( + value.to(device=self.device, dtype=self.model.dtype) + if key == "pixel_values" + else value.to(self.device) + ) + for key, value in inputs.items() + } + return inputs, temporary_path + except Exception: + _remove_temporary(temporary_path) + raise + + @staticmethod + def _gate_inputs(inputs: dict[str, Any]) -> dict[str, Any]: + return {key: inputs[key] for key in ("pixel_values", "image_grid_thw", "patch_positions") if key in inputs} + + +def _materialize_video(data: Any) -> tuple[Path, Path | None]: + if not isinstance(data, dict): + raise ValueError("video window data must be an object containing video_base64") + encoded = data.get("video_base64") + if not isinstance(encoded, str) or not encoded: + raise ValueError("video window requires non-empty video_base64") + try: + payload = base64.b64decode(encoded, validate=True) + except (ValueError, binascii.Error) as error: + raise ValueError("video_base64 is not valid base64") from error + if not payload: + raise ValueError("decoded video payload is empty") + with tempfile.NamedTemporaryFile(prefix="mage-vl-", suffix=".mp4", delete=False) as handle: + handle.write(payload) + path = Path(handle.name) + return path, path + + +def _remove_temporary(path: Path | None) -> None: + if path is not None: + with contextlib.suppress(FileNotFoundError): + path.unlink() diff --git a/vllm_omni/experimental/fullduplex/mage_vl/serving/server.py b/vllm_omni/experimental/fullduplex/mage_vl/serving/server.py new file mode 100644 index 000000000..e1cfde41b --- /dev/null +++ b/vllm_omni/experimental/fullduplex/mage_vl/serving/server.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Production WebSocket transport for Mage-VL full-duplex sessions.""" + +from __future__ import annotations + +import argparse +import asyncio +import contextlib +import importlib +import inspect +import json +import secrets +import time +import uuid +from collections.abc import AsyncIterator, Callable +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Any + +import uvicorn +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from fastapi.responses import JSONResponse + +from vllm_omni.experimental.fullduplex.core import protocol as event_protocol +from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter +from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig +from vllm_omni.experimental.fullduplex.mage_vl import MageVLDuplexAdapter, MageVLDuplexRuntime + +AdapterFactory = Callable[[], DuplexAdapter | Any] + + +@dataclass(frozen=True) +class MageVLServingConfig: + model: str = "microsoft/Mage-VL" + max_sessions: int = 32 + idle_timeout_s: float = 300.0 + max_message_bytes: int = 8 * 1024 * 1024 + auth_token: str | None = None + + def __post_init__(self) -> None: + if self.max_sessions <= 0 or self.idle_timeout_s <= 0 or self.max_message_bytes <= 0: + raise ValueError("max_sessions, idle_timeout_s, and max_message_bytes must be positive") + + +class _SessionLease: + def __init__(self, session_id: str, runtime: MageVLDuplexRuntime) -> None: + self.session_id = session_id + self.runtime = runtime + self.last_activity = time.monotonic() + + def touch(self) -> None: + self.last_activity = time.monotonic() + + +class _SessionRegistry: + def __init__(self, config: MageVLServingConfig, factory: AdapterFactory) -> None: + self.config = config + self.factory = factory + self._leases: dict[str, _SessionLease] = {} + self._lock = asyncio.Lock() + + async def acquire(self, requested_id: str | None) -> _SessionLease | None: + async with self._lock: + if len(self._leases) >= self.config.max_sessions: + return None + session_id = requested_id or f"mage-{uuid.uuid4().hex}" + if not _valid_session_id(session_id) or session_id in self._leases: + return None + adapter = self.factory() + if inspect.isawaitable(adapter): + adapter = await adapter + if not isinstance(adapter, MageVLDuplexAdapter): + raise TypeError("adapter factory must return a MageVLDuplexAdapter") + capability = adapter.capabilities() + session = DuplexSession( + session_id, + DuplexSessionConfig( + input_modalities=tuple(capability.input_modalities), + output_modalities=tuple(capability.output_modalities), + proactive=capability.proactive, + ), + ) + lease = _SessionLease(session_id, MageVLDuplexRuntime(session, adapter)) + self._leases[session_id] = lease + return lease + + async def release(self, session_id: str) -> None: + async with self._lock: + self._leases.pop(session_id, None) + + async def close(self) -> None: + async with self._lock: + leases = list(self._leases.values()) + self._leases.clear() + for lease in leases: + with contextlib.suppress(asyncio.CancelledError): + await lease.runtime.close() + + @property + def active_count(self) -> int: + return len(self._leases) + + +def _valid_session_id(value: str) -> bool: + return 0 < len(value) <= 128 and all(char.isalnum() or char in "-_." for char in value) + + +def _authorized(websocket: WebSocket, token: str | None) -> bool: + if token is None: + return True + return secrets.compare_digest(websocket.headers.get("authorization", ""), f"Bearer {token}") + + +def _event_from_text(text: str, max_bytes: int) -> dict[str, Any]: + if len(text.encode("utf-8")) > max_bytes: + raise ValueError("message exceeds configured max_message_bytes") + event = json.loads(text) + if not isinstance(event, dict) or not isinstance(event.get("type"), str): + raise ValueError("event must be a JSON object with a string type") + supported = { + event_protocol.INPUT_APPEND, + event_protocol.INPUT_COMMIT, + event_protocol.RESPONSE_CREATE, + event_protocol.RESPONSE_CANCEL, + event_protocol.PLAYBACK_ACK, + event_protocol.CLOSE, + } + if event["type"] not in supported: + raise ValueError(f"unsupported event type: {event['type']}") + if event["type"] == event_protocol.INPUT_APPEND and not isinstance(event.get("modality"), str): + raise ValueError("input.append requires a modality") + return event + + +def load_adapter_factory(path: str) -> AdapterFactory: + module_name, separator, attribute_name = path.partition(":") + if not separator: + raise ValueError("adapter path must use module:attribute syntax") + value = getattr(importlib.import_module(module_name), attribute_name) + if isinstance(value, DuplexAdapter): + return lambda: value + if not callable(value): + raise TypeError(f"adapter target {path!r} is not callable") + return value + + +def create_app(adapter_factory: AdapterFactory, config: MageVLServingConfig | None = None) -> FastAPI: + config = config or MageVLServingConfig() + registry = _SessionRegistry(config, adapter_factory) + + @asynccontextmanager + async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + try: + yield + finally: + await registry.close() + + app = FastAPI(title="Mage-VL Duplex Server", version="1", lifespan=lifespan) + app.state.session_registry = registry + + @app.get("/health") + async def health() -> dict[str, str]: + return {"status": "ok"} + + @app.get("/ready") + async def ready() -> JSONResponse: + return JSONResponse({"status": "ready", "active_sessions": registry.active_count}) + + @app.get("/v1/models") + async def models() -> dict[str, Any]: + return {"object": "list", "data": [{"id": config.model, "object": "model", "owned_by": "microsoft"}]} + + @app.websocket("/v1/mage-vl/duplex") + async def duplex(websocket: WebSocket) -> None: + if not _authorized(websocket, config.auth_token): + await websocket.close(code=4401, reason="unauthorized") + return + await websocket.accept() + lease = await registry.acquire(websocket.query_params.get("session_id")) + if lease is None: + await websocket.close(code=4429, reason="session unavailable") + return + + async def inputs() -> AsyncIterator[dict[str, Any]]: + while True: + event = _event_from_text( + await asyncio.wait_for(websocket.receive_text(), config.idle_timeout_s), config.max_message_bytes + ) + lease.touch() + yield event + if event.get("type") == event_protocol.CLOSE: + return + + async def emit(event: dict[str, Any]) -> None: + await websocket.send_json(event) + + capability = lease.runtime.capabilities + try: + await websocket.send_json( + { + "type": "session.created", + "session_id": lease.session_id, + "model": config.model, + "capabilities": { + "input_modalities": sorted(capability.input_modalities), + "output_modalities": sorted(capability.output_modalities), + "proactive": capability.proactive, + }, + } + ) + await lease.runtime.run(inputs(), emit) + except WebSocketDisconnect: + await lease.runtime.cancel_response() + except asyncio.TimeoutError: + with contextlib.suppress(Exception): + await websocket.send_json({"type": "error", "message": "idle timeout"}) + except (ValueError, json.JSONDecodeError) as error: + with contextlib.suppress(Exception): + await websocket.send_json({"type": "error", "message": str(error)}) + await lease.runtime.cancel_response() + finally: + await registry.release(lease.session_id) + with contextlib.suppress(Exception): + await websocket.close() + + return app + + +def main() -> None: + parser = argparse.ArgumentParser(description="Serve Mage-VL full-duplex WebSocket sessions.") + parser.add_argument("--adapter", help="optional module:attribute adapter factory") + parser.add_argument("--model", default="microsoft/Mage-VL") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8090) + parser.add_argument("--max-sessions", type=int, default=32) + parser.add_argument("--idle-timeout", type=float, default=300.0) + parser.add_argument("--max-message-bytes", type=int, default=8 * 1024 * 1024) + parser.add_argument("--auth-token") + parser.add_argument("--device", default="cuda") + parser.add_argument("--attn-impl", choices=("flash_attention_2", "sdpa", "eager"), default="sdpa") + parser.add_argument("--video-backend", choices=("frames", "codec"), default="frames") + parser.add_argument("--num-frames", type=int, default=8) + parser.add_argument("--target-fps", type=float, default=1.0) + parser.add_argument("--max-new-tokens", type=int, default=128) + parser.add_argument("--gate-threshold", type=float, default=0.5) + parser.add_argument("--window-size", type=int, default=4) + args = parser.parse_args() + if args.adapter: + adapter_factory = load_adapter_factory(args.adapter) + else: + from vllm_omni.experimental.fullduplex.mage_vl.serving.backend import MageVLTransformersBackend + + backend = MageVLTransformersBackend( + args.model, + device=args.device, + attn_impl=args.attn_impl, + video_backend=args.video_backend, + num_frames=args.num_frames, + target_fps=args.target_fps, + max_new_tokens=args.max_new_tokens, + gate_threshold=args.gate_threshold, + window_size=args.window_size, + ) + backend.load() + adapter_factory = backend.adapter_factory + app = create_app( + adapter_factory, + MageVLServingConfig( + model=args.model, + max_sessions=args.max_sessions, + idle_timeout_s=args.idle_timeout, + max_message_bytes=args.max_message_bytes, + auth_token=args.auth_token, + ), + ) + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 4438d2200..cb5167597 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from vllm.model_executor.models.registry import ( _VLLM_MODELS, _LazyRegisteredModel,