| 20 | self.queue = QueueStorage.get_wrapper(config) |
| 21 | |
| 22 | async def read(self, batch_size: Optional[int] = None, **kwargs) -> List[Experience]: |
| 23 | batch_size = self.read_batch_size if batch_size is None else batch_size |
| 24 | try: |
| 25 | exp_bytes = await self.queue.get_batch.remote( |
| 26 | batch_size, timeout=self.timeout, **kwargs |
| 27 | ) |
| 28 | except Exception as e: |
| 29 | if "StopAsyncIteration" in traceback.format_exc(): |
| 30 | raise StopAsyncIteration() from e |
| 31 | else: |
| 32 | raise |
| 33 | exps = Experience.deserialize_many(exp_bytes) |
| 34 | if len(exps) != batch_size: |
| 35 | raise TimeoutError( |
| 36 | f"Read incomplete batch ({len(exps)}/{batch_size}), please check your workflow." |
| 37 | ) |
| 38 | return exps |
| 39 | |
| 40 | def state_dict(self) -> Dict: |
| 41 | return {"current_index": 0} |