stratum-ai 0.0.0.dev1

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.

Design Principles

• Provide seamless and unrestricted support for arbitrary ML libraries without operator porting.
• Enable lazy evaluation and provide operator semantics that enable logical rewrites and cost-based optimizations.
• 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(
    rust_backend=True,
    scheduler=True,
    stats=True,
    debug_timing=False,
)

Example Code

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

def main():
    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()

    skrub.set_config(rust_backend=True, debug_timing=True, scheduler=True, stats=True)
    tv = skrub.TableVectorizer(high_cardinality=skrub.StringEncoder(), low_cardinality=OneHotEncoder(sparse_output=False))
    X_enc = X.skb.apply(tv)
    print(f"Encoded data shape: {X_enc.shape.skb.preview()}")

    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()

Repository Layout

stratum/
├─ pyproject.toml           # Project metadata + Python/Rust build config (maturin)
├─ README.md
├─ LICENSE
├─ _rust/                   # Rust crate (PyO3 extension)
│  ├─ Cargo.toml
│  └─ src/lib.rs            # Defines #[pymodule] fn _rust_backend_native(...)
└─ stratum/                 # Python package
   ├─ __init__.py           # Façade over skrub + automatic patching
   ├─ _config.py            # set_config/get_config + runtime/env sync
   ├─ _api.py               # High-level grid search / evaluate helpers
   ├─ _rust_backend.py      # Python <-> Rust shim (re-exports native fns)
   ├─ adapters/             # Public API (dispatch to Rust or fall back to skrub)
   │  ├─ string_encoder.py  # RustyStringEncoder
   │  └─ one_hot_encoder.py # RustyOneHotEncoder
   ├─ optimizer/
   │  ├─ ir/                # DAG representation 
   │  └─ _optimize.py       # logical rewrites
   ├─ runtime/              # Schedulers and runtime execution
   ├─ patching/             # Hooks that patch upstream skrub
   └─ tests/                # Test suite

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.