02-experiment-tracking/model-registry.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MLflow's Model Registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlflow.tracking import MlflowClient\n",
    "\n",
    "\n",
    "MLFLOW_TRACKING_URI = \"sqlite:///mlflow.db\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interacting with the MLflow tracking server\n",
    "\n",
    "The `MlflowClient` object allows us to interact with...\n",
    "- an MLflow Tracking Server that creates and manages experiments and runs.\n",
    "- an MLflow Registry Server that creates and manages registered models and model versions. \n",
    "\n",
    "To instantiate it we need to pass a tracking URI and/or a registry URI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<Experiment: artifact_location='./mlruns/0', experiment_id='0', lifecycle_stage='active', name='Default', tags={}>,\n",
       " <Experiment: artifact_location='./mlruns/1', experiment_id='1', lifecycle_stage='active', name='nyc-taxi-experiment', tags={}>,\n",
       " <Experiment: artifact_location='./mlruns/3', experiment_id='3', lifecycle_stage='active', name='my-cool-experiment', tags={}>]"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client = MlflowClient(tracking_uri=MLFLOW_TRACKING_URI)\n",
    "\n",
    "client.list_experiments()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.create_experiment(name=\"my-cool-experiment\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check the latest versions for the experiment with id `1`..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlflow.entities import ViewType\n",
    "\n",
    "runs = client.search_runs(\n",
    "    experiment_ids='1',\n",
    "    filter_string=\"metrics.rmse < 7\",\n",
    "    run_view_type=ViewType.ACTIVE_ONLY,\n",
    "    max_results=5,\n",
    "    order_by=[\"metrics.rmse ASC\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "run id: 7db08e4f93af4ee1bcbce1d8a763e23a, rmse: 6.3040\n",
      "run id: a06a6b594fff409cb0d34e203b49f33f, rmse: 6.7423\n",
      "run id: b8904012c84343b5bf8ee72aa8f0f402, rmse: 6.9047\n",
      "run id: 54493fed643c4952be5232279e309053, rmse: 6.9213\n"
     ]
    }
   ],
   "source": [
    "for run in runs:\n",
    "    print(f\"run id: {run.info.run_id}, rmse: {run.data.metrics['rmse']:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interacting with the Model Registry\n",
    "\n",
    "In this section We will use the `MlflowClient` instance to:\n",
    "\n",
    "1. Register a new version for the experiment `nyc-taxi-regressor`\n",
    "2. Retrieve the latests versions of the model `nyc-taxi-regressor` and check that a new version `4` was created.\n",
    "3. Transition the version `4` to \"Staging\" and adding annotations to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Registered model 'nyc-taxi-regressor' already exists. Creating a new version of this model...\n",
      "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\n",
      "Created version '4' of model 'nyc-taxi-regressor'.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<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>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "run_id = \"b8904012c84343b5bf8ee72aa8f0f402\"\n",
    "model_uri = f\"runs:/{run_id}/model\"\n",
    "mlflow.register_model(model_uri=model_uri, name=\"nyc-taxi-regressor\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "version: 1, stage: Staging\n",
      "version: 2, stage: Production\n",
      "version: 4, stage: None\n"
     ]
    }
   ],
   "source": [
    "model_name = \"nyc-taxi-regressor\"\n",
    "latest_versions = client.get_latest_versions(name=model_name)\n",
    "\n",
    "for version in latest_versions:\n",
    "    print(f\"version: {version.version}, stage: {version.current_stage}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<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>"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_version = 4\n",
    "new_stage = \"Staging\"\n",
    "client.transition_model_version_stage(\n",
    "    name=model_name,\n",
    "    version=model_version,\n",
    "    stage=new_stage,\n",
    "    archive_existing_versions=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<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>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "date = datetime.today().date()\n",
    "client.update_model_version(\n",
    "    name=model_name,\n",
    "    version=model_version,\n",
    "    description=f\"The model version {model_version} was transitioned to {new_stage} on {date}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparing versions and selecting the new \"Production\" model\n",
    "\n",
    "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.\n",
    "\n",
    "These are the steps:\n",
    "\n",
    "1. Load the test dataset, which corresponds to the NYC Green Taxi data from the month of March 2021.\n",
    "2. Download the `DictVectorizer` that was fitted using the training data and saved to MLflow as an artifact, and load it with pickle.\n",
    "3. Preprocess the test set using the `DictVectorizer` so we can properly feed the regressors.\n",
    "4. Make predictions on the test set using the model versions that are currently in the \"Staging\" and \"Production\" stages, and compare their performance.\n",
    "5. Based on the results, update the \"Production\" model version accordingly.\n",
    "\n",
    "\n",
    "**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.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import mean_squared_error\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "def read_dataframe(filename):\n",
    "    df = pd.read_csv(filename)\n",
    "\n",
    "    df.lpep_dropoff_datetime = pd.to_datetime(df.lpep_dropoff_datetime)\n",
    "    df.lpep_pickup_datetime = pd.to_datetime(df.lpep_pickup_datetime)\n",
    "\n",
    "    df['duration'] = df.lpep_dropoff_datetime - df.lpep_pickup_datetime\n",
    "    df.duration = df.duration.apply(lambda td: td.total_seconds() / 60)\n",
    "\n",
    "    df = df[(df.duration >= 1) & (df.duration <= 60)]\n",
    "\n",
    "    categorical = ['PULocationID', 'DOLocationID']\n",
    "    df[categorical] = df[categorical].astype(str)\n",
    "    \n",
    "    return df\n",
    "\n",
    "\n",
    "def preprocess(df, dv):\n",
    "    df['PU_DO'] = df['PULocationID'] + '_' + df['DOLocationID']\n",
    "    categorical = ['PU_DO']\n",
    "    numerical = ['trip_distance']\n",
    "    train_dicts = df[categorical + numerical].to_dict(orient='records')\n",
    "    return dv.transform(train_dicts)\n",
    "\n",
    "\n",
    "def test_model(name, stage, X_test, y_test):\n",
    "    model = mlflow.pyfunc.load_model(f\"models:/{name}/{stage}\")\n",
    "    y_pred = model.predict(X_test)\n",
    "    return {\"rmse\": mean_squared_error(y_test, y_pred, squared=False)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/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.\n",
      "  df = pd.read_csv(filename)\n"
     ]
    }
   ],
   "source": [
    "df = read_dataframe(\"data/green_tripdata_2021-03.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/Users/cristian.martinez/Repositories/mlops-zoomcamp/02-experiment-tracking/preprocessor'"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.download_artifacts(run_id=run_id, path='preprocessor', dst_path='.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open(\"preprocessor/preprocessor.b\", \"rb\") as f_in:\n",
    "    dv = pickle.load(f_in)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test = preprocess(df, dv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = \"duration\"\n",
    "y_test = df[target].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 139 ms, sys: 44.6 ms, total: 183 ms\n",
      "Wall time: 447 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'rmse': 6.659623830022514}"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time test_model(name=model_name, stage=\"Production\", X_test=X_test, y_test=y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.94 s, sys: 216 ms, total: 7.16 s\n",
      "Wall time: 7.28 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'rmse': 6.881555517147188}"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time test_model(name=model_name, stage=\"Staging\", X_test=X_test, y_test=y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<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>"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.transition_model_version_stage(\n",
    "    name=model_name,\n",
    "    version=4,\n",
    "    stage=\"Production\",\n",
    "    archive_existing_versions=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "0848c9d6c7d415ad6c477ff7ff8e98694d1a4aa96d0deee89244642e6b630036"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 ('exp-tracking-env')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}