| 303 | |
| 304 | |
| 305 | class LocalTargetCacheWriter: |
| 306 | def __init__(self, *, rank_dir: str, max_shard_bytes: int): |
| 307 | self.rank_dir = rank_dir |
| 308 | self.max_shard_bytes = int(max_shard_bytes) |
| 309 | self.local_index_path = os.path.join(rank_dir, "samples.local.idx") |
| 310 | self.index_handle = open(self.local_index_path, "wb") |
| 311 | self.current_shard_id = -1 |
| 312 | self.current_shard_handle = None |
| 313 | self.current_shard_size = 0 |
| 314 | self.local_shard_files = [] |
| 315 | self.num_local_samples = 0 |
| 316 | |
| 317 | def close(self): |
| 318 | if self.current_shard_handle is not None: |
| 319 | self.current_shard_handle.flush() |
| 320 | os.fsync(self.current_shard_handle.fileno()) |
| 321 | self.current_shard_handle.close() |
| 322 | self.current_shard_handle = None |
| 323 | if getattr(self, "index_handle", None) is not None: |
| 324 | self.index_handle.flush() |
| 325 | os.fsync(self.index_handle.fileno()) |
| 326 | self.index_handle.close() |
| 327 | self.index_handle = None |
| 328 | |
| 329 | def _open_new_shard(self): |
| 330 | if self.current_shard_handle is not None: |
| 331 | self.current_shard_handle.flush() |
| 332 | os.fsync(self.current_shard_handle.fileno()) |
| 333 | self.current_shard_handle.close() |
| 334 | self.current_shard_id += 1 |
| 335 | file_name = f"shard-local-{self.current_shard_id:05d}.bin" |
| 336 | shard_path = os.path.join(self.rank_dir, file_name) |
| 337 | self.current_shard_handle = open(shard_path, "wb") |
| 338 | self.current_shard_size = 0 |
| 339 | self.local_shard_files.append(file_name) |
| 340 | |
| 341 | def _ensure_shard(self, sample_nbytes: int): |
| 342 | if self.current_shard_handle is None: |
| 343 | self._open_new_shard() |
| 344 | return |
| 345 | if ( |
| 346 | self.current_shard_size > 0 |
| 347 | and self.current_shard_size + int(sample_nbytes) > self.max_shard_bytes |
| 348 | ): |
| 349 | self._open_new_shard() |
| 350 | |
| 351 | def write_sample_bytes(self, sample: TargetCacheSampleBytes): |
| 352 | sample_nbytes = ( |
| 353 | len(sample.input_ids) |
| 354 | + len(sample.attention_mask) |
| 355 | + len(sample.loss_mask) |
| 356 | + len(sample.target_hidden_states) |
| 357 | + len(sample.target_last_hidden_states) |
| 358 | ) |
| 359 | self._ensure_shard(sample_nbytes) |
| 360 | input_ids_offset = self.current_shard_size |
| 361 | self.current_shard_handle.write(sample.input_ids) |
| 362 | self.current_shard_size += len(sample.input_ids) |