Inference API

Inference executes a configured sequence or dependency graph of steps. The configuration describes the pipeline, builders resolve step factories from `INFERENCE_STEPS`, each step reads/writes artifacts in an `InferenceContext`, and the pipeline returns the updated context. For command line usage, see Inference.

Configuration

`InferenceConfig` exposes pipeline name, parallel mode, steps, and requested runtime device from an inference YAML mapping.

chimera_ml.inference.config.InferenceConfig dataclass

Bases: ExperimentConfig

Typed accessors for the inference YAML schema.

The wrapper keeps the raw mapping inherited from ExperimentConfig and exposes pipeline name, parallel flag, step list, and requested runtime device. It does not build steps; use build_inference_pipeline for that.

Context and graph

`InferenceContext` is the mutable artifact store passed to steps. An `InferenceGraphNode` associates one step with its `after` dependencies.

chimera_ml.inference.context.InferenceContext dataclass

Mutable state and artifact store shared by inference steps.

Attributes:

Name	Type	Description
input_path	Path	Input file or sample path for the current run.
work_dir	Path	Directory for caches and intermediate artifacts.
device	str	Resolved runtime device string.
config	dict[str, Any]	Raw inference configuration mapping.
artifacts	dict[str, Any]	Named values produced by steps; predictions is conventional.
written_artifact_keys	frozenset[str]	Keys written by the current step context and used by the pipeline when merging results.

Steps should use get_artifact and set_artifact rather than mutating the artifact ownership bookkeeping directly.

get_artifact

get_artifact(name, default=None)

Return a named artifact or a default when it is absent.

Parameters:

Name	Type	Description	Default
name	str	Artifact key.	required
default	Any	Value returned when the key is missing.	None

set_artifact

set_artifact(name, value)

Store an artifact and mark it as written by this step.

Parameters:

Name	Type	Description	Default
name	str	Artifact key.	required
value	Any	Value to store.	required

chimera_ml.inference.pipeline.InferenceGraphNode dataclass

One inference step and the node identifiers it depends on.

Attributes:

Name	Type	Description
node_id	str	Unique identifier used by after references.
step	BaseInferenceStep	Object implementing run(ctx) -> ctx.
after	tuple[str, ...]	Dependencies that must complete before this node runs.

Pipeline

`InferencePipeline` validates node identifiers and dependencies, schedules ready nodes, isolates step contexts, and merges written artifacts. In parallel mode, unrelated nodes may run concurrently; conflicting artifact writes are rejected unless the dependency graph makes the overwrite ordered.

chimera_ml.inference.pipeline.InferencePipeline

Execute inference steps as a validated sequential or parallel DAG.

The constructor rejects duplicate IDs, self-dependencies, unknown dependencies, and cycles. run copies the initial artifact mapping into step contexts and merges only artifacts explicitly written by each step. In parallel mode, unordered writes to the same artifact are rejected.

Attributes:

Name	Type	Description
nodes		Original graph nodes in configuration order.
name		Human-readable pipeline name.
parallel_mode		Whether explicit after dependencies control concurrency.

__init__

__init__(nodes, *, name='inference_pipeline', parallel_mode=False)

Validate and initialize a pipeline.

Parameters:

Name	Type	Description	Default
nodes	list[InferenceGraphNode]	Graph nodes with unique IDs and valid dependencies.	required
name	str	Human-readable pipeline name.	'inference_pipeline'
parallel_mode	bool	Allow independent nodes to run concurrently and enforce artifact ownership rules.	False

Raises: ValueError: If node IDs or dependencies are invalid.

run

run(ctx)

Run all nodes and update the supplied context's artifacts.

Parameters:

Name	Type	Description	Default
ctx	InferenceContext	Initial context containing input, runtime, config, and seed artifacts.	required

Returns: The same context object with merged artifacts. Raises: ValueError: If dependencies cannot be resolved or parallel artifact writes conflict.

Builders and utilities

Use `build_inference_step` for one registry-backed step and `build_inference_pipeline` for a complete `InferenceConfig`. Device selection accepts `auto`, `cpu`, or `cuda` through `resolve_inference_device`.

chimera_ml.inference.builders.build_inference_step

build_inference_step(cfg, *, inject=None)

Build and validate one step from INFERENCE_STEPS.

Parameters:

Name	Type	Description	Default
cfg	dict[str, Any]	Mapping with registry name and optional params.	required
inject	Mapping[str, Any] | None	Runtime values selectively passed to factories.	None

Returns: A step implementing run(ctx) -> ctx. Raises: TypeError: If the factory result has no callable run contract.

chimera_ml.inference.builders.build_inference_pipeline

build_inference_pipeline(config, *, inject=None)

Build a sequential or parallel pipeline from InferenceConfig.

Parameters:

Name	Type	Description	Default
config	InferenceConfig	Loaded inference configuration.	required
inject	Mapping[str, Any] | None	Optional runtime dependencies for step factories.	None

Returns: A validated InferencePipeline. Raises: TypeError: If a step configuration is not a mapping. ValueError: If a step has no usable name/id.

chimera_ml.inference.utils.resolve_inference_device

resolve_inference_device(device)

Resolve requested inference device to a concrete runtime value.

Step protocol and built-ins

Custom steps implement the small `run(ctx) -> ctx` protocol. The built-in checkpoint step resolves local or HTTP(S) references into a cache; JSON steps write or print the `predictions` artifact.

chimera_ml.inference.steps.base.BaseInferenceStep

Bases: Protocol

Protocol implemented by every inference step.

A step receives an InferenceContext, may read existing artifacts, and returns the context after adding or updating artifacts. Steps should keep configuration in the object instance and avoid hidden global state.

chimera_ml.inference.steps.checkpoint_steps.ResolveCheckpointsStep dataclass

Resolve local or HTTP(S) checkpoint references into a work-dir cache.

Attributes:

Name	Type	Description
checkpoints	dict[str, str]	Named checkpoint reference mapping.
cache_dir	str	Relative-to-work-dir or absolute cache directory.
force_download	bool	Redownload remote files even when cached.
chunk_size	int	Streaming HTTP chunk size in bytes.

run stores the resolved mapping under the checkpoints artifact and the cache directory under cache_dir. Local paths must point to files; remote downloads use a temporary .part file before replacement.

run

run(ctx)

Resolve configured references and publish checkpoint artifacts.

Parameters:

Name	Type	Description	Default
ctx	InferenceContext	Inference context whose work_dir and artifacts are used.	required

Returns: The context with checkpoints and cache_dir artifacts. Raises: FileNotFoundError: If a local or downloaded checkpoint is missing. TypeError: If an existing checkpoints artifact is not a mapping.

chimera_ml.inference.steps.checkpoint_steps.resolve_checkpoints_step

resolve_checkpoints_step(**params)

chimera_ml.inference.steps.output_steps.WriteJsonPredictionsStep dataclass

Write the predictions artifact and input path as JSON.

Attributes:

Name	Type	Description
output_path	str | None	Optional explicit output file. When omitted, the input suffix is replaced with .json.

chimera_ml.inference.steps.output_steps.write_json_predictions_step

write_json_predictions_step(**params)

Create WriteJsonPredictionsStep from registry parameters.

Registry key: write_json_predictions_step.

chimera_ml.inference.steps.output_steps.PrintJsonPredictionsStep dataclass

Print the predictions artifact and input path as JSON.

chimera_ml.inference.steps.output_steps.print_json_predictions_step

print_json_predictions_step(**params)

Create PrintJsonPredictionsStep from registry parameters.

Registry key: print_json_predictions_step.