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() → InferStreamSession

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

There are two ways to drive the session:

1. Data-source driven — configure a streaming dataset under [data_request.infer_stream] (e.g. KafkaStreamDataset). A DataLoader is built over it, so the returned session can be iterated directly:

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

2. Manual — leave [data_request.infer_stream] unset and feed batches yourself:

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

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

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

_context

__iter__()

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

Yields:

tuple[dict, np.ndarray | dict[str, np.ndarray]] – The collated input batch and the model results for it, in the form returned by process().

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.

_close_provider()

Tear down the streaming data source, releasing any connection it holds.

Kept separate from stop() because stopping is only a signal: a data loader left suspended mid-iteration (the caller hit an exception, or broke out of the loop) never resumes to clean up after itself, so the connection has to be closed from the outside or it lingers for the life of the process.

process(batch: dict) → numpy.typing.NDArray | dict[str, 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. When the model’s infer_batch returns a dict of tensors, this is a dict with the same keys whose values are numpy arrays; otherwise it is a single numpy array.

Return type:

np.ndarray | dict[str, np.ndarray]

Raises:

RuntimeError – If the session has already been closed, or if the model returns something other than a tensor or a dict of tensors.

close()

Commit results and return the result dataset.

Returns:

The accumulated results from all batches processed in this session.

Return type:

ResultDataset

__enter__()

__exit__(*_)