| |
| |
| |
|
|
| import logging, argparse |
| import time, os |
| from typing import Dict, Optional, Tuple |
|
|
| from typing_extensions import override |
| import websockets.sync.client |
|
|
| from . import msgpack_numpy |
|
|
|
|
| class WebsocketClientPolicy: |
| """Implements the Policy interface by communicating with a server over websocket. |
| |
| See WebsocketPolicyServer for a corresponding server implementation. |
| """ |
|
|
| def __init__(self, host: str = "127.0.0.1", port: Optional[int] = 10093, api_key: Optional[str] = None) -> None: |
| |
| self._uri = f"ws://{host}" |
| if port is not None: |
| self._uri += f":{port}" |
| self._packer = msgpack_numpy.Packer() |
| self._api_key = api_key |
| self._ws, self._server_metadata = self._wait_for_server() |
|
|
| def get_server_metadata(self) -> Dict: |
| return self._server_metadata |
|
|
| def _wait_for_server(self, timeout: float = 300) -> Tuple[websockets.sync.client.ClientConnection, Dict]: |
| logging.info(f"Waiting for server at {self._uri}...") |
| start_time = time.time() |
| |
| for k in ("HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy", "ALL_PROXY", "all_proxy"): |
| os.environ.pop(k, None) |
| |
| while True: |
| if time.time() - start_time > timeout: |
| raise TimeoutError(f"Failed to connect to server within {timeout} seconds") |
| |
| try: |
| headers = {"Authorization": f"Api-Key {self._api_key}"} if self._api_key else None |
| conn = websockets.sync.client.connect( |
| self._uri, |
| compression=None, |
| max_size=None, |
| additional_headers=headers, |
| open_timeout=150, |
| ping_interval=None, |
| ping_timeout=None, |
| ) |
| metadata = msgpack_numpy.unpackb(conn.recv()) |
| return conn, metadata |
| except ConnectionRefusedError: |
| logging.info(f"Still waiting for server {self._uri} ...") |
| time.sleep(2) |
|
|
| def close(self) -> None: |
| try: |
| self._ws.close() |
| except Exception: |
| pass |
| |
| @override |
| def predict_action(self, query_info: Dict) -> Dict: |
| data = self._packer.pack(query_info) |
| self._ws.send(data) |
| response = self._ws.recv() |
| if isinstance(response, str): |
| raise RuntimeError(f"Error in inference server:\n{response}") |
| return msgpack_numpy.unpackb(response) |
|
|
|
|
|
|
|
|
