| 143 | |
| 144 | |
| 145 | class AsyncQueue(asyncio.Queue, QueueBuffer): |
| 146 | def __init__(self, capacity: int): |
| 147 | """ |
| 148 | Initialize the async queue with a specified capacity. |
| 149 | |
| 150 | Args: |
| 151 | capacity (`int`): The maximum number of items the queue can hold. |
| 152 | """ |
| 153 | super().__init__(maxsize=capacity) |
| 154 | self._closed = False |
| 155 | self.min_model_version = 0 |
| 156 | |
| 157 | async def put(self, item: List[Experience]): |
| 158 | if len(item) == 0: |
| 159 | return |
| 160 | await super().put(item) |
| 161 | |
| 162 | async def get(self): |
| 163 | while True: |
| 164 | item = await super().get() |
| 165 | if ( |
| 166 | self.min_model_version <= 0 |
| 167 | or item[0].info["model_version"] >= self.min_model_version |
| 168 | ): |
| 169 | return item |
| 170 | |
| 171 | async def close(self) -> None: |
| 172 | """Close the queue.""" |
| 173 | self._closed = True |
| 174 | for getter in self._getters: |
| 175 | if not getter.done(): |
| 176 | getter.set_exception(StopAsyncIteration()) |
| 177 | self._getters.clear() |
| 178 | |
| 179 | def stopped(self) -> bool: |
| 180 | """Check if there is no more data to read.""" |
| 181 | return self._closed and self.empty() |
| 182 | |
| 183 | |
| 184 | class AsyncPriorityQueue(QueueBuffer): |