mmirage.core.process — Processors¶
Variables¶
Variable system for MMIRAGE pipeline with multimodal support.
- class mmirage.core.process.variables.BaseVar(name='')[source]¶
Bases:
ABCBase class for variables in the MMIRAGE pipeline.
- Parameters:
name (str)
- class mmirage.core.process.variables.InputVar(name='', key='', type='text')[source]¶
Bases:
BaseVarInput variable extracted from source datasets.
- type¶
Variable type - “text” or “image”.
- Type:
Literal[‘text’, ‘image’]
- class mmirage.core.process.variables.OutputVar(name='', type='')[source]¶
Bases:
BaseVarOutput variable generated by processors.
Output variables are created by processors (e.g., LLMs) and can depend on input variables and previously computed output variables.
- class mmirage.core.process.variables.VariableEnvironment(var_env, image_vars=None)[source]¶
Bases:
objectEnvironment for storing and accessing variables during processing.
- with_variable(key, value, is_image=False)[source]¶
Create a new environment with an additional variable.
- Parameters:
- Returns:
New environment with the added variable.
- Return type:
- to_dict()[source]¶
Get an immutable view of the variable dictionary.
- Returns:
MappingProxyType providing read-only access to variables.
- Return type:
- get_image_vars()[source]¶
Get all image variable names.
- Returns:
Copy of the set containing names of all image variables.
- Return type:
- 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:
- static from_input_variables(sample, input_vars, image_base_path=None)[source]¶
Create a variable environment from a single sample.
- Parameters:
- Returns:
Environment populated with extracted variables.
- Return type:
- Raises:
ValueError – If a required input variable is not found in the sample.
Base processor¶
Base classes and registry for processors in MMIRAGE.
- class mmirage.core.process.base.BaseProcessorConfig(type)[source]¶
Bases:
objectBase 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:
- class mmirage.core.process.base.TokenCounts(input_tokens, output_tokens)[source]¶
Bases:
objectCumulative token counts from LLM processors.
- class mmirage.core.process.base.BaseProcessor(config, shard_id=0, **kwargs)[source]¶
-
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 (BaseProcessorConfig)
shard_id (int)
- 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:
- 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:
- 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:
- class mmirage.core.process.base.ProcessorRegistry[source]¶
Bases:
objectRegistry 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:
name (str)
config_cls (Type[BaseProcessorConfig])
- 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:
- 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:
- 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:
- 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:
- class mmirage.core.process.base.AutoProcessor[source]¶
Bases:
objectFactory 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:
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:
objectMapper 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:
- rewrite_batch(batch, image_base_path=None)[source]¶
Transform a batch of samples by computing output variables.
- Parameters:
- Returns:
List of VariableEnvironments with all output variables computed.
- Raises:
RuntimeError – If an output variable type has no registered processor.
- Return type:
- get_token_counts()[source]¶
Return cumulative token counts aggregated across all LLM processors.
Sums
input_tokensandoutput_tokensfrom every processor that exposes aget_token_counts()method (i.e.,LLMProcessor).- Returns:
TokenCounts with
input_tokensandoutput_tokensfields.- Return type:
- get_load_time()[source]¶
Return total model-loading time (seconds) summed across all LLM processors.
- Return type:
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:
objectServer arguments for SGLang engine.
- Parameters:
- class mmirage.core.process.processors.llm.config.SGLangLLMConfig(type='llm', server_args=<factory>, default_sampling_params=<factory>, chat_template='')[source]¶
Bases:
BaseProcessorConfigConfiguration 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.
- server_args: SGLangServerArgs¶
- class mmirage.core.process.processors.llm.config.LLMSchemaField(type, min=None, max=None)[source]¶
Bases:
objectOne 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.
- type¶
Schema type name (“str”, “int”, “float” or “bool”, or one of their aliases “string”, “integer”, “number”, “boolean”).
- Type:
- min¶
Inclusive lower bound, numeric types only. Numeric strings (e.g. produced by ${ENV_VAR} expansion) are coerced.
- 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'>}¶
- classmethod from_spec(spec)[source]¶
Normalize any accepted output_schema entry form into a field.
- 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:
OutputVarOutput variable generated by LLM processor.
Uses Jinja2 templating for prompts and supports both plain text and structured JSON outputs.
- Parameters:
- 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.
- 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
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:
engine_args (SGLangLLMConfig)
shard_id (int)
- 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:
- 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:
- 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:
- batch_process_sample(batch, output_var)[source]¶
Process a batch of variable environments to generate LLM outputs.
- Parameters:
batch (List[VariableEnvironment]) – List of variable environments to process.
output_var (LLMOutputVar) – Output variable defining prompt and output format.
- 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:
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:
BaseProcessorConfigConfiguration 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:
type (Literal['batch_api'])
provider_config (BatchProviderConfig | None)
export_prompts_dir (str | None)
- provider_config¶
Resolved provider-specific batch configuration.
- Type:
- provider_config: BatchProviderConfig | None = None¶
- class mmirage.core.process.processors.batch_api.config.BatchApiOutputVar(name='', type='', prompt='', output_schema=<factory>, output_type='')[source]¶
Bases:
OutputVarOutput variable generated by the batch API processor.
- 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.
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:
config (BatchApiProcessorConfig)
shard_id (int)
- __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_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:
- build_multimodal_prompt(prompt_template, var_env)[source]¶
Build a prompt and extract its images.
- Returns:
(formatted_prompt, images)
- Parameters:
prompt_template (str)
var_env (VariableEnvironment)
- Return type:
- batch_process_sample(batch, output_var)[source]¶
Serialize a batch of samples into provider requests.
- Parameters:
batch (List[VariableEnvironment]) – List of variable environments to process.
output_var (BatchApiOutputVar) – Output variable defining prompt and output format.
- Returns:
The variable environments with a placeholder set for
output_var.- Return type:
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:
objectAccumulate requests across map iterations and submit full-ready chunks.
- Parameters:
adapter (BatchSubmissionAdapter)
config (BatchProviderConfig)
export_prompts_path (Optional[str])
export_batch_prefix (str)
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:
objectNormalized result returned by any provider adapter after chunk submission.
- status¶
Provider status mapped to ‘completed’, ‘failed’, ‘in_progress’ or ‘unknown’, so generic code never reads a provider vocabulary.
- Type:
- class mmirage.core.process.batch.adapter.BatchSubmissionAdapter[source]¶
Bases:
ABCAbstract 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.
- 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:
- 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.
- abstractmethod submit_chunk(chunk_id, requests, config)[source]¶
Submit one pre-chunked request group to the provider.
- Parameters:
- Returns:
Raw provider response payload as a mapping.
- Return type:
- abstractmethod parse_submission_result(raw_result)[source]¶
Normalize provider submission output into a shared result model.
- 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
BatchSubmissionResultwherestatusreflects the latest provider-reported lifecycle state for the batch.- Return type:
- 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_textso downstream collectors can consume a provider-agnostic result shape. Reported usage should likewise be exposed asinput_tokensandoutput_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:
OpenAI adapter¶
Concrete OpenAI implementation of batch submission contracts.
- class mmirage.core.process.batch.openai_adapter.OpenAIBatchAdapter[source]¶
Bases:
BatchSubmissionAdapterProvider adapter for OpenAI Batch API.
- 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:
- 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.
- submit_chunk(chunk_id, requests, config)[source]¶
Submit one pre-chunked request group to the provider.
- Parameters:
- Returns:
Raw provider response payload as a mapping.
- Return type:
- 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
BatchSubmissionResultwherestatusreflects the latest provider-reported lifecycle state for the batch.- Return type:
- 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.
- parse_submission_result(raw_result)[source]¶
Normalize provider submission output into a shared result model.
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:
objectChunk of provider-ready requests with aggregate metadata.
- class mmirage.core.process.batch.chunking.BatchRequestChunker(adapter, config)[source]¶
Bases:
objectSplit request sequences into chunks using serialized-byte limits.
- Parameters:
adapter (BatchSubmissionAdapter)
config (BatchProviderConfig)
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.
- 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_failedso one stale receipt does not hide the status of every batch after it.- Parameters:
metadata_records (Sequence[BatchMetadataRecord])
provider_configs (Mapping[str, BatchProviderConfig])
- Return type:
- 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