Repository: LMCache/LMCache
File: lmcache/v1/cache_controller/worker.py (Python)

Code before:

        assert config.controller_pull_url is not None

        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:
            controller_rep_url = config.controller_reply_url
            self.req_socket = get_zmq_socket(
                self.context,
                controller_rep_url,
                protocol="tcp",
                role=zmq.REQ,  # type: ignore[attr-defined]
                bind_or_connect="connect",
            )

        lmcache_worker_port = config.lmcache_worker_ports[self.worker_id]

        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)
        self.thread.start()
        asyncio.run_coroutine_threadsafe(self.start_all(), self.loop)

        self.msg_queue: asyncio.Queue[WorkerMsg] = asyncio.Queue()


Review comment (performance):
Is this branch redundant, since we set lmcache_worker_ids to `list(range(world_size))`? If so we can remove this branch.

Code after:

        assert config.controller_pull_url is not None

        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:
            controller_rep_url = config.controller_reply_url
            self.req_socket = get_zmq_socket(
                self.context,
                controller_rep_url,
                protocol="tcp",
                role=zmq.REQ,  # type: ignore[attr-defined]
                bind_or_connect="connect",
            )

        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
            assert len(config.lmcache_worker_ports) == metadata.world_size
            lmcache_worker_port = config.lmcache_worker_ports[self.worker_id]
        else:
            # start lmcache worker on given worker ids
            assert len(lmcache_worker_ids) == 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",