hyrax.datasets.kafka_stream_dataset
===================================

.. py:module:: hyrax.datasets.kafka_stream_dataset

.. autoapi-nested-parse::

   Streaming dataset that reads JSON messages from a Kafka topic.

   :class:`KafkaStreamDataset` is a :class:`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
   :meth:`__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**: :meth:`_decode` returns the parsed JSON
   object as a flat ``dict``. Extracting the object id and grouping model-input fields is
   the responsibility of
   :class:`~hyrax.datasets.streaming_data_provider.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:

   .. code-block:: python

       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 :meth:`stop` signal would not
       propagate.



Attributes
----------

.. autoapisummary::

   hyrax.datasets.kafka_stream_dataset.logger


Classes
-------

.. autoapisummary::

   hyrax.datasets.kafka_stream_dataset.KafkaStreamDataset


Module Contents
---------------

.. py:data:: logger

.. py:class:: KafkaStreamDataset(config: dict, data_location=None)

   Bases: :py:obj:`hyrax.datasets.dataset_registry.HyraxDataset`, :py:obj:`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. :meth:`_decode` returns it as a
   flat ``dict`` (e.g. ``{"object_id": "...", "image": [...], ...}``); the wrapping
   :class:`~hyrax.datasets.streaming_data_provider.StreamingDataProvider` turns each
   flat sample into the structured form the collation + model machinery expect.

   .. py:method:: __init__

   Overall initialization for all Datasets which saves the config

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

   .. code-block:: python

       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:

   .. code-block:: python

       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)

   :param config: The runtime configuration for hyrax
   :type config: dict, Optional
   :param metadata_table: An Astropy Table with
                          1. the metadata columns desired for visualization AND
                          2. in the order your data will be enumerated.
   :type metadata_table: Optional[Table], optional
   :param object_id_column_name: The name of the column containing object IDs. If None, uses the default
                                 from config or creates one from the ids() method.
   :type object_id_column_name: Optional[str], optional


   .. py:attribute:: bootstrap_servers


   .. py:attribute:: topics


   .. py:attribute:: group_id


   .. py:attribute:: auto_offset_reset


   .. py:attribute:: poll_timeout


   .. py:attribute:: batch_flush_timeout


   .. py:attribute:: batch_size


   .. py:attribute:: _stop


   .. py:attribute:: _consumer
      :value: None



   .. py:attribute:: _buffered
      :type:  list[dict]
      :value: []



   .. py:method:: stop()

      Signal :meth:`__iter__` to flush any pending batch and stop iterating.



   .. py:method:: __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.



   .. py:method:: _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.



   .. py:method:: _ensure_consumer()

      Return the shared consumer, creating it on first use.



   .. py:method:: _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
      :class:`~hyrax.datasets.streaming_data_provider.StreamingDataProvider`.

      :param msg: A Kafka message whose ``value()`` is JSON bytes/str.
      :type msg: object

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



   .. py:method:: peek_sample() -> dict

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

      Polls until a message arrives (or :meth:`stop` is set), decodes it, and buffers
      it so :meth:`__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.
      :rtype: dict

      :raises RuntimeError: If the stream is stopped before any message arrives.



   .. py:method:: __iter__()

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