vllm-project--vllm-omni
wehub-resource-sync
318 行
10 KiB
Python
318 行
10 KiB
Python
import asyncio
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import numpy as np
|
|
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_and_unpack_round_trip_numpy_values():
|
|
payload = {
|
|
"image": np.arange(6, dtype=np.uint8).reshape(2, 3),
|
|
"action": np.asarray([[1.0, 2.0]], dtype=np.float32),
|
|
"scalar": np.float32(3.5),
|
|
"nested": [{"done": np.bool_(True)}],
|
|
}
|
|
|
|
decoded = openpi_connection._unpack(openpi_connection._pack(payload))
|
|
|
|
np.testing.assert_array_equal(decoded["image"], payload["image"])
|
|
np.testing.assert_allclose(decoded["action"], payload["action"])
|
|
assert decoded["image"].dtype == np.uint8
|
|
assert decoded["action"].dtype == np.float32
|
|
assert decoded["scalar"] == np.float32(3.5)
|
|
assert decoded["nested"][0]["done"] == np.bool_(True)
|
|
|
|
|
|
def test_unpack_accepts_msgpack_numpy_marker_dicts():
|
|
action = np.asarray([[1.0, 2.0]], dtype=np.float32)
|
|
payload = {
|
|
b"actions": {
|
|
b"nd": True,
|
|
b"type": action.dtype.str,
|
|
b"kind": action.dtype.kind,
|
|
b"shape": action.shape,
|
|
b"data": action.tobytes(),
|
|
}
|
|
}
|
|
|
|
decoded = openpi_connection._unpack_numpy(payload)
|
|
|
|
np.testing.assert_allclose(decoded[b"actions"], action)
|
|
assert decoded[b"actions"].flags.writeable is True
|
|
decoded[b"actions"][:, -1] = 0.0
|
|
np.testing.assert_allclose(decoded[b"actions"], np.asarray([[1.0, 0.0]], dtype=np.float32))
|
|
|
|
|
|
def test_unpack_accepts_openpi_client_ndarray_markers():
|
|
image = np.arange(6, dtype=np.uint8).reshape(2, 3)
|
|
payload = {
|
|
"observation/exterior_image_0_left": {
|
|
b"__ndarray__": True,
|
|
b"dtype": image.dtype.str,
|
|
b"shape": image.shape,
|
|
b"data": image.tobytes(),
|
|
}
|
|
}
|
|
|
|
decoded = openpi_connection._unpack_numpy(payload)
|
|
|
|
np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image)
|
|
assert decoded["observation/exterior_image_0_left"].dtype == np.uint8
|
|
|
|
|
|
def test_unpack_openpi_client_packed_observation():
|
|
openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy")
|
|
image = np.zeros((180, 320, 3), dtype=np.uint8)
|
|
obs = {
|
|
"observation/exterior_image_0_left": image,
|
|
"observation/joint_position": np.zeros(7, dtype=np.float32),
|
|
}
|
|
|
|
decoded = openpi_connection._unpack(openpi_msgpack.packb(obs))
|
|
|
|
np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image)
|
|
np.testing.assert_allclose(decoded["observation/joint_position"], obs["observation/joint_position"])
|
|
|
|
|
|
def test_pack_uses_openpi_client_ndarray_markers():
|
|
openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy")
|
|
actions = np.asarray([[1.0, 2.0]], dtype=np.float32)
|
|
|
|
packed = openpi_connection._pack(actions)
|
|
decoded = openpi_msgpack.unpackb(packed)
|
|
|
|
np.testing.assert_allclose(decoded, actions)
|
|
|
|
|
|
def test_unpack_leaves_user_dict_without_numpy_kind_marker_unchanged():
|
|
payload = {
|
|
"metadata": {
|
|
"nd": True,
|
|
"type": "<f4",
|
|
"data": b"user payload",
|
|
}
|
|
}
|
|
|
|
decoded = openpi_connection._unpack_numpy(payload)
|
|
|
|
assert decoded == 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 == []
|