Skip to content

python-api.md

Latest commit

 

History

History
873 lines (632 loc) · 41.2 KB

python-api.md

File metadata and controls

873 lines (632 loc) · 41.2 KB
Raw

Python API

The Python API is intentionally thin. It describes what the user wants to cache or read, while the Rust core owns the cache format, validation, sampling, and iterator state.

CacheInput

class CacheInput(NamedTuple):
    data: bytes | bytearray | memoryview
    metadata: Mapping[str, bool | int | float | str]

data must be bytes-like. DatasetRT v0.1 does not accept arbitrary Python objects.

metadata must be a mapping with string keys and primitive values. Every emitted row must use the same metadata keys and compatible field types.

CacheSource

class CacheSource(Protocol):
    name: str

    def __iter__(self) -> Iterator[CacheInput]:
        ...

A CacheSource is a synchronous Python iterable. Rust pulls from it and owns the cache construction rules and bounded sample queues. Optional spawned source writers execute independent iterables in child processes without transferring payloads.

DatasetRuntime

runtime = DatasetRuntime(num_workers=4)

DatasetRuntime owns exactly num_workers reusable Rust threads. Cache loading, sample reading, and cache writing called through the runtime share that pool. The worker count is never inferred from the machine and is not repeated on operations.

DatasetRuntime.write_cache

results = runtime.write_cache(
    source,
    base_cache_dir,
    writer_config=WriterConfig(
        prefetch_size=64,
        max_shard_bytes=...,
        shard_compression=ShardCompression(algo="none", ratio=1.0),
        show_progress=True,
        validate_cache=False,
    ),
)
results = runtime.write_cache(
    [source_a, source_b],
    base_cache_dir,
    writer_config=WriterConfig(...),
)

DatasetRuntime.write_cache creates immutable caches and returns one result per source.

path is always the base cache directory. Python passes this through to Rust as base_cache_dir.

Each source is written by Rust under base_cache_dir / name. The returned list preserves source order and contains either CacheWriteSuccess(source_name, path) or CacheWriteError(source_name, message).

Rust writes each cache into a temporary directory under base_cache_dir / tmp first and publishes it by renaming only after the manifest is complete. If one source in a multi-source write fails, its temporary cache is cleaned up and the remaining sources still run.

Per-source failures do not stop the rest of the write. For example, an empty source returns CacheWriteError(source_name="empty", message="cache source yielded no samples") while other sources can still return CacheWriteSuccess.

for result in results:
    match result:
        case CacheWriteSuccess(source_name, path):
            print(source_name, path)
        case CacheWriteError(source_name, message):
            print(source_name, message)

prefetch_size bounds the writer's in-flight task and result window. If Python iteration is faster than writing, Rust applies backpressure when that window reaches min(prefetch_size, num_workers). For multi-source writes, the window may span source boundaries while ordered commit preserves deterministic output.

The runtime's fixed worker count limits active jobs. Calls do not create threads or resize the pool.

show_progress controls Rust-owned write progress. It is enabled by default and reports committed samples/s and MB/s for the active source. For multi-source writes, it also reports completed sources with an ETA based on wall time and completed source count. Set it to False for quiet tests, background jobs, or logging systems that do not want terminal progress output.

validate_cache controls checksum validation when the writer reuses an existing cache. It is False by default so DatasetRuntime.from_cache_sources does not hash every payload shard before handing paths to the dataset loader. Set it to True for integrity checks.

Writer configuration is a frozen pydantic model:

CompressionAlgo = Literal["none", "lz4"]

class ShardCompression(BaseModel):
    algo: CompressionAlgo = "none"
    ratio: float = 1.0

class WriterConfig(BaseModel):
    num_processes: int = 0
    process_timeout_seconds: float = 3600.0
    prefetch_size: int = 64
    max_shard_bytes: int = 64 * 1024 * 1024
    shard_compression: ShardCompression = ShardCompression()
    show_progress: bool = True
    validate_cache: bool = False
    profiler: WriterProfilerConfig = WriterProfilerConfig()

class WriterProfilerConfig(BaseModel):
    enabled: bool = False
    path: Path = Path("dataset_rt_profile.json")

Python validates the config shape with pydantic, then Rust validates it again before writing. DatasetRT v0.1 supports algo="none" with ratio=1.0 and algo="lz4" with a positive finite ratio value.

LZ4 compression is applied per payload record, not to the whole shard as one stream. This keeps random sample access direct: the index locates one compressed record, and Rust decompresses that record before returning the original bytes. The commonly named Rust zstd crate is a binding to the C zstd library, so it is not used for v0.1. A pure-Rust zstd implementation may be considered later once its compression path is accepted as stable enough for DatasetRT.

Set profiler=WriterProfilerConfig(enabled=True, path=Path("profile.json")) to write a structured timing summary after the write returns, including handled failures such as Ctrl-C. The summary separates Python iterator time (python_next), Python-to-Rust extraction (python_extract), queue backpressure (ingress_wait), Rust metadata/record work, compression, disk writes, finish steps, and publish time.

Rules:

  • Every target path must not already contain a completed cache.
  • Source names used as subdirectories must be plain path segments.
  • prefetch_size must be greater than zero.
  • validate_cache=True verifies existing cache checksums before reuse.
  • The source must yield at least one sample.
  • Payload data must be bytes-like.
  • Metadata schema is inferred from the first sample and validated for every sample.
  • Rust assigns cache_id, sample_id, shard offsets, and checksums.

DatasetRuntime.cached_dataset

dataset = runtime.cached_dataset(
    [path],
    reader_config=ReaderConfig(
        seed=42,
        prefetch_size=64,
        shuffle=True,
        validate_cache=False,
    ),
)

The method loads one or more immutable caches and returns a synchronous CachedDataset. The dataset retains the runtime pool internally, so reads remain valid for the dataset's lifetime.

Reader configuration is a frozen pydantic model:

class ReaderConfig(BaseModel):
    seed: int
    prefetch_size: int = 64
    shuffle: bool = True
    validate_cache: bool = False

validate_cache is False by default so dataset construction reads manifests, metadata, and indexes without hashing every shard payload. Set it to True when startup should verify metadata, index, and shard checksums before iteration.

Rules:

  • Dataset state is owned by Rust.
  • Multiple readers may load the same cache concurrently.
  • prefetch_size controls the bounded Rust reader queue.
  • The runtime worker count limits active read jobs.
  • validate_cache=True verifies metadata, index, and shard checksums during dataset construction.
  • With shuffle=True, each Python iterator snapshots the current weight vector and the next finite window of the deterministic draw stream.
  • With shuffle=False, samples are emitted from a cyclic window over active table order.
  • Changing weights or metadata after iterator construction affects future iterators, not existing ones.

DatasetRuntime.from_cache_sources

result = runtime.from_cache_sources(
    [source_a, source_b],
    base_cache_dir,
    reader_config=ReaderConfig(seed=42),
    writer_config=WriterConfig(),
)

match result:
    case CacheSourcesDatasetSuccess(dataset, results):
        ...
    case CacheSourcesDatasetError(results, message):
        ...

The factory asks Rust to generate cache paths, reuse valid existing caches, create missing caches, and then load every successful cache path. It returns CacheSourcesDatasetSuccess(dataset, results) when at least one cache was loaded, or CacheSourcesDatasetError(results, message) when no cache could be loaded. The results list preserves every per-source CacheWriteSuccess or CacheWriteError, so callers can see what was missing or malformed.

Samples Metadata

metadata = dataset.get_metadata()
dataset.update_metadata(metadata)

Active samples and weights are represented inside a Polars metadata DataFrame, not as a bare vector. The table is designed to be filtered and edited by metadata:

cache_id | sample_id | <metadata columns...> | weight

Required columns:

  • cache_id: persistent v3 manifest ID, or supplied path position for legacy v2 caches.
  • sample_id: physical row within that cache.
  • Metadata columns: one column per metadata field stored in the cache.
  • weight: positive finite float, default 1.0.

get_metadata returns a copy. Mutating the Polars frame has no effect until it is passed to update_metadata.

update_metadata accepts a Polars DataFrame. Rust reads and validates the authoritative identity, stored metadata, and weight columns:

  • At least one physical sample must be included.
  • Unknown (cache_id, sample_id) pairs are rejected.
  • Stored metadata columns must be present.
  • Every weight must be finite and positive.

Rows may be filtered, duplicated, or reordered by Polars operations before calling update_metadata. Removed rows are excluded from future iterators. Duplicated rows repeat the same physical sample in the active row space, and non-shuffled iterators follow the active table order exactly.

Additional columns are preserved in the active metadata table. They will be returned by the next get_metadata() call.

dataset.set_epoch_len(n) changes how many samples future iterators emit before stopping. The active metadata table remains the sampling population. Ordered reads continue through the active table cyclically across iterator boundaries; shuffled reads continue one deterministic multinomial draw stream. Updating metadata resets the cursor or draw stream and sets epoch length to the new table row count. Call set_epoch_len afterward to override it.

For a quick development run:

dataset.update_metadata(dataset.get_metadata().head(100))

Metadata columns are included for ergonomic filtering and auditing. They are not trusted for sample identity; cache_id and sample_id are the identity fields.

Generated from public docstrings behind the dataset_rt.api facade.

CacheInput

One sample yielded by a CacheSource.

data is an already-serialized payload. metadata is stored separately in Arrow-compatible columns and is available for weighting/filtering.

Fields:

  • data: BytesLike: Bytes-like payload to store in DatasetRT shards.
  • metadata: Metadata: Primitive metadata columns for this physical sample.

CacheSource

Synchronous source protocol consumed by the Rust cache writer.

Fields:

  • name: str: Plain source name used by Rust when generating base_cache_dir/name.

CacheSource.__iter__() -> Iterator[CacheInput]

Yield cache inputs synchronously.

Rust pulls from this iterator and owns bounded sample prefetching, worker threads, and commits. Optional spawned source writers execute this iterator in a child; payloads stay in that process.

CacheSourcesDatasetError

Dataset creation result when no source produced a cache.

Fields:

  • results: list[CacheWriteResult]: Per-source write outcomes explaining why no dataset was loaded.
  • message: str: Human-readable summary of the failed dataset creation.

CacheSourcesDatasetResult

CacheSourcesDatasetResult = CacheSourcesDatasetSuccess | CacheSourcesDatasetError

Best-effort result returned by DatasetRuntime.from_cache_sources.

CacheSourcesDatasetSuccess

Dataset creation result when at least one source produced a cache.

Fields:

  • dataset: CachedDataset: Dataset loaded from all successful cache writes.
  • results: list[CacheWriteResult]: Per-source write outcomes, including failures.

CacheWriteError

Failed per-source cache write result.

Fields:

  • source_name: str: Source name, or a source index label if the name could not be read.
  • message: str: Human-readable reason the source was not written.

CacheWriteResult

CacheWriteResult = CacheWriteSuccess | CacheWriteError

Per-source cache write outcome returned by DatasetRuntime.write_cache.

CacheWriteSuccess

Successful per-source cache write result.

Fields:

  • source_name: str: Source name used to derive the cache path.
  • path: Path: Published cache directory.

CachedDataset

Synchronous iterable view over Rust-owned dataset state.

Users do not construct this class directly; use DatasetRuntime.cached_dataset or DatasetRuntime.from_cache_sources. Rust owns cache validation, active metadata state, epoch planning, sampling, bounded prefetching, and iterator cancellation.

Fields:

  • cache_paths: list[Path]: Cache directories in physical traversal order; only legacy IDs depend on it.
  • reader_config: ReaderConfig: Reader configuration used when this dataset was loaded.

CachedDataset.update_manifests(version: Literal[3]) -> Result[ManifestUpdateReport, ManifestUpdateError]

Persist current cache IDs when upgrading v2 manifests to v3.

Construct legacy datasets in the original cache order used by saved metadata tables. This explicit migration preserves those IDs; it does not replace them with name hashes. Already-v3 manifests are unchanged. Rust preflights all targets, then atomically replaces each manifest. Partial failures retain completed entries; retry in the original order. Caller owns manifest mutation exclusively. Sample files, active metadata, weights and iterator state are unchanged.

CachedDataset.__init__()

Reject direct construction because every dataset requires a runtime.

CachedDataset.dump_config(path: DirectoryPath) -> FilePath

Save construction settings and active metadata Arrow into an existing directory.

Return the generated dataset.json file. Cache payloads are referenced by absolute paths, not copied. Active row order, duplicates, weights, extra columns, and epoch length are preserved; reader progress is not saved. The directory must be exclusively owned for export and contain no prior export. Publish configuration after metadata is complete; existing files are never overwritten. Keep the artifacts immutable while readers use them.

CachedDataset.from_config(cls, path: FilePath) -> CachedDataset

Construct fresh native state from an existing config and its metadata Arrow.

Validate paths and config before construction, then validate metadata in Rust. Restore the finite epoch length after metadata and start fresh cursors. No caller runtime or saved native handles are required.

CachedDataset.__iter__() -> Iterator[CachedSample]

Create an iterator from the current active metadata table and epoch length.

Iterator construction snapshots Rust runtime state. Later set_epoch_len or update_metadata calls affect future iterators, not this iterator.

With ReaderConfig.shuffle=True, Rust creates a deterministic weighted multinomial stream over active metadata rows. With shuffle=False, iteration follows active metadata table order from the current cyclic cursor, including duplicate rows.

CachedDataset.__len__() -> int

Return the number of samples emitted by each future iterator.

This value changes after set_epoch_len or update_metadata. Existing iterators keep their own snapshot even if this value changes mid-epoch.

CachedDataset.get_item(cache_id: int, sample_id: int) -> CachedSample

Read a physical sample without advancing an iterator or sampling cursor.

IDs refer to immutable cache contents, so filtered or duplicated active metadata rows do not change which sample this method returns.

Raises: TypeError: If either ID is not a plain integer. IndexError: If either ID is outside the loaded cache range.

CachedDataset.set_epoch_len(epoch_len: int) -> None

Set how many samples each future iterator emits before stopping.

epoch_len must be at least 1. The active metadata table remains the sampling population; this method changes only the finite window length. Existing iterators keep their snapshot.

With ReaderConfig.shuffle=False, future iterators continue from the current cyclic active-row cursor. With ReaderConfig.shuffle=True, future iterators continue from the current multinomial draw stream. update_metadata resets the cursor or draw stream and sets epoch_len to the new table row count. Call set_epoch_len afterward to override it.

CachedDataset.to_torch_iterable_dataset() -> SizedTorchIterableDataset

Return a sized torch.utils.data.IterableDataset view.

The adapter yields the same CachedSample objects as DatasetRT's normal iterator and implements __len__, so PyTorch consumers can use it with DataLoader while keeping domain decoding in Python.

Raises: ImportError: If PyTorch is not installed in the active environment.

CachedDataset.to_torch_dataloader(*, batch_size: PositiveInt = 1, sample_transform_fn: Callable[[CachedSample], T], collate_fn: Callable[[list[T]], BatchT], shuffle: bool = True, seed: int | Literal['random'] = 'random', samples_per_epoch: PositiveInt | None = None, worker_partition: Literal['split', 'replicate'] | None = None, work_dir: DirectoryPath | None = None, num_workers: int | Literal['auto'] = 'auto', drop_last: bool = False, pin_memory: bool = False, timeout: float = 0, native_num_workers: int | Literal['auto'] = 'auto', multiprocessing_context: Literal['inherit', 'fork', 'spawn', 'forkserver'] | BaseContext = 'spawn', worker_init_fn: Callable[[int], None] | None = None, prefetch_factor: int | Literal['auto'] = 'auto', persistent_workers: bool | Literal['auto'] = 'auto', verbose: bool = True) -> DataLoader[CachedSample | T]

Collate lists of transformed samples into a caller-defined result.

CachedDataset.to_torch_dataloader(*, batch_size: PositiveInt = 1, sample_transform_fn: None = None, collate_fn: Callable[[list[CachedSample]], BatchT], shuffle: bool = True, seed: int | Literal['random'] = 'random', samples_per_epoch: PositiveInt | None = None, worker_partition: Literal['split', 'replicate'] | None = None, work_dir: DirectoryPath | None = None, num_workers: int | Literal['auto'] = 'auto', drop_last: bool = False, pin_memory: bool = False, timeout: float = 0, native_num_workers: int | Literal['auto'] = 'auto', multiprocessing_context: Literal['inherit', 'fork', 'spawn', 'forkserver'] | BaseContext = 'spawn', worker_init_fn: Callable[[int], None] | None = None, prefetch_factor: int | Literal['auto'] = 'auto', persistent_workers: bool | Literal['auto'] = 'auto', verbose: bool = True) -> DataLoader[CachedSample | T]

Collate lists of cache records without a sample transform.

CachedDataset.to_torch_dataloader(*, batch_size: PositiveInt = 1, sample_transform_fn: Callable[[CachedSample], T] | None = None, collate_fn: None = None, shuffle: bool = True, seed: int | Literal['random'] = 'random', samples_per_epoch: PositiveInt | None = None, worker_partition: Literal['split', 'replicate'] | None = None, work_dir: DirectoryPath | None = None, num_workers: int | Literal['auto'] = 'auto', drop_last: bool = False, pin_memory: bool = False, timeout: float = 0, native_num_workers: int | Literal['auto'] = 'auto', multiprocessing_context: Literal['inherit', 'fork', 'spawn', 'forkserver'] | BaseContext = 'spawn', worker_init_fn: Callable[[int], None] | None = None, prefetch_factor: int | Literal['auto'] = 'auto', persistent_workers: bool | Literal['auto'] = 'auto', verbose: bool = True) -> DataLoader[CachedSample | T]

Delegate to Torch default collation when no custom callback is supplied.

CachedDataset.to_torch_dataloader(*, shuffle: bool = True, seed: int | Literal['random'] = 'random', samples_per_epoch: PositiveInt | None = None, worker_partition: Literal['split', 'replicate'] | None = None, work_dir: DirectoryPath | None = None, batch_size: PositiveInt = 1, num_workers: int | Literal['auto'] = 'auto', sample_transform_fn: Callable[[CachedSample], T] | None = None, collate_fn: Callable[[list[T]], BatchT] | Callable[[list[CachedSample]], BatchT] | None = None, drop_last: bool = False, pin_memory: bool = False, timeout: float = 0, native_num_workers: int | Literal['auto'] = 'auto', multiprocessing_context: Literal['inherit', 'fork', 'spawn', 'forkserver'] | BaseContext = 'spawn', worker_init_fn: Callable[[int], None] | None = None, prefetch_factor: int | Literal['auto'] = 'auto', persistent_workers: bool | Literal['auto'] = 'auto', verbose: bool = True) -> DataLoader[CachedSample | T]

Build a finite PyTorch DataLoader for training or evaluation.

samples_per_epoch counts global samples before batching, not batches or optimizer steps. None inherits len(self) at loader construction, including any set_epoch_len override. A positive integer overrides that inherited count for this loader. Sequential readers wrap around their population when necessary. Later changes to the source dataset do not change an existing loader's sample budget. The inherited or explicit count must be between 1 and min(sys.maxsize, 2**53), keeping PyTorch's batch-length calculation exact.

With one training process, len(loader) is exact: with batch size B and sample budget N it is ceil(N / B), or floor(N / B) with drop_last=True. For example, N=10 and B=3 yields four batches, or three when dropping the incomplete batch. batch_size must be a positive integer no larger than sys.maxsize; None is not supported. Worker quotas consist of whole batches plus at most one incomplete batch. Increasing num_workers neither multiplies N nor adds dropped tails. Workers without an assigned quota yield nothing.

With an initialized default distributed group of R training processes, each receives ceil(N/R) samples when drop_last=False, or Bfloor(N/(RB)) when drop_last=True. All report the same batch length. A remainder pads the global count or discards samples respectively; process zero warns with requested and effective counts. For N=1000, R=3, B=32, each gets 334 samples and 11 batches, or 320 samples and 10 batches when dropping tails. A zero effective count yields no batches.

worker_partition=None (the default) selects replicate for shuffled training and split for sequential reading. Replicate preserves global weighted frequencies while retaining full private sampling state in each consumer. Ten million f64 weights alone require 80 MB per consumer.

worker_partition="split" divides active row positions across training processes, then their loading workers. Padding may repeat row positions. With shuffle=False, metadata stays in its existing order and workers read their populations sequentially. With shuffle=True, metadata rows use one common seeded ordering before splitting, then each worker samples its own population by weight with replacement. Randomized splits improve the weight mix but do not guarantee the global weighted frequency: worker quotas follow batch counts, not partition weight sums.

worker_partition="replicate" gives every worker the full population while sharing the same total sample budget. With shuffle=True, each worker samples by weight with replacement; with shuffle=False, each starts at the beginning, so different workers can emit duplicate samples. With one training process and one consumer, both modes use the full population. Torch interleaves worker batches; multiworker output need not match global metadata order. Weights do not affect sequential reads.

An integer seed is accepted only with shuffle=True; otherwise ValueError avoids silently ignoring it. An explicit seed reproduces fresh loaders with the same worker topology. "random" (the default) selects randomness once during construction and is unused for sequential reading. None is not accepted as a seed or worker-count mode. Repeated passes reuse native state in the main process or persistent workers and continue their draw streams/cursors. Nonpersistent workers reconstruct fresh readers with the same saved seed and reset draw offset or sequential cursor each pass, repeating the same sample sequence when metadata is unchanged. This also happens with seed="random", because its randomness is chosen once per loader. A short epoch can therefore keep omitting the same population rows. With num_workers>0, use persistent_workers="auto" (the default) to continue across epochs; explicit False emits a warning. With num_workers=0 the reader stays in the main process and continues without persistent workers; auto disables the worker-only flag in that mode. Explicit persistent_workers=True requires positive loading workers. Split populations are fixed for the loader's lifetime, including when persistent workers are used.

Construction saves a config and active metadata Arrow without loading payloads or creating a consuming native reader. Parent edits do not change snapshots. Workers receive the config path and compact partition instructions, map shared snapshot buffers, and build their own native sampling state. They do not query distributed groups. With shuffle=True and seed="random", training processes agree a random partition seed during construction; construct loaders in the same order with the same global budget and batch size on every process. Each worker uses a distinct derived sampling seed.

work_dir is an existing directory for exported construction artifacts. Each loader creates its own subdirectory there, retained for inspection or reconstruction. None creates a managed temporary directory, removed when the parent loader and its active iterators release their dataset. Keep artifacts immutable and accessible to every consuming process. A hard process kill can leave temporary artifacts behind. Native setup happens on first consumption, or before worker_init_fn in each worker. num_workers="auto" selects up to four Torch processes from the CPU budget per local rank minus one; a one-CPU budget selects zero. This leaves room in the process-count heuristic for training, but does not reserve CPUs or account for other workload threads. Explicit zero keeps loading in the main process and allows non-picklable local transforms. native_num_workers="auto" then selects 1–4 Rust reader threads per consumer. Both modes share one budget observation from CPU affinity and visible cgroup v1/v2 quota and CPU-set limits, divided among consumers and local ranks. LOCAL_WORLD_SIZE supplies local rank count; absent launcher information conservatively uses distributed world size. Explicit positive integers override detection. Fractional quotas round down with a one-thread minimum; unreadable limits warn and fall back to one CPU. Automatic counts can vary across machines; fix both counts for reproducible topology. This is a resource budget, not a throughput guarantee. Queues remain bounded by ReaderConfig.prefetch_size. With num_workers>0, sample_transform_fn is checked for picklability before export or worker startup, for every start method. Other callbacks must also be picklable for spawn/forkserver. Successful serialization does not guarantee child imports will succeed; startup and callback errors propagate rather than replacing samples.

persistent_workers="auto" enables persistence exactly when loading workers exist. prefetch_factor="auto" selects two batches per loading worker, or disables prefetch for main-process loading. Explicit prefetch must be a positive integer with loading workers. None is not accepted for either selection mode.

verbose=True reports resolved sampling seeds, process/native thread counts, partition, persistence, prefetching, and reader starts/passes. Nonpersistent reader starts report sampling_reset=True and new PIDs. Events use INFO on the standard logger "dataset_rt.loader", respecting application handlers and filters; unconfigured processes use stderr. verbose=False suppresses lifecycle diagnostics but never the replay warning. Logs describe configured budgets, not measured thread activity or exact native draw offsets. Sampling seeds do not seed random user transforms.

Collation always receives a list, including when batch_size=1. The list contains transformed samples if sample_transform_fn is supplied, otherwise CachedSample values. multiprocessing_context selects how workers start: "spawn" (the default), "fork", "forkserver", or a multiprocessing BaseContext object. "inherit" uses Python's platform default. With num_workers=0, the context is omitted because no subprocesses are created. Resolved pinning, timeout, prefetch, and worker lifetime follow ordinary PyTorch semantics. PyTorch is optional until called.

CachedDataset.samples_metadata() -> pl.DataFrame

Compatibility alias for get_metadata.

Returns the same active metadata table as get_metadata. Prefer get_metadata in new code.

CachedDataset.get_metadata() -> pl.DataFrame

Return the in-memory metadata table that controls future iterators.

Contract:

  • Returns a Polars DataFrame with shared immutable buffers; editing it does not mutate the dataset until the whole frame is passed to update_metadata.
  • Columns are cache_id, sample_id, every metadata column stored in the cache, weight, and any extra columns preserved from the previous update_metadata call.
  • cache_id is persisted in v3 manifests. Legacy v2 caches use their supplied path positions until explicit manifest migration. sample_id is the physical row inside that cache.
  • Each row is one active sampling row. Duplicate (cache_id, sample_id) rows are allowed and represent repeated entries for the same physical sample.
  • len(dataset) equals the configured epoch length. By default it matches the active row count, including duplicate rows.
  • Original cache metadata is mapped lazily on first access. Later calls share the retained Arrow buffers. No cache files are rewritten.

CachedDataset.set_samples_metadata(metadata: pl.DataFrame) -> None

Compatibility alias for update_metadata.

Applies the same validation and runtime-only update semantics as update_metadata. Prefer update_metadata in new code.

CachedDataset.update_metadata(metadata: pl.DataFrame) -> None

Replace the Rust-owned active metadata table for future iterators.

Input contract:

  • metadata must be a Polars DataFrame.
  • Required columns are cache_id, sample_id, every metadata column stored in the cache, and weight.
  • cache_id and sample_id must be non-null integer columns that map to known physical cache samples.
  • Stored metadata columns must be present with the same Arrow types as the immutable cache metadata schema.
  • weight must be a non-null numeric column, and every value must be positive and finite.
  • Extra columns are allowed and preserved in runtime memory; they are returned by the next get_metadata call.

Row semantics:

  • Rows absent from metadata are removed from the active sampling space and excluded from future iterators.
  • Duplicate (cache_id, sample_id) rows are allowed. Each duplicate is a separate active row that points to the same immutable physical sample, useful for row-duplication balancing or OHEM.
  • len(dataset) becomes the new table row count, replacing any previous set_epoch_len override.
  • With ReaderConfig(shuffle=False), future iterators emit active rows exactly in table order, including duplicates.
  • With ReaderConfig.shuffle=True, future iterators sample with replacement over active rows using each row's weight.

Mutation boundary:

  • Rust validates the full table before replacing runtime state; a validation error leaves the previous active table intact.
  • The update is runtime-only and does not rewrite metadata.arrow, index.bin, shards, or manifests.
  • Iterators created before this call keep their existing snapshot; iterators created after this call use the new active table.

CachedSample

One sample emitted by CachedDataset iteration.

Fields:

  • data: bytes: Payload bytes loaded from the immutable cache.
  • metadata: dict[str, MetadataValue]: Metadata row associated with this physical sample.
  • cache_id: int: Stored v3 cache identity, or the supplied cache-path position for legacy v2.
  • sample_id: int: Physical sample row within the source cache.

DatasetRuntime

Owner of the fixed Rust worker pool used by DatasetRT operations.

Create one runtime per process or training job, then call its methods to write caches, load datasets, and iterate samples. The worker count is chosen once at construction and reused; per-operation APIs do not resize the pool.

DatasetRuntime.__init__(*, num_workers: int)

Create exactly num_workers reusable Rust worker threads.

num_workers must be positive. Rust validates the value before creating the native pool. The pool is owned by this runtime and kept alive by any datasets loaded through it.

DatasetRuntime.num_workers

Return the fixed worker count selected when this runtime was created.

DatasetRuntime.write_cache(sources: CacheSource | list[CacheSource], path: str | Path, *, writer_config: WriterConfig = DEFAULT_WRITER_CONFIG) -> list[CacheWriteResult]

Write one immutable cache per source under path.

sources may be one CacheSource or a list of sources. Each source must expose a plain name and yield CacheInput values. Rust writes source name under path / name, validates a stable metadata schema, stores payload bytes in shards, writes metadata.arrow and index.bin, then publishes a manifest only after the cache is complete.

Returns one CacheWriteSuccess or CacheWriteError per source in input order. Per-source failures are reported as values instead of exceptions when Rust can handle them cleanly.

writer_config.num_processes=0 keeps this native serial path. A positive count spawns source-writing processes, each with self.num_workers native threads. Sources must then be picklable and importable. Jobs and outcomes cross processes; sample payloads do not. Worker death or process_timeout_seconds stops the pool without retrying unfinished sources. Profiling keeps serial execution even with a positive count.

Each source has one owning worker; callers must exclude concurrent write invocations targeting the same destinations. Parallel mode refuses pre-existing temporary cache paths and reserves the tmp source name. show_progress reports completed source count from the parent; child sample progress is disabled. Use a main guard when launching processes.

DatasetRuntime.cached_dataset(paths: Sequence[str | Path], *, reader_config: ReaderConfig) -> CachedDataset

Load immutable cache directories into a CachedDataset.

V3 manifests define persistent cache_id values. Legacy v2 IDs remain positions in paths until explicit update_manifests(version=3) migration. DatasetRT validates manifests, schemas, metadata/index shape, and shard lengths while loading. Expensive checksum hashing is controlled by reader_config.validate_cache.

The returned dataset keeps this runtime's Rust worker pool alive and uses it for every future iterator.

DatasetRuntime.from_cache_sources(sources: CacheSource | list[CacheSource], path: str | Path, *, reader_config: ReaderConfig, writer_config: WriterConfig = DEFAULT_WRITER_CONFIG) -> CacheSourcesDatasetResult

Create or reuse source caches, then load all successful cache paths.

Existing complete caches under path / source.name are reused. Missing caches are written with writer_config. The method then loads every successful cache with reader_config.

Returns CacheSourcesDatasetSuccess when at least one cache is loaded; returns CacheSourcesDatasetError when every source failed or no source was provided. The result always includes per-source write outcomes so callers can audit partial success.

Parallel execution follows write_cache's process, timeout, ownership, and serialization contract. Successful paths preserve input source order, regardless of completion order; IDs come from v3 manifests or v2 positions.

ReaderConfig

Configuration for Rust-owned dataset reading and sampling.

Fields:

  • seed: int: Seed used for deterministic shuffled epoch planning.
  • prefetch_size: int: Capacity of Rust's bounded reader result queue.
  • shuffle: bool: Whether future iterators use deterministic weighted sampling.
  • validate_cache: bool: Whether cache checksums are verified while loading metadata and indexes.

ShardCompression

Compression policy requested for payload shards.

ratio is part of the explicit policy object so callers and manifests use the same structured shape. ratio is advisory metadata for compressed algorithms; Rust validates algo="none" with ratio == 1.0.

Fields:

  • algo: CompressionAlgo: Compression algorithm to apply independently to each payload record.
  • ratio: float: Expected compression ratio; algo='none' requires exactly 1.0.

SizedTorchIterableDataset

Sized PyTorch iterable view returned by to_torch_iterable_dataset.

SizedTorchIterableDataset.__iter__() -> Iterator[CachedSample]

Yield cached samples in DatasetRT iterator order.

SizedTorchIterableDataset.__len__() -> int

Return the physical sample count visible to PyTorch.

WriterConfig

Native cache writing with optional spawned source-process supervision.

Fields:

  • num_processes: int: Number of spawned source writers; zero writes in the calling process.
  • process_timeout_seconds: float: Deadline in seconds for child startup or a dispatched source, including deserialization. Timeout stops the pool without retries; ignored in serial mode.
  • prefetch_size: int: Maximum number of writer tasks/results buffered by Rust.
  • max_shard_bytes: int: Target shard byte size before Rust rotates to a new shard.
  • shard_compression: ShardCompression: Per-record payload compression policy for new shards.
  • show_progress: bool: Whether Rust renders cache write progress with samples/s and MB/s.
  • validate_cache: bool: Whether reused existing caches are checksum-validated before loading.
  • profiler: WriterProfilerConfig: Optional JSON writer-stage profiler configuration.

WriterProfilerConfig

Optional writer profiler output.

Profiling is disabled by default. When enabled, Rust writes a structured JSON summary at path after successful writes and handled failures such as Ctrl-C.

Fields:

  • enabled: bool: Whether Rust collects and writes writer-stage timing statistics.
  • path: Path: JSON summary path used when profiling is enabled.

AbsentManifestTarget

A rejected request has no individual filesystem target.

Fields:

  • reason: str: Why the rejected request cannot identify an individual manifest.

Err

Carry an expected failure without converting it into exception control flow.

Fields:

  • error: E: Typed rejection, available after narrowing this outcome to Err.

ManifestUpdateEntry

Describe a replacement or unchanged manifest; durability may be unconfirmed.

Fields:

  • path: Path: Resolved manifest file that was replaced or left unchanged.
  • cache_id: int: Dataset identity retained in the persisted manifest.
  • status: Literal['updated', 'unchanged', 'durability_unknown']: Replacement state; durability_unknown means rename succeeded but directory sync failed.

ManifestUpdateError

Retain completed replacements so callers can safely inspect and retry.

Fields:

  • path: Path | AbsentManifestTarget: Failed manifest target, or a reason-carrying absence for invalid requests.
  • message: str: Boundary diagnostic describing the expected validation or filesystem failure.
  • entries: tuple[ManifestUpdateEntry, ...]: Completed or durability-uncertain outcomes before the request stopped.

ManifestUpdateReport

One entry per cache after a complete, identity-preserving migration.

Fields:

  • entries: tuple[ManifestUpdateEntry, ...]: Immutable input-ordered outcomes for every cache in the completed request.

Ok

Carry a successful outcome that callers narrow with pattern matching.

Fields:

  • value: T: Accepted value, available after narrowing this outcome to Ok.

Result

Result = Ok[T] | Err[E]

Typed expected success/failure contract; interrupts and unexpected bugs may unwind.

Iteration

for sample in dataset:
    reveal_type(sample)  # CachedSample

CachedSample contains:

  • data: bytes
  • metadata: dict[str, bool | int | float | str]
  • cache_id: int
  • sample_id: int

The default epoch length is the active metadata row count. dataset.set_epoch_len(n) changes how many samples future iterators emit. Each iterator snapshots the active metadata table and epoch length when it is created, so set_epoch_len and update_metadata affect future iterators and do not rewrite an iterator that is already mid-epoch.

PyTorch Adapter

torch_dataset = dataset.to_torch_iterable_dataset()
loader = torch.utils.data.DataLoader(torch_dataset, batch_size=None, num_workers=0)

The adapter is available only when PyTorch is installed. It returns a sized torch.utils.data.IterableDataset view that yields CachedSample objects and implements __len__.

DatasetRT does not decode tensors in the adapter. Payload decoding remains in Python, usually inside your collate function, transform, or training loop.

Keep PyTorch DataLoader(num_workers=0) with this adapter. Use DatasetRuntime(num_workers=N) for parallel cache reads; the adapter rejects PyTorch worker processes to avoid duplicated iterable streams.