hyrax.datasets.kafka_stream_dataset

Streaming dataset that reads JSON messages from a Kafka topic.

KafkaStreamDataset is a torch.utils.data.IterableDataset intended for live, open-ended inference (e.g. a telescope alert stream). Unlike the map-style Hyrax datasets, it has no length: it polls a Kafka topic and yields batches as data arrives.

The defining feature is latency-bounded batching. A PyTorch DataLoader cannot emit a partial batch on a timeout, so the batching logic lives here in __iter__(): messages are accumulated and a batch is yielded as soon as either batch_size messages have arrived or batch_flush_timeout seconds have elapsed since the first message of the current batch. This means inference still proceeds on a short batch during quiet periods instead of blocking until the batch fills.

The stream is intentionally a dumb decoder: _decode() returns the parsed JSON object as a flat dict. Extracting the object id and grouping model-input fields is the responsibility of StreamingDataProvider, which wraps this dataset and knows the [data_request] (primary_id_field and fields). The provider is what is passed to the DataLoader; configure it the normal way:

import hyrax

hy = hyrax.Hyrax()
hy.config["data_request"] = {
    "infer_stream": {
        "data": {
            "dataset_class": "KafkaStreamDataset",
            "primary_id_field": "object_id",
            "fields": ["image"],
        }
    }
}
hy.config["data_set"]["KafkaStreamDataset"]["topic"] = "ztf-alerts"

with hy.infer_stream() as session:  # builds the provider + loader internally
    for batch, results in session:
        ...

Warning

The stream uses a single in-process Kafka consumer, so the loader must run with num_workers = 0 (the default applied by dist_data_loader). With multiple workers each would open its own consumer and the stop() signal would not propagate.

Attributes

logger

Classes

KafkaStreamDataset

Reads JSON messages from a Kafka topic and yields latency-bounded batches.

Module Contents

logger

class KafkaStreamDataset(config: dict, data_location=None)

Bases: hyrax.datasets.dataset_registry.HyraxDataset, torch.utils.data.IterableDataset

Reads JSON messages from a Kafka topic and yields latency-bounded batches.

The location of the Kafka broker(s) and topic(s) is configured in [data_set.KafkaStreamDataset]. Alternatively, the data_location of the dataset can be specified inline in the data request as the URI kafka://<host>:<port>[/<topic>]. The inline URI takes precedence over the configuration file.

Each Kafka message is expected to be a JSON object. _decode() returns it as a flat dict (e.g. {"object_id": "...", "image": [...], ...}); the wrapping StreamingDataProvider turns each flat sample into the structured form the collation + model machinery expect.

__init__()

Overall initialization for all Datasets which saves the config

Subclasses of HyraxDataset ought call this at the end of their __init__ like:

from hyrax.datasets import HyraxDataset

class MyDataset(HyraxDataset):
    def __init__(config):
        <your code>
        super().__init__(config)

If per tensor metadata is available, it is recommended that dataset authors create an astropy Table of that data, in the same order as their data and pass that metadata_table as shown below:

from hyrax.datasets import HyraxDataset
from astropy.table import Table

class MyDataset(HyraxDataset):
    def __init__(config):
        <your code>
        metadata_table = Table(<Your catalog data goes here>)
        super().__init__(config, metadata_table)

Parameters:

• config (dict, Optional) – The runtime configuration for hyrax

• metadata_table (Optional[Table], optional) – An Astropy Table with 1. the metadata columns desired for visualization AND 2. in the order your data will be enumerated.

• object_id_column_name (Optional[str], optional) – The name of the column containing object IDs. If None, uses the default from config or creates one from the ids() method.

bootstrap_servers

topics

group_id

auto_offset_reset

poll_timeout

batch_flush_timeout

batch_size

_stop

_consumer = None

_buffered: list[dict] = []

stop()

Signal __iter__() to flush any pending batch and stop iterating.

__len__()

A live stream has no length.

Defined only so HyraxDataset.__init_subclass__ (which requires a __len__ attribute) accepts the class. The iterable branch of dist_data_loader never calls it.

_make_consumer()

Create and subscribe a Kafka consumer.

Built lazily (not in __init__) because Kafka consumers are not safe to fork or pickle across DataLoader workers.

_ensure_consumer()

Return the shared consumer, creating it on first use.

_decode(msg) → dict

Decode a single Kafka message into a flat dict via JSON.

The stream is intentionally a dumb decoder: it returns the parsed JSON object as-is. Extracting the object id and grouping model-input fields is the job of StreamingDataProvider.

Parameters:

msg (object) – A Kafka message whose value() is JSON bytes/str.

Returns:

The parsed JSON object, e.g. {"object_id": "...", "image": [...], ...}.

Return type:

dict

peek_sample() → dict

Return one decoded sample without removing it from the batch stream.

Polls until a message arrives (or stop() is set), decodes it, and buffers it so __iter__() replays it as part of the first batch. Used to pre-flight the model architecture without losing a live message.

Returns:

The flat decoded sample.

Return type:

dict

Raises:

RuntimeError – If the stream is stopped before any message arrives.

__iter__()

Poll Kafka and yield list[dict] batches with latency-bounded flushing.