mlflow_sklearn_model_training_logging_and_registry.py

1import mlflow.sklearn
2from mlflow.models import infer_signature
3from sklearn.ensemble import RandomForestRegressor
4from sklearn.metrics import mean_squared_error
5from sklearn.model_selection import train_test_split
6from sklearn import datasets
7
8# Prepare training data
9iris = datasets.load_iris()
10x = iris.data
11y = iris.target
12x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)
13
14# Train a model
15params = {"n_estimators": 100, "random_state": 42}
16rf = RandomForestRegressor(**params)
17rf.fit(x_train, y_train)
18
19# Infer the model signature
20signature = infer_signature(x_train, rf.predict(x_train))
21
22# Log the model and register it in the Model Registry
23with mlflow.start_run() as run:
24    # Log parameters and metrics
25    mlflow.log_params(params)
26    y_pred = rf.predict(x_test)
27    mlflow.log_metric("mse", mean_squared_error(y_test, y_pred))
28
29    # Log the sklearn model and register as "iris-classifier"
30    mlflow.sklearn.log_model(
31        sk_model=rf,
32        artifact_path="sklearn-model",
33        signature=signature,
34        registered_model_name="iris-classifier"
35    )
36
37print(f"Model registered. Run ID: {run.info.run_id}")