LADDER RoboLab-Long for LeRobot EnvHub (LADDER v1.0.0)

01171b7 verified 3 days ago

9.49 kB

from abc import abstractmethod
from typing import Any, Dict, Optional, TypeVar

import zmq
from typing_extensions import Protocol

from ladder.components.inference.api.core.requests import RequestType, parse_request_type
from ladder.components.inference.api.core.serialization import deserialize, serialize
from ladder.core.common.logger import get_logger
from ladder.core.databib.config import Config, Configurable

LOGGER = get_logger(__name__)


class CallbackProtocol(Protocol):
    def __call__(self, message: Dict) -> Dict:
        ...


class RequestReplyServerConfig(Config):
    pass


RequestReplyServerConfigT = TypeVar('RequestReplyServerConfigT', bound=RequestReplyServerConfig)


class RequestReplyServer(Configurable[RequestReplyServerConfigT]):
    """
    Abstract base class for a request-reply server that listens
    for requests, dispatches them to the provided callback, and
    sends the output back to the client.
    Methods:
        start(on_request_cb: CallbackProtocol): Begins listening for requests
    """

    @abstractmethod
    def start(self, on_request_cb: CallbackProtocol) -> None:
        pass


class RequestReplyClientConfig(Config):
    pass


RequestReplyClientConfigT = TypeVar('RequestReplyClientConfigT', bound=RequestReplyClientConfig)


class RequestReplyClient(Configurable[RequestReplyClientConfigT]):
    """
    Abstract base class for a request-reply client that sends
    requests to the server and receives the output from the server.
    Methods:
        send(request: Dict) -> Optional[Dict]: Send a request to the server.
    """

    @abstractmethod
    def send(self, request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        pass


class ZmqRequestReplyServerConfig(RequestReplyServerConfig):
    port: int  # Port number to listen on
    recv_timeout_ms: int = 5000  # Timeout in milliseconds for calls to socket.recv


class ZmqRequestReplyServer(RequestReplyServer[ZmqRequestReplyServerConfig]):
    """
    ZeroMQ-based implementation of RequestReplyServer.
    Methods:
        start(on_request_cb: CallbackProtocol): Begins listening for requests.
    Notes:
        Designed to be used in conjunction with the ZmqRequestReplyClient class.
    """

    def __init__(self, config: ZmqRequestReplyServerConfig):
        super().__init__(config)
        self._context = zmq.Context()
        self._init_socket()

    def _init_socket(self):
        self.socket = self._context.socket(zmq.REP)
        self.socket.bind(f"tcp://*:{self.config.port}")
        self.socket.setsockopt(zmq.SNDHWM, 5)
        self.socket.setsockopt(zmq.LINGER, 0)
        self.socket.setsockopt(zmq.RCVTIMEO, self.config.recv_timeout_ms)

    def start(self, on_request_cb: CallbackProtocol) -> None:
        """
        Starts listening for requests over the socket, blocking on the current thread.
        When a request is received, the on_request_cb function is called with the request
        as an argument, and the output is sent back as a reply to the client.
        Args:
            on_request_cb (CallbackProtocol): Function to be called when a request is received.
        Notes: this function is blocking
        """
        LOGGER.info("Starting ZMQ request reply server")
        while True:
            LOGGER.debug("Waiting for new request")
            try:
                message = self.socket.recv()
            except zmq.Again:
                LOGGER.warn("ZMQ receive timeout. Continuing.")
                continue
            except zmq.ZMQError as e:
                LOGGER.error(f"A ZMQ error occured while receiving. Error: {e}")
                raise e
            LOGGER.debug(f"Received new request")

            LOGGER.debug(f"Processing new request")
            message = deserialize(message)
            res = on_request_cb(message)
            res = serialize(res)
            LOGGER.debug(f"Request processed")

            LOGGER.debug(f"Sending reply")
            try:
                self.socket.send(res)
            except zmq.ZMQError as e:
                LOGGER.error(f"A ZMQ error occured while sending. Error: {e}")
                raise e
            LOGGER.debug(f"Reply sent")

            # If the request is an exit request, we stop the loop and
            # let the process exit gracefully.
            if parse_request_type(message) == RequestType.EXIT:
                LOGGER.info("Received EXIT request. Stopping server.")
                break

    def __del__(self):
        if self.socket:
            self.socket.close()
        self._context.term()


class ZmqRequestReplyClientConfig(RequestReplyClientConfig):
    ip: str  # IP address of the server
    port: int  # Port number of the server
    recv_timeout_ms: int = 15000  # Timeout in milliseconds for calls to socket.recv
    send_num_retries: int = 3  # Number of times to retry sending a message
    # If True, recv timeouts are raised as ``zmq.Again`` to the caller after the broken
    # REQ socket has been reset. Use when a timeout indicates the worker should fail and
    # be respawned by its supervisor (e.g. a remote inference client). Default False
    # keeps the "log and return None" behaviour.
    propagate_timeout: bool = False


class ZmqRequestReplyClient(RequestReplyClient[ZmqRequestReplyClientConfig]):
    """
    ZeroMQ-based implementation of RequestReplyClient.
    It communicates with the server using a ZeroMQ in REQ mode.
    Methods:
        send(request: dict) -> Optional[Dict]: Send a request to the server and
            receive a reply.
    Notes:
        This class is designed to be used in conjunction with the
        ZmqRequestReplyServer class.
    """

    def __init__(self, config: ZmqRequestReplyClientConfig):
        super().__init__(config)
        self._closed = False
        self.context = zmq.Context()
        self.socket = None
        self._init_socket()

    def _init_socket(self):
        """
        Reset the socket connection, this is needed when REQ is in a broken state.
        """
        if self.socket:
            LOGGER.info(f"Closing and deleting existing socket")
            self.socket.close()
            del self.socket
        self.socket = self.context.socket(zmq.REQ)
        self.socket.setsockopt(zmq.RCVTIMEO, self.config.recv_timeout_ms)
        # this option is needed, among other things, to prevent context.term()
        # from handing if socket.recv() is interrupted e.g. by a KeyboardInterrupt
        self.socket.setsockopt(zmq.LINGER, 0)

        LOGGER.info(f"Connecting to {self.config.ip}:{self.config.port}")
        self.socket.connect(f"tcp://{self.config.ip}:{self.config.port}")
        LOGGER.info(f"Connected to {self.config.ip}:{self.config.port}")

    def send(self, request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        if self.socket is None or self.socket.closed:
            LOGGER.error(f"The socket is closed, request not sent.")
            return None

        try:
            LOGGER.debug(f"Sending request")
            sent = False
            request_serialized = serialize(request)
            for i in range(self.config.send_num_retries):
                try:
                    self.socket.send(request_serialized)
                    sent = True
                    LOGGER.debug(f"Request sent.")
                    break
                except zmq.ZMQError as e:
                    LOGGER.error(
                        f"Failed to send message to {self.config.ip}:{self.config.port}. Error: {e}."
                    )

                LOGGER.debug(f"Attempt {i} failed. Retrying.")

            if not sent:
                LOGGER.warn("Number of retries exceeded. Aborting.")
                return None

            # If the request is an exit request, we don't wait for a reply.
            if request.get("type") == RequestType.EXIT:
                return None

            LOGGER.debug(f"Waiting for reply")
            try:
                message = self.socket.recv()
            except zmq.Again:
                # REQ socket is now stuck in "waiting for reply", so any subsequent send
                # would fail with ZMQError. Reset before deciding how to surface the timeout.
                self._init_socket()
                if self.config.propagate_timeout:
                    raise
                LOGGER.warn("ZMQ receive timeout. Aborting.")
                return None
            except zmq.ZMQError as e:
                LOGGER.error(f"A ZMQ error occured while receiving. Error: {e}.")
                raise e

            return deserialize(message)

        except KeyboardInterrupt:
            # Ctrl+C might be sent mid-episode. In this case, if the exception is raised inside the try
            # block, the socket will fall into a broken state. Here we reinitialize it to keep the
            # communication working.
            self._init_socket()
            raise

    def close(self) -> None:
        """Idempotently tear down the socket and ZMQ context. Safe to call multiple times."""
        if self._closed:
            return
        if self.socket:
            try:
                self.socket.disconnect(f"tcp://{self.config.ip}:{self.config.port}")
            except zmq.ZMQError:
                # Socket already disconnected (e.g. _init_socket re-created it after a timeout).
                pass
            self.socket.close()
        self.context.term()
        self._closed = True

    def __del__(self):
        self.close()

	from abc import abstractmethod
	from typing import Any, Dict, Optional, TypeVar

	import zmq
	from typing_extensions import Protocol

	from ladder.components.inference.api.core.requests import RequestType, parse_request_type
	from ladder.components.inference.api.core.serialization import deserialize, serialize
	from ladder.core.common.logger import get_logger
	from ladder.core.databib.config import Config, Configurable

	LOGGER = get_logger(__name__)


	class CallbackProtocol(Protocol):
	def __call__(self, message: Dict) -> Dict:
	...


	class RequestReplyServerConfig(Config):
	pass


	RequestReplyServerConfigT = TypeVar('RequestReplyServerConfigT', bound=RequestReplyServerConfig)


	class RequestReplyServer(Configurable[RequestReplyServerConfigT]):
	"""
	Abstract base class for a request-reply server that listens
	for requests, dispatches them to the provided callback, and
	sends the output back to the client.

	Methods:
	start(on_request_cb: CallbackProtocol): Begins listening for requests
	"""

	@abstractmethod
	def start(self, on_request_cb: CallbackProtocol) -> None:
	pass


	class RequestReplyClientConfig(Config):
	pass


	RequestReplyClientConfigT = TypeVar('RequestReplyClientConfigT', bound=RequestReplyClientConfig)


	class RequestReplyClient(Configurable[RequestReplyClientConfigT]):
	"""
	Abstract base class for a request-reply client that sends
	requests to the server and receives the output from the server.

	Methods:
	send(request: Dict) -> Optional[Dict]: Send a request to the server.
	"""

	@abstractmethod
	def send(self, request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
	pass


	class ZmqRequestReplyServerConfig(RequestReplyServerConfig):
	port: int # Port number to listen on
	recv_timeout_ms: int = 5000 # Timeout in milliseconds for calls to socket.recv


	class ZmqRequestReplyServer(RequestReplyServer[ZmqRequestReplyServerConfig]):
	"""
	ZeroMQ-based implementation of RequestReplyServer.

	Methods:
	start(on_request_cb: CallbackProtocol): Begins listening for requests.

	Notes:
	Designed to be used in conjunction with the ZmqRequestReplyClient class.
	"""

	def __init__(self, config: ZmqRequestReplyServerConfig):
	super().__init__(config)
	self._context = zmq.Context()
	self._init_socket()

	def _init_socket(self):
	self.socket = self._context.socket(zmq.REP)
	self.socket.bind(f"tcp://*:{self.config.port}")
	self.socket.setsockopt(zmq.SNDHWM, 5)
	self.socket.setsockopt(zmq.LINGER, 0)
	self.socket.setsockopt(zmq.RCVTIMEO, self.config.recv_timeout_ms)

	def start(self, on_request_cb: CallbackProtocol) -> None:
	"""
	Starts listening for requests over the socket, blocking on the current thread.
	When a request is received, the on_request_cb function is called with the request
	as an argument, and the output is sent back as a reply to the client.

	Args:
	on_request_cb (CallbackProtocol): Function to be called when a request is received.

	Notes: this function is blocking
	"""
	LOGGER.info("Starting ZMQ request reply server")
	while True:
	LOGGER.debug("Waiting for new request")
	try:
	message = self.socket.recv()
	except zmq.Again:
	LOGGER.warn("ZMQ receive timeout. Continuing.")
	continue
	except zmq.ZMQError as e:
	LOGGER.error(f"A ZMQ error occured while receiving. Error: {e}")
	raise e
	LOGGER.debug(f"Received new request")

	LOGGER.debug(f"Processing new request")
	message = deserialize(message)
	res = on_request_cb(message)
	res = serialize(res)
	LOGGER.debug(f"Request processed")

	LOGGER.debug(f"Sending reply")
	try:
	self.socket.send(res)
	except zmq.ZMQError as e:
	LOGGER.error(f"A ZMQ error occured while sending. Error: {e}")
	raise e
	LOGGER.debug(f"Reply sent")

	# If the request is an exit request, we stop the loop and
	# let the process exit gracefully.
	if parse_request_type(message) == RequestType.EXIT:
	LOGGER.info("Received EXIT request. Stopping server.")
	break

	def __del__(self):
	if self.socket:
	self.socket.close()
	self._context.term()


	class ZmqRequestReplyClientConfig(RequestReplyClientConfig):
	ip: str # IP address of the server
	port: int # Port number of the server
	recv_timeout_ms: int = 15000 # Timeout in milliseconds for calls to socket.recv
	send_num_retries: int = 3 # Number of times to retry sending a message
	# If True, recv timeouts are raised as ``zmq.Again`` to the caller after the broken
	# REQ socket has been reset. Use when a timeout indicates the worker should fail and
	# be respawned by its supervisor (e.g. a remote inference client). Default False
	# keeps the "log and return None" behaviour.
	propagate_timeout: bool = False


	class ZmqRequestReplyClient(RequestReplyClient[ZmqRequestReplyClientConfig]):
	"""
	ZeroMQ-based implementation of RequestReplyClient.

	It communicates with the server using a ZeroMQ in REQ mode.

	Methods:
	send(request: dict) -> Optional[Dict]: Send a request to the server and
	receive a reply.

	Notes:
	This class is designed to be used in conjunction with the
	ZmqRequestReplyServer class.
	"""

	def __init__(self, config: ZmqRequestReplyClientConfig):
	super().__init__(config)
	self._closed = False
	self.context = zmq.Context()
	self.socket = None
	self._init_socket()

	def _init_socket(self):
	"""
	Reset the socket connection, this is needed when REQ is in a broken state.
	"""
	if self.socket:
	LOGGER.info(f"Closing and deleting existing socket")
	self.socket.close()
	del self.socket
	self.socket = self.context.socket(zmq.REQ)
	self.socket.setsockopt(zmq.RCVTIMEO, self.config.recv_timeout_ms)
	# this option is needed, among other things, to prevent context.term()
	# from handing if socket.recv() is interrupted e.g. by a KeyboardInterrupt
	self.socket.setsockopt(zmq.LINGER, 0)

	LOGGER.info(f"Connecting to {self.config.ip}:{self.config.port}")
	self.socket.connect(f"tcp://{self.config.ip}:{self.config.port}")
	LOGGER.info(f"Connected to {self.config.ip}:{self.config.port}")

	def send(self, request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
	if self.socket is None or self.socket.closed:
	LOGGER.error(f"The socket is closed, request not sent.")
	return None

	try:
	LOGGER.debug(f"Sending request")
	sent = False
	request_serialized = serialize(request)
	for i in range(self.config.send_num_retries):
	try:
	self.socket.send(request_serialized)
	sent = True
	LOGGER.debug(f"Request sent.")
	break
	except zmq.ZMQError as e:
	LOGGER.error(
	f"Failed to send message to {self.config.ip}:{self.config.port}. Error: {e}."
	)

	LOGGER.debug(f"Attempt {i} failed. Retrying.")

	if not sent:
	LOGGER.warn("Number of retries exceeded. Aborting.")
	return None

	# If the request is an exit request, we don't wait for a reply.
	if request.get("type") == RequestType.EXIT:
	return None

	LOGGER.debug(f"Waiting for reply")
	try:
	message = self.socket.recv()
	except zmq.Again:
	# REQ socket is now stuck in "waiting for reply", so any subsequent send
	# would fail with ZMQError. Reset before deciding how to surface the timeout.
	self._init_socket()
	if self.config.propagate_timeout:
	raise
	LOGGER.warn("ZMQ receive timeout. Aborting.")
	return None
	except zmq.ZMQError as e:
	LOGGER.error(f"A ZMQ error occured while receiving. Error: {e}.")
	raise e

	return deserialize(message)

	except KeyboardInterrupt:
	# Ctrl+C might be sent mid-episode. In this case, if the exception is raised inside the try
	# block, the socket will fall into a broken state. Here we reinitialize it to keep the
	# communication working.
	self._init_socket()
	raise

	def close(self) -> None:
	"""Idempotently tear down the socket and ZMQ context. Safe to call multiple times."""
	if self._closed:
	return
	if self.socket:
	try:
	self.socket.disconnect(f"tcp://{self.config.ip}:{self.config.port}")
	except zmq.ZMQError:
	# Socket already disconnected (e.g. _init_socket re-created it after a timeout).
	pass
	self.socket.close()
	self.context.term()
	self._closed = True

	def __del__(self):
	self.close()