"""Lazy random access to one immutable ordered dataset sequence."""
from __future__ import annotations
import bisect
import threading
from collections import OrderedDict
from collections.abc import Iterator, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, TypeVar, overload
from signal_dataset._internal.constants import LEGACY_ROOT_DOCUMENT, ROOT_DOCUMENT
from signal_dataset._internal.json import frozen_mapping
from signal_dataset.config import AccessMode, CachePolicy, StorageOptions
from signal_dataset.dataset.control_io import read_document
from signal_dataset.dataset.instructions import ReadInstruction
from signal_dataset.dataset.layout import (
DatasetRoot,
Manifest,
ManifestReference,
ShardEntry,
Snapshot,
)
from signal_dataset.dataset.versions import require_reader
from signal_dataset.errors import CorruptDatasetError
from signal_dataset.record.codec import SafeTensorsRecordCodec
from signal_dataset.record.contracts import RecordCodec
from signal_dataset.record.metadata import RecordMetadataCodec
from signal_dataset.record.model import Record, RecordMetadata
from signal_dataset.storage import ArrayRecordShardStore, is_absolute_uri, object_store_for
from signal_dataset.storage.contracts import ObjectStore, ShardStore
from signal_dataset.storage.paths import join_uri, validate_root_uri
from signal_dataset.storage.reference import ObjectReference
from signal_dataset.storage.staging import default_staging
T = TypeVar("T")
if TYPE_CHECKING:
from signal_dataset.dataset.view import DatasetView
def _truncated_shard_message(uri: str, index: int) -> str:
"""Name the object, not just the index that failed.
A shard whose upload was truncated reads back short, and the first sign of
it is an out-of-range index here rather than at publication time. Naming
the object and the remedy is the whole cost of publishing without
`verify_record_counts`, and it is a better error than the bare index the
reader used to raise.
"""
return (
f"shard {uri} does not contain record index {index}. "
"The shard may be truncated or may not match the manifest that names it; "
"republish with PublicationOptions(verify_record_counts=True) to catch "
"this at publication time instead."
)
def _resolve(root: str, path: str) -> str:
return path if is_absolute_uri(path) else join_uri(root, path)
class _ManifestCache:
def __init__(self, capacity: int) -> None:
self._capacity = capacity
self._values: OrderedDict[str, Manifest] = OrderedDict()
# A reader is shared across threads: Grain calls __getitem__ from a
# pool of sixteen by default, and a PyTorch DataLoader may too. A read
# here mutates the ordering, so it needs the lock as much as a write.
self._lock = threading.Lock()
def get(self, key: str) -> Manifest | None:
with self._lock:
value = self._values.get(key)
if value is not None:
self._values.move_to_end(key)
return value
def put(self, key: str, value: Manifest) -> None:
with self._lock:
self._values[key] = value
self._values.move_to_end(key)
while len(self._values) > self._capacity:
self._values.popitem(last=False)
@dataclass
class _Index:
root: str
snapshot: Snapshot
objects: ObjectStore
shards: ShardStore
options: StorageOptions
def __post_init__(self) -> None:
self._starts = tuple(item.first_ordinal for item in self.snapshot.manifests)
self._cache = _ManifestCache(self.options.manifest_cache_pages)
# Shard URIs already checked against the size their manifest declares.
# See _check_shard. Guarded because a Dataset is read concurrently; the
# worst a race costs is a duplicate HEAD, but the set itself must not be
# mutated from several threads at once.
self._checked: set[str] = set()
self._checked_lock = threading.Lock()
self._record_codec = SafeTensorsRecordCodec(policy=self.options.resources)
self._metadata_codec = RecordMetadataCodec(
control=self.options.control,
resources=self.options.resources,
)
def locate(self, ordinal: int) -> tuple[ShardEntry, int]:
page_position = bisect.bisect_right(self._starts, ordinal) - 1
reference = self.snapshot.manifests[page_position]
manifest = self._manifest(reference)
starts = tuple(item.first_ordinal for item in manifest.shards)
shard_position = bisect.bisect_right(starts, ordinal) - 1
shard = manifest.shards[shard_position]
offset = ordinal - shard.first_ordinal
local = offset if shard.record_indices is None else shard.record_indices[offset]
return shard, local
def read_instructions(
self, start: int, stop: int, *, metadata: bool
) -> list[ReadInstruction]:
"""Say where ordinals `start` to `stop` live, without reading them."""
if start < 0 or stop > self.snapshot.record_count or start > stop:
raise ValueError(
f"range [{start}, {stop}) is outside this dataset of "
f"{self.snapshot.record_count} records"
)
results: list[ReadInstruction] = []
ordinal = start
while ordinal < stop:
shard, _ = self.locate(ordinal)
if shard.record_indices is not None:
raise ValueError(
"read_instructions is not available for a dataset whose shards "
"are referenced with explicit record indices, because the "
"physical shard's own record count is not recorded there. "
"Call materialize() to copy the view into a dataset that "
"stores its own shards."
)
reference = shard.metadata if metadata else shard.data
offset = ordinal - shard.first_ordinal
take = min(shard.record_count - offset, stop - ordinal)
results.append(
ReadInstruction(
filename=_resolve(self.root, reference.path),
skip=offset,
take=take,
examples_in_shard=shard.record_count,
first_ordinal=ordinal,
)
)
ordinal += take
return results
def _manifest(self, reference: ManifestReference) -> Manifest:
key = f"{reference.object.path}#{reference.object.generation}"
cached = self._cache.get(key)
if cached is not None:
return cached
uri = _resolve(self.root, reference.object.path)
info = self.objects.info(uri, generation=reference.object.generation)
if info.size != reference.object.stored_bytes:
raise CorruptDatasetError("manifest stored_bytes does not match object")
try:
value = Manifest.from_dict(
read_document(
self.objects,
uri,
generation=reference.object.generation,
policy=self.options.control,
)
)
except (KeyError, TypeError, AttributeError, ValueError) as exc:
if isinstance(exc, CorruptDatasetError):
raise
raise CorruptDatasetError(f"invalid manifest: {exc}") from exc
if (
value.first_ordinal != reference.first_ordinal
or value.record_count != reference.record_count
):
raise CorruptDatasetError("manifest range does not match snapshot reference")
if len(value.shards) != reference.shard_count:
raise CorruptDatasetError("manifest shard count does not match snapshot reference")
stored_bytes = sum(
item.data.stored_bytes + item.metadata.stored_bytes for item in value.shards
)
if stored_bytes != reference.stored_bytes:
raise CorruptDatasetError("manifest stored bytes do not match snapshot reference")
self._cache.put(key, value)
return value
def _check_shard(self, uri: str, reference: ObjectReference) -> None:
"""Verify a shard object is the size its manifest says, once per shard.
A shard's object name is derived from its producer's work id and
attempt, not from its content, so the same name can be reused for
different bytes. Nothing in this library replaces a referenced object --
every write is create-only -- but an operator or a lifecycle rule can
delete one, and a later run with the same work id will then write
different content to the same path. Without a check, an older snapshot
that still names it would quietly return the wrong records.
Once per shard rather than once per record: the annotation read path
makes the same comparison, and paying it per element is what makes that
path expensive.
Skipped for a reference dataset's absolute paths, which name objects in
another dataset that this dataset's store is not confined to and cannot
address. Reading that dataset directly still checks them.
"""
if is_absolute_uri(reference.path):
return
with self._checked_lock:
if uri in self._checked:
return
observed = self.objects.info(uri, generation=reference.generation)
if observed.size != reference.stored_bytes:
raise CorruptDatasetError(
f"shard {uri} holds {observed.size} bytes where the manifest that "
f"names it declares {reference.stored_bytes}; the object it points "
"at is not the one that was published"
)
with self._checked_lock:
self._checked.add(uri)
def read(self, ordinal: int, *, metadata: bool) -> Record | RecordMetadata:
shard, local = self.locate(ordinal)
reference = shard.metadata if metadata else shard.data
uri = _resolve(self.root, reference.path)
self._check_shard(uri, reference)
try:
payload = self.shards.read(
uri,
local,
generation=reference.generation,
options=self.options.reader_options,
file_reader_buffer_size=self.options.file_reader_buffer_size,
)
except (IndexError, ValueError) as exc:
raise CorruptDatasetError(_truncated_shard_message(uri, local)) from exc
codec: RecordCodec[Record] | RecordCodec[RecordMetadata]
codec = self._metadata_codec if metadata else self._record_codec
return codec.decode(payload)
def _iter_payloads(self, *, metadata: bool) -> Iterator[bytes]:
"""Yield encoded entries in bounded batches across physical-shard runs."""
pending_reference = None
pending_indices: list[int] = []
def read_pending() -> tuple[bytes, ...]:
if pending_reference is None or not pending_indices:
return ()
uri = _resolve(self.root, pending_reference.path)
self._check_shard(uri, pending_reference)
try:
payloads = self.shards.read_many(
uri,
pending_indices,
generation=pending_reference.generation,
options=self.options.reader_options,
file_reader_buffer_size=self.options.file_reader_buffer_size,
)
except (IndexError, ValueError) as exc:
raise CorruptDatasetError(
_truncated_shard_message(uri, pending_indices[0])
) from exc
return payloads
for manifest_reference in self.snapshot.manifests:
manifest = self._manifest(manifest_reference)
for shard in manifest.shards:
reference = shard.metadata if metadata else shard.data
indices = (
tuple(range(shard.record_count))
if shard.record_indices is None
else shard.record_indices
)
if pending_reference is not None and reference != pending_reference:
yield from read_pending()
pending_indices.clear()
pending_reference = reference
for index in indices:
pending_indices.append(index)
if len(pending_indices) == self.options.records_per_read_batch:
yield from read_pending()
pending_indices.clear()
yield from read_pending()
def iter_records(self) -> Iterator[Record]:
"""Yield full records through bounded physical-shard requests."""
for payload in self._iter_payloads(metadata=False):
yield self._record_codec.decode(payload)
def iter_metadata(self) -> Iterator[RecordMetadata]:
"""Yield metadata through bounded physical-shard requests."""
for payload in self._iter_payloads(metadata=True):
yield self._metadata_codec.decode(payload)
class RecordMetadataSequence(Sequence[RecordMetadata]):
def __init__(self, index: _Index) -> None:
self._index = index
def __len__(self) -> int:
return self._index.snapshot.record_count
@overload
def __getitem__(self, ordinal: int) -> RecordMetadata: ...
@overload
def __getitem__(self, ordinal: slice) -> list[RecordMetadata]: ...
def __getitem__(self, ordinal: int | slice) -> RecordMetadata | list[RecordMetadata]:
if isinstance(ordinal, slice):
return [self[index] for index in range(*ordinal.indices(len(self)))]
normalized = _normalize(ordinal, len(self))
value = self._index.read(normalized, metadata=True)
assert isinstance(value, RecordMetadata)
return value
[docs]
class Dataset(Sequence[Record]):
def __init__(
self,
root: str,
dataset_root: DatasetRoot,
snapshot: Snapshot,
index: _Index,
) -> None:
self.root = root
self.id = dataset_root.id
self.snapshot_id = snapshot.id
self.metadata = frozen_mapping(dataset_root.metadata)
self.record_metadata = RecordMetadataSequence(index)
self.access_mode = AccessMode.LAZY_RANDOM_ACCESS
self.cache_policy = CachePolicy.NONE
self._index = index
from signal_dataset.annotation.catalog import open_catalog
# Opened eagerly, deliberately. Deferring it would save one request per
# open -- and would let a dataset whose annotation catalog is corrupt
# open successfully and fail later, somewhere else. That is a bad trade
# for one request.
self.annotations = open_catalog(self)
def __repr__(self) -> str:
"""Stable across processes.
Grain's DataLoader writes repr(source) into a checkpoint and compares it
on restore, so a default repr carrying a memory address fails every
restore. Naming the dataset rather than the object is also simply more
useful in a log line.
"""
return (
f"{type(self).__name__}(root={self.root!r}, id={self.id!r}, "
f"snapshot_id={self.snapshot_id!r}, records={len(self)})"
)
[docs]
def read_instructions(
self, start: int = 0, stop: int | None = None, *, metadata: bool = False
) -> list[ReadInstruction]:
"""Where a range of records lives, as shard URIs and offsets.
Performs no I/O beyond manifest pages already loaded, so it is safe to
call in a `DataLoader.__init__` or inside a worker.
This is the seam for a caller who wants their own fetching. The library
owns addressing and decoding because those *are* the format; ordering,
concurrency, prefetch and caching belong to whatever is feeding the
trainer. Pair it with :func:`signal_dataset.decode`::
for instruction in dataset.read_instructions(0, 1024):
payloads = my_reader(instruction.filename,
instruction.skip, instruction.take)
records = [sds.decode(payload) for payload in payloads]
`metadata=True` addresses the aligned metadata shards instead, which
carry every record's identity and shape without its samples.
"""
return self._index.read_instructions(
start, len(self) if stop is None else stop, metadata=metadata
)
@property
def object_store(self) -> ObjectStore:
return self._index.objects
@property
def shard_store(self) -> ShardStore:
return self._index.shards
@property
def storage_options(self) -> StorageOptions:
return self._index.options
def __len__(self) -> int:
return self._index.snapshot.record_count
@overload
def __getitem__(self, ordinal: int) -> Record: ...
@overload
def __getitem__(self, ordinal: slice) -> list[Record]: ...
def __getitem__(self, ordinal: int | slice) -> Record | list[Record]:
if isinstance(ordinal, slice):
return [self[index] for index in range(*ordinal.indices(len(self)))]
value = self._index.read(_normalize(ordinal, len(self)), metadata=False)
assert isinstance(value, Record)
return value
def __iter__(self) -> Iterator[Record]:
"""Yield every record, shard by shard.
Routed through the same batched path as `iter_records`. Reading one
ordinal at a time opens a reader per record, which is the obvious thing
to write and was the expensive one.
"""
return self._index.iter_records()
[docs]
def iter_records(self) -> Iterator[Record]:
"""Yield all full records with bounded, shard-batched I/O."""
return self._index.iter_records()
[docs]
def select(self, indices: Sequence[int]) -> DatasetView:
from signal_dataset.dataset.view import DatasetView
return DatasetView(self, indices)
def _normalize(ordinal: int, length: int) -> int:
normalized = ordinal + length if ordinal < 0 else ordinal
if normalized < 0 or normalized >= length:
raise IndexError("dataset index out of range")
return normalized
[docs]
def open(
root: str,
*,
object_store: ObjectStore | None = None,
options: StorageOptions | None = None,
shard_store: ShardStore | None = None,
) -> Dataset:
root = validate_root_uri(root)
objects = (
object_store_for(root, root_is_directory=True) if object_store is None else object_store
)
settings = StorageOptions() if options is None else options
shards = (
ArrayRecordShardStore(
staging=default_staging(
objects, root=root, max_open_shards=settings.max_open_shards
)
)
if shard_store is None
else shard_store
)
try:
root_uri = join_uri(root, ROOT_DOCUMENT)
try:
root_value = read_document(objects, root_uri, policy=settings.control)
except FileNotFoundError as exc:
legacy_uri = join_uri(root, LEGACY_ROOT_DOCUMENT)
try:
objects.info(legacy_uri)
except FileNotFoundError:
raise exc from None
raise CorruptDatasetError(
"the earlier split-based 0.1 prototype is unsupported; "
"publish a new ordered dataset"
) from exc
dataset_root = DatasetRoot.from_dict(root_value)
require_reader(dataset_root.min_reader_version, root=root)
snapshot_reference = dataset_root.snapshot
snapshot_uri = join_uri(root, snapshot_reference.path)
snapshot_info = objects.info(snapshot_uri, generation=snapshot_reference.generation)
if snapshot_info.size != snapshot_reference.stored_bytes:
raise CorruptDatasetError("snapshot stored_bytes does not match object")
snapshot = Snapshot.from_dict(
read_document(
objects,
snapshot_uri,
generation=snapshot_reference.generation,
policy=settings.control,
)
)
except (KeyError, TypeError, AttributeError, ValueError) as exc:
if isinstance(exc, CorruptDatasetError):
raise
raise CorruptDatasetError(f"invalid dataset control document: {exc}") from exc
index = _Index(root, snapshot, objects, shards, settings)
return Dataset(root, dataset_root, snapshot, index)