distributed/shm_broadcast.py: MessageQueue

    @staticmethod
    def create_from_process_group(
        pg: ProcessGroup,
        max_chunk_bytes: int,
        max_chunks: int,
        writer_rank: int = 0,
    ) -> "MessageQueue":
        group_rank = dist.get_rank(pg)
        group_world_size = dist.get_world_size(pg)
        global_ranks = dist.get_process_group_ranks(pg)

        # assume the group does not cross nodes, so every reader is a local reader
        n_reader = group_world_size - 1
        local_reader_ranks = [i for i in range(group_world_size) if i != writer_rank]
        buffer_io: MessageQueue
        if group_rank == writer_rank:
            buffer_io = MessageQueue(
                n_reader=n_reader,
                n_local_reader=n_reader,
                local_reader_ranks=local_reader_ranks,
                max_chunk_bytes=max_chunk_bytes,
                max_chunks=max_chunks,
            )
            dist.broadcast_object_list([buffer_io.export_handle()], src=global_ranks[writer_rank], group=pg)
        else:
            recv = [None]
            dist.broadcast_object_list(recv, src=global_ranks[writer_rank], group=pg)
            handle = recv[0]  # type: ignore
            buffer_io = MessageQueue.create_from_handle(handle, group_rank)
        buffer_io.wait_until_ready()
        return buffer_io

        if n_local_reader > 0:
            # for local readers, we will:
            # 1. create a shared memory ring buffer to communicate small data
            # 2. create a publish-subscribe socket to communicate large data
            self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks)

            # XPUB is very similar to PUB,
            # except that it can receive subscription messages
            # to confirm the number of subscribers
            self.local_socket = context.socket(XPUB)
            # set the verbose option so that we can receive every subscription
            # message. otherwise, we will only receive the first subscription
            # see http://api.zeromq.org/3-3:zmq-setsockopt for more details
            self.local_socket.setsockopt(XPUB_VERBOSE, True)
            local_subscribe_port = get_open_port()
            socket_addr = f"tcp://127.0.0.1:{local_subscribe_port}"
            logger.debug("Binding to %s", socket_addr)
            self.local_socket.bind(socket_addr)

            self.current_idx = 0

        else:
            self.buffer = None  # type: ignore
            local_subscribe_port = None
            self.local_socket = None
            self.current_idx = -1

    @contextmanager
    def acquire_write(self, timeout: float | None = None) -> Generator[memoryview, None, None]:
        assert self._is_writer, "Only writers can acquire write"
        start_time = time.monotonic()
        n_warning = 1
        while True:
            with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
                read_count = sum(metadata_buffer[1:])
                written_flag = metadata_buffer[0]
                if written_flag and read_count != self.buffer.n_reader:
                    # this block is written and not read by all readers
                    # for writers, `self.current_idx` is the next block to write
                    # if this block is not ready to write,
                    # we need to wait until it is read by all readers
                    n_warning = self._spin_wait(start_time, n_warning, timeout)
                    continue
                # found a block that is either
                # (1) not written
                # (2) read by all readers

                # mark the block as not written
                metadata_buffer[0] = 0
                # let caller write to the buffer
                with self.buffer.get_data(self.current_idx) as buf:
                    yield buf

                # caller has written to the buffer
                # NOTE: order is important here
                # first set the read flags to 0
                # then set the written flag to 1
                # otherwise, the readers may think they already read the block
                for i in range(1, self.buffer.n_reader + 1):
                    # set read flag to 0, meaning it is not read yet
                    metadata_buffer[i] = 0
                # mark the block as written
                metadata_buffer[0] = 1
                self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
                break

    @contextmanager
    def acquire_read(self, timeout: float | None = None) -> Generator[memoryview, None, None]:
        assert self._is_local_reader, "Only readers can acquire read"
        start_time = time.monotonic()
        n_warning = 1
        while True:
            with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
                read_flag = metadata_buffer[self.local_reader_rank + 1]
                written_flag = metadata_buffer[0]
                if not written_flag or read_flag:
                    # this block is either
                    # (1) not written
                    # (2) already read by this reader

                    # for readers, `self.current_idx` is the next block to read
                    # if this block is not ready,
                    # we need to wait until it is written
                    n_warning = self._spin_wait(start_time, n_warning, timeout)
                    continue
                # found a block that is not read by this reader
                # let caller read from the buffer
                with self.buffer.get_data(self.current_idx) as buf:
                    yield buf

                # caller has read from the buffer
                # set the read flag
                metadata_buffer[self.local_reader_rank + 1] = 1
                self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
                break

    def enqueue(self, obj: Any, timeout: float | None = None) -> None:
        """Write to message queue with optional timeout (in seconds)"""
        assert self._is_writer, "Only writers can enqueue"
        serialized_obj = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
        if self.n_local_reader > 0:
            # first byte of the chunk is the overflow flag; the payload is stored
            # inline if it fits, otherwise it is sent through the local socket
            overflow = len(serialized_obj) >= self.buffer.max_chunk_bytes
            with self.acquire_write(timeout) as buf:
                buf[0] = overflow
                if not overflow:
                    buf[1 : len(serialized_obj) + 1] = serialized_obj
            if overflow:
                self.local_socket.send(serialized_obj)
        if self.n_remote_reader > 0:
            self.remote_socket.send(serialized_obj)

    def dequeue(self, timeout: float | None = None) -> Any:
        """Read from message queue with optional timeout (in seconds)"""
        if self._is_local_reader:
            with self.acquire_read(timeout) as buf:
                overflow = buf[0] == 1
                if not overflow:
                    # no need to know the size of serialized object
                    # pickle format contains the size information internally
                    # see https://docs.python.org/3/library/pickle.html
                    obj = pickle.loads(buf[1:])
            if overflow:
                obj = pickle.loads(self.local_socket.recv())
        elif self._is_remote_reader:
            obj = pickle.loads(self.remote_socket.recv())
        else:
            raise RuntimeError("Only readers can dequeue")
        return obj

    def wait_until_ready(self) -> None:
        """This is a collective operation. All processes (including the
        readers and the writer) should call this function.
        """
        if self._is_writer:
            # wait for subscription messages from all readers, then send a
            # message to confirm the publish channel is working
            for n_reader, socket in (
                (self.n_local_reader, self.local_socket),
                (self.n_remote_reader, self.remote_socket),
            ):
                for _ in range(n_reader):
                    socket.recv()
                if n_reader > 0:
                    socket.send(b"READY")
        elif self._is_local_reader:
            # wait for the writer to send a message
            assert self.local_socket.recv() == b"READY"
        elif self._is_remote_reader:
            # wait for the writer to send a message
            assert self.remote_socket.recv() == b"READY"

    @staticmethod
    def _spin_wait(start_time: float, n_warning: int, timeout: float | None = None) -> int:
        """自旋等待一轮：让出 CPU、按间隔打 debug 日志、超时抛 TimeoutError。返回更新后的 n_warning。"""
        os.sched_yield()
        elapsed = time.monotonic() - start_time
        if elapsed > envs.DEEPINFER_RINGBUFFER_WARNING_INTERVAL * n_warning:
            logger.debug("No available shared memory block after %.1f seconds", elapsed)
            n_warning += 1
        if timeout is not None and elapsed > timeout:
            raise TimeoutError
        return n_warning