(self, block_idx)
| 165 | ) |
| 166 | |
| 167 | def _wait_blocks_move(self, block_idx): |
| 168 | if block_idx not in self.futures: |
| 169 | return |
| 170 | |
| 171 | if self.debug: |
| 172 | print(f"[{self.block_type}] Wait for block {block_idx}") |
| 173 | start_time = time.perf_counter() |
| 174 | |
| 175 | future = self.futures.pop(block_idx) |
| 176 | _, bidx_to_cuda = future.result() |
| 177 | |
| 178 | assert block_idx == bidx_to_cuda, f"Block index mismatch: {block_idx} != {bidx_to_cuda}" |
| 179 | |
| 180 | if self.debug: |
| 181 | print(f"[{self.block_type}] Waited for block {block_idx}: {time.perf_counter()-start_time:.2f}s") |
| 182 | |
| 183 | |
| 184 | class ModelOffloader(Offloader): |
no outgoing calls
no test coverage detected