[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>

tests/entrypoints/openai_api/test_openpi_connection.py

@@ -0,0 +1,270 @@
+import asyncio
+import builtins
+import sys
+import types
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from vllm_omni.entrypoints.openpi import connection as openpi_connection
+from vllm_omni.entrypoints.openpi.serving import PolicyServerConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeWebSocket:
+    def __init__(self, messages):
+        self._messages = list(messages)
+        self.sent_bytes = []
+        self.sent_texts = []
+        self.accepted = False
+        self.closed = False
+
+    async def accept(self):
+        self.accepted = True
+
+    async def send_bytes(self, data):
+        self.sent_bytes.append(data)
+
+    async def send_text(self, data):
+        self.sent_texts.append(data)
+
+    async def receive(self):
+        return self._messages.pop(0)
+
+    async def close(self):
+        self.closed = True
+
+
+def _serving_mock():
+    serving = MagicMock()
+    serving.policy_server_config = PolicyServerConfig(
+        {
+            "image_resolution": (180, 320),
+            "n_external_cameras": 2,
+            "needs_wrist_camera": True,
+            "needs_stereo_camera": False,
+            "needs_session_id": True,
+            "action_space": "joint_position",
+        }
+    )
+    serving.infer = AsyncMock(return_value=[0.0])
+    return serving
+
+
+def test_pack_reports_clear_error_when_openpi_client_is_missing(monkeypatch):
+    real_import = builtins.__import__
+
+    def import_without_openpi_client(name, globals=None, locals=None, fromlist=(), level=0):
+        if name == "openpi_client":
+            raise ModuleNotFoundError("No module named 'openpi_client'", name="openpi_client")
+        return real_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", import_without_openpi_client)
+
+    with pytest.raises(ImportError) as exc_info:
+        openpi_connection._pack({"prompt": "pick up the object"})
+
+    message = str(exc_info.value)
+    assert "/v1/realtime/robot/openpi" in message
+    assert "pip install openpi-client" in message
+
+
+def test_pack_and_unpack_delegate_to_openpi_msgpack_numpy(monkeypatch):
+    calls = []
+
+    class FakeMsgpackNumpy:
+        @staticmethod
+        def packb(obj):
+            calls.append(("packb", obj))
+            return b"packed"
+
+        @staticmethod
+        def unpackb(data):
+            calls.append(("unpackb", data))
+            return {"unpacked": data}
+
+    fake_openpi_client = types.ModuleType("openpi_client")
+    fake_openpi_client.msgpack_numpy = FakeMsgpackNumpy
+    monkeypatch.setitem(sys.modules, "openpi_client", fake_openpi_client)
+
+    assert openpi_connection._pack({"x": 1}) == b"packed"
+    assert openpi_connection._unpack(b"payload") == {"unpacked": b"payload"}
+    assert calls == [
+        ("packb", {"x": 1}),
+        ("unpackb", b"payload"),
+    ]
+
+
+def test_handle_connection_returns_structured_error_for_invalid_payload(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+    monkeypatch.setattr(
+        openpi_connection,
+        "_unpack",
+        lambda _data: (_ for _ in ()).throw(ValueError("bad payload traceback")),
+    )
+
+    websocket = FakeWebSocket(
+        [
+            {"type": "websocket.receive", "bytes": b"bad"},
+            {"type": "websocket.disconnect"},
+        ]
+    )
+    serving = MagicMock()
+    connection = openpi_connection.RobotRealtimeConnection(websocket, serving)
+
+    asyncio.run(connection.handle_connection())
+
+    assert websocket.accepted is True
+    assert websocket.sent_bytes[1] == {"type": "error", "message": "Invalid request payload"}
+    assert "traceback" not in str(websocket.sent_bytes[1]).lower()
+    assert websocket.sent_texts == []
+    serving.infer.assert_not_called()
+    serving.reset.assert_not_called()
+
+
+def test_handle_connection_rejects_oversized_payload_before_unpack(monkeypatch):
+    unpack_mock = MagicMock(side_effect=AssertionError("_unpack should not be called"))
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+    monkeypatch.setattr(openpi_connection, "_unpack", unpack_mock)
+    monkeypatch.setattr(openpi_connection, "MAX_OPENPI_PAYLOAD_BYTES", 4)
+
+    websocket = FakeWebSocket(
+        [
+            {"type": "websocket.receive", "bytes": b"too-large"},
+            {"type": "websocket.disconnect"},
+        ]
+    )
+    serving = MagicMock()
+    connection = openpi_connection.RobotRealtimeConnection(websocket, serving)
+
+    asyncio.run(connection.handle_connection())
+
+    assert websocket.sent_bytes[1] == {"type": "error", "message": "Invalid request payload"}
+    unpack_mock.assert_not_called()
+    serving.infer.assert_not_called()
+    serving.reset.assert_not_called()
+
+
+def test_handle_connection_returns_structured_error_for_infer_exception(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+    monkeypatch.setattr(
+        openpi_connection,
+        "_unpack",
+        lambda _data: {"prompt": "pick up the object"},
+    )
+
+    websocket = FakeWebSocket(
+        [
+            {"type": "websocket.receive", "bytes": b"request"},
+            {"type": "websocket.disconnect"},
+        ]
+    )
+    serving = MagicMock()
+    serving.infer = AsyncMock(side_effect=RuntimeError("secret traceback text"))
+    connection = openpi_connection.RobotRealtimeConnection(websocket, serving)
+
+    asyncio.run(connection.handle_connection())
+
+    assert websocket.sent_bytes[1] == {"type": "error", "message": "Internal inference error"}
+    assert "secret traceback text" not in str(websocket.sent_bytes[1])
+    assert websocket.sent_texts == []
+    serving.infer.assert_awaited_once_with(
+        {"prompt": "pick up the object"},
+        session_id="default",
+        reset=True,
+    )
+
+
+def test_handle_connection_closes_websocket_on_idle_timeout(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+
+    websocket = FakeWebSocket([])
+
+    async def never_receives():
+        await asyncio.sleep(1)
+
+    websocket.receive = never_receives
+    serving = MagicMock()
+    serving.policy_server_config = PolicyServerConfig(
+        {
+            "image_resolution": (180, 320),
+            "n_external_cameras": 2,
+            "needs_wrist_camera": True,
+            "needs_stereo_camera": False,
+            "needs_session_id": True,
+            "action_space": "joint_position",
+        }
+    )
+    connection = openpi_connection.RobotRealtimeConnection(
+        websocket,
+        serving,
+        idle_timeout=0.01,
+    )
+
+    asyncio.run(connection.handle_connection())
+
+    assert websocket.accepted is True
+    assert websocket.sent_bytes[0]["action_space"] == "joint_position"
+    assert websocket.closed is True
+    assert websocket.sent_texts == []
+    serving.infer.assert_not_called()
+
+
+def test_handle_connection_keeps_session_state_per_websocket(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+    requests = {
+        b"a1": {"prompt": "first", "session_id": "session-a"},
+        b"a2": {"prompt": "second", "session_id": "session-a"},
+        b"b1": {"prompt": "other", "session_id": "session-b"},
+    }
+    monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data]))
+    serving = _serving_mock()
+
+    websocket_a = FakeWebSocket(
+        [
+            {"type": "websocket.receive", "bytes": b"a1"},
+            {"type": "websocket.receive", "bytes": b"a2"},
+            {"type": "websocket.disconnect"},
+        ]
+    )
+    websocket_b = FakeWebSocket(
+        [
+            {"type": "websocket.receive", "bytes": b"b1"},
+            {"type": "websocket.disconnect"},
+        ]
+    )
+
+    asyncio.run(openpi_connection.RobotRealtimeConnection(websocket_a, serving).handle_connection())
+    asyncio.run(openpi_connection.RobotRealtimeConnection(websocket_b, serving).handle_connection())
+
+    calls = serving.infer.await_args_list
+    assert calls[0].kwargs == {"session_id": "session-a", "reset": True}
+    assert calls[1].kwargs == {"session_id": "session-a", "reset": False}
+    assert calls[2].kwargs == {"session_id": "session-b", "reset": True}
+
+
+def test_handle_connection_reset_endpoint_resets_next_infer(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+    requests = {
+        b"a1": {"prompt": "first", "session_id": "session-a"},
+        b"reset": {"endpoint": "reset"},
+        b"a2": {"prompt": "second", "session_id": "session-a"},
+    }
+    monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data]))
+    serving = _serving_mock()
+    websocket = FakeWebSocket(
+        [
+            {"type": "websocket.receive", "bytes": b"a1"},
+            {"type": "websocket.receive", "bytes": b"reset"},
+            {"type": "websocket.receive", "bytes": b"a2"},
+            {"type": "websocket.disconnect"},
+        ]
+    )
+
+    asyncio.run(openpi_connection.RobotRealtimeConnection(websocket, serving).handle_connection())
+
+    assert [call.kwargs["reset"] for call in serving.infer.await_args_list] == [True, True]
+    serving.reset.assert_called_once_with({})
+    assert websocket.sent_bytes[2] == {"status": "reset successful"}
+    assert websocket.sent_texts == []

tests/entrypoints/openai_api/test_openpi_serving.py

@@ -0,0 +1,338 @@
+import asyncio
+import json
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from fastapi import FastAPI, WebSocket
+from omegaconf import OmegaConf
+from starlette.testclient import TestClient
+
+from vllm_omni.entrypoints.openpi import connection as openpi_connection
+from vllm_omni.entrypoints.openpi import serving as openpi_serving
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TEST_POLICY_SERVER_CONFIG = {
+    "image_resolution": (180, 320),
+    "n_external_cameras": 2,
+    "needs_wrist_camera": True,
+    "needs_stereo_camera": False,
+    "needs_session_id": True,
+    "action_space": "joint_position",
+}
+
+
+def _json_default(obj):
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
+
+
+def _json_pack(obj):
+    return json.dumps(obj, default=_json_default).encode()
+
+
+def _json_unpack(data):
+    return json.loads(data.decode())
+
+
+def _engine_with_policy_config(policy_config=None):
+    od_config = SimpleNamespace(model_config={"policy_server_config": policy_config or TEST_POLICY_SERVER_CONFIG})
+    return SimpleNamespace(get_diffusion_od_config=lambda: od_config)
+
+
+class RecordingEngine:
+    def __init__(self):
+        self.od_config = SimpleNamespace(model_config={"policy_server_config": TEST_POLICY_SERVER_CONFIG})
+        self.generate_calls = []
+
+    def get_diffusion_od_config(self):
+        return self.od_config
+
+    def generate(self, *, prompt, request_id, sampling_params_list):
+        async def _generate():
+            self.generate_calls.append(
+                {
+                    "prompt": prompt,
+                    "request_id": request_id,
+                    "sampling_params_list": sampling_params_list,
+                }
+            )
+            yield SimpleNamespace(multimodal_output={"actions": [0.0]})
+
+        return _generate()
+
+
+class ConcurrentRecordingEngine(RecordingEngine):
+    def __init__(self, *, expected_calls: int):
+        super().__init__()
+        self.expected_calls = expected_calls
+        self.condition = threading.Condition()
+        self.saw_overlap = False
+
+    def _wait_for_expected_calls(self):
+        with self.condition:
+            completed = self.condition.wait_for(
+                lambda: len(self.generate_calls) >= self.expected_calls,
+                timeout=5.0,
+            )
+            self.saw_overlap = self.saw_overlap or completed
+
+    def generate(self, *, prompt, request_id, sampling_params_list):
+        async def _generate():
+            with self.condition:
+                self.generate_calls.append(
+                    {
+                        "prompt": prompt,
+                        "request_id": request_id,
+                        "sampling_params_list": sampling_params_list,
+                    }
+                )
+                if len(self.generate_calls) >= self.expected_calls:
+                    self.saw_overlap = True
+                    self.condition.notify_all()
+
+            await asyncio.to_thread(self._wait_for_expected_calls)
+            yield SimpleNamespace(multimodal_output={"actions": [0.0]})
+
+        return _generate()
+
+
+def test_policy_server_config_reads_diffusion_model_config():
+    policy_config = {
+        "image_resolution": [64, 64],
+        "n_external_cameras": 1,
+        "custom_model_key": {"nested": True},
+    }
+    od_config = SimpleNamespace(model_config={"policy_server_config": policy_config})
+    engine_client = SimpleNamespace(get_diffusion_od_config=lambda: od_config)
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+    assert serving.policy_server_config.to_dict() == policy_config
+
+
+def test_policy_server_config_reads_stage_config_model_config():
+    policy_config = {"custom_model_key": "from-stage-config"}
+    engine_client = SimpleNamespace(
+        get_diffusion_od_config=lambda: None,
+        stage_configs=[
+            SimpleNamespace(
+                stage_type="diffusion",
+                engine_args=SimpleNamespace(model_config={"policy_server_config": policy_config}),
+            )
+        ],
+    )
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+    assert serving.policy_server_config.to_dict() == policy_config
+
+
+def test_policy_server_config_reads_omegaconf_stage_config():
+    engine_client = SimpleNamespace(
+        get_diffusion_od_config=lambda: None,
+        stage_configs=[
+            SimpleNamespace(
+                stage_type="diffusion",
+                engine_args=SimpleNamespace(
+                    model_config=OmegaConf.create({"policy_server_config": {"custom_model_key": "from-omegaconf"}})
+                ),
+            )
+        ],
+    )
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+    assert serving.policy_server_config.to_dict() == {"custom_model_key": "from-omegaconf"}
+
+
+def test_policy_server_config_is_required():
+    od_config = SimpleNamespace(model_config={})
+    engine_client = SimpleNamespace(get_diffusion_od_config=lambda: od_config)
+
+    with pytest.raises(ValueError) as exc_info:
+        openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+    assert "policy_server_config" in str(exc_info.value)
+
+
+def test_create_policy_server_returns_none_without_policy_config():
+    od_config = SimpleNamespace(model_config={})
+    engine_client = SimpleNamespace(get_diffusion_od_config=lambda: od_config)
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI.create_policy_server(
+        engine_client=engine_client,
+        model_name="generic-model",
+    )
+
+    assert serving is None
+
+
+def test_policy_server_config_reads_engine_model_config():
+    policy_config = {"custom_model_key": "custom-value"}
+    engine_client = SimpleNamespace(model_config=SimpleNamespace(policy_server_config=policy_config))
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+    assert serving.policy_server_config.to_dict() == policy_config
+
+
+def test_build_request_uses_unique_engine_request_id_per_inference():
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config())
+
+    request_a = serving._build_request(
+        {"prompt": "pick up the object"},
+        session_id="session-a",
+        reset=True,
+    )
+    request_b = serving._build_request(
+        {"prompt": "pick up the object"},
+        session_id="session-a",
+        reset=False,
+    )
+
+    assert request_a.sampling_params.extra_args["reset"] is True
+    assert request_b.sampling_params.extra_args["reset"] is False
+    assert request_a.sampling_params.extra_args["session_id"] == "session-a"
+    assert request_b.sampling_params.extra_args["session_id"] == "session-a"
+    assert request_a.sampling_params.extra_args["robot_obs"]["prompt"] == "pick up the object"
+    assert request_b.sampling_params.extra_args["robot_obs"]["prompt"] == "pick up the object"
+
+    assert request_a.request_id == "robot-session-a-0"
+    assert request_b.request_id == "robot-session-a-1"
+    assert request_a.request_id != request_b.request_id
+
+
+def test_infer_keeps_session_state_but_uses_unique_engine_request_ids():
+    engine = RecordingEngine()
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine)
+
+    async def run_requests():
+        await serving.infer({"prompt": "pick up the object"}, session_id="session-a", reset=True)
+        await serving.infer({"prompt": "pick up the object"}, session_id="session-a", reset=False)
+
+    asyncio.run(run_requests())
+
+    assert [call["request_id"] for call in engine.generate_calls] == [
+        "robot-session-a-0",
+        "robot-session-a-1",
+    ]
+    assert engine.generate_calls[0]["request_id"] != engine.generate_calls[1]["request_id"]
+
+    sampling_params_a = engine.generate_calls[0]["sampling_params_list"][0]
+    sampling_params_b = engine.generate_calls[1]["sampling_params_list"][0]
+    assert sampling_params_a.extra_args["session_id"] == "session-a"
+    assert sampling_params_b.extra_args["session_id"] == "session-a"
+    assert sampling_params_a.extra_args["reset"] is True
+    assert sampling_params_b.extra_args["reset"] is False
+
+
+def test_two_websocket_clients_without_session_id_do_not_conflict(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", _json_pack)
+    monkeypatch.setattr(openpi_connection, "_unpack", _json_unpack)
+
+    engine = ConcurrentRecordingEngine(expected_calls=2)
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine)
+    app = FastAPI()
+
+    @app.websocket("/v1/realtime/robot/openpi")
+    async def openpi_endpoint(websocket: WebSocket):
+        connection = openpi_connection.RobotRealtimeConnection(websocket, serving)
+        await connection.handle_connection()
+
+    def run_client(prompt: str):
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/robot/openpi") as websocket:
+                metadata = _json_unpack(websocket.receive_bytes())
+                assert metadata["needs_session_id"] is True
+
+                websocket.send_bytes(_json_pack({"prompt": prompt}))
+                actions = _json_unpack(websocket.receive_bytes())
+                np.testing.assert_array_equal(
+                    np.asarray(actions, dtype=np.float32),
+                    np.asarray([0.0], dtype=np.float32),
+                )
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        futures = [
+            executor.submit(run_client, "first client"),
+            executor.submit(run_client, "second client"),
+        ]
+        for future in futures:
+            future.result(timeout=10.0)
+
+    request_ids = [call["request_id"] for call in engine.generate_calls]
+    assert len(request_ids) == 2
+    assert len(set(request_ids)) == 2
+    assert all(request_id.startswith("robot-default-") for request_id in request_ids)
+    assert engine.saw_overlap is True
+
+    sampling_params = [call["sampling_params_list"][0] for call in engine.generate_calls]
+    assert [params.extra_args["session_id"] for params in sampling_params] == ["default", "default"]
+    assert [params.extra_args["reset"] for params in sampling_params] == [True, True]
+
+
+def test_infer_extracts_actions_from_generic_multimodal_output():
+    class FakeEngineClient:
+        def get_diffusion_od_config(self):
+            return SimpleNamespace(model_config={"policy_server_config": TEST_POLICY_SERVER_CONFIG})
+
+        async def generate(self, **kwargs):
+            self.generate_kwargs = kwargs
+            yield SimpleNamespace(multimodal_output={"actions": [[1.0, 2.0, 3.0]]})
+
+    engine_client = FakeEngineClient()
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+    actions = asyncio.run(serving.infer({"prompt": "pick up"}, session_id="session-a", reset=True))
+
+    np.testing.assert_allclose(actions, np.array([[1.0, 2.0, 3.0]], dtype=np.float32))
+    assert engine_client.generate_kwargs["prompt"] == "pick up"
+    assert engine_client.generate_kwargs["request_id"] == "robot-session-a-0"
+
+
+def test_infer_preserves_dict_actions_from_multimodal_output():
+    class FakeEngineClient:
+        def get_diffusion_od_config(self):
+            return SimpleNamespace(model_config={"policy_server_config": TEST_POLICY_SERVER_CONFIG})
+
+        async def generate(self, **kwargs):
+            self.generate_kwargs = kwargs
+            yield SimpleNamespace(
+                multimodal_output={
+                    "actions": {
+                        "left_arm": [[1.0, 2.0]],
+                        "right_arm": np.array([[3.0, 4.0]], dtype=np.float64),
+                    }
+                }
+            )
+
+    engine_client = FakeEngineClient()
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+    actions = asyncio.run(serving.infer({"prompt": "pick up"}, session_id="session-a", reset=True))
+
+    assert isinstance(actions, dict)
+    assert set(actions) == {"left_arm", "right_arm"}
+    np.testing.assert_allclose(actions["left_arm"], np.array([[1.0, 2.0]], dtype=np.float32))
+    np.testing.assert_allclose(actions["right_arm"], np.array([[3.0, 4.0]], dtype=np.float32))
+    assert actions["left_arm"].dtype == np.float32
+    assert actions["right_arm"].dtype == np.float32
+
+
+def test_extract_actions_does_not_iterate_result_object():
+    class IterableResult:
+        multimodal_output = {"actions": [[1.0, 2.0, 3.0]]}
+
+        def __iter__(self):
+            raise AssertionError("result object should not be iterated")
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config())
+
+    actions = serving._extract_actions(IterableResult())
+
+    np.testing.assert_allclose(actions, np.array([[1.0, 2.0, 3.0]], dtype=np.float32))