"""Single-process and distributed annotation publication."""
from __future__ import annotations
import hashlib
import re
import time
from collections.abc import Callable, Mapping, Sequence
from typing import Any
from signal_dataset._internal.json import compact_json_bytes
from signal_dataset._internal.retry import RetryPolicy, retrying
from signal_dataset.annotation.catalog import (
CATALOG,
AnnotationSet,
DatasetContext,
_catalog_pointer,
_catalog_revision,
_parse,
)
from signal_dataset.annotation.codec import SafeTensorsAnnotationCodec
from signal_dataset.annotation.model import (
SET_NAME as _SET_NAME,
)
from signal_dataset.annotation.model import (
AnnotationRoot,
AnnotationShard,
PublishedAnnotationShard,
)
from signal_dataset.annotation.model import (
published_shard_from_dict as _published_shard_from_dict,
)
from signal_dataset.annotation.model import (
published_shard_to_dict as _published_shard_to_dict,
)
from signal_dataset.annotation.model import (
validated_published_shard as _validated_published_shard,
)
from signal_dataset.config import AnnotationOptions, StorageOptions
from signal_dataset.errors import CorruptDatasetError, PublicationCollisionError, StorageError
from signal_dataset.record.model import AnnotationRecord
from signal_dataset.storage.contracts import ObjectStore
from signal_dataset.storage.paths import join_uri, relative_uri
from signal_dataset.storage.reference import ObjectReference
from signal_dataset.storage.resume import create_file_or_resume
def require_compare_and_swap(objects: ObjectStore) -> None:
"""Refuse before writing anything, not at the catalog update.
Annotation shards are written first and the catalog is swung last, so a
store that cannot swap would otherwise upload every shard and then fail --
leaving objects nothing references and no way to delete them. DatasetView
calls this earlier still, before it publishes the dataset those annotations
would belong to.
"""
if not getattr(objects, "supports_compare_and_swap", True):
raise StorageError(
"annotations cannot be published through a mounted bucket: the "
"annotation catalog is the only object in the format that is "
"modified in place, and updating it safely needs a compare-and-swap "
"that neither gcsfuse nor Mountpoint for Amazon S3 supports. "
"Publish annotations against the native URI instead -- "
'sds.publish_annotations("gs://bucket/path.sds", ...) -- which '
"addresses the same dataset."
)
[docs]
def publish_annotations(
dataset: DatasetContext,
name: str,
records: Sequence[AnnotationRecord],
*,
metadata: Mapping[str, Any] | None = None,
options: AnnotationOptions | None = None,
retry: int | RetryPolicy | None = None,
) -> AnnotationSet:
"""Publish a dense annotation set for every record in the dataset.
Safe to re-run. With `publication_id` left unset it is derived from the
snapshot and set name, so a repeat writes to the same places and converges
rather than orphaning a new set of objects under a fresh identity.
"""
require_compare_and_swap(dataset.object_store)
policy = RetryPolicy.coerce(retry)
def attempt() -> AnnotationSet:
return _publish_annotations_once(
dataset, name, records, metadata=metadata, options=options
)
return retrying(policy, attempt, description=f"publishing annotation set {name!r}")
def _publish_annotations_once(
dataset: DatasetContext,
name: str,
records: Sequence[AnnotationRecord],
*,
metadata: Mapping[str, Any] | None,
options: AnnotationOptions | None,
) -> AnnotationSet:
if _SET_NAME.fullmatch(name) is None:
raise ValueError("annotation set name must be a portable identifier")
if len(records) != len(dataset):
raise ValueError("annotation set must be dense and dataset-aligned")
for index, record in enumerate(records):
if (
record.source_index != index
or record.source_record_id != dataset.record_metadata[index].id
):
raise ValueError("annotation source identity does not match dataset ordinal")
objects = dataset.object_store
shards = dataset.shard_store
storage = dataset.storage_options
settings = AnnotationOptions() if options is None else options
# Derived rather than random. A fresh UUID made every re-run write a whole
# new set of objects under a new prefix and then collide on the catalog
# name, orphaning the lot -- so publish_annotations was idempotent only for
# a caller who knew to pin AnnotationOptions.publication_id.
publication_id = settings.publication_id or hashlib.sha256(
f"{dataset.snapshot_id}:{name}".encode()
).hexdigest()[:32]
if _SET_NAME.fullmatch(publication_id) is None:
raise ValueError("annotation publication_id must be a portable identifier")
codec = SafeTensorsAnnotationCodec(options=dataset.storage_options)
published: list[AnnotationShard] = []
import tempfile
from pathlib import Path
for first in range(0, len(records), settings.records_per_shard):
values = records[first : first + settings.records_per_shard]
encoded = tuple(codec.encode(value) for value in values)
content_hash = hashlib.sha256()
for payload in encoded:
content_hash.update(len(payload).to_bytes(8, "little"))
content_hash.update(payload)
shard_id = hashlib.sha256(
f"{publication_id}:{first}:{content_hash.hexdigest()}".encode()
).hexdigest()
path = (
f"annotations/sets/{name}/{publication_id}/objects/"
f"{shard_id[:2]}/{shard_id}.arrayrecord"
)
with tempfile.TemporaryDirectory(prefix="signal-dataset-annotation-") as directory:
temporary = Path(directory) / "annotation.arrayrecord"
count = shards.write(
temporary,
encoded,
options=storage.writer_options,
)
size = temporary.stat().st_size
shard_uri = join_uri(dataset.root, path)
create_file_or_resume(
objects,
shard_uri,
temporary,
content_type="application/octet-stream",
size=size,
expected_count=count,
shards=shards,
reader_options=storage.reader_options,
file_reader_buffer_size=storage.file_reader_buffer_size,
)
published.append(AnnotationShard(first, count, ObjectReference(path, None, size)))
root_specification = AnnotationRoot(
name,
publication_id,
dataset.snapshot_id,
len(records),
tuple(published),
metadata or {},
)
root_path = f"annotations/sets/{name}/{publication_id}/root.json"
root_value = _root_to_dict(root_specification)
root_payload = compact_json_bytes(root_value)
root_uri = join_uri(dataset.root, root_path)
try:
objects.create(root_uri, root_payload, content_type="application/json")
except PublicationCollisionError as exc:
existing_root = objects.read_version(root_uri)
if existing_root.data != root_payload:
raise PublicationCollisionError(f"annotation set name already exists: {name}") from exc
_update_catalog(
dataset.root,
objects,
name,
root_path,
len(root_payload),
settings.catalog_retry_limit,
storage,
)
return AnnotationSet(
dataset.root,
root_specification,
objects,
shards,
storage,
dataset.record_metadata,
)
def write_annotation_shard(
dataset: DatasetContext,
name: str,
publication_id: str,
first_ordinal: int,
records: Sequence[AnnotationRecord],
*,
logical_shard_id: str,
retry: int | RetryPolicy | None = None,
) -> PublishedAnnotationShard:
"""Write one immutable annotation shard without updating the catalog."""
policy = RetryPolicy.coerce(retry)
def attempt() -> PublishedAnnotationShard:
return _write_annotation_shard_once(
dataset,
name,
publication_id,
first_ordinal,
records,
logical_shard_id=logical_shard_id,
)
return retrying(policy, attempt, description=f"writing annotation shard {logical_shard_id!r}")
def _write_annotation_shard_once(
dataset: DatasetContext,
name: str,
publication_id: str,
first_ordinal: int,
records: Sequence[AnnotationRecord],
*,
logical_shard_id: str,
) -> PublishedAnnotationShard:
if _SET_NAME.fullmatch(name) is None or _SET_NAME.fullmatch(publication_id) is None:
raise ValueError("annotation name and publication_id must be portable identifiers")
if _SET_NAME.fullmatch(logical_shard_id) is None:
raise ValueError("logical_shard_id must be a portable identifier")
if not records:
raise ValueError("annotation shards must be non-empty")
record_limit = dataset.storage_options.max_records_per_shard
if record_limit is not None and len(records) > record_limit:
raise ValueError("annotation shard exceeds max_records_per_shard")
if first_ordinal < 0 or first_ordinal + len(records) > len(dataset):
raise ValueError("annotation shard range is outside the source dataset")
for offset, record in enumerate(records):
ordinal = first_ordinal + offset
if (
record.source_index != ordinal
or record.source_record_id != dataset.record_metadata[ordinal].id
):
raise ValueError("annotation shard identity does not match source ordinal")
import tempfile
from pathlib import Path
with tempfile.TemporaryDirectory(prefix="signal-dataset-annotation-") as directory:
temporary = Path(directory) / "annotation.arrayrecord"
codec = SafeTensorsAnnotationCodec(options=dataset.storage_options)
encoded = tuple(codec.encode(value) for value in records)
count = dataset.shard_store.write(
temporary,
encoded,
options=dataset.storage_options.writer_options,
)
if count != len(records):
raise RuntimeError("annotation shard count changed while writing")
size = temporary.stat().st_size
byte_limit = dataset.storage_options.max_encoded_shard_bytes
if byte_limit is not None and size > byte_limit:
raise ValueError("annotation shard exceeds max_encoded_shard_bytes")
# Addressed by content, so an identical retry writes to the same place
# and resumes instead of orphaning an object under a fresh UUID. It is
# deliberately *not* addressed by identity: two workers racing the same
# logical shard with a nondeterministic annotator must both succeed and
# let the pointer below pick a winner, which identity-addressing would
# turn into a corruption error.
content_hash = hashlib.sha256()
for payload in encoded:
content_hash.update(len(payload).to_bytes(8, "little"))
content_hash.update(payload)
attempt_id = hashlib.sha256(
f"{name}:{publication_id}:{logical_shard_id}:{content_hash.hexdigest()}".encode()
).hexdigest()[:32]
path = (
f"annotations/sets/{name}/{publication_id}/objects/"
f"{attempt_id[:2]}/{attempt_id}.arrayrecord"
)
uri = join_uri(dataset.root, path)
# The shard pointer written below is a persisted control document, so
# the version this returns is deliberately dropped: a pointer written
# on GCS and then replayed against the same root through a mounted
# bucket would otherwise hand a generation to a local store, which
# rejects one outright.
create_file_or_resume(
dataset.object_store,
uri,
temporary,
content_type="application/octet-stream",
size=size,
)
candidate = PublishedAnnotationShard(
name,
publication_id,
dataset.snapshot_id,
logical_shard_id,
first_ordinal,
tuple(item.source_record_id for item in records),
uri,
size,
None,
tuple(item.status.value for item in records),
tuple(item.detail_status or item.status.value for item in records),
)
pointer_id = hashlib.sha256(logical_shard_id.encode("utf-8")).hexdigest()
pointer_path = (
f"annotations/sets/{name}/{publication_id}/shards/{pointer_id[:2]}/{pointer_id}.json"
)
pointer_uri = join_uri(dataset.root, pointer_path)
payload = compact_json_bytes(_published_shard_to_dict(candidate))
try:
dataset.object_store.create(pointer_uri, payload, content_type="application/json")
return candidate
except PublicationCollisionError:
existing = dataset.object_store.read_version(pointer_uri)
value = _parse(existing.data, dataset.storage_options, "annotation shard pointer")
return _published_shard_from_dict(value)
def publish_annotation_shards(
dataset: DatasetContext,
name: str,
publication_id: str,
published: Sequence[PublishedAnnotationShard],
*,
metadata: Mapping[str, Any] | None = None,
options: AnnotationOptions | None = None,
retry: int | RetryPolicy | None = None,
) -> AnnotationSet:
"""Validate dense shard coverage, then publish the set root and catalog."""
require_compare_and_swap(dataset.object_store)
policy = RetryPolicy.coerce(retry)
def attempt() -> AnnotationSet:
return _publish_annotation_shards_once(
dataset, name, publication_id, published, metadata=metadata, options=options
)
return retrying(policy, attempt, description=f"publishing annotation set {name!r}")
def _publish_annotation_shards_once(
dataset: DatasetContext,
name: str,
publication_id: str,
published: Sequence[PublishedAnnotationShard],
*,
metadata: Mapping[str, Any] | None,
options: AnnotationOptions | None,
) -> AnnotationSet:
if _SET_NAME.fullmatch(name) is None or _SET_NAME.fullmatch(publication_id) is None:
raise ValueError("annotation name and publication_id must be portable identifiers")
settings = AnnotationOptions() if options is None else options
winners: dict[str, PublishedAnnotationShard] = {}
for candidate in published:
item = _validated_published_shard(candidate)
previous = winners.get(item.logical_shard_id)
if previous is not None and previous != item:
raise ValueError("annotation logical shard has conflicting winners")
winners[item.logical_shard_id] = item
ordered = sorted(winners.values(), key=lambda item: item.first_ordinal)
ordinal = 0
shards: list[AnnotationShard] = []
for item in ordered:
if (
item.name != name
or item.publication_id != publication_id
or item.source_snapshot_id != dataset.snapshot_id
):
raise ValueError("annotation shard identity does not match publication")
if not item.record_ids:
raise ValueError("annotation shard must be non-empty")
if item.first_ordinal != ordinal:
raise ValueError("annotation shards do not provide contiguous dense coverage")
expected_ids = tuple(
dataset.record_metadata[index].id
for index in range(ordinal, ordinal + len(item.record_ids))
)
if item.record_ids != expected_ids:
raise ValueError("annotation shard record IDs do not match source ordinals")
try:
path = relative_uri(dataset.root, item.uri)
except ValueError as exc:
raise ValueError("annotation shard URI is outside its publication prefix") from exc
parts = path.split("/")
if (
len(parts) != 7
or parts[:5] != ["annotations", "sets", name, publication_id, "objects"]
or not re.fullmatch(r"[0-9a-f]{2}", parts[5])
or not re.fullmatch(r"[0-9a-f]{32}\.arrayrecord", parts[6])
):
raise ValueError("annotation shard URI is outside its publication prefix")
info = dataset.object_store.info(item.uri, generation=item.generation)
if info.size != item.stored_bytes:
raise CorruptDatasetError("annotation shard size does not match descriptor")
if settings.verify_record_counts:
count = dataset.shard_store.count(
item.uri,
generation=item.generation,
options=dataset.storage_options.reader_options,
file_reader_buffer_size=dataset.storage_options.file_reader_buffer_size,
)
if count != len(item.record_ids):
raise CorruptDatasetError("annotation shard count does not match descriptor")
shards.append(
AnnotationShard(
ordinal,
len(item.record_ids),
ObjectReference(path, None, item.stored_bytes),
)
)
ordinal += len(item.record_ids)
if ordinal != len(dataset):
raise ValueError("annotation shards do not cover the source dataset")
specification = AnnotationRoot(
name,
publication_id,
dataset.snapshot_id,
len(dataset),
tuple(shards),
metadata or {},
)
root_path = f"annotations/sets/{name}/{publication_id}/root.json"
root_payload = compact_json_bytes(_root_to_dict(specification))
root_uri = join_uri(dataset.root, root_path)
try:
dataset.object_store.create(root_uri, root_payload, content_type="application/json")
except PublicationCollisionError as exc:
existing = dataset.object_store.read_version(root_uri)
if existing.data != root_payload:
raise PublicationCollisionError(f"annotation set name already exists: {name}") from exc
_update_catalog(
dataset.root,
dataset.object_store,
name,
root_path,
len(root_payload),
settings.catalog_retry_limit,
dataset.storage_options,
)
return AnnotationSet(
dataset.root,
specification,
dataset.object_store,
dataset.shard_store,
dataset.storage_options,
dataset.record_metadata,
)
def _update_catalog(
root: str,
objects: ObjectStore,
name: str,
path: str,
stored_bytes: int,
retry_limit: int,
options: StorageOptions,
sleep: Callable[[float], None] = time.sleep,
) -> None:
uri = join_uri(root, CATALOG)
# Losing a compare-and-swap means another publisher won the race, so
# retrying in lockstep reproduces the collision. The schedule is the shared
# one, tuned short: catalog contention clears in milliseconds, unlike the
# storage failures RetryPolicy's defaults are built for.
# `attempts` is deliberately not passed; see S3ObjectStore.backoff. The
# loop below bounds itself by `retry_limit`, which is validated separately.
contention = RetryPolicy(base_seconds=0.01, max_seconds=0.5)
for attempt in range(retry_limit):
value: dict[str, Any]
try:
observed = objects.read_version(uri)
except FileNotFoundError:
value, expected = {"sets": {}}, None
else:
expected = observed.generation
try:
pointer = _parse(observed.data, options, "annotation catalog pointer")
revision_path, revision_generation = _catalog_pointer(pointer)
revision_data = objects.read(
join_uri(root, revision_path), generation=revision_generation
)
revision_id = revision_path.rsplit("/", 1)[-1].removesuffix(".json")
if hashlib.sha256(revision_data).hexdigest() != revision_id:
raise CorruptDatasetError("annotation catalog revision hash does not match")
value = _catalog_revision(
_parse(revision_data, options, "annotation catalog revision")
)
except FileNotFoundError as exc:
raise CorruptDatasetError("annotation catalog revision is missing") from exc
except (KeyError, TypeError, ValueError) as exc:
raise CorruptDatasetError(f"invalid annotation catalog: {exc}") from exc
existing = value["sets"].get(name)
reference = {
"path": path,
"generation": None,
"stored_bytes": stored_bytes,
}
if existing is not None and existing != reference:
raise PublicationCollisionError(f"annotation set name already exists: {name}")
value["sets"][name] = reference
revision_payload = compact_json_bytes(value)
revision_id = hashlib.sha256(revision_payload).hexdigest()
revision_path = f"annotations/catalogs/{revision_id}.json"
revision_uri = join_uri(root, revision_path)
try:
objects.create(revision_uri, revision_payload, content_type="application/json")
except PublicationCollisionError as exc:
existing_revision = objects.read_version(revision_uri)
if existing_revision.data != revision_payload:
raise CorruptDatasetError("catalog revision path contains different bytes") from exc
pointer = {
"revision": {
"path": revision_path,
"generation": None,
}
}
try:
objects.compare_and_swap(
uri,
compact_json_bytes(pointer),
expected_generation=expected,
content_type="application/json",
)
return
except PublicationCollisionError:
if attempt < retry_limit - 1:
sleep(contention.delay(attempt))
continue
raise PublicationCollisionError("annotation catalog update retry limit exceeded")
def _root_to_dict(value: AnnotationRoot) -> dict[str, Any]:
return {
"name": value.name,
"publication_id": value.publication_id,
"source_snapshot_id": value.source_snapshot_id,
"record_count": value.record_count,
"metadata": value.metadata,
"shards": [
{
"first_ordinal": item.first_ordinal,
"record_count": item.record_count,
"object": item.object.to_dict(),
}
for item in value.shards
],
}
write_shard = write_annotation_shard
publish_shards = publish_annotation_shards
__all__ = ["publish_annotations", "publish_shards", "write_shard"]