Add a custom store

Warning

Pre-implementation. APIs describe the target plugin surface.

Goal

Add a Python store backend that persists Record objects in a new on-disk or cloud format alongside ZarrStore and WebDatasetStore.

A store backend is the runtime component that opens a dataset location, writes records, reads them back by stable sample_id, iterates available records, and preserves append behavior. Generation creates the records; the store owns their serialization and lookup contract.

When to use this

Use this when no existing store backend matches your access pattern, on-disk format, or cloud target.

Target format examples:

  • TFRecord mirror for a TensorFlow trainer.

  • Parquet plus in-phase/quadrature (IQ) sample blobs for an analytics warehouse.

  • SQLite-backed store for local debugging.

Mature library examples to wrap: zarr-python for Zarr, webdataset for tar-sharded training data, h5py for HDF5, and SigMF readers for Signal Metadata Format (SigMF) capture interop. See Reference / Storage Layout for the role each format plays in the canonical dataset layout.

If you only need to change the codec, chunk size, or bucket of an existing backend, configure ZarrStore instead; custom-store scope is a new on-disk format, not a tuned variant.

Prerequisites

Read Concepts / Storage, particularly the Store and Handle Interfaces. The store layer uses two ABCs because the configured factory and the open I/O session have different lifetimes:

A custom backend subclasses both. The IQ tensor crosses the boundary as (2, N) float32 in memory, where axis 0 is real then imaginary components and N is the number of time samples. The handle converts that tensor to storage [N] complex64 so storage libraries can use their native complex-array support; see Reference / IQ Layout Policy for the exact conversion contract.

Minimal implementation

from collections.abc import Iterable, Iterator, Mapping
from pathlib import Path

from rfgen.config import SceneAssetsConfig
from rfgen.core_types import GeometryAssetRef, Record
from rfgen.enums import StoreMode
from rfgen.storage import BaseStore, StoreHandle, canonical_sample_id


class TFRecordStoreHandle(StoreHandle):
    def __init__(self, uri: str, mode: StoreMode | str) -> None:
        self._uri = uri
        self._mode = mode
        # Open the TensorFlow reader/writer and load the sample/shard indexes.

    def write(self, record: Record, *, sample_id: str | None = None) -> str:
        sid = sample_id or canonical_sample_id(record)
        # If sid exists, verify identical canonical content and return sid.
        # Otherwise serialize the record with TensorFlow's TFRecord APIs.
        return sid

    def write_shard(
        self,
        shard_id: str,
        records: Iterable[Record],
    ) -> tuple[str, ...]:
        if self._completion_marker_exists(shard_id):
            return self._committed_sample_ids(shard_id)
        sample_ids = []
        for record in records:          # consume lazily; iteration can raise
            sample_ids.append(self.write(record))
        self._publish_completion_marker(shard_id, tuple(sample_ids))
        return tuple(sample_ids)

    def read(self, sample_id: str) -> Record:
        ...

    def append_annotation_row(
        self,
        sample_id: str,
        *,
        annotation_type: str,
        template_id: str,
        run_id: str,
        row: Mapping[str, object],
    ) -> Record:
        ...

    def iter_sample_ids(self) -> Iterator[str]:
        ...

    def read_siblings(self, record: Record) -> Iterable[Record]:
        ...

    def write_asset_bundle(
        self,
        ref: GeometryAssetRef,
        files: Mapping[str, bytes],
    ) -> GeometryAssetRef:
        ...

    def resolve_asset(
        self,
        ref: GeometryAssetRef,
        *,
        worker_cache_dir: Path | None = None,
    ) -> Path | str:
        ...

    def prepare_worker_assets(
        self,
        scene_assets: SceneAssetsConfig,
        cache_root: Path,
    ) -> SceneAssetsConfig:
        ...

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

    def __len__(self) -> int:
        ...

    def close(self) -> None:
        ...


class TFRecordStore(BaseStore):
    name = "tfrecord"

    def __init__(self, uri: str) -> None:
        self._uri = uri

    def open(self, uri: str, mode: StoreMode | str) -> TFRecordStoreHandle:
        return TFRecordStoreHandle(uri, mode)

    def has_shard(self, shard_id: str) -> bool:
        # Probe only the durable completion marker under self._uri.
        ...

    def dataset_uri(self) -> str:
        return self._uri

The skeleton includes every abstract method on the shipped BaseStore and StoreHandle surfaces. Replace each ellipsis and format-specific helper with calls into the mature storage library selected for the backend. sample_ids() is inherited from StoreHandle and materializes iter_sample_ids().

write() returns the caller-supplied sample_id, or calls canonical_sample_id(record) for the same canonical CBOR and SHA-256 identity used by shipped stores. The handle is a context manager; close runs on __exit__ and must flush pending writes and release library resources.

Shard write lifecycle

write_shard(…) receives the lazy iterator created by the Layer 9 worker. Keep two kinds of stored state separate: record data and a small completion marker. The marker is the authoritative flag that every requested record in the shard finished. In this guide, commit means publishing that marker; it does not mean rolling back record writes if a later operation fails.

A failed attempt may leave records but must not leave the marker. A retry regenerates the same content-derived sample_id for each identical record. When an identical record already exists, write(...) returns that ID without creating a duplicate; this safe no-op is the idempotent replay behavior used to repair partial shards.

Implement it in this order:

  1. If the durable completion marker already exists, return its stored sample IDs without consuming records.

  2. Iterate records once and call the idempotent write(...) for each value. Do not convert the iterable to a list before this loop.

  3. If iterator consumption or write(...) raises, let the exception propagate. Do not publish the completion marker. Already-written content-addressed records may remain and must be safe to replay.

  4. Only after the iterator is exhausted normally, durably store the ordered sample IDs and publish the backend’s completion marker.

  5. Return the ordered sample-ID tuple.

The Layer 9 worker intentionally raises its private abort from inside records when any sample-local generation or labeling failure occurred. The ordering above lets that exception leave write_shard(…) before step 4, so BaseStore.has_shard(shard_id) remains False and a retry repairs the shard.

Register

Declare the entry-point group rfgen.stores in the third-party package metadata. The shipped registry discovers the class by its entry-point name; there is no register_store decorator. See Reference / Plugin Metadata § Entry-point group names for the canonical group list.

[project.entry-points."rfgen.stores"]
tfrecord = "my_pkg.stores:TFRecordStore"

The third-party package also declares a top-level PLUGIN: PluginMetadata attribute with plugin_kind="store"; see Reference / Plugin Metadata § Plugin registration for the full schema.

Configure

storage:
  _target_: my_pkg.stores.TFRecordStore
  path: file://./out/demo/samples.tfrecord

_target_ is Hydra object-instantiation syntax: it tells the config layer which Python class to construct for the storage block. Use a local file:// path for the first smoke test. A cloud path such as gs://rf-fm-datasets-synth/demo/v1/samples.tfrecord is the same store contract with cloud credentials resolved outside the store.

For URLs, use the same scheme convention the canonical ZarrStore accepts (file://, gs://, s3://, az://); the credentials chain is resolved by the framework’s credentials providers, not by the store.

Verify

rfgen list-stores | grep tfrecord
rfgen validate storage=tfrecord_smoke
rfgen generate +preset=narrowband_classifier_baseline_xs storage=tfrecord_smoke

storage=tfrecord_smoke swaps the whole Hydra storage config group. A minimal smoke config would live at configs/storage/tfrecord_smoke.yaml:

_target_: my_pkg.stores.TFRecordStore
path: file://./out/tfrecord-smoke/samples.tfrecord

Hydra uses +preset=<name> to add a named scenario preset, and bare assignments such as storage=tfrecord_smoke swap config groups. See Config Schema / CLI override patterns for the broader syntax, and run rfgen show-preset narrowband_classifier_baseline_xs before generating if you need to inspect the preset.

Append mode is only for Phase 2 annotation or text metadata attached to records that already exist from Phase 1 generation. It must not rewrite generated IQ or structured labels such as bounding boxes, segmentation masks, or per-emitter metadata.

Confirm:

  • Round-trip: store.open(path, "w"), write(record), reopen with mode="r", read back, and assert record.iq, record.labels, and record.metadata match field-for-field.

  • Idempotency: writing the same Record twice with the same sample_id produces a single on-disk entry, not two.

  • IQ layout: the in-memory (2, N) float32 tensor round-trips through the store as storage-side complex64 without dtype drift.

  • Append (if your backend supports it): mode="a" opens an existing dataset and appends Phase 2 text-run subgroups without rewriting the IQ or structured-label arrays. If your backend cannot support append, raise StorageError for mode="a" rather than silently truncating.

  • Add contract tests for shape, dtype, sample_id stability, deterministic iteration order, and concurrent-writer safety. See Reference / Contract Tests for the store-test fixtures.

Troubleshoot

Symptom

Fix

rfgen list-stores does not show the backend

Check that the rfgen.stores entry-point name matches the class-level name.

Round-trip changes IQ values

Confirm the (2, N) float32 to complex64 conversion uses the policy in IQ Layout Policy; do not coerce dtypes ad hoc.

Re-runs are not idempotent

Use the content-hash sample_id from the Record; do not generate fresh UUIDs per write.

Cloud writes fail with credentials errors

The store does not parse credentials; resolve them through the framework’s credentials providers before opening.

Next steps

See Also