oumi.core.inference

Inference module for the Oumi (Open Universal Machine Intelligence) library.

This module provides base classes for model inference in the Oumi framework.

class oumi.core.inference.BaseInferenceEngine(model_params: ModelParams, *, generation_params: GenerationParams | None = None)

Bases: ABC

Base class for running model inference.

apply_chat_template(conversation: Conversation, **tokenizer_kwargs) → str

Applies the chat template to the conversation.

Parameters:

• conversation – The conversation to apply the chat template to.

• tokenizer_kwargs – Additional keyword arguments to pass to the tokenizer.

Returns:

The conversation with the chat template applied.

Return type:

str

get_batch_results_partial(batch_id: str, conversations: list[Conversation]) → BatchResult

Gets partial results of a completed batch job.

Engines that support batch inference should override this method.

Parameters:

• batch_id – The batch job ID.

• conversations – Original conversations used to create the batch.

Returns:

BatchResult with successful conversations and failure details.

Raises:

NotImplementedError – If the engine does not support batch.

abstractmethod get_supported_params() → set[str]

Returns a set of supported generation parameters for this engine.

Override this method in derived classes to specify which parameters are supported.

Returns:

A set of supported parameter names.

Return type:

Set[str]

infer(input: list[Conversation] | None = None, inference_config: InferenceConfig | None = None) → list[Conversation]

Runs model inference.

Parameters:

• input – A list of conversations to run inference on. Optional.

• inference_config – Parameters for inference. If not specified, a default config is inferred.

Returns:

Inference output.

Return type:

List[Conversation]

infer_partial(input: list[Conversation] | None = None, inference_config: InferenceConfig | None = None, *, progress_path: str | None = None) → InferenceResult

Runs model inference, tolerating per-row failures.

Unlike infer(), which raises if any conversation fails, this returns an InferenceResult pairing each successful conversation with its original input index and reporting per-row failures separately, so callers can keep successes and decide which failures to resubmit.

Parameters:

• input – A list of conversations to run inference on. Optional.

• inference_config – Parameters for inference.

• progress_path – Path to a JSON file where progress counters are periodically written for external pollers. Takes precedence over inference_config.progress_path.

Returns:

Successful conversations and per-row failures.

Return type:

InferenceResult

list_models(chat_only: bool = True) → list[str]

Returns a list of model IDs supported by this engine.

Override this method in derived classes to query the provider’s API for available models. The default implementation returns the model name this engine was initialized with.

Parameters:

chat_only – If True (default), only return models that support chat completions. If False, return all models.

Returns:

A list of supported model ID strings.

Return type:

list[str]

class oumi.core.inference.BatchResult(successful: list[tuple[int, Conversation]], failed_indices: list[int], error_messages: dict[int, str])

Bases: object

Result of a partial batch retrieval, separating successes from failures.

error_messages: dict[int, str]

Mapping of failed index to error message.

failed_indices: list[int]

Indices of requests that failed.

property has_failures: bool

Return True if any requests failed.

successful: list[tuple[int, Conversation]]

List of (original_index, conversation) for successful requests.

class oumi.core.inference.FailureDetail(error_message: str, status_code: int | None = None, is_retryable: bool = True, error_type: InferenceErrorType = InferenceErrorType.UNKNOWN)

Bases: object

Details about a single failed inference request.

error_message: str

Human-readable description of the failure.

error_type: InferenceErrorType = 'unknown'

Failure category for this request.

is_retryable: bool = True

Whether resubmitting this request could plausibly succeed.

status_code: int | None = None

HTTP status code of the final failed attempt, if applicable.

class oumi.core.inference.InferenceErrorType(value)

Bases: str, Enum

Failure category for an inference request.

API_STATUS = 'api_status'

The API returned a non-success HTTP status code.

CONFIG = 'config'

The request failed due to invalid configuration.

CONNECTION = 'connection'

A network-level error prevented reaching the API.

ENGINE_FAILURE = 'engine_failure'

The inference engine itself failed.

PARSE_ERROR = 'parse_error'

The API response could not be parsed into a conversation.

RUNTIME = 'runtime'

An error was raised while running inference.

UNKNOWN = 'unknown'

the failure fits no other category.

Type:

Default and catch-all

class oumi.core.inference.InferenceResult(successful: list[tuple[int, Conversation]], failures: dict[int, FailureDetail])

Bases: object

Result of partial online inference, separating successes from failures.

property error_messages: dict[int, str]

Mapping of failed index to error message.

property failed_indices: list[int]

Sorted indices of requests that failed.

failures: dict[int, FailureDetail]

Mapping of failed index to structured failure info.

property has_failures: bool

Return True if any requests failed.

successful: list[tuple[int, Conversation]]

List of (original_index, conversation) for successful requests.

class oumi.core.inference.ProgressFileReporter(path: str, total: int, min_write_interval: float = 1.0)

Bases: object

Writes inference progress counters to a JSON file for external pollers.

The snapshot format is:

{"total": N, "completed": n, "failed": f, "updated_at": "<iso8601 utc>"}

The run is complete when completed + failed == total. Writes are atomic (temp file + os.replace), so a polling process never observes partial JSON, and are throttled to at most one per min_write_interval seconds (start() and finalize() always write).

Filesystem failures are logged and swallowed: a broken progress path must never kill inference. Counter updates are thread-safe.

finalize() → None

Writes a final snapshot, bypassing the throttle.

record_completed(n: int = 1) → None

Records n successfully completed rows.

record_failed(n: int = 1) → None

Records n failed rows.

start(completed: int = 0, failed: int = 0) → None

Initializes counters and writes the first snapshot.