[Entrypoint] Add realtime OpenPI robot serving API (#3673)

Signed-off-by: Yangshen Deng <yangshen.d@outlook.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Meng <meng_chen99@163.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

vllm_omni/entrypoints/openai/api_server.py

@@ -133,6 +133,7 @@ from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER
 from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import decode_input_reference
+from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 
 logger = init_logger(__name__)
@@ -633,6 +634,10 @@ async def omni_init_app_state(
         )
         state.openai_streaming_speech = None
         state.openai_streaming_video = None
+        state.openai_serving_realtime_robot = ServingRealtimeRobotOpenPI.create_policy_server(
+            engine_client=engine_client,
+            model_name=model_name,
+        )
 
         state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False)
         state.server_load_metrics = 0
@@ -963,6 +968,7 @@ async def omni_init_app_state(
         model_name=served_model_names[0] if served_model_names else None,
         stage_configs=state.stage_configs,
     )
+    state.openai_serving_realtime_robot = None
 
     state.enable_server_load_tracking = args.enable_server_load_tracking
     state.server_load_metrics = 0
@@ -1406,6 +1412,23 @@ async def realtime_websocket(websocket: WebSocket):
     await connection.handle_connection()
 
 
+@router.websocket("/v1/realtime/robot/openpi")
+async def realtime_robot_openpi(websocket: WebSocket):
+    """WebSocket endpoint for robot policy inference via OpenPI messages."""
+    from vllm_omni.entrypoints.openpi.connection import (
+        RobotRealtimeConnection,
+    )
+
+    serving = getattr(websocket.app.state, "openai_serving_realtime_robot", None)
+    if serving is None:
+        await websocket.accept()
+        await websocket.send_json({"type": "error", "error": "Robot policy not available", "code": "unsupported"})
+        await websocket.close()
+        return
+    connection = RobotRealtimeConnection(websocket, serving)
+    await connection.handle_connection()
+
+
 # Health and Model endpoints for diffusion mode
 
 

vllm_omni/entrypoints/openpi/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

vllm_omni/entrypoints/openpi/connection.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""WebSocket connection for robot policy inference (OpenPI protocol).
+
+Protocol (compatible with OpenPI policy clients):
+    Connect  -> server sends msgpack(PolicyServerConfig fields)
+    Infer    -> client sends msgpack(obs), server sends msgpack(ndarray)
+    Reset    -> client sends msgpack({endpoint:reset}), server sends msgpack(status)
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from fastapi import WebSocket
+from starlette.websockets import WebSocketDisconnect
+from vllm.logger import init_logger
+
+from vllm_omni.entrypoints.openpi.serving import (
+    ServingRealtimeRobotOpenPI,
+)
+
+logger = init_logger(__name__)
+_DEFAULT_IDLE_TIMEOUT = 30.0
+MAX_OPENPI_PAYLOAD_BYTES = 64 * 1024 * 1024
+
+
+def _get_msgpack_numpy() -> Any:
+    try:
+        from openpi_client import msgpack_numpy
+    except ImportError as exc:
+        raise ImportError(
+            "The `/v1/realtime/robot/openpi` endpoint requires the optional "
+            "`openpi-client` dependency. Install it with `pip install openpi-client`."
+        ) from exc
+
+    return msgpack_numpy
+
+
+def _pack(obj: Any) -> bytes:
+    return _get_msgpack_numpy().packb(obj)
+
+
+def _unpack(data: bytes) -> Any:
+    return _get_msgpack_numpy().unpackb(data)
+
+
+class RobotRealtimeConnection:
+    """WebSocket connection for robot policy inference."""
+
+    def __init__(
+        self,
+        websocket: WebSocket,
+        serving: ServingRealtimeRobotOpenPI,
+        idle_timeout: float = _DEFAULT_IDLE_TIMEOUT,
+    ) -> None:
+        self.websocket = websocket
+        self.serving = serving
+        self._idle_timeout = idle_timeout
+        self._current_session_id: str | None = None
+        self._call_count = 0
+
+    def reset(self) -> None:
+        self._current_session_id = None
+        self._call_count = 0
+
+    async def _send_error(self, message: str) -> None:
+        await self.websocket.send_bytes(_pack({"type": "error", "message": message}))
+
+    def _unpack_request(self, data: bytes) -> dict[str, Any]:
+        if len(data) > MAX_OPENPI_PAYLOAD_BYTES:
+            raise ValueError("OpenPI request payload too large")
+        obs = _unpack(data)
+        if not isinstance(obs, dict):
+            raise ValueError("Invalid request payload")
+        return obs
+
+    async def handle_connection(self) -> None:
+        """Main loop for OpenPI-compatible policy serving."""
+        await self.websocket.accept()
+
+        try:
+            # Send model-specific PolicyServerConfig resolved by serving from
+            # diffusion od_config.model_config.
+            metadata = self.serving.policy_server_config.to_dict()
+            await self.websocket.send_bytes(_pack(metadata))
+
+            while True:
+                try:
+                    msg = await asyncio.wait_for(
+                        self.websocket.receive(),
+                        timeout=self._idle_timeout,
+                    )
+                except asyncio.TimeoutError:
+                    logger.info("Robot OpenPI connection idle timeout after %.1f seconds", self._idle_timeout)
+                    try:
+                        await self.websocket.close()
+                    except Exception:
+                        logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True)
+                    return
+
+                if msg.get("type") == "websocket.disconnect":
+                    break
+
+                if "bytes" not in msg or not msg["bytes"]:
+                    continue
+
+                try:
+                    obs = self._unpack_request(msg["bytes"])
+                except Exception:
+                    logger.exception("Invalid robot OpenPI request payload")
+                    try:
+                        await self._send_error("Invalid request payload")
+                    except Exception:
+                        break
+                    continue
+
+                try:
+                    endpoint = obs.pop("endpoint", "infer")
+
+                    if endpoint == "reset":
+                        self.reset()
+                        self.serving.reset(obs)
+                        await self.websocket.send_bytes(_pack({"status": "reset successful"}))
+                    else:
+                        session_id = str(obs.get("session_id") or self._current_session_id or "default")
+                        if session_id != self._current_session_id:
+                            if self._current_session_id is not None:
+                                logger.info(
+                                    "Robot OpenPI session changed %s -> %s",
+                                    self._current_session_id,
+                                    session_id,
+                                )
+                            self._current_session_id = session_id
+                            self._call_count = 0
+
+                        self._call_count += 1
+                        actions = await self.serving.infer(
+                            obs,
+                            session_id=session_id,
+                            reset=self._call_count <= 1,
+                        )
+                        await self.websocket.send_bytes(_pack(actions))
+                except Exception:
+                    logger.exception("Error handling request")
+                    try:
+                        await self._send_error("Internal inference error")
+                    except Exception:
+                        break
+
+        except WebSocketDisconnect:
+            pass
+        except Exception:
+            logger.exception("Connection error")

vllm_omni/entrypoints/openpi/serving.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Serving layer for robot policy inference via `/v1/realtime/robot/openpi`.
+
+Flow: raw obs → engine request → actions.
+The loaded policy model owns dataset transforms inside its pipeline.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from itertools import count
+from typing import Any
+
+import numpy as np
+from omegaconf import OmegaConf
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+ActionOutput = np.ndarray | dict[str, np.ndarray]
+
+
+def _to_builtin_container(value: Any) -> Any:
+    if OmegaConf.is_config(value):
+        return OmegaConf.to_container(value, resolve=True)
+    if isinstance(value, Mapping):
+        return {key: _to_builtin_container(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_to_builtin_container(item) for item in value]
+    return value
+
+
+@dataclass(frozen=True)
+class PolicyServerConfig:
+    """OpenPI policy server handshake config.
+
+    Values are model-specific and must be provided by the loaded policy model.
+    """
+
+    values: dict[str, Any]
+
+    @classmethod
+    def from_model_config(cls, model_config: Any) -> PolicyServerConfig:
+        if isinstance(model_config, Mapping):
+            raw_config = model_config.get("policy_server_config")
+        else:
+            raw_config = getattr(model_config, "policy_server_config", None)
+
+        if raw_config is None:
+            raise ValueError("Robot OpenPI serving requires policy_server_config.")
+        if isinstance(raw_config, cls):
+            return raw_config
+        if not isinstance(raw_config, Mapping):
+            raise ValueError("Robot OpenPI serving requires policy_server_config.")
+        return cls(_to_builtin_container(raw_config))
+
+    def to_dict(self) -> dict[str, Any]:
+        return _to_builtin_container(self.values)
+
+
+class ServingRealtimeRobotOpenPI:
+    """Robot policy serving layer for OpenPI protocol.
+
+    Model-specific transform/state lives in the diffusion pipeline.
+    """
+
+    def __init__(
+        self,
+        engine_client: Any,
+        model_name: str | None = None,
+    ) -> None:
+        self.engine_client = engine_client
+        self.model_name = model_name
+        self.policy_server_config = self._get_policy_server_config(engine_client)
+        self._request_counter = count()
+
+    @classmethod
+    def create_policy_server(
+        cls,
+        engine_client: Any,
+        model_name: str | None = None,
+    ) -> ServingRealtimeRobotOpenPI | None:
+        try:
+            return cls(engine_client=engine_client, model_name=model_name)
+        except ValueError as exc:
+            if "policy_server_config" not in str(exc):
+                raise
+            logger.info("Robot OpenPI serving disabled for model %s", model_name)
+            return None
+
+    @staticmethod
+    def _get_policy_server_config(engine_client: Any) -> PolicyServerConfig:
+        model_config = None
+        get_od_config = getattr(engine_client, "get_diffusion_od_config", None)
+        if callable(get_od_config):
+            od_config = get_od_config()
+            model_config = getattr(od_config, "model_config", None)
+
+        if model_config is None:
+            for stage_config in getattr(engine_client, "stage_configs", []) or []:
+                if getattr(stage_config, "stage_type", None) != "diffusion":
+                    continue
+                engine_args = getattr(stage_config, "engine_args", None)
+                model_config = getattr(engine_args, "model_config", None)
+                if model_config is not None:
+                    break
+
+        if model_config is None:
+            od_config = getattr(engine_client, "od_config", None)
+            model_config = getattr(od_config, "model_config", None)
+
+        if model_config is None:
+            model_config = getattr(engine_client, "model_config", None)
+        return PolicyServerConfig.from_model_config(model_config)
+
+    def reset(self, obs: dict) -> None:
+        """Compatibility hook; per-connection state lives in RobotRealtimeConnection."""
+
+    async def infer(self, obs: dict, *, session_id: str, reset: bool) -> ActionOutput:
+        """raw obs → engine → actions."""
+        # Build request, run inference through AsyncOmni
+        request = self._build_request(obs, session_id=session_id, reset=reset)
+        result = None
+        # OpenPI policy serving is one request -> one action reply. AsyncOmni
+        # exposes an async iterator, so consume it to completion and use the
+        # final output, matching other non-streaming OpenAI serving paths.
+        async for output in self.engine_client.generate(
+            prompt=request.prompts[0],
+            request_id=request.request_id,
+            sampling_params_list=[request.sampling_params],
+        ):
+            result = output
+        if result is None:
+            raise RuntimeError("Robot OpenPI request produced no output.")
+
+        return self._extract_actions(result)
+
+    def _next_request_id(self, session_id: str) -> str:
+        return f"robot-{session_id}-{next(self._request_counter)}"
+
+    def _build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any:
+        """Build engine request from raw robot obs.
+
+        Returns an `OmniDiffusionRequest` payload consumed by
+        `AsyncOmni.generate()` and routed to the diffusion stage.
+        """
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        extra_args = {
+            "reset": reset,
+            "session_id": session_id,
+            "robot_obs": obs,
+        }
+
+        prompt = obs.get("prompt", "")
+        sampling_params = OmniDiffusionSamplingParams(extra_args=extra_args)
+        return OmniDiffusionRequest(
+            prompts=[prompt],
+            sampling_params=sampling_params,
+            request_id=self._next_request_id(session_id),
+        )
+
+    def _extract_actions(self, result: Any) -> ActionOutput:
+        """Extract actions from engine result."""
+        multimodal_output = getattr(result, "multimodal_output", None)
+        if not isinstance(multimodal_output, Mapping):
+            raise RuntimeError("Missing multimodal_output in robot policy result")
+
+        actions = multimodal_output.get("actions")
+        if actions is None:
+            raise RuntimeError("Missing multimodal_output['actions'] in robot policy result")
+        if isinstance(actions, Mapping):
+            return {str(key): np.asarray(value, dtype=np.float32) for key, value in actions.items()}
+        return np.asarray(actions, dtype=np.float32)