An asynchronous priority queue that manages a fixed-size buffer of experience items. Items are prioritized using a user-defined function and reinserted after a cooldown period. Attributes: capacity (int): Maximum number of items the queue can hold. This value is automatically
| 182 | |
| 183 | |
| 184 | class AsyncPriorityQueue(QueueBuffer): |
| 185 | """ |
| 186 | An asynchronous priority queue that manages a fixed-size buffer of experience items. |
| 187 | Items are prioritized using a user-defined function and reinserted after a cooldown period. |
| 188 | |
| 189 | Attributes: |
| 190 | capacity (int): Maximum number of items the queue can hold. This value is automatically |
| 191 | adjusted to be at most twice the read batch size. |
| 192 | reuse_cooldown_time (float): Delay before reusing an item (set to infinity to disable). |
| 193 | priority_fn (callable): Function used to determine the priority of an item. |
| 194 | priority_groups (SortedDict): Maps priorities to deques of items with the same priority. |
| 195 | """ |
| 196 | |
| 197 | def __init__( |
| 198 | self, |
| 199 | capacity: int, |
| 200 | reuse_cooldown_time: Optional[float] = None, |
| 201 | priority_fn: str = "linear_decay", |
| 202 | priority_fn_args: Optional[dict] = None, |
| 203 | ): |
| 204 | """ |
| 205 | Initialize the async priority queue. |
| 206 | |
| 207 | Args: |
| 208 | capacity (`int`): The maximum number of items the queue can store. |
| 209 | reuse_cooldown_time (`float`): Time to wait before reusing an item. Set to None to disable reuse. |
| 210 | priority_fn (`str`): Name of the function to use for determining item priority. |
| 211 | kwargs: Additional keyword arguments for the priority function. |
| 212 | """ |
| 213 | from trinity.buffer.storage import PRIORITY_FUNC |
| 214 | |
| 215 | self.capacity = capacity |
| 216 | self.item_count = 0 |
| 217 | self.priority_groups = SortedDict() # Maps priority -> deque of items |
| 218 | priority_fn_cls = PRIORITY_FUNC.get(priority_fn) |
| 219 | kwargs = priority_fn_cls.default_config() |
| 220 | kwargs.update(priority_fn_args or {}) |
| 221 | self.priority_fn = priority_fn_cls(**kwargs) |
| 222 | self.reuse_cooldown_time = reuse_cooldown_time |
| 223 | self._condition = asyncio.Condition() # For thread-safe operations |
| 224 | self._closed = False |
| 225 | self.min_model_version = 0 |
| 226 | |
| 227 | async def _put(self, item: List[Experience], delay: float = 0) -> None: |
| 228 | """ |
| 229 | Insert an item into the queue, replacing the lowest-priority item if full. |
| 230 | |
| 231 | Args: |
| 232 | item (`List[Experience]`): A list of experiences to add. |
| 233 | delay (`float`): Optional delay before insertion (for simulating timing behavior). |
| 234 | """ |
| 235 | if delay > 0: |
| 236 | await asyncio.sleep(delay) |
| 237 | if len(item) == 0: |
| 238 | return |
| 239 | |
| 240 | priority, put_into_queue = self.priority_fn(item=item) |
| 241 | if not put_into_queue: |