File size: 5,813 Bytes
8101696 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | # Copyright 2025 starVLA community. All rights reserved.
# Licensed under the MIT License, Version 1.0 (the "License");
import asyncio
import logging
import traceback
import time
import websockets.asyncio.server
import websockets.frames
# from openpi_client import base_policy as _base_policy
from . import msgpack_numpy
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,
host: str = "0.0.0.0",
port: int = 10093,
idle_timeout: int = -1, # Negative values disable the idle timeout.
metadata: dict | None = None,
) -> None:
self._policy = policy
self._host = host
self._port = port
self._metadata = metadata or {}
self._idle_timeout = idle_timeout
self._last_active = time.time()
logging.getLogger("websockets.server").setLevel(logging.INFO)
def serve_forever(self) -> None:
asyncio.run(self.run())
async def run(self):
async with websockets.asyncio.server.serve(
self._handler,
self._host,
self._port,
compression=None,
max_size=None,
) as server:
if self._idle_timeout > 0:
await self._idle_watchdog(server)
else:
await server.serve_forever()
async def _idle_watchdog(self, server):
"""Close the server after the configured idle period."""
while True:
await asyncio.sleep(5)
if time.time() - self._last_active > self._idle_timeout:
logging.info(f"Idle timeout ({self._idle_timeout}s) reached, shutting down server.")
server.close()
await server.wait_closed()
break
async def _handler(self, websocket: websockets.asyncio.server.ServerConnection):
logging.info(f"Connection from {websocket.remote_address} opened")
packer = msgpack_numpy.Packer()
await websocket.send(packer.pack(self._metadata))
while True:
try:
msg = msgpack_numpy.unpackb(await websocket.recv())
self._last_active = time.time()
ret = self._route_message(msg)
await websocket.send(packer.pack(ret))
except websockets.ConnectionClosed:
logging.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
# route logic: recognize request from client
def _route_message(self, msg: dict) -> dict:
"""
Route rules (fault-tolerant):
- Supports messages of form:
{"type": "ping|init|infer|reset", "request_id": "...", "payload": {...}}
or a flat dict (will be treated as payload).
- Does NOT raise inside this function: all exceptions are caught and encoded in response.
"""
if not isinstance(msg, dict):
return {
"status": "error",
"ok": False,
"type": "inference_result",
"request_id": "default",
"error": {"message": "Payload must be a dict", "payload_type": str(type(msg))},
}
req_id = msg.get("request_id", "default")
mtype = msg.get("type", "infer") # default = infer
# The envelope fields are addressed to this router, not to the policy.
payload = {key: value for key, value in msg.items() if key not in ("type", "request_id")}
# ping
if mtype == "ping":
return {"status": "ok", "ok": True, "type": "ping", "request_id": req_id}
# infer --> framework.predict_action
elif mtype == "infer" or mtype == "predict_action":
try:
ouput_dict = self._policy.predict_action(**payload)
except Exception as e:
logging.exception("Policy inference error (request_id=%s)", req_id)
logging.exception(e)
return {
"status": "error",
"ok": False,
"type": "inference_result",
"request_id": req_id,
"error": {
"message": str(e),
# "traceback": traceback.format_exc(),
},
}
data = ouput_dict
return {
"status": "ok",
"ok": True,
"type": "inference_result",
"request_id": req_id,
"data": data,
}
# unknow request type
else:
return {
"status": "error",
"ok": False,
"type": "unknown",
"request_id": req_id,
"error": {"message": f"Unsupported message type '{mtype}'"},
}
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO, force=True)
# Example usage:
# policy = YourPolicyClass() # Replace with your actual policy class
# server = WebsocketPolicyServer(policy, host="localhost", port=10091)
# server.serve_forever()
raise NotImplementedError("This module is not intended to be run directly.")
#
# Instead, it should be imported and used in a server context.
|