Class QueueReader

trinity/buffer/reader/queue_reader.py:13–44

Reader of the Queue buffer.

class QueueReader(BufferReader):
    """Reader of the Queue buffer."""

    def __init__(self, config: StorageConfig):
        assert config.storage_type == StorageType.QUEUE.value
        self.timeout = config.max_read_timeout
        self.read_batch_size = config.batch_size
        self.queue = QueueStorage.get_wrapper(config)

    async def read(self, batch_size: Optional[int] = None, **kwargs) -> List[Experience]:
        batch_size = self.read_batch_size if batch_size is None else batch_size
        try:
            exp_bytes = await self.queue.get_batch.remote(
                batch_size, timeout=self.timeout, **kwargs
            )
        except Exception as e:
            if "StopAsyncIteration" in traceback.format_exc():
                raise StopAsyncIteration() from e
            else:
                raise
        exps = Experience.deserialize_many(exp_bytes)
        if len(exps) != batch_size:
            raise TimeoutError(
                f"Read incomplete batch ({len(exps)}/{batch_size}), please check your workflow."
            )
        return exps

    def state_dict(self) -> Dict:
        return {"current_index": 0}

    def load_state_dict(self, state_dict):
        return None