Read a batch of tasks according to the current schedule. For each step: - Checks if a new epoch has started; rebuilds order if so - Determines which tasksets contribute to this batch - Uses each taskset's selector to pick specific samples - Annotates
(self)
| 182 | return self.step >= self.max_steps |
| 183 | |
| 184 | async def read(self) -> List: |
| 185 | """Read a batch of tasks according to the current schedule. |
| 186 | |
| 187 | For each step: |
| 188 | - Checks if a new epoch has started; rebuilds order if so |
| 189 | - Determines which tasksets contribute to this batch |
| 190 | - Uses each taskset's selector to pick specific samples |
| 191 | - Annotates each task with its source taskset_id |
| 192 | - Returns combined list of tasks |
| 193 | |
| 194 | Raises: |
| 195 | StopAsyncIteration: When total_epochs is reached |
| 196 | |
| 197 | Returns: |
| 198 | List[Task]: A batch of tasks from potentially multiple tasksets |
| 199 | """ |
| 200 | if self._should_stop(): |
| 201 | raise StopAsyncIteration |
| 202 | |
| 203 | batch_size = self.read_batch_size |
| 204 | start = self.step * batch_size % len(self.base_taskset_ids) |
| 205 | end = start + batch_size |
| 206 | if end <= len(self.base_taskset_ids): |
| 207 | taskset_ids = self.orders[start:end] |
| 208 | if end == len(self.base_taskset_ids): |
| 209 | self.epoch += 1 |
| 210 | self.orders = self.build_orders(self.epoch) |
| 211 | else: |
| 212 | taskset_ids = self.orders[start:] |
| 213 | self.epoch += 1 |
| 214 | self.orders = self.build_orders(self.epoch) |
| 215 | taskset_ids += self.orders[: (end - len(self.base_taskset_ids))] |
| 216 | |
| 217 | counter = Counter(taskset_ids) |
| 218 | batch = [] |
| 219 | for taskset_id, count in counter.items(): |
| 220 | tasks = await self.tasksets[taskset_id].read(batch_size=count) |
| 221 | # Annotate each task with its origin |
| 222 | for task in tasks: |
| 223 | task.index["taskset_id"] = taskset_id |
| 224 | batch.extend(tasks) |
| 225 | |
| 226 | self.step += 1 |
| 227 | return batch |
| 228 | |
| 229 | def state_dict(self) -> List[Dict]: |
| 230 | """ |