lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/v1/cache_controller/worker.py

@@ -0,0 +1,665 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import TYPE_CHECKING, Optional
+import asyncio
+import threading
+
+# Third Party
+import msgspec
+import zmq
+import zmq.asyncio
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.cache_controller.full_sync_sender import FullSyncSender
+from lmcache.v1.cache_controller.message import (
+    BatchedP2PLookupMsg,
+    BatchedP2PLookupRetMsg,
+    ClearWorkerMsg,
+    ClearWorkerRetMsg,
+    CompressWorkerMsg,
+    CompressWorkerRetMsg,
+    DecompressWorkerMsg,
+    DecompressWorkerRetMsg,
+    DeRegisterMsg,
+    ErrorMsg,
+    FullSyncStartMsg,
+    FullSyncStartRetMsg,
+    FullSyncStatusMsg,
+    FullSyncStatusRetMsg,
+    HealthWorkerMsg,
+    HealthWorkerRetMsg,
+    HeartbeatMsg,
+    HeartbeatRetMsg,
+    MoveWorkerMsg,
+    MoveWorkerRetMsg,
+    Msg,
+    PinWorkerMsg,
+    PinWorkerRetMsg,
+    RegisterMsg,
+    RegisterRetMsg,
+    WorkerMsg,
+    WorkerReqMsg,
+    WorkerReqRetMsg,
+)
+from lmcache.v1.config import LMCacheEngineConfig
+from lmcache.v1.metadata import LMCacheMetadata
+from lmcache.v1.rpc_utils import (
+    DEFAULT_SOCKET_RECV_TIMEOUT_MS,
+    DEFAULT_SOCKET_SEND_TIMEOUT_MS,
+    close_zmq_socket,
+    get_ip,
+    get_zmq_context,
+    get_zmq_socket,
+    get_zmq_socket_with_timeout,
+)
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.cache_engine import LMCacheEngine
+
+logger = init_logger(__name__)
+
+
+class LMCacheWorker:
+    """
+    LMCache Worker class to handle the execution of cache operations.
+    This class is responsible for receiving requests from the executor and
+    executing the corresponding operations on the LMCache engine.
+    Each worker is associated with a specific LMCache instance and a worker id.
+    """
+
+    def __init__(
+        self,
+        config: LMCacheEngineConfig,
+        metadata: LMCacheMetadata,
+        lmcache_engine: "LMCacheEngine",
+    ):
+        # TODO (Jiayi): "instance_id" might not be needed anymore.
+        # Please consider removing it.
+        self.config = config
+        self.lmcache_instance_id = config.lmcache_instance_id
+        if self.lmcache_instance_id is None:
+            raise ValueError(
+                "lmcache_instance_id is required when enable_controller=True"
+            )
+        self.lmcache_engine = lmcache_engine
+        self.worker_id = metadata.worker_id
+
+        self.context = get_zmq_context()
+
+        # Load timeout configurations from extra_config (in milliseconds)
+        self.socket_recv_timeout_ms = config.get_extra_config_value(
+            "worker_socket_recv_timeout_ms", DEFAULT_SOCKET_RECV_TIMEOUT_MS
+        )
+        self.socket_send_timeout_ms = config.get_extra_config_value(
+            "worker_socket_send_timeout_ms", DEFAULT_SOCKET_SEND_TIMEOUT_MS
+        )
+
+        if config.controller_pull_url is None:
+            raise ValueError(
+                "controller_pull_url is required when enable_controller=True"
+            )
+
+        controller_pull_url = config.controller_pull_url
+        self.push_socket = get_zmq_socket(
+            self.context,
+            controller_pull_url,
+            protocol="tcp",
+            role=zmq.PUSH,  # type: ignore[attr-defined]
+            bind_or_connect="connect",
+        )
+
+        if config.controller_reply_url is not None:
+            self.controller_rep_url = config.controller_reply_url
+            self._create_req_socket()
+
+        # Heartbeat socket will be created dynamically after register
+        # based on heartbeat_url returned from controller
+        self.heartbeat_socket: Optional[zmq.asyncio.Socket] = None
+        self.controller_heartbeat_url: Optional[str] = None
+
+        # metadata.world_size comes from vLLM's parallel_config.world_size.
+        # For MLA models, vLLM divides this by tp_size (e.g. TP=8 PP=1 on
+        # 8 GPUs → world_size=1), so it may be much smaller than the total
+        # GPU count.  For non-MLA models it equals TP × PP.
+        #
+        # get_lmcache_worker_ids() decides which workers run an LMCache
+        # instance: [0] for MLA (only one worker needed since KV caches
+        # are not TP-sharded), or range(world_size) for non-MLA.
+        #
+        # We use >= because extra ports are harmless — port selection
+        # indexes by worker_id or lmcache_worker_ids position, so
+        # trailing entries are never bound to sockets.
+        lmcache_worker_ids = config.get_lmcache_worker_ids(
+            metadata.use_mla, metadata.world_size
+        )
+        if not lmcache_worker_ids:
+            # start lmcache worker on all ranks;
+            # need at least one port per rank (world_size)
+            if len(config.lmcache_worker_ports) < metadata.world_size:
+                raise ValueError(
+                    f"lmcache_worker_ports must have at least {metadata.world_size} "
+                    f"port(s) (world_size), got {len(config.lmcache_worker_ports)}"
+                )
+            lmcache_worker_port = config.lmcache_worker_ports[self.worker_id]
+        else:
+            # start lmcache worker on given worker ids;
+            # need at least one port per explicitly listed worker
+            if len(config.lmcache_worker_ports) < len(lmcache_worker_ids):
+                raise ValueError(
+                    f"lmcache_worker_ports must have at least "
+                    f"{len(lmcache_worker_ids)} "
+                    f"port(s) (one per lmcache worker), "
+                    f"got {len(config.lmcache_worker_ports)}"
+                )
+            index = lmcache_worker_ids.index(self.worker_id)
+            lmcache_worker_port = config.lmcache_worker_ports[index]
+
+        self.lmcache_worker_internal_url = f"*:{lmcache_worker_port}"
+        self.lmcache_worker_ip = get_ip()
+        self.lmcache_worker_port = lmcache_worker_port
+
+        self.p2p_init_url = None
+        if config.enable_p2p:
+            self.p2p_host = config.p2p_host
+            self.p2p_init_port = config.p2p_init_ports[self.worker_id]
+            self.p2p_init_url = f"{self.p2p_host}:{self.p2p_init_port}"
+
+        self.reply_socket = get_zmq_socket(
+            self.context,
+            self.lmcache_worker_internal_url,
+            protocol="tcp",
+            role=zmq.REP,  # type: ignore[attr-defined]
+            bind_or_connect="bind",
+        )
+
+        logger.info(f"Reply socket established at {self.lmcache_worker_internal_url}")
+
+        self.loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(
+            target=self.loop.run_forever, daemon=True, name="lmcache-worker-thread"
+        )
+        self.thread.start()
+        asyncio.run_coroutine_threadsafe(self.start_all(), self.loop)
+
+        self.msg_queue: asyncio.Queue[WorkerMsg] = asyncio.Queue()
+
+        # Full sync sender (initialized lazily when needed)
+        self._full_sync_sender: Optional["FullSyncSender"] = None
+
+    async def register(self):
+        """
+        Register the lmcache worker with the controller via DEALER-ROUTER.
+
+        This method sends a RegisterMsg and waits for RegisterRetMsg
+        which contains extra_config (e.g., heartbeat_url).
+        """
+        assert self.lmcache_instance_id is not None
+        logger.info(
+            "Registering lmcache instance-worker: %s",
+            (self.lmcache_instance_id, self.worker_id),
+        )
+
+        register_msg = RegisterMsg(
+            instance_id=self.lmcache_instance_id,
+            worker_id=self.worker_id,
+            ip=self.lmcache_worker_ip,
+            port=self.lmcache_worker_port,
+            peer_init_url=self.p2p_init_url,
+        )
+
+        # Send via DEALER socket (empty frame + payload) and wait for response
+        try:
+            await self.req_socket.send_multipart(
+                [b"", msgspec.msgpack.encode(register_msg)]
+            )
+            # DEALER receives: [empty_frame, payload]
+            frames = await self.req_socket.recv_multipart()
+            serialized_ret_msg = frames[-1]
+            ret_msg = msgspec.msgpack.decode(serialized_ret_msg, type=Msg)
+
+            if isinstance(ret_msg, RegisterRetMsg):
+                self._process_register_response(ret_msg)
+            else:
+                logger.warning("Unexpected register response type: %s", type(ret_msg))
+        except zmq.ZMQError as e:
+            logger.error("Failed to register with controller: %s", e)
+            raise
+
+    def _process_register_response(self, ret_msg: RegisterRetMsg):
+        """Process RegisterRetMsg and initialize components based on extra_config."""
+        extra_config = ret_msg.extra_config
+
+        # Initialize heartbeat socket if heartbeat_url is provided
+        heartbeat_url = extra_config.get("heartbeat_url")
+        if heartbeat_url:
+            logger.info("Received heartbeat_url from controller: %s", heartbeat_url)
+            self.controller_heartbeat_url = heartbeat_url
+            self._create_heartbeat_socket()
+        else:
+            logger.info("No dedicated heartbeat_url provided by controller")
+
+    def deregister(self):
+        """
+        De-register the lmcache worker from the controller.
+        """
+        assert self.lmcache_instance_id is not None
+        self.put_msg(
+            DeRegisterMsg(
+                instance_id=self.lmcache_instance_id,
+                worker_id=self.worker_id,
+                ip=self.lmcache_worker_ip,
+                port=self.lmcache_worker_port,
+            )
+        )
+
+    async def async_put_and_wait_msg(
+        self,
+        msg: WorkerReqMsg,
+    ) -> WorkerReqRetMsg:
+        """
+        Send a message to the controller and wait for the response.
+
+        This method handles different types of WorkerReqMsg using appropriate sockets:
+        - HeartbeatMsg: Uses dedicated heartbeat socket
+        - Other messages (RegisterMsg, BatchedP2PLookupMsg, FullSyncStartMsg,
+          FullSyncStatusMsg): Uses DEALER socket (req_socket)
+
+        Note: With DEALER-ROUTER mode, we no longer need to separate FullSync
+        messages to heartbeat socket since DEALER supports async concurrent requests.
+        """
+        # Send heartbeat via dedicated heartbeat socket
+        if isinstance(msg, HeartbeatMsg):
+            return await self._send_heartbeat_msg(msg)
+
+        # Send other messages via DEALER socket
+        try:
+            # DEALER socket: send [empty_frame, payload]
+            await self.req_socket.send_multipart([b"", msgspec.msgpack.encode(msg)])
+            frames = await self.req_socket.recv_multipart()
+            # DEALER receives: [empty_frame, payload]
+            serialized_ret_msg = frames[-1]
+            ret_msg = msgspec.msgpack.decode(serialized_ret_msg, type=Msg)
+            return ret_msg
+        except zmq.Again as e:
+            logger.error("Timeout occurred, recreating socket. Error: %s", e)
+            self._recreate_req_socket()
+            return self._on_request_failure(msg)
+        except zmq.ZMQError as e:
+            logger.error("ZMQ error occurred, recreating socket. Error: %s", e)
+            self._recreate_req_socket()
+            return self._on_request_failure(msg)
+        except Exception as e:
+            logger.error("Error happens in lmcache worker req_socket. Error: %s", e)
+            return self._on_request_failure(msg)
+
+    def _create_req_socket(self):
+        self.req_socket = get_zmq_socket_with_timeout(
+            self.context,
+            self.controller_rep_url,
+            "tcp",
+            zmq.DEALER,  # type: ignore[attr-defined]
+            "connect",
+            self.socket_recv_timeout_ms,
+            self.socket_send_timeout_ms,
+        )
+
+    def _create_heartbeat_socket(self):
+        logger.info(
+            "Creating heartbeat socket to connect to: %s, "
+            "recv_timeout: %dms, send_timeout: %dms",
+            self.controller_heartbeat_url,
+            self.socket_recv_timeout_ms,
+            self.socket_send_timeout_ms,
+        )
+        self.heartbeat_socket = get_zmq_socket_with_timeout(
+            self.context,
+            self.controller_heartbeat_url,
+            "tcp",
+            zmq.DEALER,  # type: ignore[attr-defined]
+            "connect",
+            self.socket_recv_timeout_ms,
+            self.socket_send_timeout_ms,
+        )
+
+    def _recreate_heartbeat_socket(self):
+        try:
+            self.heartbeat_socket.close(linger=0)
+        except Exception as e:
+            logger.error("Error closing heartbeat socket: %s", e)
+        self._create_heartbeat_socket()
+
+    def _recreate_req_socket(self):
+        try:
+            self.req_socket.close(linger=0)
+        except Exception as e:
+            logger.error("Error closing req socket: %s", e)
+        self._create_req_socket()
+
+    def _get_full_sync_sender(self):
+        """Lazy initialization of FullSyncSender"""
+        if self._full_sync_sender is None:
+            # Get the local_cpu_backend from lmcache_engine
+            local_cpu_backend = self.lmcache_engine.storage_manager.local_cpu_backend
+            self._full_sync_sender = FullSyncSender(
+                config=self.config,
+                worker=self,
+                lmcache_engine=self.lmcache_engine,
+                local_cpu_backend=local_cpu_backend,
+            )
+        return self._full_sync_sender
+
+    def _on_request_failure(self, msg: WorkerReqMsg) -> WorkerReqRetMsg:
+        """
+        Create a default return message when worker -> controller
+        request encounters an error (e.g., timeout, ZMQ error).
+        """
+        if isinstance(msg, BatchedP2PLookupMsg):
+            return BatchedP2PLookupRetMsg(layout_info=[("", "", 0, "")])
+        elif isinstance(msg, HeartbeatMsg):
+            return HeartbeatRetMsg()  # No command by default
+        elif isinstance(msg, FullSyncStartMsg):
+            return FullSyncStartRetMsg(
+                sync_id=msg.sync_id,
+                accepted=False,
+                error_msg="Communication error",
+            )
+        elif isinstance(msg, FullSyncStatusMsg):
+            return FullSyncStatusRetMsg(
+                sync_id=msg.sync_id,
+                is_complete=False,
+                global_progress=0.0,
+                can_exit_freeze=False,
+            )
+        else:
+            raise ValueError(f"Unknown message type: {type(msg)}")
+
+    def put_msg(self, msg: WorkerMsg):
+        """
+        Put a message into the message queue.
+        """
+        # TODO(Jiayi): This might introduce ~0.05ms latency than
+        # a normal function call.
+        # Not sure how much overhead is blocking though.
+        self.loop.call_soon_threadsafe(self.msg_queue.put_nowait, msg)
+
+    async def batched_get_msg(self, max_bsz: int = 50) -> list[WorkerMsg]:
+        """
+        Get a batch of messages from the message queue.
+        """
+        batch = []
+
+        # use blocking get for the first msg
+        try:
+            item = await self.msg_queue.get()
+            batch.append(item)
+        except asyncio.CancelledError:
+            return batch  # shutdown path
+
+        for _ in range(max_bsz - 1):
+            try:
+                item = self.msg_queue.get_nowait()
+                batch.append(item)
+            except asyncio.QueueEmpty:
+                break
+        return batch
+
+    async def _send_heartbeat_msg(self, msg: HeartbeatMsg) -> HeartbeatRetMsg:
+        """
+        Send heartbeat message via dedicated heartbeat DEALER socket.
+        This is separate from async_put_and_wait_msg to keep heartbeat independent.
+        """
+        if self.heartbeat_socket is None:
+            logger.warning("Heartbeat socket is not initialized")
+            return HeartbeatRetMsg()
+        try:
+            # DEALER socket: send [empty_frame, payload]
+            logger.info("Sending heartbeat message to controller...")
+            await self.heartbeat_socket.send_multipart(
+                [b"", msgspec.msgpack.encode(msg)]
+            )
+            logger.info("Heartbeat message sent, waiting for response...")
+            frames = await self.heartbeat_socket.recv_multipart()
+            logger.info("Received heartbeat response with %d frames", len(frames))
+            # DEALER receives: [empty_frame, payload]
+            serialized_ret_msg = frames[-1]
+            ret_msg = msgspec.msgpack.decode(serialized_ret_msg, type=Msg)
+            return ret_msg
+        except zmq.Again as e:
+            logger.error("Heartbeat timeout occurred, recreating socket. Error: %s", e)
+            self._recreate_heartbeat_socket()
+            return HeartbeatRetMsg()
+        except zmq.ZMQError as e:
+            logger.error(
+                "Heartbeat ZMQ error occurred, recreating socket. Error: %s", e
+            )
+            self._recreate_heartbeat_socket()
+            return HeartbeatRetMsg()
+        except Exception as e:
+            logger.error("Error happens in heartbeat socket. Error: %s", e)
+            return HeartbeatRetMsg()
+
+    async def heartbeat(self):
+        """
+        Send periodic heartbeats to the controller (DEALER-ROUTER mode).
+
+        Process any commands received in the heartbeat response.
+        Uses dedicated heartbeat socket to avoid blocking from other requests.
+        """
+        enable_heartbeat = (
+            self.config.lmcache_worker_heartbeat_time is not None
+            and self.config.lmcache_worker_heartbeat_time > 0
+            and self.heartbeat_socket is not None
+        )
+        if enable_heartbeat:
+            await asyncio.sleep(self.config.lmcache_worker_heartbeat_delay_time)
+            logger.info(
+                "Start heartbeat in %s : %s, delay time: %ss, heartbeat time: %ss",
+                self.lmcache_instance_id,
+                self.worker_id,
+                self.config.lmcache_worker_heartbeat_delay_time,
+                self.config.lmcache_worker_heartbeat_time,
+            )
+            while True:
+                # Send heartbeat via dedicated heartbeat socket
+                heartbeat_msg = HeartbeatMsg(
+                    instance_id=self.lmcache_instance_id,
+                    worker_id=self.worker_id,
+                    ip=self.lmcache_worker_ip,
+                    port=self.lmcache_worker_port,
+                    peer_init_url=self.p2p_init_url,
+                )
+
+                try:
+                    ret_msg = await self._send_heartbeat_msg(heartbeat_msg)
+
+                    if isinstance(ret_msg, HeartbeatRetMsg):
+                        self._handle_heartbeat_commands(ret_msg)
+                    else:
+                        logger.warning(
+                            "Unexpected heartbeat response type: %s", type(ret_msg)
+                        )
+                except Exception as e:
+                    logger.error("Error during heartbeat: %s", e)
+
+                await asyncio.sleep(self.config.lmcache_worker_heartbeat_time)
+
+    def _handle_heartbeat_commands(self, ret_msg: HeartbeatRetMsg):
+        """
+        Handle commands received in heartbeat response.
+
+        Uses polymorphic dispatch - each command class implements its own
+        execute() method. Commands are executed sequentially.
+        """
+        if not ret_msg.has_commands():
+            return
+
+        for command in ret_msg.commands:
+            logger.info(
+                "Executing heartbeat command: %s",
+                command.describe(),
+            )
+            try:
+                command.execute(self)
+            except NotImplementedError:
+                logger.warning(
+                    "Command %s.execute() not implemented yet",
+                    type(command).__name__,
+                )
+            except Exception as e:
+                logger.error(
+                    "Error executing command %s: %s",
+                    type(command).__name__,
+                    e,
+                )
+
+    async def push(self):
+        while True:
+            try:
+                msgs = await self.batched_get_msg()
+                logger.debug(f"Sending {len(msgs)} messages")
+                self.push_socket.send_multipart(
+                    [msgspec.msgpack.encode(msg) for msg in msgs]
+                )
+
+            except Exception as e:
+                logger.error(f"Push error: {e}")
+
+    async def handle_request(self):
+        """
+        Handle incoming requests (control msgs) from the controller.
+        """
+        while True:
+            try:
+                serialized_request = await self.reply_socket.recv()
+                request = msgspec.msgpack.decode(serialized_request, type=Msg)
+                logger.debug(f"Received message: {request}")
+                if isinstance(request, MoveWorkerMsg):
+                    tokens = request.tokens
+                    old_position = request.old_position
+                    new_position = request.new_position
+                    do_copy = request.copy
+                    worker_event_id = request.worker_event_id
+
+                    # Intra node move
+                    if new_position[0] == self.lmcache_worker_internal_url:
+                        # TODO(Jiayi): currently we only support moving from
+                        # local disk to local cpu.
+                        assert old_position[1] == "LocalDiskBackend"
+                        assert new_position[1] == "LocalCPUBackend"
+                        assert do_copy
+
+                        # TODO(Jiayi): We need to align prefetch and move.
+                        logger.debug("Executing prefetch operation.")
+                        raise NotImplementedError(
+                            "Prefetch from controller is not implemented yet."
+                        )
+                    else:
+                        assert new_position[1] == "LocalCPUBackend", (
+                            "Only support moving to cpu for now."
+                        )
+                        logger.debug("Executing cross-node move operation.")
+                        num_tokens = self.lmcache_engine.move(
+                            tokens=tokens,
+                            old_position=old_position,
+                            new_position=new_position,
+                            event_id=worker_event_id,
+                            do_copy=do_copy,
+                        )
+
+                    # TODO(Jiayi): LMCache needs to have an event tracking
+                    # pool to enable more advanced control-plane optims.
+                    # For now, we use a dummy `event_id`.
+                    serialized_ret_msg = msgspec.msgpack.encode(
+                        MoveWorkerRetMsg(num_tokens=num_tokens)
+                    )
+                elif isinstance(request, CompressWorkerMsg):
+                    num_compressed_tokens = self.lmcache_engine.compress(
+                        tokens=request.tokens,
+                        method=request.method,
+                        location=request.location,
+                        event_id=request.worker_event_id,
+                    )
+                    serialized_ret_msg = msgspec.msgpack.encode(
+                        CompressWorkerRetMsg(num_tokens=num_compressed_tokens)
+                    )
+                elif isinstance(request, DecompressWorkerMsg):
+                    num_decompressed_tokens = self.lmcache_engine.decompress(
+                        tokens=request.tokens,
+                        method=request.method,
+                        location=request.location,
+                        event_id=request.worker_event_id,
+                    )
+                    serialized_ret_msg = msgspec.msgpack.encode(
+                        DecompressWorkerRetMsg(num_tokens=num_decompressed_tokens)
+                    )
+                elif isinstance(request, PinWorkerMsg):
+                    num_pinned_tokens = self.lmcache_engine.lookup(
+                        tokens=request.tokens,
+                        search_range=[request.location],
+                        lookup_id=request.worker_event_id,
+                        pin=True,
+                    )
+                    serialized_ret_msg = msgspec.msgpack.encode(
+                        PinWorkerRetMsg(num_tokens=num_pinned_tokens)
+                    )
+                elif isinstance(request, ClearWorkerMsg):
+                    num_cleared_tokens = self.lmcache_engine.clear(
+                        locations=[request.location],
+                    )
+                    serialized_ret_msg = msgspec.msgpack.encode(
+                        ClearWorkerRetMsg(num_tokens=num_cleared_tokens)
+                    )
+                elif isinstance(request, HealthWorkerMsg):
+                    error_code = self.lmcache_engine.health()
+                    serialized_ret_msg = msgspec.msgpack.encode(
+                        HealthWorkerRetMsg(error_code=error_code)
+                    )
+                else:
+                    logger.error(f"Unknown message: {request}")
+                    serialized_ret_msg = msgspec.msgpack.encode(
+                        ErrorMsg(error=f"Unknown message: {request}")
+                    )
+
+                await self.reply_socket.send(serialized_ret_msg)
+            except Exception as e:
+                logger.error(f"Worker error: {e}")
+                serialized_ret_msg = msgspec.msgpack.encode(
+                    ErrorMsg(error=f"Worker error: {e}")
+                )
+                await self.reply_socket.send(serialized_ret_msg)
+
+    async def start_all(self):
+        try:
+            # Register first to get heartbeat_url before starting heartbeat task
+            await self.register()
+
+            logger.info(
+                f"Starting lmcache worker {self.worker_id}"
+                f"for instance {self.lmcache_instance_id}"
+            )
+            await asyncio.gather(
+                self.push(),
+                self.handle_request(),
+                self.heartbeat(),
+            )
+        except Exception as e:
+            logger.error(
+                f"Instance {self.lmcache_instance_id}, "
+                f"worker {self.worker_id} error: {e}"
+            )
+
+    def close(self):
+        self.deregister()
+        if self.loop.is_running():
+            self.loop.call_soon_threadsafe(self.loop.stop)
+        if self.thread.is_alive():
+            self.thread.join()
+        self.loop.close()
+        close_zmq_socket(self.push_socket)
+        close_zmq_socket(self.reply_socket)
+        if self.heartbeat_socket is not None:
+            close_zmq_socket(self.heartbeat_socket)
+        if hasattr(self, "req_socket"):
+            close_zmq_socket(self.req_socket)