ReTime / deployment /model_server /tools /websocket_policy_server.py

Add Florence-Bridge code and handoff documentation

977b06e verified 7 days ago

5.89 kB

# 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.

	# 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.