hyrax.verbs.infer_stream
========================

.. py:module:: hyrax.verbs.infer_stream


Attributes
----------

.. autoapisummary::

   hyrax.verbs.infer_stream.logger


Classes
-------

.. autoapisummary::

   hyrax.verbs.infer_stream.InferStream
   hyrax.verbs.infer_stream.InferStreamSession


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

.. py:data:: logger

.. py:class:: InferStream(config)

   Bases: :py:obj:`hyrax.verbs.verb_registry.Verb`


   Streaming inference verb — loads model once, processes batches on demand.

   .. py:method:: __init__

   Overall initialization for all verbs that saves the config


   .. py:attribute:: cli_name
      :value: 'infer_stream'



   .. py:attribute:: add_parser_kwargs


   .. py:attribute:: description
      :value: 'Run streaming inference: load model once and process batches interactively.'



   .. py:attribute:: REQUIRED_DATA_GROUPS
      :value: ('infer_stream',)



   .. py:attribute:: OPTIONAL_DATA_GROUPS
      :value: ()



   .. py:method:: setup_parser(parser)
      :staticmethod:


      No CLI arguments needed.



   .. py:method:: run_cli(args=None)
      :abstractmethod:


      CLI stub — infer_stream is a programmatic API only.



   .. py:method:: run(sample_batch: dict | None = None) -> InferStreamSession

      Set up the model and return a session for streaming inference.

      There are two ways to drive the session:

      1. **Data-source driven** (``sample_batch=None``) — configure a streaming
         dataset under ``[data_request.infer_stream]`` (e.g. ``KafkaStreamDataset``).
         The model is pre-flighted from the stream itself and a DataLoader is built,
         so the returned session can be iterated directly::

             with hy.infer_stream() as session:
                 for batch, results in session:
                     ...

      2. **Manual** — pass a representative ``sample_batch`` and feed batches yourself::

             with hy.infer_stream(sample_batch=batch) as session:
                 results = session.process(batch)

      :param sample_batch: A representative batch dict with ``"object_id"`` and model-specific data
                           fields, used to pre-flight the model architecture. When ``None``, the model
                           is pre-flighted from a ``[data_request.infer_stream]`` streaming dataset
                           instead.
      :type sample_batch: dict | None

      :returns: A context manager / session object. Iterate it (data-source driven) or call
                ``session.process(batch)`` (manual); call ``session.close()`` when done.
      :rtype: InferStreamSession

      :raises ValueError: If ``sample_batch`` is None and no ``[data_request.infer_stream]`` is configured.



.. py:class:: InferStreamSession(process_func, save_batch_callback, config, results_dir, close_logger_fn, load_dataset_fn, data_loader=None, provider=None)

   Context manager for streaming inference.

   Holds a loaded model and Lance writer. When constructed with a ``data_loader``
   (the data-source-driven path), the session is iterable and yields
   ``(batch, results)`` pairs as data arrives; otherwise feed batches yourself with
   :meth:`process`.

   .. warning::
       ``process()`` is **not** thread-safe. Do not call it concurrently.


   .. py:attribute:: _process_func


   .. py:attribute:: _save_batch


   .. py:attribute:: _config


   .. py:attribute:: _results_dir


   .. py:attribute:: _close_logger


   .. py:attribute:: _load_dataset


   .. py:attribute:: data_loader
      :value: None



   .. py:attribute:: _provider
      :value: None



   .. py:attribute:: _closed
      :value: False



   .. py:method:: __iter__()

      Iterate the configured data source, processing each batch as it arrives.

      :Yields: *tuple[dict, np.ndarray]* -- The collated input ``batch`` and the model ``results`` for it.

      :raises RuntimeError: If the session was created without a data source (no
          ``[data_request.infer_stream]`` configuration).



   .. py:method:: stop()

      Signal the underlying streaming data source to stop iterating.



   .. py:method:: process(batch: dict) -> numpy.typing.NDArray

      Run inference on a single batch and save results.

      :param batch: Must contain ``"object_id"`` (list of str) and model-specific data fields.
      :type batch: dict

      :returns: Model output on CPU, detached from the computation graph.
      :rtype: np.ndarray

      :raises RuntimeError: If the session has already been closed.



   .. py:method:: close()

      Commit results and return the result dataset.

      :returns: The accumulated results from all batches processed in this session.
      :rtype: ResultDataset



   .. py:method:: __enter__()


   .. py:method:: __exit__(*_)