| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| from dataclasses import dataclass |
| import functools |
| import io |
| import json |
| from typing import Any, Callable |
|
|
| import msgpack |
| import msgpack_numpy as mnp |
| import numpy as np |
| import zmq |
|
|
| from gr00t.data.types import ModalityConfig |
| from gr00t.data.utils import to_json_serializable |
|
|
| from .policy import BasePolicy |
|
|
|
|
| class MsgSerializer: |
| """msgpack_numpy serializer with a hard ``allow_pickle=False`` boundary. |
| |
| Implementation note: msgpack_numpy's ``Packer``/``Unpacker`` wire any |
| user-provided ``default``/``object_hook`` *behind* their own |
| ``mnp.encode``/``mnp.decode`` (``functools.partial(encode, chain=user_fn)``). |
| That means an object-dtype ndarray is serialised by ``mnp.encode`` (and a |
| forged ``{nd: True, kind: 'O', ...}`` payload is fed to ``pickle.loads`` by |
| ``mnp.decode``) *before* a hook installed via ``mnp.packb``/``mnp.unpackb`` |
| can intervene. We therefore drive ``msgpack`` directly and chain |
| ``_safe_{encode,decode} → mnp.{encode,decode} → custom`` ourselves, so the |
| refusal runs first. |
| """ |
|
|
| @staticmethod |
| def to_bytes(data: Any) -> bytes: |
| default = functools.partial(MsgSerializer._safe_encode, chain=MsgSerializer._encode_custom) |
| return msgpack.packb(data, default=default) |
|
|
| @staticmethod |
| def from_bytes(data: bytes) -> Any: |
| object_hook = functools.partial( |
| MsgSerializer._safe_decode, chain=MsgSerializer._decode_custom |
| ) |
| return msgpack.unpackb(data, object_hook=object_hook, raw=False) |
|
|
| @staticmethod |
| def _safe_encode(obj, chain=None): |
| |
| |
| |
| |
| if isinstance(obj, np.ndarray): |
| if obj.dtype.kind == "O": |
| raise TypeError( |
| f"Refusing to encode object-dtype ndarray (shape={obj.shape}); " |
| f"msgpack_numpy would invoke pickle. Convert to a concrete " |
| f"numeric dtype before sending." |
| ) |
| return mnp.encode(obj, chain=chain) |
|
|
| @staticmethod |
| def _safe_decode(obj, chain=None): |
| if isinstance(obj, dict): |
| marker = obj.get("__ndarray_class__", obj.get(b"__ndarray_class__")) |
| if marker: |
| payload = obj.get("as_npy", obj.get(b"as_npy")) |
| if payload is None: |
| raise ValueError( |
| "Malformed ndarray payload: marker present but 'as_npy' missing" |
| ) |
| return np.load(io.BytesIO(payload), allow_pickle=False) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| if isinstance(obj, dict): |
| nd_val = obj.get(b"nd", obj.get("nd")) |
| kind_val = obj.get(b"kind", obj.get("kind")) |
| if nd_val and kind_val in (b"O", "O"): |
| raise ValueError( |
| "Refusing to decode object-dtype ndarray payload (pickle-bearing); " |
| "the allow_pickle=False contract is enforced by MsgSerializer." |
| ) |
| return mnp.decode(obj, chain=chain) |
|
|
| @staticmethod |
| def _encode_custom(obj): |
| if isinstance(obj, ModalityConfig): |
| return {"__ModalityConfig__": True, "as_json": to_json_serializable(obj)} |
| return obj |
|
|
| @staticmethod |
| def _decode_custom(obj): |
| if not isinstance(obj, dict): |
| return obj |
| |
| |
| has_modality_marker = ( |
| "__ModalityConfig__" in obj |
| or b"__ModalityConfig__" in obj |
| or "__ModalityConfig_class__" in obj |
| or b"__ModalityConfig_class__" in obj |
| ) |
| if has_modality_marker: |
| key = next((k for k in ("as_json", b"as_json") if k in obj), None) |
| if key is None: |
| raise ValueError( |
| f"Malformed ModalityConfig payload: marker present but " |
| f"'as_json' missing. keys={sorted(repr(k) for k in obj.keys())}" |
| ) |
| payload = obj[key] |
| if isinstance(payload, bytes): |
| payload = payload.decode() |
| if isinstance(payload, str): |
| payload = json.loads(payload) |
| return ModalityConfig(**payload) |
| return obj |
|
|
|
|
| @dataclass |
| class EndpointHandler: |
| handler: Callable |
| requires_input: bool = True |
|
|
|
|
| class PolicyServer: |
| """ |
| An inference server that spin up a ZeroMQ socket and listen for incoming requests. |
| Can add custom endpoints by calling `register_endpoint`. |
| """ |
|
|
| |
| |
| _CLOSE_LINGER_MS = 1000 |
|
|
| def __init__( |
| self, |
| policy: BasePolicy, |
| host: str = "*", |
| port: int = 5555, |
| api_token: str = None, |
| ): |
| self.policy = policy |
| self.host = host |
| self.port = port |
| self.running = True |
| self._closed = False |
| self.context = zmq.Context() |
| self.socket = self.context.socket(zmq.REP) |
| self.socket.bind(f"tcp://{host}:{port}") |
| self._endpoints: dict[str, EndpointHandler] = {} |
| self.api_token = api_token |
|
|
| |
| self.register_endpoint("ping", self._handle_ping, requires_input=False) |
| self.register_endpoint("kill", self._kill_server, requires_input=False) |
| self.register_endpoint("get_action", self.policy.get_action) |
| self.register_endpoint("reset", self.policy.reset) |
| self.register_endpoint( |
| "get_modality_config", |
| getattr(self.policy, "get_modality_config", lambda: {}), |
| requires_input=False, |
| ) |
|
|
| def _kill_server(self): |
| """Stop the run loop. Does not release the socket / context — use ``close()``.""" |
| self.running = False |
|
|
| def close(self) -> None: |
| """Release the bound socket and ZMQ context. Idempotent.""" |
| if getattr(self, "_closed", True): |
| return |
| self._closed = True |
| self.running = False |
| socket = getattr(self, "socket", None) |
| if socket is not None: |
| try: |
| socket.close(linger=self._CLOSE_LINGER_MS) |
| except Exception: |
| pass |
| context = getattr(self, "context", None) |
| if context is not None: |
| try: |
| context.term() |
| except Exception: |
| pass |
|
|
| def __enter__(self): |
| |
| |
| print(f"\n✓ Server ready — listening on {self.host}:{self.port}\n") |
| return self |
|
|
| def __exit__(self, exc_type, exc, tb): |
| self.close() |
|
|
| def _handle_ping(self) -> dict: |
| """ |
| Simple ping handler that returns a success message. |
| """ |
| return {"status": "ok", "message": "Server is running"} |
|
|
| def register_endpoint(self, name: str, handler: Callable, requires_input: bool = True): |
| """ |
| Register a new endpoint to the server. |
| |
| Args: |
| name: The name of the endpoint. |
| handler: The handler function that will be called when the endpoint is hit. |
| requires_input: Whether the handler requires input data. |
| """ |
| self._endpoints[name] = EndpointHandler(handler, requires_input) |
|
|
| def _validate_token(self, request: dict) -> bool: |
| """ |
| Validate the API token in the request. |
| """ |
| if self.api_token is None: |
| return True |
| return request.get("api_token") == self.api_token |
|
|
| def run(self): |
| addr = self.socket.getsockopt_string(zmq.LAST_ENDPOINT) |
| print(f"Server is ready and listening on {addr}") |
| while self.running: |
| try: |
| message = self.socket.recv() |
| request = MsgSerializer.from_bytes(message) |
|
|
| |
| if not self._validate_token(request): |
| self.socket.send( |
| MsgSerializer.to_bytes({"error": "Unauthorized: Invalid API token"}) |
| ) |
| continue |
|
|
| endpoint = request.get("endpoint", "get_action") |
|
|
| if endpoint not in self._endpoints: |
| raise ValueError(f"Unknown endpoint: {endpoint}") |
|
|
| handler = self._endpoints[endpoint] |
| result = ( |
| handler.handler(**request.get("data", {})) |
| if handler.requires_input |
| else handler.handler() |
| ) |
| self.socket.send(MsgSerializer.to_bytes(result)) |
| except Exception as e: |
| print(f"Error in server: {e}") |
| import traceback |
|
|
| print(traceback.format_exc()) |
| self.socket.send(MsgSerializer.to_bytes({"error": str(e)})) |
|
|
| @staticmethod |
| def start_server(policy: BasePolicy, port: int, host: str = "*", api_token: str = None): |
| with PolicyServer(policy, host=host, port=port, api_token=api_token) as server: |
| server.run() |
|
|
|
|
| class PolicyClient(BasePolicy): |
| def __init__( |
| self, |
| host: str = "localhost", |
| port: int = 5555, |
| timeout_ms: int = 15000, |
| api_token: str = None, |
| strict: bool = False, |
| ): |
| super().__init__(strict=strict) |
| self._closed = False |
| self.context = zmq.Context() |
| self.host = host |
| self.port = port |
| self.timeout_ms = timeout_ms |
| self.api_token = api_token |
| self._init_socket() |
|
|
| def _init_socket(self): |
| """Initialize or reinitialize the socket with current settings""" |
| self.socket = self.context.socket(zmq.REQ) |
| self.socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms) |
| self.socket.setsockopt(zmq.SNDTIMEO, self.timeout_ms) |
| self.socket.connect(f"tcp://{self.host}:{self.port}") |
|
|
| def ping(self) -> bool: |
| try: |
| self.call_endpoint("ping", requires_input=False) |
| return True |
| except zmq.error.ZMQError: |
| self._init_socket() |
| return False |
|
|
| def kill_server(self): |
| """ |
| Kill the server. |
| """ |
| self.call_endpoint("kill", requires_input=False) |
|
|
| def call_endpoint( |
| self, endpoint: str, data: dict | None = None, requires_input: bool = True |
| ) -> Any: |
| """ |
| Call an endpoint on the server. |
| |
| Args: |
| endpoint: The name of the endpoint. |
| data: The input data for the endpoint. |
| requires_input: Whether the endpoint requires input data. |
| """ |
| request: dict = {"endpoint": endpoint} |
| if requires_input: |
| request["data"] = data |
| if self.api_token: |
| request["api_token"] = self.api_token |
|
|
| try: |
| self.socket.send(MsgSerializer.to_bytes(request)) |
| message = self.socket.recv() |
| except zmq.error.Again: |
| |
| |
| |
| self._init_socket() |
| raise |
| if message == b"ERROR": |
| raise RuntimeError("Server error. Make sure we are running the correct policy server.") |
| response = MsgSerializer.from_bytes(message) |
|
|
| if isinstance(response, dict) and "error" in response: |
| raise RuntimeError(f"Server error: {response['error']}") |
| return response |
|
|
| def close(self) -> None: |
| """Release the REQ socket and ZMQ context. Idempotent.""" |
| if getattr(self, "_closed", True): |
| return |
| self._closed = True |
| socket = getattr(self, "socket", None) |
| if socket is not None: |
| try: |
| socket.close(linger=0) |
| except Exception: |
| pass |
| context = getattr(self, "context", None) |
| if context is not None: |
| try: |
| context.term() |
| except Exception: |
| pass |
|
|
| def __enter__(self): |
| return self |
|
|
| def __exit__(self, exc_type, exc, tb): |
| self.close() |
|
|
| def __del__(self): |
| |
| |
| |
| try: |
| self.close() |
| except Exception: |
| pass |
|
|
| def _get_action( |
| self, observation: dict[str, Any], options: dict[str, Any] | None = None |
| ) -> tuple[dict[str, Any], dict[str, Any]]: |
| response = self.call_endpoint( |
| "get_action", {"observation": observation, "options": options} |
| ) |
| return tuple(response) |
|
|
| def reset(self, options: dict[str, Any] | None = None) -> dict[str, Any]: |
| return self.call_endpoint("reset", {"options": options}) |
|
|
| def get_modality_config(self) -> dict[str, ModalityConfig]: |
| return self.call_endpoint("get_modality_config", requires_input=False) |
|
|
| def check_observation(self, observation: dict[str, Any]) -> None: |
| raise NotImplementedError( |
| "check_observation is not implemented. Please use `strict=False` to disable strict mode or implement this method in the subclass." |
| ) |
|
|
| def check_action(self, action: dict[str, Any]) -> None: |
| raise NotImplementedError( |
| "check_action is not implemented. Please use `strict=False` to disable strict mode or implement this method in the subclass." |
| ) |
|
|
