"""Bounded publication of aligned data and metadata shards."""
from __future__ import annotations
import hashlib
import os
import struct
import tempfile
from collections.abc import Iterable
from pathlib import Path
from signal_dataset._internal.retry import RetryPolicy, retrying
from signal_dataset.config import StorageOptions
from signal_dataset.record.codec import SafeTensorsRecordCodec
from signal_dataset.record.metadata import RecordMetadataCodec
from signal_dataset.record.model import PublishedShard, Record, RecordMetadata
from signal_dataset.storage import ArrayRecordShardStore, default_registry
from signal_dataset.storage.contracts import ObjectStore, ShardStore
from signal_dataset.storage.paths import validate_root_uri
from signal_dataset.storage.resume import create_file_or_resume
from signal_dataset.storage.staging import default_staging
class RecordShardWriter:
def __init__(
self,
*,
object_store: ObjectStore | None = None,
shard_store: ShardStore | None = None,
options: StorageOptions | None = None,
) -> None:
self._object_store = object_store
self._options = StorageOptions() if options is None else options
# write() never stages, but a default built here must still carry the
# caller's store rather than resolving its own.
self._shard_store = (
ArrayRecordShardStore(
staging=default_staging(
object_store, max_open_shards=self._options.max_open_shards
)
)
if shard_store is None
else shard_store
)
self._records = SafeTensorsRecordCodec(policy=self._options.resources)
self._metadata = RecordMetadataCodec(
control=self._options.control,
resources=self._options.resources,
)
def write(
self,
records: Iterable[Record],
root: str,
*,
work_id: str,
attempt: int = 0,
retry: int | RetryPolicy | None = None,
) -> PublishedShard:
"""Write an aligned data and metadata shard pair.
Safe to re-run with the same `work_id` and `attempt`: an identical
object already in place is accepted rather than treated as a conflict,
so a producer that died between the two writes can simply run again.
"""
policy = RetryPolicy.coerce(retry)
# Every other entry point validates the root; without this a directory
# literally named "gs:notascheme" would resolve to a GCS client.
root = validate_root_uri(root)
shard_id = hashlib.sha256(f"{work_id}:{attempt}".encode()).hexdigest()
fanout = shard_id[:2]
data_uri = f"{root}/objects/data/{fanout}/{shard_id}.arrayrecord"
metadata_uri = f"{root}/objects/metadata/{fanout}/{shard_id}.arrayrecord"
object_store = self._object_store
if object_store is None:
# A filesystem backend is confined to a directory spanning both
# objects; a bucket backend is not bound to a location at all.
backend = default_registry().backend_for(data_uri)
store_root = (
os.path.commonpath((data_uri, metadata_uri)) if backend.filesystem else data_uri
)
object_store = backend.factory(store_root, root_is_directory=backend.filesystem)
with tempfile.TemporaryDirectory(prefix="signal-dataset-") as directory:
base = Path(directory)
data_path = base / "data.arrayrecord"
metadata_path = base / "metadata.arrayrecord"
count = self._write_pair(records, data_path, metadata_path)
data_bytes = data_path.stat().st_size
metadata_bytes = metadata_path.stat().st_size
byte_limit = self._options.max_encoded_shard_bytes
if byte_limit is not None and max(data_bytes, metadata_bytes) > byte_limit:
raise ValueError("encoded shard exceeds max_encoded_shard_bytes")
# A producer that died between these two writes used to poison its
# own (work_id, attempt) forever: the retry collided on the data
# object, and there is no delete to clean up with.
generations: list[int | None] = []
for uri, local, stored in (
(data_uri, data_path, data_bytes),
(metadata_uri, metadata_path, metadata_bytes),
):
def store_one(
uri: str = uri, local: Path = local, stored: int = stored
) -> int | None:
return create_file_or_resume(
object_store,
uri,
local,
content_type="application/octet-stream",
size=stored,
expected_count=count,
shards=self._shard_store,
reader_options=self._options.reader_options,
file_reader_buffer_size=self._options.file_reader_buffer_size,
)
generations.append(
retrying(policy, store_one, description=f"writing {uri}")
)
data_generation, metadata_generation = generations
return PublishedShard(
data_uri,
metadata_uri,
work_id,
attempt,
count,
data_bytes,
metadata_bytes,
data_generation,
metadata_generation,
)
def _write_pair(self, records: Iterable[Record], data: Path, metadata: Path) -> int:
spool = metadata.with_suffix(".records")
def encoded_data() -> Iterable[bytes]:
with spool.open("wb") as stream:
for index, record in enumerate(records):
record_limit = self._options.max_records_per_shard
if record_limit is not None and index >= record_limit:
raise ValueError("shard exceeds max_records_per_shard")
value = self._metadata.encode(RecordMetadata.from_record(record))
stream.write(struct.pack("<Q", len(value)))
stream.write(value)
yield self._records.encode(record)
def encoded_metadata() -> Iterable[bytes]:
with spool.open("rb") as stream:
while header := stream.read(8):
if len(header) != 8:
raise RuntimeError("truncated metadata spool length")
length = struct.unpack("<Q", header)[0]
value = stream.read(length)
if len(value) != length:
raise RuntimeError("truncated metadata spool record")
yield value
count = self._shard_store.write(data, encoded_data(), options=self._options.writer_options)
metadata_count = self._shard_store.write(
metadata, encoded_metadata(), options=self._options.writer_options
)
if metadata_count != count:
raise RuntimeError("aligned metadata shard record count changed")
return count
[docs]
def write_shard(
records: Iterable[Record],
root: str,
*,
work_id: str = "0",
attempt: int = 0,
object_store: ObjectStore | None = None,
shard_store: ShardStore | None = None,
options: StorageOptions | None = None,
retry: int | RetryPolicy | None = None,
) -> PublishedShard:
writer = RecordShardWriter(object_store=object_store, shard_store=shard_store, options=options)
return writer.write(records, root, work_id=work_id, attempt=attempt, retry=retry)