(
self,
*,
rank_dir: str,
max_shard_bytes: int,
max_queue_size: int = 128,
)
| 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): |
nothing calls this directly
no test coverage detected