| 408 | |
| 409 | |
| 410 | class AsyncTargetCacheWriter: |
| 411 | def __init__( |
| 412 | self, |
| 413 | *, |
| 414 | rank_dir: str, |
| 415 | max_shard_bytes: int, |
| 416 | max_queue_size: int = 128, |
| 417 | ): |
| 418 | self.writer = LocalTargetCacheWriter( |
| 419 | rank_dir=rank_dir, |
| 420 | max_shard_bytes=max_shard_bytes, |
| 421 | ) |
| 422 | # Queue CPU byte records only; never hold CUDA tensor references here. |
| 423 | self.queue = queue.Queue(maxsize=int(max_queue_size)) |
| 424 | self.sentinel = object() |
| 425 | self.num_local_samples = 0 |
| 426 | self._closed = False |
| 427 | self._exception = None |
| 428 | self.thread = threading.Thread( |
| 429 | target=self._run, |
| 430 | name=f"target-cache-writer-{os.path.basename(rank_dir)}", |
| 431 | ) |
| 432 | self.thread.start() |
| 433 | |
| 434 | @property |
| 435 | def local_shard_files(self): |
| 436 | return self.writer.local_shard_files |
| 437 | |
| 438 | def _run(self): |
| 439 | try: |
| 440 | while True: |
| 441 | item = self.queue.get() |
| 442 | try: |
| 443 | if item is self.sentinel: |
| 444 | break |
| 445 | self.writer.write_sample_bytes(item) |
| 446 | finally: |
| 447 | self.queue.task_done() |
| 448 | except BaseException as exc: |
| 449 | self._exception = exc |
| 450 | finally: |
| 451 | try: |
| 452 | self.writer.close() |
| 453 | except BaseException as exc: |
| 454 | if self._exception is None: |
| 455 | self._exception = exc |
| 456 | |
| 457 | def _raise_if_failed(self): |
| 458 | if self._exception is not None: |
| 459 | raise RuntimeError("Async target cache writer failed.") from self._exception |
| 460 | |
| 461 | def _put(self, item): |
| 462 | while True: |
| 463 | self._raise_if_failed() |
| 464 | try: |
| 465 | self.queue.put(item, timeout=1.0) |
| 466 | return |
| 467 | except queue.Full: |