Stratum is an ML system for efficiently executing large-scale agentic pipeline search. It integrates with MLE agents by representing batches of agent-generated pipelines as lazily evaluated DAGs, applying logical and runtime optimizations, and executing them across heterogeneous backends, including a Rust-based runtime. Stratum builds on skrub's operator abstraction and is under active development.

📄 Paper: The motivation and vision behind stratum are described in our VLDB 2026 paper.

---

Design Principles

• Provide seamless and unrestricted support for arbitrary ML libraries without operator porting.
• A semantic abstraction built on a minimal set of logical operators that enables rewrites and lazy evaluation with physical operator independence.
• Implement a runtime with efficient operator kernels (in Rust), scheduling across CPUs, GPUs, and distributed backends, plus runtime optimizations such as buffer pools, reuse of intermediates, and inter- and intra-operator parallelization.

---

Installation

Install the current development build from PyPI:

python -m pip install "stratum-ai==0.0.0.dev1"

Pre-built Rust wheels are provided for the supported CPython 3.11+ platforms, so a Rust toolchain is not required for a normal installation. The source-build instructions below are for contributors developing Stratum itself.

For source development, you need Python 3.11+, a Rust toolchain (nightly not required; stable is fine), and maturin (python -m pip install maturin).

From the repository root, install the extension in editable (development) mode:

maturin develop --release

For more details (including building wheels), see the Developer Instructions section below.

---

Usage

To leverage stratum, agent prompts or pipelines need minor changes. Prompts should be modified to generate code following skrub DataOps syntax.

Stratum can also significantly speed up human-written skrub code.

The following flags enable different features of Stratum. These flags can be set via environment variables or directly in code:

import stratum

stratum.set_config(
    scheduler=True,
    implementation_selector="greedy",
    explain=True,
    stats=True,
)

Example Code

import stratum as skrub #drop-in replacement
from sklearn.preprocessing import OneHotEncoder
from sklearn.linear_model import LinearRegression

def main():
    # Collect and prepare datasets
    dataset = skrub.datasets.fetch_employee_salaries()
    df = skrub.as_data_op(dataset.employee_salaries).skb.subsample()
    df_clean = df.dropna()
    y = df_clean["current_annual_salary"].skb.mark_as_y()
    X = df_clean.drop(columns=["current_annual_salary"]).skb.mark_as_X()

    # Apply feature transformations
    skrub.set_config(scheduler=True, implementation_selector="greedy", explain=True, stats=True)
    tv = skrub.TableVectorizer()
    X_enc = X.skb.apply(tv)
    print(f"Encoded data shape: {X_enc.shape.skb.preview()}")

    # Training and cross-validation
    pred = X_enc.skb.apply(LinearRegression(), y=y)
    search = pred.skb.make_grid_search(cv=3, fitted=True, scoring="r2", refit=False)
    print(search.results_)

if __name__ == "__main__":
    main()

---

Developer Instructions

Running the Tests

Install all extras and run the full test suite:

uv sync --all-extras
pytest -v stratum/tests

Or, more concisely:

uv run pytest

---

Local Dev Install (Editable, without uv)

maturin develop				# Debug mode
maturin develop --release	# Optimized dev build

Building Wheels

This produces redistributable .whl files under dist/.

# Linux / macOS
maturin build --release -o dist --interpreter python3.11 --compatibility linux

# Windows
maturin build --release -o dist

Then install with:

python -m pip install ./dist/stratum_ai-*.whl

---

License

Apache License 2.0. See LICENSE for details.