Class ZmqTcpServer

fastdeploy/inter_communicator/zmq_server.py:395–474

ZmqTcpServer, used when FD_ENABLE_INTERNAL_ADAPTER=1

class ZmqTcpServer(ZmqServerBase):
    """
    ZmqTcpServer, used when FD_ENABLE_INTERNAL_ADAPTER=1
    """

    def __init__(self, port, mode):
        super(ZmqTcpServer, self).__init__()
        self.mode = mode
        self.port = port
        self.cached_results = defaultdict(list)
        self.ZMQ_SNDHWM = int(envs.FD_ZMQ_SNDHWM)
        self.aggregate_send = envs.FD_USE_AGGREGATE_SEND

        self.mutex = threading.Lock()
        self.req_dict = dict()
        self.running = True
        self.context = zmq.Context()
        self._create_socket()
        self.response_token_lock = threading.Lock()

    def _create_socket(self):
        """create and return a ZeroMQ socket."""
        self.socket = self.context.socket(self.mode)
        self.socket.setsockopt(zmq.SNDHWM, self.ZMQ_SNDHWM)
        self.socket.setsockopt(zmq.SNDTIMEO, -1)
        self.address = f"tcp://*:{self.port}"
        self.socket.bind(self.address)
        return self.socket

    def recv_control_cmd(self):
        """
        Recieve control command from client
        """
        self._ensure_socket()
        try:
            client, _, task_data = self.socket.recv_multipart(flags=zmq.NOBLOCK)
            task = msgpack.unpackb(task_data)
            task_id_str = task["task_id"]
        except zmq.Again:
            return None
        with self.mutex:
            self.req_dict[task_id_str] = client
        return task

    def response_for_control_cmd(self, task_id, result):
        """
        Send command result back to client.
        """
        self._ensure_socket()
        if self.socket is None:
            raise RuntimeError("Router socket not created.")
        try:
            result = msgpack.packb(result)
            self.socket.send_multipart([self.req_dict[task_id], b"", result])

        except Exception as e:
            llm_logger.error(f"Send result to zmq client failed: {e}")