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): if policy_config is None: policy_config = TEST_POLICY_SERVER_CONFIG od_config = SimpleNamespace(model_config={"policy_server_config": policy_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_allows_explicit_empty_config(): serving = openpi_serving.ServingRealtimeRobotOpenPI( engine_client=_engine_with_policy_config(policy_config={}), model_name="nvidia/Cosmos3-Nano-Policy-DROID", ) assert serving.policy_server_config.to_dict() == {} 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))