mmirage.core.process — Processors

Variables

Variable system for MMIRAGE pipeline with multimodal support.

class mmirage.core.process.variables.BaseVar(name='')[source]

Bases: ABC

Base class for variables in the MMIRAGE pipeline.

Parameters:

name (str)

name

Name of the variable.

Type:

str

name: str = ''
class mmirage.core.process.variables.InputVar(name='', key='', type='text')[source]

Bases: BaseVar

Input variable extracted from source datasets.

Parameters:
key

JMESPath query to extract the variable from a sample.

Type:

str

type

Variable type - “text” or “image”.

Type:

Literal[‘text’, ‘image’]

key: str = ''
type: Literal['text', 'image'] = 'text'
is_image()[source]

Check if this input variable represents an image.

Returns:

True if the variable type is “image”, False otherwise.

Return type:

bool

class mmirage.core.process.variables.OutputVar(name='', type='')[source]

Bases: BaseVar

Output variable generated by processors.

Output variables are created by processors (e.g., LLMs) and can depend on input variables and previously computed output variables.

Parameters:
name

Name of the variable.

Type:

str

type

Type identifier for the processor that generates this variable.

Type:

str

type: str = ''
abstractmethod is_computable(vars)[source]
Parameters:

vars (Sequence[BaseVar])

Return type:

bool

class mmirage.core.process.variables.VariableEnvironment(var_env, image_vars=None)[source]

Bases: object

Environment for storing and accessing variables during processing.

Parameters:
  • var_env (Dict[str, Any])

  • image_vars (Optional[set])

__init__(var_env, image_vars=None)[source]

Initialize a variable environment.

Parameters:
  • var_env (Dict[str, Any]) – Dictionary mapping variable names to their values.

  • image_vars (set | None) – Set of variable names that represent images. Defaults to empty set.

Return type:

None

with_variable(key, value, is_image=False)[source]

Create a new environment with an additional variable.

Parameters:
  • key (str) – Name of the variable to add.

  • value (Any) – Value of the variable.

  • is_image (bool) – Whether the variable represents an image.

Returns:

New environment with the added variable.

Return type:

VariableEnvironment

to_dict()[source]

Get an immutable view of the variable dictionary.

Returns:

MappingProxyType providing read-only access to variables.

Return type:

MappingProxyType

get(key, default=None)[source]

Get a variable value by name.

Parameters:
  • key (str) – Name of the variable to retrieve.

  • default (Any) – Default value to return if variable is not found.

Returns:

Variable value, or default if not found.

Return type:

Any

is_image_var(key)[source]

Check if a variable represents an image.

Parameters:

key (str) – Name of the variable to check.

Returns:

True if the variable is an image variable, False otherwise.

Return type:

bool

get_image_vars()[source]

Get all image variable names.

Returns:

Copy of the set containing names of all image variables.

Return type:

set

get_images()[source]

Get image values in a deterministic order.

Return type:

List[Any]

has_images()[source]

Check if the environment contains any image variables.

Returns:

True if at least one image variable is present, False otherwise.

Return type:

bool

static from_input_variables(sample, input_vars, image_base_path=None)[source]

Create a variable environment from a single sample.

Parameters:
  • sample (Dict[str, Any]) – Dictionary containing the data for one sample.

  • input_vars (List[InputVar]) – List of input variable definitions to extract.

  • image_base_path (str | None) – Optional base directory for resolving relative image paths.

Returns:

Environment populated with extracted variables.

Return type:

VariableEnvironment

Raises:

ValueError – If a required input variable is not found in the sample.

static from_batch_input_variables(batch, input_vars, image_base_path=None)[source]

Extract input variables from a batch of samples.

Parameters:
  • batch (Dict[str, List[Any]]) – Dictionary mapping column names to lists of values.

  • input_vars (List[InputVar]) – List of input variable definitions.

  • image_base_path (str | None) – Optional base directory for resolving relative image paths.

Returns:

List of VariableEnvironments, one for each sample in the batch.

Return type:

List[VariableEnvironment]

Base processor

Base classes and registry for processors in MMIRAGE.

class mmirage.core.process.base.BaseProcessorConfig(type)[source]

Bases: object

Base configuration class for processors.

All processor configurations must inherit from this class.

Parameters:

type (str)

type

Registry key of the processor. Concrete configs narrow this to their own literal (e.g. Literal["llm"] = "llm").

Type:

str

type: str
class mmirage.core.process.base.TokenCounts(input_tokens, output_tokens)[source]

Bases: object

Cumulative token counts from LLM processors.

Parameters:
  • input_tokens (int)

  • output_tokens (int)

input_tokens: int
output_tokens: int
class mmirage.core.process.base.BaseProcessor(config, shard_id=0, **kwargs)[source]

Bases: ABC, Generic[C]

Abstract base class for data processors.

Processors are responsible for transforming data by generating new output variables from existing variables.

Type Parameters:

C: The output variable type this processor works with.

Parameters:
config

Configuration object for this processor.

__init__(config, shard_id=0, **kwargs)[source]

Initialize the processor with configuration.

Parameters:
  • config (BaseProcessorConfig) – Configuration object for this processor.

  • shard_id (int) – Optional shard identifier accepted for compatibility with callers that forward it during processor construction.

  • **kwargs – Additional keyword arguments. Any unexpected keyword arguments will raise TypeError.

Raises:

TypeError – If unexpected keyword arguments are provided.

Return type:

None

abstractmethod batch_process_sample(batch, output_var)[source]

Process a batch of variable environments.

Parameters:
  • batch (List[VariableEnvironment]) – List of variable environments to process.

  • output_var (C) – Output variable definition to generate.

Returns:

List of updated variable environments with the new output variable.

Raises:

NotImplementedError – If not implemented by subclass.

Return type:

List[VariableEnvironment]

finalize()[source]

Optional lifecycle hook; override when a processor buffers state.

Return type:

None

abstractmethod get_token_counts()[source]

Get cumulative token counts from this processor.

Returns:

TokenCounts object containing input and output token counts.

Raises:

NotImplementedError – If not implemented by subclass.

Return type:

TokenCounts

abstractmethod get_load_time()[source]

Get the time taken to load any necessary resources (e.g., models).

Returns:

Time in seconds taken to load resources.

Raises:

NotImplementedError – If not implemented by subclass.

Return type:

float

shutdown()[source]

Release any resources held by this processor.

Override in subclasses that hold GPU memory, open file handles, or network connections. The default implementation is a no-op.

Return type:

None

class mmirage.core.process.base.ProcessorRegistry[source]

Bases: object

Registry for managing and accessing available processors.

Provides a centralized registry for processor classes, their configuration classes, and their output variable classes.

_registry

Mapping from processor name to registered processor class.

_config_registry

Mapping from processor name to its configuration class.

_output_var_registry

Mapping from processor name to its output variable class.

classmethod register_types(name, config_cls, output_var_cls)[source]

Register config/output-var types without importing processor implementations.

Parameters:
Return type:

None

classmethod register(name, config_cls, output_var_cls)[source]

Register a processor class with its associated classes.

Parameters:
  • name (str) – String identifier for the processor.

  • config_cls (Type[BaseProcessorConfig]) – Configuration class associated with this processor.

  • output_var_cls (Type[OutputVar]) – Output variable class associated with this processor.

Returns:

Decorator function to register the processor class.

Return type:

Callable

classmethod get_processor(name)[source]

Get a registered processor class by name.

Parameters:

name (str) – String identifier of the processor.

Returns:

The registered processor class.

Raises:

ValueError – If no processor is registered under the given name.

Return type:

Type[BaseProcessor]

classmethod get_config_cls(name)[source]

Get a registered configuration class by processor name.

Parameters:

name (str) – String identifier of the processor.

Returns:

The registered configuration class.

Raises:

ValueError – If no processor is registered under the given name.

Return type:

Type[BaseProcessorConfig]

classmethod get_output_var_cls(name)[source]

Get a registered output variable class by processor name.

Parameters:

name (str) – String identifier of the processor.

Returns:

The registered output variable class.

Raises:

ValueError – If no processor is registered under the given name.

Return type:

Type[OutputVar]

class mmirage.core.process.base.AutoProcessor[source]

Bases: object

Factory class for instantiating processors by name.

classmethod from_name(name)[source]

Retrieve a processor class by its registered name.

Parameters:

name (str) – The registry name of the processor.

Returns:

The registered processor class.

Raises:

ValueError – If no processor is registered under the given name.

Return type:

Type[BaseProcessor]

Mapper

Mapper for orchestrating variable transformations.

class mmirage.core.process.mapper.MMIRAGEMapper(processor_configs, input_vars, output_vars, export_prompts_dir=None, shard_id=0)[source]

Bases: object

Mapper for orchestrating variable transformations in the MMIRAGE pipeline.

Manages processors, validates variable dependencies, and applies transformations to batches of data. Supports multimodal inputs.

Parameters:
processors

Dictionary mapping processor types to processor instances.

output_vars

List of output variables to generate.

input_vars

List of input variables to extract.

__init__(processor_configs, input_vars, output_vars, export_prompts_dir=None, shard_id=0)[source]

Initialize the MMIRAGE mapper.

Parameters:
  • processor_configs (List[BaseProcessorConfig]) – List of processor configurations.

  • input_vars (List[InputVar]) – List of input variable definitions.

  • output_vars (List[OutputVar]) – List of output variable definitions.

  • export_prompts_dir (str | None) – Value of –export-prompts.

  • shard_id (int) – Shard index for this worker, forwarded to processors.

Return type:

None

validate_vars()[source]

Validate that all output variables are computable.

Checks that each output variable can be computed given the available variables (inputs and previously computed outputs).

Returns:

True if all variables are computable, False otherwise.

Return type:

bool

rewrite_batch(batch, image_base_path=None)[source]

Transform a batch of samples by computing output variables.

Parameters:
  • batch (Dict[str, List[Any]]) – Dictionary mapping column names to lists of values.

  • image_base_path (str | None) – Optional base directory for resolving relative image paths.

Returns:

List of VariableEnvironments with all output variables computed.

Raises:

RuntimeError – If an output variable type has no registered processor.

Return type:

List[VariableEnvironment]

get_token_counts()[source]

Return cumulative token counts aggregated across all LLM processors.

Sums input_tokens and output_tokens from every processor that exposes a get_token_counts() method (i.e., LLMProcessor).

Returns:

TokenCounts with input_tokens and output_tokens fields.

Return type:

TokenCounts

get_load_time()[source]

Return total model-loading time (seconds) summed across all LLM processors.

Return type:

float

finalize_processors()[source]

Finalize processors that expose a finalize lifecycle hook.

Return type:

None

shutdown()[source]

Shut down all processors and release their resources.

Return type:

None

LLM processor

Configuration

Configuration for LLM processor in MMIRAGE.

class mmirage.core.process.processors.llm.config.SGLangServerArgs(model_path='none', tp_size=<factory>, trust_remote_code=True, disable_custom_all_reduce=False, extra_engine_args=<factory>)[source]

Bases: object

Server arguments for SGLang engine.

Parameters:
  • model_path (str)

  • tp_size (int)

  • trust_remote_code (bool)

  • disable_custom_all_reduce (bool)

  • extra_engine_args (Dict[str, Any])

model_path

Path to the model or HuggingFace model ID.

Type:

str

tp_size

Tensor parallelism size.

Type:

int

trust_remote_code

Whether to trust remote code from HuggingFace.

Type:

bool

disable_custom_all_reduce

Whether to disable custom all reduce.

Type:

bool

extra_engine_args

Any additional keyword arguments forwarded verbatim to sgl.Engine. Use this to pass SGLang-specific options that are not listed above, e.g.:

extra_engine_args:
  max_running_requests: 512
  chunked_prefill_size: 32768
  mem_fraction_static: 0.88
Type:

Dict[str, Any]

model_path: str = 'none'
tp_size: int
trust_remote_code: bool = True
disable_custom_all_reduce: bool = False
extra_engine_args: Dict[str, Any]
class mmirage.core.process.processors.llm.config.SGLangLLMConfig(type='llm', server_args=<factory>, default_sampling_params=<factory>, chat_template='')[source]

Bases: BaseProcessorConfig

Configuration for LLM processor using SGLang.

Supports both text-only and multimodal (vision-language) models.

Parameters:
server_args

SGLang server arguments including model path and TP size.

Type:

mmirage.core.process.processors.llm.config.SGLangServerArgs

default_sampling_params

Default sampling parameters for generation.

Type:

Dict[str, Any]

chat_template

Chat template name for vision-language models (e.g., “qwen2-vl”).

Type:

str

type: Literal['llm'] = 'llm'
server_args: SGLangServerArgs
default_sampling_params: Dict[str, Any]
chat_template: str = ''
class mmirage.core.process.processors.llm.config.LLMSchemaField(type, min=None, max=None)[source]

Bases: object

One field of a JSON output_schema: a type name plus optional bounds.

Validates itself on construction, so an invalid spec fails as soon as it is built.

Parameters:
type

Schema type name (“str”, “int”, “float” or “bool”, or one of their aliases “string”, “integer”, “number”, “boolean”).

Type:

str

min

Inclusive lower bound, numeric types only. Numeric strings (e.g. produced by ${ENV_VAR} expansion) are coerced.

Type:

int | float | str | None

max

Inclusive upper bound, same rules as min.

Type:

int | float | str | None

TYPE_MAP: ClassVar[dict[str, type]] = {'bool': <class 'bool'>, 'boolean': <class 'bool'>, 'float': <class 'float'>, 'int': <class 'int'>, 'integer': <class 'int'>, 'number': <class 'float'>, 'str': <class 'str'>, 'string': <class 'str'>}
type: str
min: int | float | str | None = None
max: int | float | str | None = None
classmethod from_spec(spec)[source]

Normalize any accepted output_schema entry form into a field.

Parameters:

spec (str | dict[str, str | int | float | None] | LLMSchemaField)

Return type:

LLMSchemaField

python_type()[source]

Map the schema type name (“int”, “str”, …) to its Python type.

Return type:

type

field_type()[source]

The field’s type for create_model, carrying any declared bounds.

Return type:

object

property has_bounds: bool

Whether the field declares a min/max bound.

mmirage.core.process.processors.llm.config.SchemaFieldSpec = str | dict[str, str | int | float | None] | mmirage.core.process.processors.llm.config.LLMSchemaField

a type-name shorthand, a raw type/min/max mapping (kept raw so unknown keys are still rejected), or an already-built LLMSchemaField.

Type:

Accepted forms for one output_schema entry

class mmirage.core.process.processors.llm.config.LLMOutputVar(name='', type='', prompt='', output_schema=<factory>, output_type='')[source]

Bases: OutputVar

Output variable generated by LLM processor.

Uses Jinja2 templating for prompts and supports both plain text and structured JSON outputs.

Parameters:
name

Name of the variable.

Type:

str

type

Type identifier (must be “llm”).

Type:

str

prompt

Jinja2 template for the LLM prompt.

Type:

str

output_schema

JSON output fields, either a list of field names (all typed as strings) or a mapping of field name to an LLMSchemaField spec: a type name (“str”, “int”, “float” or “bool”) or a nested mapping with keys type (required) and, for numeric types, optional min/max bounds enforced during constrained decoding. Empty for plain text.

Type:

list[str] | dict[str, str | dict[str, str | int | float | None] | mmirage.core.process.processors.llm.config.LLMSchemaField]

output_type

Output format - “JSON” or “plain”.

Type:

str

prompt: str = ''
output_schema: list[str] | dict[str, str | dict[str, str | int | float | None] | LLMSchemaField]
output_type: str = ''
has_schema_constraints()[source]

Whether any field declares a min/max bound.

Return type:

bool

get_output_schema()[source]

Generate a Pydantic model for JSON output validation.

Returns:

A Pydantic BaseModel class if output_type is “JSON” and output_schema is non-empty, otherwise None.

Raises:

ValueError – If output_schema maps a field to an unsupported type name or to an invalid constraint mapping.

Return type:

Type[pydantic.BaseModel] | None

is_computable(vars)[source]

Check if all variables referenced in the prompt are available.

Parameters:

vars (Sequence[BaseVar]) – Sequence of currently available variables.

Returns:

True if all template variables are declared, False otherwise.

Return type:

bool

Implementation

LLM processor implementation using SGLang with multimodal support.

class mmirage.core.process.processors.llm.llm_processor.LLMProcessor(engine_args, shard_id=0, **kwargs)[source]

Bases: BaseProcessor[LLMOutputVar]

LLM processor for generating text using SGLang.

Supports both plain text and JSON output formats, with automatic chat template formatting and structured output validation.

Also supports multimodal (vision-language) inputs. For SGLang, image_data is expected to be aligned with prompts: a list where each element is either None (text-only), a single image (path/URL/PIL), or (optionally) a list of images for that prompt.

Parameters:
llm

SGLang engine for text generation.

tokenizer

Hugging Face tokenizer for chat template formatting.

sampling_params

Default sampling parameters for generation.

__init__(engine_args, shard_id=0, **kwargs)[source]

Initialize the LLM processor.

Parameters:
  • engine_args (SGLangLLMConfig) – SGLang runtime configuration.

  • shard_id (int) – Shard index for this worker.

Return type:

None

get_load_time()[source]

Return the wall-clock seconds spent initializing the SGLang engine.

Return type:

float

get_token_counts()[source]

Return cumulative token counts for this processor.

Returns:

TokenCounts object containing input and output token counts accumulated since this processor was created.

Return type:

TokenCounts

build_prompt(prompt_template, vars_samples)[source]

Build formatted prompts from a Jinja2 template and variable environments.

Parameters:
  • prompt_template (str) – Jinja2 template string for the prompt.

  • vars_samples (List[VariableEnvironment]) – List of variable environments containing values.

Returns:

List of formatted prompts with chat template applied.

Return type:

List[str]

batch_process_sample(batch, output_var)[source]

Process a batch of variable environments to generate LLM outputs.

Parameters:
Returns:

List of updated variable environments with LLM-generated values.

Raises:
  • ValueError – If output_type is JSON but no output_schema is defined.

  • RuntimeError – If output batch size doesn’t match input batch size.

Return type:

List[VariableEnvironment]

shutdown()[source]

Shutdown the LLM engine.

Return type:

None

Batch API processor

Configuration

Configuration for the batch API processor in MMIRAGE.

class mmirage.core.process.processors.batch_api.config.BatchApiProcessorConfig(type='batch_api', provider_config=None, export_prompts_dir=None)[source]

Bases: BaseProcessorConfig

Configuration for the batch API processor.

Provider settings are declared inline in YAML and resolved to the matching provider config class:

processors:
  - type: batch_api
    provider: openai
    model: gpt-4o-mini
Parameters:
provider_config

Resolved provider-specific batch configuration.

Type:

mmirage.config.batch_provider.BatchProviderConfig | None

export_prompts_dir

Value of –export-prompts.

Type:

str | None

type: Literal['batch_api'] = 'batch_api'
provider_config: BatchProviderConfig | None = None
export_prompts_dir: str | None = None
classmethod from_raw(data)[source]

Build the config from a raw YAML block, dispatching on provider.

Parameters:

data (Dict[str, Any])

Return type:

BatchApiProcessorConfig

class mmirage.core.process.processors.batch_api.config.BatchApiOutputVar(name='', type='', prompt='', output_schema=<factory>, output_type='')[source]

Bases: OutputVar

Output variable generated by the batch API processor.

Parameters:
prompt

Jinja2 template for the request prompt.

Type:

str

output_schema

Expected JSON field names, sent to the provider as a hint. Unlike LLMOutputVar, typed/bounded field specs are not supported here: providers are not given schema constraints, only field names.

Type:

list[str]

output_type

Output format - “JSON” or “plain”.

Type:

str

prompt: str = ''
output_schema: list[str]
output_type: str = ''
is_computable(vars)[source]

Check if all variables referenced in the prompt are available.

Parameters:

vars (Sequence[BaseVar])

Return type:

bool

Implementation

Batch API processor implementation for provider batch submission.

class mmirage.core.process.processors.batch_api.batch_api_processor.BatchApiProcessor(config, shard_id=0, **kwargs)[source]

Bases: BaseProcessor[BatchApiOutputVar]

Processor that submits generation requests to a provider batch API.

No model runs locally: each sample is serialized into a provider request and accumulated by an orchestrator, which uploads chunks and writes metadata receipts. Processed samples receive a __BATCH_SUBMITTED__ placeholder that the receiver utilities later replace with the provider results.

Text-only and multimodal requests are accumulated separately so each batch job stays homogeneous.

Parameters:
__init__(config, shard_id=0, **kwargs)[source]

Initialize the batch API processor.

Parameters:
  • config (BatchApiProcessorConfig) – Batch API configuration holding the resolved provider config.

  • shard_id (int) – Shard index for this worker.

Return type:

None

get_load_time()[source]

Return 0: no model is loaded in batch submission mode.

Return type:

float

get_token_counts()[source]

Return zero counts: no generation happens at submission time.

Provider usage is reported with the batch results, so it is read by the receiver (mmirage.core.process.batch.collector) instead.

Return type:

TokenCounts

build_multimodal_prompt(prompt_template, var_env)[source]

Build a prompt and extract its images.

Returns:

(formatted_prompt, images)

Parameters:
Return type:

Tuple[str, List[Image.Image | str]]

batch_process_sample(batch, output_var)[source]

Serialize a batch of samples into provider requests.

Parameters:
Returns:

The variable environments with a placeholder set for output_var.

Return type:

List[VariableEnvironment]

finalize()[source]

Flush both accumulators, submitting any remaining requests.

Return type:

None

Batch processing

The batch subsystem handles asynchronous, provider-backed inference (OpenAI Batch API, Anthropic Message Batches). It is activated by declaring a batch_api processor.

Orchestrator

Stateful provider-agnostic orchestration for batch submission.

class mmirage.core.process.batch.orchestrator.BatchSubmissionOrchestrator(adapter, config, export_prompts_path=None, export_batch_prefix='')[source]

Bases: object

Accumulate requests across map iterations and submit full-ready chunks.

Parameters:
property pending_count: int
add_requests(requests, source_indices, model_params_snapshot=None)[source]

Append requests and submit only chunks that are ready mid-stream.

Parameters:
Return type:

List[BatchSubmissionResult]

finalize(model_params_snapshot=None)[source]

Flush all remaining requests at end-of-dataset lifecycle.

Parameters:

model_params_snapshot (Mapping[str, Any] | None)

Return type:

List[BatchSubmissionResult]

Adapter (provider-neutral interface)

Provider-agnostic batch submission adapter contracts.

Adapters implement translation from internal request payloads into provider request formats and normalize submission responses into a shared result shape.

class mmirage.core.process.batch.adapter.BatchSubmissionResult(provider_batch_id, status, raw_response=<factory>)[source]

Bases: object

Normalized result returned by any provider adapter after chunk submission.

Parameters:
provider_batch_id

Provider-side identifier for the submitted job/batch.

Type:

str

status

Provider status mapped to ‘completed’, ‘failed’, ‘in_progress’ or ‘unknown’, so generic code never reads a provider vocabulary.

Type:

str

raw_response

Original provider response payload for traceability.

Type:

Dict[str, Any]

provider_batch_id: str
status: str
raw_response: Dict[str, Any]
class mmirage.core.process.batch.adapter.BatchSubmissionAdapter[source]

Bases: ABC

Abstract interface for provider-specific batch submission adapters.

Implementations should be deterministic for request building and byte estimation so chunk boundaries can be reproduced across retries.

required_credentials: Tuple[str, ...] = ()
abstractmethod build_request(custom_id, payload, config)[source]

Build a single provider-ready request object.

Parameters:
  • custom_id (str) – Stable request identifier used to map async results back to source rows.

  • payload (Dict[str, Any]) – Provider-neutral request payload assembled by the core processing layer.

  • config (BatchProviderConfig) – Provider configuration contract that may influence request shaping.

Returns:

A provider-specific request object represented as a mapping.

Return type:

Dict[str, Any]

abstractmethod estimate_request_bytes(request)[source]

Estimate serialized UTF-8 bytes for a request payload.

The estimate must match or safely upper-bound the size produced by the serializer used for submission so chunk boundaries are enforced correctly.

Parameters:

request (Dict[str, Any]) – Provider request object returned by build_request.

Returns:

Estimated byte size for the serialized request.

Return type:

int

abstractmethod submit_chunk(chunk_id, requests, config)[source]

Submit one pre-chunked request group to the provider.

Parameters:
  • chunk_id (str) – Internal chunk identifier generated by orchestration.

  • requests (Sequence[Dict[str, Any]]) – Provider-ready request objects belonging to this chunk.

  • config (BatchProviderConfig) – Provider config containing credentials and submission knobs.

Returns:

Raw provider response payload as a mapping.

Return type:

Dict[str, Any]

abstractmethod parse_submission_result(raw_result)[source]

Normalize provider submission output into a shared result model.

Parameters:

raw_result (Dict[str, Any]) – Raw payload returned by submit_chunk.

Returns:

A normalized BatchSubmissionResult for provider-neutral orchestration and metadata persistence.

Return type:

BatchSubmissionResult

abstractmethod check_batch_status(provider_batch_id, config)[source]

Retrieve and normalize the latest status for a provider batch job.

Parameters:
  • provider_batch_id (str) – Provider-side batch/job identifier to query.

  • config (BatchProviderConfig) – Provider configuration containing credentials and endpoint overrides.

Returns:

A normalized BatchSubmissionResult where status reflects the latest provider-reported lifecycle state for the batch.

Return type:

BatchSubmissionResult

abstractmethod retrieve_results(provider_batch_id, config)[source]

Download and parse completed batch results from the provider.

Implementations should normalize each returned row into a plain mapping and, when a text payload is available, expose it as generated_text so downstream collectors can consume a provider-agnostic result shape. Reported usage should likewise be exposed as input_tokens and output_tokens, omitted when the provider reports none.

Parameters:
  • provider_batch_id (str) – Provider-side batch/job identifier.

  • config (BatchProviderConfig) – Provider configuration containing credentials and endpoint overrides.

Returns:

Sequence of parsed result rows (provider JSONL records normalized to dictionaries) preserving provider output order.

Return type:

Sequence[Dict[str, Any]]

OpenAI adapter

Concrete OpenAI implementation of batch submission contracts.

class mmirage.core.process.batch.openai_adapter.OpenAIBatchAdapter[source]

Bases: BatchSubmissionAdapter

Provider adapter for OpenAI Batch API.

required_credentials: Tuple[str, ...] = ('api_key',)
build_request(custom_id, payload, config)[source]

Build a single provider-ready request object.

Parameters:
  • custom_id (str) – Stable request identifier used to map async results back to source rows.

  • payload (Dict[str, Any]) – Provider-neutral request payload assembled by the core processing layer.

  • config (BatchProviderConfig) – Provider configuration contract that may influence request shaping.

Returns:

A provider-specific request object represented as a mapping.

Return type:

Dict[str, Any]

estimate_request_bytes(request)[source]

Estimate serialized UTF-8 bytes for a request payload.

The estimate must match or safely upper-bound the size produced by the serializer used for submission so chunk boundaries are enforced correctly.

Parameters:

request (Dict[str, Any]) – Provider request object returned by build_request.

Returns:

Estimated byte size for the serialized request.

Return type:

int

submit_chunk(chunk_id, requests, config)[source]

Submit one pre-chunked request group to the provider.

Parameters:
  • chunk_id (str) – Internal chunk identifier generated by orchestration.

  • requests (Sequence[Dict[str, Any]]) – Provider-ready request objects belonging to this chunk.

  • config (BatchProviderConfig) – Provider config containing credentials and submission knobs.

Returns:

Raw provider response payload as a mapping.

Return type:

Dict[str, Any]

check_batch_status(provider_batch_id, config)[source]

Retrieve and normalize the latest status for a provider batch job.

Parameters:
  • provider_batch_id (str) – Provider-side batch/job identifier to query.

  • config (BatchProviderConfig) – Provider configuration containing credentials and endpoint overrides.

Returns:

A normalized BatchSubmissionResult where status reflects the latest provider-reported lifecycle state for the batch.

Return type:

BatchSubmissionResult

retrieve_results(provider_batch_id, config)[source]

Download completed OpenAI batch rows and normalize text into generated_text.

OpenAI batch outputs can surface the assistant payload in nested response bodies, so this method flattens the provider-specific shape before returning rows to the provider-agnostic collector.

Parameters:
Return type:

Sequence[Dict[str, Any]]

parse_submission_result(raw_result)[source]

Normalize provider submission output into a shared result model.

Parameters:

raw_result (Dict[str, Any]) – Raw payload returned by submit_chunk.

Returns:

A normalized BatchSubmissionResult for provider-neutral orchestration and metadata persistence.

Return type:

BatchSubmissionResult

Anthropic adapter

Chunking

Provider-agnostic request chunking utilities for batch submission.

class mmirage.core.process.batch.chunking.RequestChunk(requests, total_bytes, has_oversized_request=False)[source]

Bases: object

Chunk of provider-ready requests with aggregate metadata.

Parameters:
requests: List[Dict[str, Any]]
total_bytes: int
has_oversized_request: bool = False
property total_requests: int
class mmirage.core.process.batch.chunking.BatchRequestChunker(adapter, config)[source]

Bases: object

Split request sequences into chunks using serialized-byte limits.

Parameters:
chunk_requests(requests)[source]

Chunk requests according to max bytes, max requests, and oversize policy.

Parameters:

requests (Sequence[Dict[str, Any]])

Return type:

List[RequestChunk]

Status checker

Receiver-side helper to check provider batch status from metadata receipts.

Designed for CLI use against JSONL receipt files. Skips malformed lines and missing keys to keep status checks resilient to partial metadata corruption.

mmirage.core.process.batch.status_checker.extract_unique_provider_batches(metadata_records)[source]

Return unique (provider, provider_batch_id) pairs.

Provider names are already lowercased and records missing either key are already dropped by _read_metadata_records, so both are assumed here.

Parameters:

metadata_records (Sequence[BatchMetadataRecord])

Return type:

List[Tuple[str, str]]

mmirage.core.process.batch.status_checker.run_status_checker(metadata_records, provider_configs)[source]

Check batch status for each referenced provider batch.

Prints a per-batch line and a per-provider summary. A batch the provider cannot resolve is counted as lookup_failed so one stale receipt does not hide the status of every batch after it.

Parameters:
Return type:

List[BatchSubmissionResult]

mmirage.core.process.batch.status_checker.check_batches(cfg, metadata_paths=None)[source]

Check every batch referenced by the receipts of cfg.

Parameters:
  • cfg (MMirageConfig) – Loaded MMIRAGE configuration declaring the batch_api processor(s).

  • metadata_paths (Optional[Sequence[str]]) – Explicit receipt paths; resolved from the config when omitted.

Returns:

1 when a batch-level failure occurred, a lookup failed or the provider configs cannot be built, 0 otherwise. Batches still running are not a failure, and per request errors only show up when the results are read.

Return type:

Exit code

mmirage.core.process.batch.status_checker.main(argv=None)[source]

CLI entry point that returns a process-style status code.

Returns 0 on success or no batches found, and 1 on configuration or provider resolution failures.

Parameters:

argv (Sequence[str] | None)

Return type:

int