Source code for signal_dataset.config

"""Operational policies independent of the wire format."""

from dataclasses import dataclass, field
from enum import StrEnum

from signal_dataset._internal.constants import (
    DEFAULT_ARRAY_RECORD_OPTIONS,
    DEFAULT_MAX_OPEN_SHARDS,
)
from signal_dataset._internal.policy import ControlPolicy, ResourcePolicy
from signal_dataset._internal.retry import RetryPolicy

__all__ = [
    "AccessMode",
    "AnnotationOptions",
    "CachePolicy",
    "PublicationOptions",
    "RetryPolicy",
    "StorageOptions",
]




@dataclass(frozen=True, slots=True)
class StorageOptions:
    writer_options: str = DEFAULT_ARRAY_RECORD_OPTIONS
    reader_options: str = "readahead_buffer_size:0,max_parallelism:0"
    file_reader_buffer_size: int | None = None
    resources: ResourcePolicy = field(default_factory=ResourcePolicy)
    control: ControlPolicy = field(default_factory=ControlPolicy)
    manifest_cache_pages: int = 16
    records_per_read_batch: int = 1_024
    max_records_per_shard: int | None = None
    max_encoded_shard_bytes: int | None = None
    #: How many staged shards to keep open at once, on a backend that stages.
    #:
    #: Only `s3://` stages; local and `gs://` hand the URI straight to the
    #: reader. A staging miss re-downloads a whole object, so under an access
    #: order that alternates between more shards than are retained, every read
    #: costs a full shard. One is the right answer for a single sequential
    #: scan and the wrong one for a shuffled reader, which is the shape a
    #: training loop usually has -- so this is a choice the caller has to be
    #: able to make, and until now could not.
    max_open_shards: int = DEFAULT_MAX_OPEN_SHARDS

    def __post_init__(self) -> None:
        size = self.file_reader_buffer_size
        if size is not None and (isinstance(size, bool) or not isinstance(size, int) or size < 0):
            raise ValueError("file_reader_buffer_size must be a nonnegative integer or None")
        if (
            isinstance(self.manifest_cache_pages, bool)
            or not isinstance(self.manifest_cache_pages, int)
            or self.manifest_cache_pages < 1
        ):
            raise ValueError("manifest_cache_pages must be positive")
        if (
            isinstance(self.records_per_read_batch, bool)
            or not isinstance(self.records_per_read_batch, int)
            or self.records_per_read_batch < 1
        ):
            raise ValueError("records_per_read_batch must be positive")
        if self.max_records_per_shard is not None and (
            isinstance(self.max_records_per_shard, bool)
            or not isinstance(self.max_records_per_shard, int)
            or self.max_records_per_shard < 1
        ):
            raise ValueError("max_records_per_shard must be positive")
        if (
            isinstance(self.max_open_shards, bool)
            or not isinstance(self.max_open_shards, int)
            or self.max_open_shards < 1
        ):
            raise ValueError("max_open_shards must be positive")
        if self.max_encoded_shard_bytes is not None and (
            isinstance(self.max_encoded_shard_bytes, bool)
            or not isinstance(self.max_encoded_shard_bytes, int)
            or self.max_encoded_shard_bytes < 1
        ):
            raise ValueError("max_encoded_shard_bytes must be positive")



class AccessMode(StrEnum):
    LAZY_RANDOM_ACCESS = "lazy_random_access"


class CachePolicy(StrEnum):
    NONE = "none"




@dataclass(frozen=True, slots=True)
class PublicationOptions:
    """How a publication is laid out, and how hard it checks its inputs.

    `verify_record_counts` re-opens every shard a producer reported and counts
    its records. That is exact, and on object storage it costs a full read of
    everything being published -- for a terabyte dataset, a terabyte
    downloaded, twice, for data and metadata. It is off by default. The byte
    size of each shard is always compared, which costs one HEAD and catches the
    failure that actually happens: a truncated or half-written upload.
    """

    shards_per_manifest: int = 10_000
    verify_record_counts: bool = False

    def __post_init__(self) -> None:
        if (
            isinstance(self.shards_per_manifest, bool)
            or not isinstance(self.shards_per_manifest, int)
            or self.shards_per_manifest < 1
        ):
            raise ValueError("shards_per_manifest must be a positive integer")





@dataclass(frozen=True, slots=True)
class AnnotationOptions:
    records_per_shard: int = 10_000
    catalog_retry_limit: int = 8
    publication_id: str | None = None
    #: See PublicationOptions.verify_record_counts; same trade, same default.
    verify_record_counts: bool = False

    def __post_init__(self) -> None:
        for name in ("records_per_shard", "catalog_retry_limit"):
            value = getattr(self, name)
            if isinstance(value, bool) or not isinstance(value, int) or value < 1:
                raise ValueError(f"{name} must be a positive integer")
        if self.publication_id is not None and not self.publication_id:
            raise ValueError("publication_id must be non-empty or None")