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:
        ...

Offsets are committed after a batch has been consumed downstream, never before; see KafkaStreamDataset.__iter__(). The consumer must therefore be closed deterministically when iteration ends – including when it ends because something raised – which is what KafkaStreamDataset.close() is for. InferStreamSession calls it on every exit path, so the with hy.infer_stream() as session form above is the supported way to drive the stream.

Recovering / replaying a stream

An unclosed consumer stays a member of its consumer group and keeps its partition assignment, so the next run joins the same group, loses the rebalance, is assigned no partitions and blocks forever. If you see “no partitions are assigned” in the logs, find the lingering member with:

kafka-consumer-groups --bootstrap-server localhost:9092 \
    --group hyrax-infer-stream --describe

Separately, auto_offset_reset only applies to a group with no committed offset. Once a group has consumed a topic, re-running resumes after the last commit rather than replaying. To read a topic from the beginning again, either set a fresh group_id or reset the existing group’s offsets while no consumer is running:

kafka-consumer-groups --bootstrap-server localhost:9092 \
    --group hyrax-infer-stream --reset-offsets --to-earliest --all-topics --execute

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.

Functions

_format_partitions(→ list[str])

Render a list of TopicPartition as ["topic[0]", "topic[1]"] for logging.

Module Contents

logger

_format_partitions(partitions) → list[str]

Render a list of TopicPartition as ["topic[0]", "topic[1]"] for logging.

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.

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

commit_after_batch

batch_size

_stop

_consumer = None

consumer_config

stop()

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

This is only a signal. It does not release the Kafka connection – use close() for that.

close()

Stop iteration and close the Kafka consumer. Idempotent.

Closing is what releases the partition assignment and removes this consumer from its group. Skipping it leaves a zombie group member: its librdkafka thread keeps running, it keeps the partitions, and the next run joining the same group_id can be assigned nothing at all and block forever. A suspended __iter__() generator cannot clean itself up (nothing resumes it), and in a notebook the retained traceback keeps it from ever being collected, so the owner of the stream has to call this explicitly on error paths.

_consumer is None is the guard, so a late second call – the generator’s finally running at collection time, long after the session closed – is a no-op rather than a double close.

Deliberately does not commit. Closing is reached on failure paths just as often as on clean ones, and librdkafka has already stored the offsets of every message it handed out – including a batch the caller never finished. Committing here would advance past exactly the alerts that still need reprocessing. Since enable.auto.commit is off, Consumer.close() commits nothing on its own; it only flushes commits __iter__() already issued for completed batches.

__enter__()

__exit__(*_)

__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.

_on_assign(consumer, partitions)

Log the partitions handed to this consumer on each rebalance.

_on_revoke(consumer, partitions)

Log the partitions taken from this consumer on each rebalance.

_connection_diagnostic(prefix: str) → str

Describe what this consumer is connected to, and why it may see no messages.

static _log_message_error(msg)

Log a message-level Kafka error instead of dropping it silently.

_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

__iter__()

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

Offsets are committed only after the caller comes back for the next batch. Because this is a generator, that resumption is the signal that the previous batch was fully processed downstream – inferred and written – so it needs no cooperation from the caller. If the caller raises instead, the generator is never resumed, the commit never happens, and the batch is redelivered on the next run. That makes delivery at-least-once; duplicates after a crash are possible, dropped alerts are not.

static _commit(consumer)

Commit the offsets of the batch the caller has just finished processing.