hyrax.datasets.streaming_data_provider

hyrax.datasets.streaming_data_provider#

A DataProvider-like wrapper that encapsulates a streaming dataset.

DataProvider is map-style — it assumes indexed random access (resolve_data(idx), joins, caching, augmentation, splits), none of which a live stream supports. StreamingDataProvider is its iterable sibling: it wraps a single streaming dataset (e.g. KafkaStreamDataset) and presents the same surface the rest of Hyrax expects — collate (shared via CollationMixin) and sample_data for model pre-flighting — while delegating iteration to the wrapped stream.

The wrapped stream is a dumb decoder that yields list[dict] batches of flat sample dicts. This provider owns the structuring: it extracts the object id (via primary_id_field) and groups model inputs (via fields) under the request friendly-name, producing the per-sample shape collate expects:

{"object_id": str, "<friendly_name>": {field: np.ndarray, ...}}

The single source of truth for primary_id_field and fields is [data_request].

Attributes#

logger

Classes#

StreamingDataProvider

Wrap a single streaming IterableDataset behind a DataProvider-like surface.

Module Contents#

logger[source]#

class StreamingDataProvider(config: dict, request: dict)[source]#

Bases: hyrax.datasets.data_provider.CollationMixin, torch.utils.data.IterableDataset

Wrap a single streaming IterableDataset behind a DataProvider-like surface.

Parameters:

config (dict) – The Hyrax runtime configuration.
request (dict) – A single-entry data request group (friendly_name -> definition). The definition must specify dataset_class (a registered IterableDataset) and primary_id_field; fields is optional (derived from the first sample when omitted).

config[source]#

data_request[source]#

friendly_name[source]#

primary_dataset[source]#

primary_dataset_id_field_name[source]#

primary_id_field[source]#

fields[source]#

_stream[source]#

prepped_datasets[source]#

custom_collate_functions: dict[source]#

field_collate_functions: dict[source]#

_register_field_collate_hooks()[source]#: Detect collate_<field> methods on the wrapped stream for each field.

_structure(sample: dict) → dict[source]#: Turn a flat decoded sample into the per-sample shape collate expects.

__iter__()[source]#: Yield list[dict] batches of structured samples for collate_fn.

sample_data() → dict[source]#

Return one structured sample for model pre-flighting (setup_model).

Peeks a single message from the stream without losing it (it is replayed in the first batch) and structures it like any other sample.

stop()[source]#: Stop the underlying stream’s iteration.