Skip to content

io

io

Schedule dataset I/O at the level where each loader knows what it needs.

Opening or mapping a file gives access to its bytes; scattered requests can still incur repeated remote waits. These helpers let callers describe multiple reads together so independent I/O can overlap. They serve two different access patterns, using the same process-wide pool of eight workers:

  • PMD uses read_ranges(source, destination, ranges). It already owns a contiguous token-window cache and knows which 4 MiB blocks are missing. Supply disjoint, contiguous array slices: each copies source[span] into destination[span]. Slice offsets are array elements, not byte offsets. The call waits for all copies; unrequested destination regions are untouched. PMD marks those blocks loaded afterward and gathers samples from its cached strided view. Fetching individual rows through RowReader would bypass this block-reuse strategy. The caller owns allocation, cache validity, and eviction; read_ranges only fills the requested regions and keeps no cache of its own.

  • PaddedDataset and ContrastivePaddedDataset use RowReader. Their sampler requests arbitrary rows from paired token/mask files, without a PMD window cache. Map the files once, then pass the mappings and each batch's row IDs to read. It returns gathered arrays in the requested order, choosing between direct mapped indexing and sorted, coalesced positional reads based on observed cost. Its class docstring gives a complete usage example.

Use read_ranges when you already manage a buffer and its missing regions; use RowReader when you need a batch of rows assembled from mapped files. Neither helper chooses samples or changes their logical order. Files must remain immutable; ordinary sequential file reads are sufficient for a full-file scan.

RowReader()

Fetch batches of scattered rows from large, immutable token/mask files.

Think of each file as a table and each training batch as a list of row IDs. Opening the file gives you a byte stream; assembling a random batch yourself means translating those IDs into offsets and issuing reads. On a remote filesystem, scattered reads can spend far more time waiting than copying bytes. Memory mapping makes indexing convenient, but uncached rows still require filesystem reads.

This reader sees the whole batch request, so it can sort physical accesses, combine nearby reads, and overlap independent requests across token and mask files. It restores the requested row order afterward, including duplicates. It starts with ordinary mapped indexing and tries parallel I/O after slow reads, keeping it only when observed timings justify the overhead. For a sequential file scan, ordinary file reads are sufficient; use this for repeated random batches whose I/O benefits from being scheduled together.

Reuse one instance per sequential loader worker. Call map(path, dtype, shape) once for each raw, row-major file, using its on-disk dtype and two-dimensional shape. Then call samples, padding = reader.read((tokens, masks), indices) for each batch. Sources may also be contiguous column slices of those mappings. Mapping does not eagerly load the whole file; the OS manages cached file pages.

Example for paired token and mask files, each with 100,000 rows and 1,025 columns. These are the stored dimensions, including any shifted-label token. Create the reader and mappings outside the batch loop::

reader = RowReader()
shape = (
    100_000,
    1025,
)
tokens = reader.map(
    "train.bin",
    np.uint32,
    shape,
)
masks = reader.map(
    "train.bin.mask",
    np.bool_,
    shape,
)
sources = (
    tokens,
    masks,
)

# Supply row IDs from your sampler for each batch.
indices = np.array(
    [9, 2, 9],
    dtype=np.int64,
)
samples, padding = (
    reader.read(
        sources,
        indices,
    )
)

Both outputs have shape (3, 1025), in row order 9, 2, 9. For subsequent batches, call reader.read(sources, next_indices) with the same reader and mappings; do not reopen the files or construct a reader per batch.

The caller supplies valid, nonnegative row indices and owns sampling/RNG state. Returned arrays are writable copies, in source-tuple and index order, with the original dtypes. Keep files unchanged and the reader alive across batches to reuse open files and timing measurements. Do not share its scheduling state between concurrent consumers. Scheduling can change how bytes are fetched, but never which samples are returned.

map(path: str, dtype: DTypeLike, shape: tuple[int, ...]) -> np.memmap

Keep the mapping and parallel reads on the same open file.

read_ranges(source: np.ndarray, destination: np.ndarray, ranges: list[slice]) -> None

Fill disjoint contiguous ranges of a caller-owned buffer.