hyrax.verbs.infer_stream

Attributes

logger

Classes

InferStream	Streaming inference verb — loads model once, processes batches on demand.
InferStreamSession	Context manager for streaming inference.

Module Contents

logger

class InferStream(config)

Bases: hyrax.verbs.verb_registry.Verb

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

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'infer_stream'

add_parser_kwargs

description = 'Run streaming inference: load model once and process batches interactively.'

REQUIRED_DATA_GROUPS = ('infer_stream',)

OPTIONAL_DATA_GROUPS = ()

static setup_parser(parser)

No CLI arguments needed.

abstractmethod run_cli(args=None)

CLI stub — infer_stream is a programmatic API only.

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)
   

Parameters:

sample_batch (dict | None) – 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.

Returns:

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

Return type:

InferStreamSession

Raises:

ValueError – If sample_batch is None and no [data_request.infer_stream] is configured.

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 process().

Warning

process() is not thread-safe. Do not call it concurrently.

_process_func

_save_batch

_config

_results_dir

_close_logger

_load_dataset

data_loader = None

_provider = None

_closed = False

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

stop()

Signal the underlying streaming data source to stop iterating.

process(batch: dict) → numpy.typing.NDArray

Run inference on a single batch and save results.

Parameters:

batch (dict) – Must contain "object_id" (list of str) and model-specific data fields.

Returns:

Model output on CPU, detached from the computation graph.

Return type:

np.ndarray

Raises:

RuntimeError – If the session has already been closed.

close()

Commit results and return the result dataset.

Returns:

The accumulated results from all batches processed in this session.

Return type:

ResultDataset

__enter__()

__exit__(*_)