|
|
import asyncio |
|
|
import http |
|
|
import logging |
|
|
import time |
|
|
import traceback |
|
|
|
|
|
from openpi_client import base_policy as _base_policy |
|
|
from openpi_client import msgpack_numpy |
|
|
import websockets.asyncio.server as _server |
|
|
import websockets.frames |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
class WebsocketPolicyServer: |
|
|
"""Serves a policy using the websocket protocol. See websocket_client_policy.py for a client implementation. |
|
|
|
|
|
Currently only implements the `load` and `infer` methods. |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
policy: _base_policy.BasePolicy, |
|
|
host: str = "0.0.0.0", |
|
|
port: int | None = None, |
|
|
metadata: dict | None = None, |
|
|
) -> None: |
|
|
self._policy = policy |
|
|
self._host = host |
|
|
self._port = port |
|
|
self._metadata = metadata or {} |
|
|
logging.getLogger("websockets.server").setLevel(logging.INFO) |
|
|
|
|
|
def serve_forever(self) -> None: |
|
|
asyncio.run(self.run()) |
|
|
|
|
|
async def run(self): |
|
|
async with _server.serve( |
|
|
self._handler, |
|
|
self._host, |
|
|
self._port, |
|
|
compression=None, |
|
|
max_size=None, |
|
|
process_request=_health_check, |
|
|
) as server: |
|
|
await server.serve_forever() |
|
|
|
|
|
async def _handler(self, websocket: _server.ServerConnection): |
|
|
logger.info(f"Connection from {websocket.remote_address} opened") |
|
|
packer = msgpack_numpy.Packer() |
|
|
|
|
|
await websocket.send(packer.pack(self._metadata)) |
|
|
|
|
|
prev_total_time = None |
|
|
while True: |
|
|
try: |
|
|
start_time = time.monotonic() |
|
|
obs = msgpack_numpy.unpackb(await websocket.recv()) |
|
|
|
|
|
infer_time = time.monotonic() |
|
|
action = self._policy.infer(obs) |
|
|
infer_time = time.monotonic() - infer_time |
|
|
|
|
|
action["server_timing"] = { |
|
|
"infer_ms": infer_time * 1000, |
|
|
} |
|
|
if prev_total_time is not None: |
|
|
|
|
|
action["server_timing"]["prev_total_ms"] = prev_total_time * 1000 |
|
|
|
|
|
await websocket.send(packer.pack(action)) |
|
|
prev_total_time = time.monotonic() - start_time |
|
|
|
|
|
except websockets.ConnectionClosed: |
|
|
logger.info(f"Connection from {websocket.remote_address} closed") |
|
|
break |
|
|
except Exception: |
|
|
await websocket.send(traceback.format_exc()) |
|
|
await websocket.close( |
|
|
code=websockets.frames.CloseCode.INTERNAL_ERROR, |
|
|
reason="Internal server error. Traceback included in previous frame.", |
|
|
) |
|
|
raise |
|
|
|
|
|
|
|
|
def _health_check(connection: _server.ServerConnection, request: _server.Request) -> _server.Response | None: |
|
|
if request.path == "/healthz": |
|
|
return connection.respond(http.HTTPStatus.OK, "OK\n") |
|
|
|
|
|
return None |
|
|
|
