File size: 5,891 Bytes
977b06e | 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");
# Implemented by [Jinhui YE / HKUST University] in [2025].
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, # 新增参数，单位秒，-1表示永不关闭
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):
"""监控空闲时间，超时则关闭服务器"""
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) # route message
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.
"""
req_id = msg.get("request_id", "default")
mtype = msg.get("type", "infer") # default = infer
msg # when no explicit payload, treat top-level as payload
# 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":
# Basic payload sanity
if not isinstance(msg, dict):
return {
"status": "error",
"ok": False,
"type": "inference_result",
"request_id": req_id,
"error": {"message": "Payload must be a dict", "payload_type": str(type(payload))}
}
try:
ouput_dict = self._policy.predict_action(**msg)
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. |