Chapter 03
MLflow's Model Registry
NotebookPython 3.9.12 ('exp-tracking-env')31 cells
MLflow's Model Registry
In [44]python · cell 2
python
from mlflow.tracking import MlflowClient
MLFLOW_TRACKING_URI = "sqlite:///mlflow.db"Interacting with the MLflow tracking server
The MlflowClient object allows us to interact with...
- an MLflow Tracking Server that creates and manages experiments and runs.
- an MLflow Registry Server that creates and manages registered models and model versions.
To instantiate it we need to pass a tracking URI and/or a registry URI
In [45]python · cell 4
python
client = MlflowClient(tracking_uri=MLFLOW_TRACKING_URI)
client.list_experiments()Output
[<Experiment: artifact_location='./mlruns/0', experiment_id='0', lifecycle_stage='active', name='Default', tags={}>,
<Experiment: artifact_location='./mlruns/1', experiment_id='1', lifecycle_stage='active', name='nyc-taxi-experiment', tags={}>,
<Experiment: artifact_location='./mlruns/3', experiment_id='3', lifecycle_stage='active', name='my-cool-experiment', tags={}>]In [11]python · cell 5
python
client.create_experiment(name="my-cool-experiment")Output
'3'
Let's check the latest versions for the experiment with id 1...
In [19]python · cell 7
python
from mlflow.entities import ViewType
runs = client.search_runs(
experiment_ids='1',
filter_string="metrics.rmse < 7",
run_view_type=ViewType.ACTIVE_ONLY,
max_results=5,
order_by=["metrics.rmse ASC"]
)In [20]python · cell 8
python
for run in runs:
print(f"run id: {run.info.run_id}, rmse: {run.data.metrics['rmse']:.4f}")Output
run id: 7db08e4f93af4ee1bcbce1d8a763e23a, rmse: 6.3040 run id: a06a6b594fff409cb0d34e203b49f33f, rmse: 6.7423 run id: b8904012c84343b5bf8ee72aa8f0f402, rmse: 6.9047 run id: 54493fed643c4952be5232279e309053, rmse: 6.9213
In [ ]python · cell 9
python
Interacting with the Model Registry
In this section We will use the MlflowClient instance to:
- Register a new version for the experiment
nyc-taxi-regressor - Retrieve the latests versions of the model
nyc-taxi-regressorand check that a new version4was created. - Transition the version
4to "Staging" and adding annotations to it.
In [46]python · cell 11
python
import mlflow
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)In [22]python · cell 12
python
run_id = "b8904012c84343b5bf8ee72aa8f0f402"
model_uri = f"runs:/{run_id}/model"
mlflow.register_model(model_uri=model_uri, name="nyc-taxi-regressor")Output
Registered model 'nyc-taxi-regressor' already exists. Creating a new version of this model... 2022/05/19 16:47:17 INFO mlflow.tracking._model_registry.client: Waiting up to 300 seconds for model version to finish creation. Model name: nyc-taxi-regressor, version 4 Created version '4' of model 'nyc-taxi-regressor'.
<ModelVersion: creation_timestamp=1652971637398, current_stage='None', description=None, last_updated_timestamp=1652971637398, name='nyc-taxi-regressor', run_id='b8904012c84343b5bf8ee72aa8f0f402', run_link=None, source='./mlruns/1/b8904012c84343b5bf8ee72aa8f0f402/artifacts/model', status='READY', status_message=None, tags={}, user_id=None, version=4>In [ ]python · cell 13
python
In [26]python · cell 14
python
model_name = "nyc-taxi-regressor"
latest_versions = client.get_latest_versions(name=model_name)
for version in latest_versions:
print(f"version: {version.version}, stage: {version.current_stage}")Output
version: 1, stage: Staging version: 2, stage: Production version: 4, stage: None
In [32]python · cell 15
python
model_version = 4
new_stage = "Staging"
client.transition_model_version_stage(
name=model_name,
version=model_version,
stage=new_stage,
archive_existing_versions=False
)Output
<ModelVersion: creation_timestamp=1652971637398, current_stage='Staging', description='The model version 4 was transitioned to Staging on 2022-05-19', last_updated_timestamp=1652972141519, name='nyc-taxi-regressor', run_id='b8904012c84343b5bf8ee72aa8f0f402', run_link=None, source='./mlruns/1/b8904012c84343b5bf8ee72aa8f0f402/artifacts/model', status='READY', status_message=None, tags={}, user_id=None, version=4>In [33]python · cell 16
python
from datetime import datetime
date = datetime.today().date()
client.update_model_version(
name=model_name,
version=model_version,
description=f"The model version {model_version} was transitioned to {new_stage} on {date}"
)Output
<ModelVersion: creation_timestamp=1652971637398, current_stage='Staging', description='The model version 4 was transitioned to Staging on 2022-05-19', last_updated_timestamp=1652972142779, name='nyc-taxi-regressor', run_id='b8904012c84343b5bf8ee72aa8f0f402', run_link=None, source='./mlruns/1/b8904012c84343b5bf8ee72aa8f0f402/artifacts/model', status='READY', status_message=None, tags={}, user_id=None, version=4>In [ ]python · cell 17
python
Comparing versions and selecting the new "Production" model
In the last section, we will retrieve models registered in the model registry and compare their performance on an unseen test set. The idea is to simulate the scenario in which a deployment engineer has to interact with the model registry to decide whether to update the model version that is in production or not.
These are the steps:
- Load the test dataset, which corresponds to the NYC Green Taxi data from the month of March 2021.
- Download the
DictVectorizerthat was fitted using the training data and saved to MLflow as an artifact, and load it with pickle. - Preprocess the test set using the
DictVectorizerso we can properly feed the regressors. - Make predictions on the test set using the model versions that are currently in the "Staging" and "Production" stages, and compare their performance.
- Based on the results, update the "Production" model version accordingly.
Note: the model registry doesn't actually deploy the model to production when you transition a model to the "Production" stage, it just assign a label to that model version. You should complement the registry with some CI/CD code that does the actual deployment.
In [34]python · cell 19
python
from sklearn.metrics import mean_squared_error
import pandas as pd
def read_dataframe(filename):
df = pd.read_csv(filename)
df.lpep_dropoff_datetime = pd.to_datetime(df.lpep_dropoff_datetime)
df.lpep_pickup_datetime = pd.to_datetime(df.lpep_pickup_datetime)
df['duration'] = df.lpep_dropoff_datetime - df.lpep_pickup_datetime
df.duration = df.duration.apply(lambda td: td.total_seconds() / 60)
df = df[(df.duration >= 1) & (df.duration <= 60)]
categorical = ['PULocationID', 'DOLocationID']
df[categorical] = df[categorical].astype(str)
return df
def preprocess(df, dv):
df['PU_DO'] = df['PULocationID'] + '_' + df['DOLocationID']
categorical = ['PU_DO']
numerical = ['trip_distance']
train_dicts = df[categorical + numerical].to_dict(orient='records')
return dv.transform(train_dicts)
def test_model(name, stage, X_test, y_test):
model = mlflow.pyfunc.load_model(f"models:/{name}/{stage}")
y_pred = model.predict(X_test)
return {"rmse": mean_squared_error(y_test, y_pred, squared=False)}In [ ]python · cell 20
python
In [36]python · cell 21
python
df = read_dataframe("data/green_tripdata_2021-03.csv")Output
/var/folders/42/f9s_rgk15078ym2w50_xtc180000gq/T/ipykernel_5486/3050441246.py:6: DtypeWarning: Columns (3) have mixed types. Specify dtype option on import or set low_memory=False. df = pd.read_csv(filename)
In [37]python · cell 22
python
client.download_artifacts(run_id=run_id, path='preprocessor', dst_path='.')Output
'/Users/cristian.martinez/Repositories/mlops-zoomcamp/02-experiment-tracking/preprocessor'
In [38]python · cell 23
python
import pickle
with open("preprocessor/preprocessor.b", "rb") as f_in:
dv = pickle.load(f_in)In [39]python · cell 24
python
X_test = preprocess(df, dv)In [40]python · cell 25
python
target = "duration"
y_test = df[target].valuesIn [ ]python · cell 26
python
In [41]python · cell 27
python
%time test_model(name=model_name, stage="Production", X_test=X_test, y_test=y_test)Output
CPU times: user 139 ms, sys: 44.6 ms, total: 183 ms Wall time: 447 ms
{'rmse': 6.659623830022514}In [42]python · cell 28
python
%time test_model(name=model_name, stage="Staging", X_test=X_test, y_test=y_test)Output
CPU times: user 6.94 s, sys: 216 ms, total: 7.16 s Wall time: 7.28 s
{'rmse': 6.881555517147188}In [43]python · cell 29
python
client.transition_model_version_stage(
name=model_name,
version=4,
stage="Production",
archive_existing_versions=True
)Output
<ModelVersion: creation_timestamp=1652971637398, current_stage='Production', description='The model version 4 was transitioned to Staging on 2022-05-19', last_updated_timestamp=1652972763255, name='nyc-taxi-regressor', run_id='b8904012c84343b5bf8ee72aa8f0f402', run_link=None, source='./mlruns/1/b8904012c84343b5bf8ee72aa8f0f402/artifacts/model', status='READY', status_message=None, tags={}, user_id=None, version=4>In [ ]python · cell 30
python
In [ ]python · cell 31
python