Tracking experimentation in ML flow via Notebooks

The notebook trains a machine learning model on the diabetes dataset while using MLflow to track experiments and manage the model lifecycle through registration and logging.
Here’s a step-by-step explanation of what it does:

1. Importing Libraries

  • MLflow: A popular open-source platform to manage the ML lifecycle, including experimentation, reproducibility, deployment, and a central model registry.

  • Scikit-learn Modules:

    • train_test_split to split data into training and test sets.

    • load_diabetes to load the diabetes dataset.

    • RandomForestRegressor to build a regression model using a random forest algorithm.

import mlflow
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_diabetes
from sklearn.ensemble import RandomForestRegressor

2. MLflow Experiment Setup and Autologging

  • Setting the Experiment:
    mlflow.set_experiment("UAT Test") groups all runs under an experiment named "UAT Test".

  • Autologging:
    The mlflow.autolog() call configures MLflow to automatically log parameters, metrics, and the model artifact during the training process. This means you don't have to manually specify logging for these components.

mlflow.set_experiment("UAT Test")
mlflow.autolog()

3. Starting an MLflow Run

To ensure that logging and model registration occur within an active run, the notebook explicitly starts an MLflow run using a context manager (which was added to prevent the AttributeError seen when no run is active):

with mlflow.start_run() as run:

4. Data Loading and Splitting

  • Data Loading:
    The notebook loads the diabetes dataset using load_diabetes().

  • Splitting Data:
    The dataset is divided into training and testing sets with train_test_split(), ensuring that the model can later be evaluated on unseen data.

 db = load_diabetes()
 X_train, X_test, y_train, y_test = train_test_split(db.data, db.target)

5. Model Initialization, Training, and Prediction

  • Model Configuration:
    A  RandomForestRegressor it is created with 100 trees (n_estimators=100), a maximum depth of 6 for each tree (max_depth=6), and only 3 features considered at each split (max_features=3).

  • Training the Model:
    The model is fit to the training data (X_train and y_train).

  • Making Predictions:
    After training, predictions are generated on the test set (X_test).

 rf = RandomForestRegressor(n_estimators=100, max_depth=6, max_features=3)
 rf.fit(X_train, y_train)
 predictions = rf.predict(X_test)
During this process, thanks to autologging, MLflow automatically records important parameters (e.g., number of estimators, depth), metrics (such as performance scores on the test data), and the trained model artifact.
 

6. Retrieving the Run ID and Registering the Model 

  • Obtaining the Run ID:
    Inside the active run context, the run's unique identifier is accessed via run.info.run_id. This ID is essential for referencing the run later when registering the model.

  • Model Registration:
    The model artifact produced by the run is registered in MLflow’s Model Registry. The model URI is constructed using the run ID (e.g., "runs:/<run_id>/model"), and the model is registered under the name "DiabetesRandomForest-2025". Registration allows for organized version control and easy deployment of the model.

     

run_id = run.info.run_id # Ensured to be non-null within the active run
print(f"Active run ID: {run_id}")

model_uri = f"runs:/{run_id}/model"
model_name = "DiabetesRandomForest-2025"
mlflow.register_model(model_uri, model_name)
 

Below, you can track the experimentation in MLflow and look at articles on how to get started that are outside this article's scope. 

Conclusion

This notebook provides a complete end-to-end solution for training a regression model on the diabetes dataset. It sets up an MLflow experiment with autologging, ensures an active run for proper tracking, and performs data splitting, model training, and prediction. Finally, it leverages MLflow’s Model Registry for streamlined versioning and deployment, demonstrating a robust framework for reproducible, production-ready machine learning workflows.

 

The best way to understand and learn how to perform this function is through hands-on experience. Follow the steps below to create the sample notebook in your Syntasa environment:

  1. Download the sample notebook .ipynb file from this article.
  2. Create a new notebook in your Syntasa environment using the import notebook option.