Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / tools / io.py: 54%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15# Adapted from the indexed dataset implementations in: 

16# - NVIDIA NeMo AutoModel 

17# - NVIDIA Megatron-LM 

18 

19"""Load indexed datasets from paired binary and index files. 

20 

21Supports the original mmap and file-pointer readers for local *.bin / *.idx 

22pairs, plus optional streaming readers for object storage (S3 and MSC). 

23 

24All three calls below are equivalent for local data: 

25 

26 from hyper_parallel.data.indexed.io import IndexedDataset 

27 

28 ds = IndexedDataset("/path/to/shard_00_text_document") 

29 print(len(ds), ds[0][:20]) 

30 

31 ds = IndexedDataset("/path/to/shard_00_text_document.bin") 

32 print(len(ds), ds[0][:20]) 

33 

34 ds = IndexedDataset("/path/to/shard_00_text_document.idx") 

35 print(len(ds), ds[0][:20]) 

36 

37For object-storage data, pass an :class:`ObjectStorageConfig`: 

38 

39 cfg = ObjectStorageConfig(path_to_idx_cache="/tmp/idx_cache") 

40 ds = IndexedDataset("s3://bucket/path/shard_00_text_document", object_storage_config=cfg) 

41""" 

42 

43from __future__ import annotations 

44 

45import gc 

46import os 

47import shutil 

48import struct 

49from abc import ABC, abstractmethod 

50from dataclasses import dataclass 

51from enum import Enum 

52from types import TracebackType 

53from typing import Any, Dict, List, Optional, Tuple, Type, Union 

54import importlib 

55import traceback 

56 

57import numpy 

58import torch 

59 

60from hyper_parallel.data.constants import INDEX_HEADER, INDEX_VERSION, MSC_PREFIX, S3_PREFIX 

61from hyper_parallel.data.dataset_logging import get_dataset_logger 

62 

63logger = get_dataset_logger(__name__) 

64 

65 

66def safe_import( 

67 module: str, 

68 *, 

69 msg: Optional[str] = None, 

70 alt: Optional[Any] = None, 

71) -> Tuple[bool, Optional[Any]]: 

72 """ 

73 A function used to import modules that may not be available. 

74 

75 This function will attempt to import a module with the given name, but it 

76 will not throw an ImportError if the module is not found. Instead, it will 

77 return the configured alternate object or ``None``. 

78 

79 Args: 

80 module (str): The name of the module to import. 

81 msg (str or None): An optional error message to be displayed if this module is used 

82 after a failed import. 

83 alt (object): An optional module to be used in place of the given module if it 

84 fails to import 

85 

86 Returns: 

87 Tuple(bool, object): Whether the intended import succeeded and the imported module, 

88 configured alternate, or ``None``. 

89 """ 

90 try: 

91 return True, importlib.import_module(module) 

92 except ImportError: 

93 exception_text = traceback.format_exc() 

94 logger.debug(f"Import of {module} failed with: {exception_text}") 

95 if msg is None: 

96 msg = f"{module} could not be imported" 

97 if alt is None: 

98 logger.debug(msg) 

99 return False, None 

100 return False, alt 

101 

102 

103# Optional dependencies for object-storage support. Install via 

104# Install ``boto3`` for S3 or ``multi-storage-client`` for MSC support. 

105HAS_BOTO3, boto3 = safe_import("boto3") 

106HAS_MSC, multi_storage_client = safe_import("multi_storage_client") 

107 

108@dataclass 

109class ObjectStorageConfig: 

110 """Configuration for reading ``.bin``/``.idx`` files from object storage. 

111 

112 Attributes: 

113 path_to_idx_cache: Local directory where the ``.idx`` file is cached on 

114 first use. Re-used across ranks via a per-host directory layout. 

115 bin_chunk_nbytes: Size in bytes of each chunked range read against the 

116 ``.bin`` object. Defaults to 256 MiB. Larger values reduce request 

117 count but increase per-rank memory footprint. 

118 """ 

119 

120 path_to_idx_cache: str 

121 bin_chunk_nbytes: int = 256 * 1024 * 1024 

122 

123 

124def _is_object_storage_path(path: str) -> bool: 

125 """Return ``True`` if ``path`` is an ``s3://`` or ``msc://`` URI.""" 

126 return path.startswith(S3_PREFIX) or path.startswith(MSC_PREFIX) 

127 

128 

129def _parse_s3_path(path: str) -> Tuple[str, str]: 

130 """Split an ``s3://bucket/key`` URI into ``(bucket, key)``.""" 

131 if not path.startswith(S3_PREFIX): 

132 raise ValueError(f"Not an S3 path: {path}") 

133 parts = path[len(S3_PREFIX) :].split("/") 

134 bucket = parts[0] 

135 key = "/".join(parts[1:]) if len(parts) > 1 else "" 

136 return bucket, key 

137 

138 

139def _get_index_cache_path(idx_path: str, object_storage_config: ObjectStorageConfig) -> str: 

140 """Return the local cache path for ``idx_path`` under ``path_to_idx_cache``.""" 

141 if idx_path.startswith(S3_PREFIX): 

142 stripped = idx_path[len(S3_PREFIX) :] 

143 elif idx_path.startswith(MSC_PREFIX): 

144 stripped = idx_path[len(MSC_PREFIX) :] 

145 else: 

146 raise ValueError(f"Not an object storage path: {idx_path}") 

147 return os.path.join(object_storage_config.path_to_idx_cache, stripped) 

148 

149 

150def _cache_index_file(remote_path: str, local_path: str) -> None: 

151 """Download ``.idx`` from object storage to ``local_path``. 

152 

153 Rank 0 performs the download and other ranks wait on a ``torch.distributed`` 

154 barrier. If the local file already exists this is a no-op. 

155 

156 Raises: 

157 ImportError: If the relevant client library (``boto3`` for ``s3://`` or 

158 ``multi_storage_client`` for ``msc://``) is not installed. 

159 ValueError: If ``remote_path`` is neither an ``s3://`` nor an 

160 ``msc://`` URI. 

161 """ 

162 torch_dist_enabled = torch.distributed.is_initialized() 

163 rank = torch.distributed.get_rank() if torch_dist_enabled else 0 

164 

165 if remote_path.startswith(S3_PREFIX): 

166 if not HAS_BOTO3: 

167 raise ImportError("boto3 is required to read s3:// datasets. Install via `pip install boto3`.") 

168 if not os.path.exists(local_path): 

169 if not torch_dist_enabled or rank == 0: 

170 os.makedirs(os.path.dirname(local_path), exist_ok=True) 

171 bucket, key = _parse_s3_path(remote_path) 

172 client = boto3.client("s3") 

173 logger.info("Downloading %s -> %s", remote_path, local_path) 

174 client.download_file(bucket, key, local_path) 

175 client.close() 

176 elif remote_path.startswith(MSC_PREFIX): 

177 if not HAS_MSC: 

178 raise ImportError( 

179 "multi_storage_client is required to read msc:// datasets. " 

180 "Install via `pip install multi-storage-client`." 

181 ) 

182 if not os.path.exists(local_path): 

183 if not torch_dist_enabled or rank == 0: 

184 os.makedirs(os.path.dirname(local_path), exist_ok=True) 

185 multi_storage_client.download_file(remote_path, local_path) 

186 else: 

187 raise ValueError(f"Unsupported object storage path: {remote_path}") 

188 

189 if torch_dist_enabled: 

190 torch.distributed.barrier() 

191 if not os.path.exists(local_path): 

192 raise RuntimeError(f"Index cache file not found after download: {local_path}") 

193 

194 

195class DType(Enum): 

196 """The NumPy data type Enum for reading the IndexedDataset indices""" 

197 

198 UINT8 = 1 

199 INT8 = 2 

200 INT16 = 3 

201 INT32 = 4 

202 INT64 = 5 

203 FLOAT64 = 6 

204 FLOAT32 = 7 

205 UINT16 = 8 

206 

207 @classmethod 

208 def code_from_dtype(cls, value: Type[numpy.number]) -> int: 

209 """Get the code from the dtype 

210 

211 Args: 

212 value (Type[numpy.number]): The dtype 

213 

214 Returns: 

215 int: The code 

216 """ 

217 return cls[value.__name__.upper()].value 

218 

219 @classmethod 

220 def dtype_from_code(cls, value: int) -> Type[numpy.number]: 

221 """Get the dtype from the code 

222 

223 Args: 

224 value (int): The code 

225 

226 Returns: 

227 Type[numpy.number]: The dtype 

228 """ 

229 return getattr(numpy, cls(value).name.lower()) 

230 

231 @classmethod 

232 def size(cls, key: Union[int, Type[numpy.number]]) -> int: 

233 """Get the size of the dtype/code in bytes 

234 

235 Args: 

236 key (Union[int, Type[numpy.number]]): The dtype or code 

237 

238 Raises: 

239 ValueError: If the key is neither dtype nor integer code 

240 

241 Returns: 

242 int: The size of the dtype/code in bytes 

243 """ 

244 if isinstance(key, int): 

245 return cls.dtype_from_code(key)().itemsize 

246 if numpy.number in key.__mro__: 

247 return key().itemsize 

248 raise ValueError("Invalid key passed to DType.size()") 

249 

250 @classmethod 

251 def optimal_dtype(cls, cardinality: Optional[int]) -> Type[numpy.number]: 

252 """Get the dtype to use for an index of a certain cardinality 

253 

254 Args: 

255 cardinality (Optional[int]): The number of elements to be indexed 

256 

257 Returns: 

258 Type[numpy.number]: The dtype to use for the index 

259 """ 

260 if cardinality is not None and cardinality < 65500: 

261 return numpy.uint16 

262 return numpy.int32 

263 

264 

265class _IndexWriter: 

266 """Object class to write the index (.idx) file 

267 

268 Args: 

269 idx_path (str): The path to the index file 

270 

271 dtype (Type[numpy.number]): The dtype of the index file 

272 """ 

273 

274 def __init__(self, idx_path: str, dtype: Type[numpy.number]) -> None: 

275 """Initialize the _IndexWriter 

276 

277 Args: 

278 idx_path (str): The path to the index file. 

279 

280 dtype (Type[numpy.number]): The dtype of the index file. 

281 """ 

282 self.idx_path = idx_path 

283 self.dtype = dtype 

284 self.idx_writer = None 

285 

286 def __enter__(self) -> "_IndexWriter": 

287 """Enter the context introduced by the 'with' keyword 

288 

289 Returns: 

290 _IndexWriter: The instance 

291 """ 

292 self.idx_writer = open(self.idx_path, "wb") 

293 # fixed, vestigial practice 

294 self.idx_writer.write(INDEX_HEADER) 

295 # fixed, vestigial practice 

296 self.idx_writer.write(struct.pack("<Q", INDEX_VERSION)) 

297 # the numeric code for the dtype 

298 self.idx_writer.write(struct.pack("<B", DType.code_from_dtype(self.dtype))) 

299 return self 

300 

301 def __exit__( 

302 self, 

303 exc_type: Optional[Type[BaseException]], 

304 exc_val: Optional[BaseException], 

305 exc_tb: Optional[TracebackType], 

306 ) -> Optional[bool]: 

307 """Exit the context introduced by the 'with' keyword 

308 

309 Args: 

310 exc_type (Optional[Type[BaseException]]): Exception type 

311 

312 exc_val (Optional[BaseException]): Exception value 

313 

314 exc_tb (Optional[TracebackType]): Exception traceback object 

315 

316 Returns: 

317 Optional[bool]: Whether to silence the exception 

318 """ 

319 self.idx_writer.close() 

320 

321 def write( 

322 self, 

323 sequence_lengths: List[int], 

324 sequence_modes: Optional[List[int]], 

325 document_indices: List[int], 

326 ) -> None: 

327 """Write the index (.idx) file 

328 

329 Args: 

330 sequence_lengths (List[int]): The length of each sequence 

331 

332 sequence_modes (Optional[List[int]]): The mode of each sequences 

333 

334 document_indices (List[int]): The seqyebce indices demarcating the end of each document 

335 """ 

336 sequence_pointers = self._sequence_pointers(sequence_lengths) 

337 

338 # the number of sequences in the dataset 

339 sequence_count = len(sequence_lengths) 

340 self.idx_writer.write(struct.pack("<Q", sequence_count)) 

341 

342 # the number of documents in the dataset 

343 document_count = len(document_indices) 

344 self.idx_writer.write(struct.pack("<Q", document_count)) 

345 

346 # the number of tokens per sequence 

347 self.idx_writer.write(numpy.array(sequence_lengths, dtype=numpy.int32).tobytes(order="C")) 

348 

349 # the byte offsets for all sequences 

350 self.idx_writer.write(numpy.array(sequence_pointers, dtype=numpy.int64).tobytes(order="C")) 

351 

352 # the sequence indices marking the end of each document 

353 self.idx_writer.write(numpy.array(document_indices, dtype=numpy.int64).tobytes(order="C")) 

354 

355 # the mode per sequence 

356 if sequence_modes is not None: 

357 self.idx_writer.write(numpy.array(sequence_modes, dtype=numpy.int8).tobytes(order="C")) 

358 

359 def _sequence_pointers(self, sequence_lengths: List[int]) -> List[int]: 

360 """Build the sequence pointers per the sequence lengths and dtype size 

361 

362 Args: 

363 sequence_lengths (List[int]): The length of each sequence 

364 

365 Returns: 

366 List[int]: The pointer to the beginning of each sequence 

367 """ 

368 itemsize = DType.size(self.dtype) 

369 curr_ptr = 0 

370 list_ptr = [] 

371 for length in sequence_lengths: 

372 list_ptr.append(curr_ptr) 

373 curr_ptr += length * itemsize 

374 return list_ptr 

375 

376 

377class _IndexReader: 

378 """Object class to read the index (.idx) file 

379 

380 Args: 

381 idx_path (str): The path to the index file 

382 

383 multimodal (bool): Whether the dataset is multimodal 

384 """ 

385 

386 def __init__(self, idx_path: str, multimodal: bool) -> None: 

387 """Initialize the _IndexReader 

388 

389 Args: 

390 idx_path (str): The path to the index file. 

391 

392 multimodal (bool): Whether the dataset is multimodal. 

393 """ 

394 logger.info("Loading index file %s", idx_path) 

395 

396 with open(idx_path, "rb") as f: 

397 header = f.read(len(INDEX_HEADER)) 

398 assert header == INDEX_HEADER, f"Bad header in {idx_path}" 

399 

400 version = struct.unpack("<Q", f.read(8))[0] 

401 assert version == INDEX_VERSION, f"Unsupported index version {version} in {idx_path}" 

402 

403 code = struct.unpack("<B", f.read(1))[0] 

404 self.dtype = DType.dtype_from_code(code) 

405 self.dtype_size = DType.size(self.dtype) 

406 

407 self.sequence_count = struct.unpack("<Q", f.read(8))[0] 

408 self.document_count = struct.unpack("<Q", f.read(8))[0] 

409 payload_offset = f.tell() 

410 

411 # memory-map the whole file for fast zero-copy slicing 

412 self._mmap = numpy.memmap(idx_path, mode="r", order="C") 

413 self._buffer = memoryview(self._mmap) 

414 

415 # extract views 

416 logger.info("Extracting sequence lengths") 

417 self.sequence_lengths = numpy.frombuffer( 

418 self._buffer, dtype=numpy.int32, count=self.sequence_count, offset=payload_offset 

419 ) 

420 logger.info("Extracting sequence pointers") 

421 self.sequence_pointers = numpy.frombuffer( 

422 self._buffer, 

423 dtype=numpy.int64, 

424 count=self.sequence_count, 

425 offset=payload_offset + self.sequence_lengths.nbytes, 

426 ) 

427 logger.info("Extracting document indices") 

428 self.document_indices = numpy.frombuffer( 

429 self._buffer, 

430 dtype=numpy.int64, 

431 count=self.document_count, 

432 offset=payload_offset + self.sequence_lengths.nbytes + self.sequence_pointers.nbytes, 

433 ) 

434 

435 self.sequence_modes: Optional[numpy.ndarray] = None 

436 if multimodal: 

437 logger.info("Extracting sequence modes") 

438 self.sequence_modes = numpy.frombuffer( 

439 self._buffer, 

440 dtype=numpy.int8, 

441 count=self.sequence_count, 

442 offset=payload_offset 

443 + self.sequence_lengths.nbytes 

444 + self.sequence_pointers.nbytes 

445 + self.document_indices.nbytes, 

446 ) 

447 

448 sequence_length_count = self.sequence_lengths.shape[0] 

449 if sequence_length_count != len(self): 

450 raise ValueError("Sequence length count does not match the dataset length") 

451 if sequence_length_count != self.sequence_count: 

452 raise ValueError("Sequence length count does not match the index sequence count") 

453 if sequence_length_count != self.document_indices[-1]: 

454 raise ValueError("Sequence length count does not match the final document index") 

455 

456 logger.info("Sequences: %d | Documents: %d", len(self), self.document_indices.shape[0] - 1) 

457 

458 def close(self) -> None: 

459 """Release memory-mapped index resources.""" 

460 if getattr(self, "_mmap", None) is None: 

461 return 

462 

463 self.sequence_lengths = None 

464 self.sequence_pointers = None 

465 self.document_indices = None 

466 self.sequence_modes = None 

467 if getattr(self, "_buffer", None) is not None: 

468 self._buffer.release() 

469 self._buffer = None 

470 

471 self._mmap = None 

472 

473 def __del__(self) -> None: 

474 """Clean up the object.""" 

475 try: 

476 self.close() 

477 except (AttributeError, BufferError, ValueError): 

478 pass 

479 

480 def __len__(self) -> int: 

481 """Get the number of sequences in the dataset 

482 

483 Returns: 

484 int: The number of sequences in the dataset 

485 """ 

486 return self.sequence_count 

487 

488 def __getitem__(self, idx: int) -> Tuple[numpy.int32, numpy.int64, Optional[numpy.int8]]: 

489 """Return the pointer, length, and mode at the index 

490 

491 Args: 

492 idx (int): The index into the dataset 

493 

494 Returns: 

495 Tuple[numpy.int32, numpy.int64, Optional[numpy.int8]]: The pointer, length and mode 

496 at the index 

497 """ 

498 return ( 

499 self.sequence_pointers[idx], 

500 self.sequence_lengths[idx], 

501 self.sequence_modes[idx] if self.sequence_modes is not None else None, 

502 ) 

503 

504 

505class _BinReader(ABC): 

506 """Abstract class to read the data (.bin) file""" 

507 

508 @abstractmethod 

509 def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray: 

510 """Read bytes into a numpy array. 

511 

512 Args: 

513 dtype (Type[numpy.number]): Data-type of the returned array. 

514 

515 count (int): Number of items to read. 

516 

517 offset (int): Start reading from this offset (in bytes). 

518 

519 Returns: 

520 numpy.ndarray: An array with `count` items and data-type `dtype` constructed from 

521 reading bytes from the data file starting at `offset`. 

522 """ 

523 

524 

525class _MMapBinReader(_BinReader): 

526 """A _BinReader that memory maps the data (.bin) file""" 

527 

528 def __init__(self, bin_path: str) -> None: 

529 """Initialize the _MMapBinReader 

530 

531 Args: 

532 bin_path (str): The path to the data (.bin) file. 

533 """ 

534 # The file must stay open for the reader's lifetime; it is closed in __del__. 

535 self._file = open(bin_path, "rb") # pylint: disable=R1732 

536 self._mmap = numpy.memmap(self._file, mode="r", order="C") 

537 self._buffer = memoryview(self._mmap.data) 

538 

539 def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray: 

540 """Read bytes into a numpy array. 

541 

542 Args: 

543 dtype (Type[numpy.number]): Data-type of the returned array. 

544 

545 count (int): Number of items to read. 

546 

547 offset (int): Start reading from this offset (in bytes). 

548 

549 Returns: 

550 numpy.ndarray: An array with `count` items and data-type `dtype` constructed from 

551 reading bytes from the data file starting at `offset`. 

552 """ 

553 return numpy.frombuffer(self._buffer, dtype=dtype, count=count, offset=offset) 

554 

555 def close(self) -> None: 

556 """Close the backing file without invalidating arrays returned from mmap.""" 

557 file_obj = getattr(self, "_file", None) 

558 if file_obj is not None: 

559 file_obj.close() 

560 self._file = None 

561 

562 def __del__(self) -> None: 

563 """Clean up the object.""" 

564 try: 

565 self.close() 

566 except (AttributeError, BufferError, ValueError): 

567 pass 

568 

569 

570class _FileBinReader(_BinReader): 

571 """A _BinReader that reads from the data (.bin) file using a file pointer""" 

572 

573 def __init__(self, bin_path: str) -> None: 

574 """Initialize the _FileBinReader 

575 

576 Args: 

577 bin_path (str): The path to the data (.bin) file. 

578 """ 

579 self._bin_path = bin_path 

580 

581 def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray: 

582 """Read bytes into a numpy array. 

583 

584 Args: 

585 dtype (Type[numpy.number]): Data-type of the returned array. 

586 

587 count (int): Number of items to read. 

588 

589 offset (int): Start reading from this offset (in bytes). 

590 

591 Returns: 

592 numpy.ndarray: An array with `count` items and data-type `dtype` constructed from 

593 reading bytes from the data file starting at `offset`. 

594 """ 

595 out = numpy.empty(count, dtype=dtype) 

596 with open(self._bin_path, "rb", buffering=0) as f: 

597 f.seek(offset) 

598 byte_view = out.view(numpy.uint8) 

599 nbytes = f.readinto(byte_view) 

600 if nbytes != out.nbytes: 

601 raise ValueError( 

602 f"Short read from {self._bin_path}: requested {out.nbytes} bytes at offset {offset}, " 

603 f"read {nbytes or 0} bytes" 

604 ) 

605 return out 

606 

607 

608class _S3BinReader(_BinReader): 

609 """Stream ``.bin`` data from S3 via chunked ranged ``GetObject`` calls. 

610 

611 A single in-memory chunk (sized by 

612 :attr:`ObjectStorageConfig.bin_chunk_nbytes`) is cached so consecutive 

613 reads within the same chunk avoid network round-trips. Random-access 

614 reads outside the current chunk trigger a new ranged ``GetObject``. 

615 """ 

616 

617 def __init__(self, bin_path: str, object_storage_config: ObjectStorageConfig) -> None: 

618 """Initialize the _S3BinReader 

619 

620 Args: 

621 bin_path (str): The ``s3://bucket/key`` URI of the data (.bin) object. 

622 

623 object_storage_config (ObjectStorageConfig): Chunked-read configuration. 

624 """ 

625 if not HAS_BOTO3: 

626 raise ImportError("boto3 is required to read s3:// datasets. Install via `pip install boto3`.") 

627 if object_storage_config.bin_chunk_nbytes <= 0: 

628 raise ValueError(f"bin_chunk_nbytes must be positive, got {object_storage_config.bin_chunk_nbytes}") 

629 self._client = boto3.client("s3") 

630 self._s3_bucket, self._s3_key = _parse_s3_path(bin_path) 

631 self._cache_nbytes = object_storage_config.bin_chunk_nbytes 

632 self._cache_bytes_start: int = 0 

633 self._cache_bytes_end: int = 0 

634 self._cache: Optional[bytes] = None 

635 

636 def _extract_from_cache(self, offset: int, size: int) -> bytes: 

637 """Return a requested byte range from the active S3 cache window.""" 

638 if self._cache is None: 

639 raise RuntimeError("Cache is empty; cannot extract before first read") 

640 start = offset - self._cache_bytes_start 

641 end = start + size 

642 if start < 0 or end > len(self._cache): 

643 raise IndexError( 

644 f"Cache window [{self._cache_bytes_start}, {self._cache_bytes_end}) " 

645 f"does not contain requested range [{offset}, {offset + size})" 

646 ) 

647 return self._cache[start:end] 

648 

649 def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray: 

650 """Read ``count`` elements of ``dtype`` starting at byte ``offset``.""" 

651 size = count * DType.size(dtype) 

652 if self._cache is not None and offset >= self._cache_bytes_start and offset + size <= self._cache_bytes_end: 

653 return numpy.frombuffer(self._extract_from_cache(offset, size), dtype=dtype) 

654 

655 bytes_start = (offset // self._cache_nbytes) * self._cache_nbytes 

656 bytes_end = max(bytes_start + self._cache_nbytes, offset + size) 

657 self._cache = self._client.get_object( 

658 Bucket=self._s3_bucket, 

659 Key=self._s3_key, 

660 Range=f"bytes={bytes_start}-{bytes_end - 1}", 

661 )["Body"].read() 

662 self._cache_bytes_start = bytes_start 

663 self._cache_bytes_end = bytes_start + len(self._cache) 

664 return numpy.frombuffer(self._extract_from_cache(offset, size), dtype=dtype) 

665 

666 def __del__(self) -> None: 

667 try: 

668 self._client.close() 

669 except (AttributeError, RuntimeError): 

670 pass 

671 

672 

673class _MultiStorageClientBinReader(_BinReader): 

674 """Read ``.bin`` data via NVIDIA's :mod:`multi_storage_client`.""" 

675 

676 def __init__(self, bin_path: str, object_storage_config: ObjectStorageConfig) -> None: 

677 """Initialize the _MultiStorageClientBinReader 

678 

679 Args: 

680 bin_path (str): The ``msc://`` URI of the data (.bin) object. 

681 

682 object_storage_config (ObjectStorageConfig): Accepted for constructor 

683 parity with the other object-storage bin readers; unused here. 

684 """ 

685 del object_storage_config 

686 if not HAS_MSC: 

687 raise ImportError( 

688 "multi_storage_client is required to read msc:// datasets. " 

689 "Install via `pip install multi-storage-client`." 

690 ) 

691 self._client, self._bin_path = multi_storage_client.resolve_storage_client(bin_path) 

692 

693 def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray: 

694 """Read ``count`` elements of ``dtype`` starting at byte ``offset``.""" 

695 size = count * DType.size(dtype) 

696 buffer = self._client.read( 

697 path=self._bin_path, 

698 byte_range=multi_storage_client.types.Range(offset=offset, size=size), 

699 ) 

700 return numpy.frombuffer(buffer, dtype=dtype) 

701 

702 

703OBJECT_STORAGE_BIN_READERS: Dict[str, Type[_BinReader]] = { 

704 "s3": _S3BinReader, 

705 "msc": _MultiStorageClientBinReader, 

706} 

707 

708 

709class IndexedDataset(torch.utils.data.Dataset): 

710 """A fast, on-disk dataset backed by paired index and binary files.""" 

711 

712 def __init__( 

713 self, 

714 path_prefix: str, 

715 multimodal: bool = False, 

716 mmap: bool = True, 

717 object_storage_config: Optional[ObjectStorageConfig] = None, 

718 ) -> None: 

719 """Initialize the IndexedDataset 

720 

721 Args: 

722 path_prefix (str): The index (.idx) and data (.bin) prefix. May be an S3 URI 

723 (``s3://bucket/key``) when ``object_storage_config`` is provided. 

724 multimodal (bool): Whether the dataset is multimodal. Defaults to False. 

725 mmap (bool): Whether to mmap the .bin files. Defaults to True. Must be False 

726 for object-storage paths. 

727 object_storage_config (Optional[ObjectStorageConfig]): When provided and 

728 ``path_prefix`` is an S3/MSC URI, the .idx file is downloaded to 

729 ``object_storage_config.path_to_idx_cache`` and the .bin file is streamed 

730 via chunked GETs. 

731 """ 

732 super().__init__() 

733 normalized_prefix = _normalize_prefix(path_prefix) 

734 if _is_object_storage_path(normalized_prefix) and object_storage_config is not None: 

735 if mmap: 

736 raise ValueError( 

737 "mmap must be False for object-storage prefixes; " 

738 "set mmap=False (or set `mmap_bin_files: false` in the recipe)." 

739 ) 

740 idx_path = get_idx_path(normalized_prefix) 

741 cache_idx_path = _get_index_cache_path(idx_path, object_storage_config) 

742 _cache_index_file(idx_path, cache_idx_path) 

743 self.initialize(normalized_prefix, multimodal, mmap, object_storage_config) 

744 

745 def initialize( 

746 self, 

747 path_prefix: str, 

748 multimodal: bool, 

749 mmap: bool, 

750 object_storage_config: Optional[ObjectStorageConfig] = None, 

751 ) -> None: 

752 """Open the index and data readers for a dataset prefix. 

753 

754 Args: 

755 path_prefix (str): The index (.idx) and data (.bin) prefix. 

756 multimodal (bool): Whether the dataset is multimodal. 

757 mmap (bool): Whether to mmap the .bin files. 

758 object_storage_config (Optional[ObjectStorageConfig]): Object-storage 

759 configuration used for ``s3://``/``msc://`` prefixes. 

760 """ 

761 idx_path = get_idx_path(path_prefix) 

762 bin_path = get_bin_path(path_prefix) 

763 

764 if _is_object_storage_path(path_prefix) and object_storage_config is not None: 

765 # .idx is already cached locally; determine local path for _IndexReader 

766 local_idx_path = _get_index_cache_path(idx_path, object_storage_config) 

767 if not os.path.exists(local_idx_path): 

768 raise RuntimeError(f"Cached .idx not found: {local_idx_path}") 

769 access = "s3" if path_prefix.startswith(S3_PREFIX) else "msc" 

770 bin_reader: _BinReader = OBJECT_STORAGE_BIN_READERS[access](bin_path, object_storage_config) 

771 index_reader = _IndexReader(local_idx_path, multimodal) 

772 else: 

773 if not os.path.exists(idx_path) or not os.path.exists(bin_path): 

774 raise FileNotFoundError(f"Missing .idx or .bin at prefix {path_prefix}") 

775 bin_reader = _MMapBinReader(bin_path) if mmap else _FileBinReader(bin_path) 

776 index_reader = _IndexReader(idx_path, multimodal) 

777 

778 self.path_prefix = path_prefix 

779 self.multimodal = multimodal 

780 self.mmap = mmap 

781 self.object_storage_config = object_storage_config 

782 self.bin_reader = bin_reader 

783 self.index = index_reader 

784 

785 def __len__(self) -> int: 

786 """Get the number of sequences in the dataset 

787 

788 Returns: 

789 int: The number of sequences in the dataset 

790 """ 

791 return len(self.index) 

792 

793 def __getitem__( 

794 self, idx: Union[int, numpy.integer, slice] 

795 ) -> Union[ 

796 numpy.ndarray, 

797 Tuple[numpy.ndarray, Any], # mode attached 

798 List[numpy.ndarray], 

799 Tuple[List[numpy.ndarray], numpy.ndarray], 

800 ]: 

801 """Read one sequence or a contiguous slice of sequences. 

802 

803 Args: 

804 idx (Union[int, numpy.integer, slice]): The index or contiguous 

805 (step=1) slice into the dataset. 

806 

807 Returns: 

808 Union[numpy.ndarray, Tuple[numpy.ndarray, Any], List[numpy.ndarray], 

809 Tuple[List[numpy.ndarray], numpy.ndarray]]: The sequence(s), each 

810 paired with its mode when the dataset is multimodal. 

811 """ 

812 if isinstance(idx, (int, numpy.integer)): 

813 ptr, length, mode = self.index[idx] 

814 seq = self.bin_reader.read(self.index.dtype, length, ptr) 

815 return (seq, mode) if mode is not None else seq 

816 

817 if isinstance(idx, slice): 

818 start, _, step = idx.indices(len(self)) 

819 if step != 1: 

820 raise ValueError("Slices into IndexedDataset must be contiguous (step=1)") 

821 lengths = self.index.sequence_lengths[idx] 

822 modes = self.index.sequence_modes[idx] if self.multimodal else None 

823 offsets = numpy.cumsum(lengths, dtype=numpy.int64) 

824 token_count = int(offsets[-1]) if offsets.size else 0 

825 buffer = self.bin_reader.read( 

826 self.index.dtype, 

827 token_count, 

828 int(self.index.sequence_pointers[start]), 

829 ) 

830 sequences = numpy.split(buffer, offsets[:-1]) 

831 return (sequences, modes) if modes is not None else sequences 

832 

833 raise TypeError(f"Unexpected index type {type(idx)}") 

834 

835 def get( 

836 self, idx: int, offset: int = 0, length: Optional[int] = None 

837 ) -> Union[numpy.ndarray, Tuple[numpy.ndarray, Any]]: 

838 """Read a sub-range of one sequence. 

839 

840 Args: 

841 idx (int): The index into the dataset. 

842 offset (int): Number of leading elements to skip. Defaults to 0. 

843 length (Optional[int]): Number of elements to read; defaults to the 

844 remainder of the sequence. 

845 

846 Returns: 

847 Union[numpy.ndarray, Tuple[numpy.ndarray, Any]]: The sequence 

848 sub-range, paired with its mode for multimodal datasets. 

849 """ 

850 ptr, seq_len, mode = self.index[idx] 

851 length = seq_len - offset if length is None else length 

852 ptr += offset * DType.size(self.index.dtype) 

853 seq = self.bin_reader.read(self.index.dtype, length, ptr) 

854 return (seq, mode) if mode is not None else seq 

855 

856 @property 

857 def sequence_lengths(self) -> numpy.ndarray: 

858 """The length of each sequence in the dataset (numpy.int32 array).""" 

859 return self.index.sequence_lengths 

860 

861 @property 

862 def document_indices(self) -> numpy.ndarray: 

863 """The sequence indices demarcating documents (numpy.int64 array).""" 

864 return self.index.document_indices 

865 

866 @staticmethod 

867 def exists(path_prefix: str) -> bool: