Sklearn Pipeline Template

scikit-learn pipeline template

The reference layout every tabular bundle in ManagerPack copies. The goal is consistency: same project structure, same MLflow conventions, same load/predict path, same notebook style. Whatever the actual problem, the plumbing is identical.

The worked example is "is this coin fair?" — a logistic regression on (flip_index, outcome). The model is trivial, the plumbing is the point.

Project layout

<bundle>/
├── README.md             # what this bundle does + how to run it
├── SKILL.md              # this file (or specialized for the bundle)
├── src/
│   ├── train.py          # train + log to MLflow
│   ├── predict.py        # load model from MLflow and predict
│   └── plots.py          # plot helpers, logged as MLflow artifacts
├── notebooks/
│   └── <name>_demo.py    # marimo notebook with mo.ui.slider
└── mlruns/               # MLflow tracking store (gitignored)

The data lives outside the bundle, in studio/data/<problem>.parquet, generated by datagen <problem>. Bundles never carry their own data — they consume parquet from the studio's shared data directory.

Conventions

Pipeline shape

Always wrap preprocessing inside the sklearn Pipeline so it travels with the model on save/load:

Pipeline([
    ("preprocess", ColumnTransformer([
        ("scaled", StandardScaler(), numeric_cols),
        ("encoded", OneHotEncoder(handle_unknown="ignore"), categorical_cols),
    ])),
    ("clf", LogisticRegression(max_iter=1000)),
])

Never separate the preprocessing step from the model — the loaded artifact must work standalone with raw input.

MLflow logging

Every run logs:

Recovery error against ground truth is the most important metric for template runs because it answers "did the model recover what we know to be true?"

Model serialization

Use mlflow.sklearn.log_model(sk_model=pipeline, name="model", input_example=X_train.head(5)).

Never use bare joblib.dump() or pickle.dump(). The MLflow path:

• Serializes the model + signature + conda env + requirements + Python version into a self-describing artifact directory
• Loads from anywhere with mlflow.sklearn.load_model("runs:/<id>/model")
• Side note: MLflow will warn that pickle is unsafe and recommend skops. For internal templates this is fine; for security-sensitive deployments, use the skops format instead.

Loading the model elsewhere

import mlflow
import mlflow.sklearn
mlflow.set_tracking_uri(f"file:{template_dir / 'mlruns'}")
model = mlflow.sklearn.load_model(f"runs:/{run_id}/model")
predictions = model.predict_proba(new_data)

The Pipeline (preprocessing + classifier) comes back as one object. No need to re-import the training code.

Coefficient interpretation

When the model is linear, log the interpretable quantities as metrics so they show up in the MLflow UI:

clf = pipeline.named_steps["clf"]
mlflow.log_metric("intercept_logit", float(clf.intercept_[0]))
mlflow.log_metric("coef_some_feature", float(clf.coef_[0][feature_idx]))

For the fair-coin template, the intercept (in logit space) maps directly to $P(\text{heads})$ at the mean of the standardized index, and the slope on flip_index is a non-stationarity detector. Always interpret what the coefficients mean in domain terms.

Plots

Generate matplotlib figures in src/plots.py, save them to a temp directory, and log them as MLflow artifacts under plots/. Always include:

• An empirical-vs-predicted view (rolling-window or per-bin)
• A calibration plot for classifiers
• A coefficient bar chart for linear models

Marimo notebook (always marimo, never Jupyter)

Every bundle ships a notebooks/<name>_demo.py marimo notebook that:

1. Loads the latest run from mlruns/
2. Loads the model via mlflow.sklearn.load_model
3. Shows the recovered coefficients
4. Has a mo.ui.slider (or anywidget/wigglystuff widget) for interactive prediction
5. Compares recovered values to ground truth from the sidecar

Validate with marimo check notebooks/<name>_demo.py. Test rendering headlessly with marimo export html notebooks/<name>_demo.py -o /tmp/x.html to confirm all cells execute without errors before shipping.

Running the worked example

# 1. Generate the dataset
datagen coin-flip --n 500 --p 0.7 --seed 42

# 2. Train and log to MLflow
cd studio/templates/sklearn-pipeline
python src/train.py

# 3. Verify recovery
# stdout shows |recovery error| ≈ 0.01 — the model recovers true_p well

# 4. Load the model and predict
python src/predict.py --run-id <id> --flip-index 100

# 5. Open the demo notebook
marimo edit notebooks/coin_flip_demo.py

Drift detection (bonus diagnostic)

datagen coin-flip --n 500 --p 0.3 --drift 0.4 --seed 7 --output /tmp/drift.parquet
python src/train.py --data /tmp/drift.parquet

The standardized slope on flip_index should be strongly positive, correctly detecting that $P(\text{heads})$ increases over the sequence. This is a free non-stationarity check that any tabular bundle with a time-like feature can borrow.

Specializing this template for a new bundle

1. Copy studio/templates/sklearn-pipeline/ to studio/scratch/<bundle-name>/ (or directly to bundles/<bundle-name>/ once the design is settled)
2. Replace the dataset reference: change coin-flip to your problem name
3. Update build_pipeline() to use the right preprocessing + estimator
4. Update the metrics: log-loss for classification, RMSE/MAE for regression, ROC-AUC for ranking, etc.
5. Update the plots in src/plots.py for the new domain
6. Update the marimo notebook to show the right interactive widgets
7. Update SKILL.md with the bundle-specific story