Training Pipeline Practice: AutoML with Optuna

What You'll Build

In this exercise you will build an AutoML training pipeline that automatically searches across multiple model families and their hyperparameters to find the best cricket in-form classifier, then tracks every trial and promotes the winner. Rather than hand-picking one algorithm and tuning it manually, you will let Optuna choose between, say, a gradient-boosted tree and a random forest, and tune the selected model's parameters within the same study, logging each trial to MLflow. This combines the orchestration mindset of Module 2 with the adaptive search of Optuna into a single automated experiment. By the end you will have a pipeline that turns the open-ended question 'which model and settings work best here' into a systematic, reproducible search whose entire history is recorded and whose best model is saved ready for deployment, the kind of automation that lets a small team evaluate far more options than they ever could by hand.

Analogy🏏Cricket

🏏 Think of it like cricket: a selection committee does not pick a squad on a hunch; they run structured trial matches under recorded conditions, log every player's scores, compare candidates on identical criteria, and keep the records so a selection can be justified later. Just as the trial matches are your tracked experiment runs, each player's logged scores are your run metrics. Just as the committee compares candidates on the same pitch to be fair, you compare model configurations on the same data. Just as a defensible selection can be reproduced from the records if challenged, your pipeline can reproduce any run from its logged inputs. The insight is that disciplined, recorded trials, not gut feel, are what make both squad selection and model selection defensible.

Prerequisites

Completion of lesson 02 (MLflow tracking) and lesson 08 (Optuna), or equivalent familiarity with both.
Python 3.10 or newer with pip, and the ability to install optuna, mlflow, and scikit-learn.
Comfort defining a Python function, returning a value, and reading a small dataset.
Understanding of the difference between hyperparameters and learned parameters, and what a cross-validation score means.
A terminal where you can run scripts and launch the MLflow UI to inspect tracked runs.

Setup & Project Structure

You will create a focused project that separates the dataset, the AutoML search logic, and the tracked outputs, mirroring how a real automated-training service is laid out. The search logic lives in src, MLflow writes runs to a local store, and the promoted best model lands in outputs. Keeping the search definition separate from configuration and results matters because it lets you rerun the same search with a different trial budget or dataset without touching the core logic, and it keeps the tuning history cleanly recorded for later inspection. Install the three dependencies and lay out the folders before writing any code.

Analogy🏏Cricket

🏏 Think of it like cricket: a well-run academy never dumps kit, players, and match records into one heap; it keeps the practice ground, the coaching staff's playbook, and the scorers' logbook in separate, clearly labelled areas so anyone can find what they need and nothing gets mixed up. Just as you separate the data directory, the src pipeline logic, and the mlruns tracking store, the academy separates its pitches, its coaching manuals, and its performance ledgers. Just as keeping these concerns apart makes a project's structure obvious at a glance, a tidy academy lets a new coach walk in and immediately know where drills, plans, and records live. Just as MLflow writes runs to a local store beside, not inside, the code, the scorers keep the logbook outside the coaching manual so results never overwrite strategy. The payoff: a clean layout means a stranger, or you in six months, can pick up the project and understand exactly how data becomes a model.

bash

# Create the AutoML project skeleton and install dependencies.
mkdir -p cricket-automl/{data,src,outputs}
cd cricket-automl

python -m venv .venv
source .venv/bin/activate        # on Windows: .venv\Scripts\activate

pip install optuna==3.* mlflow==2.* scikit-learn==1.* pandas==2.*

# Resulting structure:
# cricket-automl/
# |-- data/
# |   `-- innings.csv          # dataset (created by src/data.py)
# |-- src/
# |   |-- data.py              # Step 1: dataset + features
# |   |-- search.py            # Step 2: Optuna objective over model families
# |   `-- automl.py            # Step 3: run study, track, promote best
# |-- outputs/                 # best_model.pkl lands here
# `-- mlruns/                  # MLflow local tracking store

echo 'AutoML project skeleton ready.'

Step 1 — Foundation

Step 1 builds the dataset and a deterministic feature loader, the fixed base every trial will be evaluated against. You will generate a labelled cricket innings dataset and a function that returns features and labels reproducibly. The concept behind this step is fair comparison: an AutoML search only produces a meaningful winner if every candidate model is scored on identical, leak-free data, so the data layer must be deterministic and shared. By fixing the dataset and the split seed here, you guarantee that differences in trial scores reflect differences in models and hyperparameters, not random variation in the data each trial happened to see.

Analogy🏏Cricket

🏏 Think of it like cricket: before any trial match, the groundsman prepares one agreed pitch and the same set of balls, so every candidate is judged on identical conditions rather than a surface that changes underfoot. Just as the fixed pitch is your versioned dataset, the standard ball is your deterministic feature function. Just as changing the pitch mid-trial would make scores incomparable, a feature function with hidden randomness would make runs incomparable. Just as recording the pitch and ball used lets you recreate the conditions, recording the data version lets you recreate the inputs. The insight is that fair, repeatable comparison demands a fixed, documented starting surface.

python

# src/data.py  -- Step 1: deterministic dataset + feature loader.
import csv, os, random

SEED = 1983
DATA_VERSION = '[email protected]'

def generate_dataset(path='data/innings.csv', n=400):
    """Synthetic but reproducible: same seed -> same file."""
    rng = random.Random(SEED)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, 'w', newline='') as f:
        w = csv.writer(f)
        w.writerow(['batting_average', 'strike_rate', 'boundary_pct', 'match_winner'])
        for _ in range(n):
            avg = rng.uniform(10, 65)
            sr = rng.uniform(60, 160)
            bpct = rng.uniform(0.1, 0.75)
            # label: strong innings tend to win matches (with some noise)
            score = 0.03 * avg + 0.02 * sr + 1.5 * bpct + rng.gauss(0, 1)
            w.writerow([round(avg, 2), round(sr, 2), round(bpct, 3),
                        1 if score > 4.0 else 0])
    return path

def load_features(path='data/innings.csv'):
    X, y = [], []
    for r in csv.DictReader(open(path)):
        X.append([float(r['batting_average']), float(r['strike_rate']),
                  float(r['boundary_pct'])])
        y.append(int(r['match_winner']))
    return X, y

if __name__ == '__main__':
    generate_dataset()
    X, y = load_features()
    print(f'{DATA_VERSION}: {len(X)} innings, {sum(y)} winners')

Step 2 — Core Logic

Step 2 builds the heart of the AutoML pipeline: an Optuna objective that first chooses a model family and then tunes that family's hyperparameters, all within a single trial. This is where define-by-run shines, because a categorical suggestion for the model type branches into different parameter spaces depending on which model was picked, something a static grid cannot express cleanly. Each trial trains the chosen model with its suggested parameters, scores it with cross-validation on the fixed dataset, and returns that score for Optuna to optimise. This conditional search is what lets one study explore fundamentally different algorithms on equal footing rather than tuning a single predetermined model.

Analogy🏏Cricket

🏏 Think of it like cricket: when a candidate bats in a trial, an official scorer records the conditions they batted in, the pitch, the bowling, and their exact score, all under one innings entry. Just as the recorded conditions are your logged params, the score is your logged metric, and the innings entry is the MLflow run. Just as the same conditions and the same player should reproduce a comparable knock, the same params and seed reproduce the same accuracy. Just as a scorer who noted only the score but not the conditions leaves selectors guessing, logging a metric without its params leaves you unable to explain a result. The insight is that a result is only useful when recorded together with the conditions that produced it.

python

# src/search.py  -- Step 2: an Optuna objective over multiple model families.
import optuna
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score
from data import load_features, SEED

def build_model(trial):
    """Choose a family, then tune it -> conditional, define-by-run space."""
    family = trial.suggest_categorical('family', ['rf', 'gb', 'logreg'])
    if family == 'rf':
        return RandomForestClassifier(
            n_estimators=trial.suggest_int('rf_n_estimators', 50, 400, step=50),
            max_depth=trial.suggest_int('rf_max_depth', 2, 12),
            random_state=SEED)
    if family == 'gb':
        return GradientBoostingClassifier(
            n_estimators=trial.suggest_int('gb_n_estimators', 50, 400, step=50),
            learning_rate=trial.suggest_float('gb_lr', 1e-3, 0.3, log=True),
            max_depth=trial.suggest_int('gb_max_depth', 2, 6),
            random_state=SEED)
    return LogisticRegression(
        C=trial.suggest_float('lr_C', 1e-3, 100, log=True), max_iter=1000)

def objective(trial):
    X, y = load_features()
    model = build_model(trial)
    return cross_val_score(model, X, y, cv=3).mean()   # fair, fixed-data scoring

if __name__ == '__main__':
    study = optuna.create_study(direction='maximize',
                                sampler=optuna.samplers.TPESampler(seed=SEED))
    study.optimize(objective, n_trials=10)
    print('Best so far:', study.best_params, round(study.best_value, 3))

Step 3 — Integration & Enhancement

Step 3 integrates the search with tracking and promotion: you run the full study, log every trial to MLflow as a nested run, then retrain the best configuration on all the data and save it as the promoted model. This brings the pipeline together into something deployable, because a search that finds a great configuration is useless if the result is not recorded and the winning model is not persisted. By logging each trial under one parent run you get a complete, comparable tuning history in the MLflow UI, and by refitting and saving the best params you produce a single artifact ready to serve, with a clear record of exactly which search produced it.

Analogy🏏Cricket

🏏 Think of it like cricket: after a round of trial matches, the committee lays every candidate's recorded scores side by side, picks the top performer on the agreed metric, and formally names them to the squad while keeping the rest on record. Just as laying out all scorecards is your sweep across configs, choosing the highest scorer is your selection by logged accuracy. Just as the named player is promoted while others remain in the pool for recall, the best model is promoted while others stay tracked. Just as the selection memo cites the scores that justified it, your promotion cites the winning run ID. The insight is that evidence-based selection across recorded candidates produces a defensible, reversible choice.

python

# src/automl.py  -- Step 3: run study, track every trial, promote the best.
import pickle
import optuna, mlflow
from data import load_features, generate_dataset, SEED, DATA_VERSION
from search import build_model
from sklearn.model_selection import cross_val_score

mlflow.set_experiment('cricket-automl')

def tracked_objective(trial):
    X, y = load_features()
    model = build_model(trial)
    score = cross_val_score(model, X, y, cv=3).mean()
    with mlflow.start_run(nested=True, run_name=f'trial-{trial.number}'):
        mlflow.log_params(trial.params)         # family + its tuned params
        mlflow.log_param('data_version', DATA_VERSION)
        mlflow.log_metric('cv_accuracy', score)
    return score

def run_automl(n_trials=30):
    generate_dataset()
    with mlflow.start_run(run_name='automl-search') as parent:
        study = optuna.create_study(
            direction='maximize',
            sampler=optuna.samplers.TPESampler(seed=SEED),
            pruner=optuna.pruners.MedianPruner())
        study.optimize(tracked_objective, n_trials=n_trials)
        mlflow.log_metric('best_cv_accuracy', study.best_value)
        mlflow.log_params({f'best_{k}': v for k, v in study.best_params.items()})

        # Refit the winner on ALL data and promote it.
        X, y = load_features()
        best = build_model(optuna.trial.FixedTrial(study.best_params))
        best.fit(X, y)
        with open('outputs/best_model.pkl', 'wb') as f:
            pickle.dump(best, f)
        print(f'Best: {study.best_params}  acc={study.best_value:.3f}')
        print(f'Promoted -> outputs/best_model.pkl  (parent run {parent.info.run_id[:8]})')
    return study

if __name__ == '__main__':
    run_automl()

Step 4 — Testing & Verification

Verify two things: the search runs end to end and produces a tracked history with a promoted model, and the result is reproducible because the dataset, sampler, and model seeds are all fixed. Run the AutoML script, confirm a best configuration and saved model, then open the MLflow UI to inspect every trial nested under the parent run. Finally rerun the search and confirm the best parameters and score are identical, proving the seeds delivered reproducibility across the whole pipeline.

Analogy🏏Cricket

🏏 Think of it like cricket: before trusting a match plan you verify two things, that it actually works start to finish under real conditions and produces a clear winner, and that it is repeatable, run the same trial again in the same conditions and you get the same result rather than a fluke. Just as you confirm the pipeline runs end to end and produces a tracked leaderboard plus a promoted model, a coach confirms the full session runs from warm-up to selection and yields a ranked shortlist with a clear pick. Just as you rerun the same configuration and check it yields the same accuracy, a coach reruns the identical trial and expects the same player to top the table, proving it was skill, not chance. Just as the MLflow UI records every run with its params, metric, and artifact, the scorebook records every trial so the result can be inspected and defended. The payoff: verifying both correctness and reproducibility, in a pipeline or a trial, means you trust the winner because it holds up when repeated.

bash

# Run the AutoML search and verify reproducibility.
cd cricket-automl && source .venv/bin/activate

python src/automl.py        # generates data, runs study, promotes best

# Expected (illustrative) output:
# Best: {'family': 'gb', 'gb_n_estimators': 200, 'gb_lr': 0.043, 'gb_max_depth': 3} acc=0.812
# Promoted -> outputs/best_model.pkl  (parent run 7f3a9c1d)

# Inspect every trial nested under the parent run:
mlflow ui --port 5000       # open http://localhost:5000

# Reproducibility check: rerun and compare the best result.
python -c "from src.automl import run_automl; \
s1=run_automl(15); s2=run_automl(15); \
print('reproducible:', s1.best_params==s2.best_params, \
round(s1.best_value,3), round(s2.best_value,3))"
# Expected: reproducible: True 0.81x 0.81x

Warning: A common error is forgetting to seed one of the three randomness sources, the data generator, the Optuna sampler, or the model's random_state, and then being puzzled when two runs of the search disagree on the best configuration. AutoML compounds randomness across the dataset, the search, and each model fit, so every source must be seeded for the pipeline as a whole to be reproducible. Seed all three from a single SEED constant.

Extension Challenge: Add a fourth model family (such as a support vector classifier) and a held-out test set so the promoted model is evaluated on data the search never saw, reporting both the cross-validation score and the honest test score. For a harder stretch, persist the Optuna study to a SQLite database with load_if_exists so the search is resumable, and add a guard that refuses to overwrite outputs/best_model.pkl unless the new best beats the incumbent's recorded test score.

AutoML searches across model families and their hyperparameters together, finding a better answer than committing to one algorithm upfront.
Optuna's define-by-run objective lets a categorical model choice branch into different, conditional hyperparameter spaces within one trial.
Every candidate must be scored on the same fixed, leak-free data so trial differences reflect models and settings, not data variation.
Logging each trial as a nested MLflow run yields a complete, comparable tuning history under one parent run for inspection.
Promotion means refitting the best configuration on all data and persisting it as a single deployable artifact tied to the search that produced it.
Reproducible AutoML requires seeding every randomness source, the data generator, the sampler, and each model fit, from one shared seed.