sdk/python/using-mlflow/model-management/model_management.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model management with MLflow\n",
    "\n",
    "Model management can be done using both MLflow and Azure ML SDK/CLI v2. If you are familiar with MLflow and the capabilities it exposes, we support the entire model lifecycle using the MLFlow client. If you rather use Azure ML specific features or do model management using the CLI, in the same way you can manage the lifecycle using the Azure ML CLI/SDK v2."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites to run this notebook"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ensure you have the dependencies for this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r model_management.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Required imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "from sklearn import linear_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following notebook, we will explore an example that uses the following naming convention:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_name = \"model-management-examples\"\n",
    "model_name = \"mlflow-sample-model\"\n",
    "artifact_path = \"regressor\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before starting"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### If you are working in a Compute Instance in Azure Machine Learning\n",
    "\n",
    "If you are working in Azure Machine Learning Compute Instances, you MLflow installation is automatically connected to Azure Machine Learning, and you don't need to do anything."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### If you are working in your local machine, or in a cloud outside Azure Machine Learning\n",
    "\n",
    "You will need to connect MLflow to the Azure Machine Learning workspace you want to work on. MLflow uses the tracking URI to indicate the MLflow server you want to connect to. There are multiple ways to get the Azure Machine Learning MLflow Tracking URI. In this tutorial we will use the Azure ML SDK for Python, but you can check [Set up tracking environment - Azure Machine Learning Docs](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-use-mlflow-cli-runs#set-up-tracking-environment) for more alternatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace = \"<AML_WORKSPACE_NAME>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient\n",
    "from azure.identity import DefaultAzureCredential\n",
    "\n",
    "ml_client = MLClient(\n",
    "    DefaultAzureCredential(), subscription_id, resource_group, workspace\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use the workspace object to get the tracking URI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "azureml_tracking_uri = ml_client.workspaces.get(\n",
    "    ml_client.workspace_name\n",
    ").mlflow_tracking_uri\n",
    "mlflow.set_tracking_uri(azureml_tracking_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> To get the URI, please navigate to Azure ML Studio and select the workspace you are working on > Click on the name of the workspace at the upper right corner of the page > Click “View all properties in Azure Portal” on the pane popup > Copy the MLflow tracking URI value from the properties section."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating models from an existing run\n",
    "\n",
    "If you have an Mlflow model logged inside of a run and you want to register it in a registry, you can do that by using the experiment and run ID information from the run. Let's create a simple experiment and run to demonstrate it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.set_experiment(experiment_name=experiment_name)\n",
    "with mlflow.start_run(run_name=\"sample-run\"):\n",
    "    reg = linear_model.LinearRegression()\n",
    "    reg.fit([[0, 0], [1, 1], [2, 2]], [0, 1, 2])\n",
    "\n",
    "    mlflow.sklearn.log_model(reg, artifact_path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's search for the last run of the experiment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp = mlflow.get_experiment_by_name(experiment_name)\n",
    "last_run = mlflow.search_runs(exp.experiment_id, output_format=\"list\")[-1]\n",
    "print(last_run.info.run_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have the run identified, we can register the model using Mlflow client:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.register_model(f\"runs:/{last_run.info.run_id}/{artifact_path}\", model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating models from assets\n",
    "\n",
    "If you have a folder with an MLModel MLflow model, then you can register it directly. There is no need for the model to be always in the context of a run. To do that you can use the URI schema `file://path/to/model` to register it. Let's create a simple model and save it in MLModel format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reg = linear_model.LinearRegression()\n",
    "reg.fit([[0, 0], [1, 1], [2, 2]], [0, 1, 2])\n",
    "\n",
    "mlflow.sklearn.save_model(reg, artifact_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the files in the folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls regressor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now register the model from the local path:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.register_model(f\"file://{artifact_path}\", model_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Notice how the model URI schema `file://`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying models"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most MLflow model management operations require the MLflow client:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = mlflow.tracking.MlflowClient()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Querying all the models in the registry"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can query all the registered models in the registry using the MLflow client with the method `search_registered_models`.\n",
    "\n",
    "> __MLflow 2.0 advisory__: In older versions of Mlflow (<2.0), please use `list_registered_models` instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for model in client.search_registered_models():\n",
    "    print(f\"{model.name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are not sure of the name of the model you are looking for, you can search for it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.search_registered_models(f\"name='{model_name}'\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting model versions from registered models"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use `search_model_versions` to search for specific model versions across one or multiple registered models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.search_model_versions(f\"name='{model_name}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting specific versions of the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The command above will retrieve the model object which contains all the model versions. However, if you want to get the last registered model version of a given model, you can use `get_registered_model`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.get_registered_model(model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you need an specific version of the model, you can indicate so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.get_model_version(model_name, version=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model stages\n",
    "\n",
    "MLflow supports model's stages to manage model's lifecycle. Stage are assigned to model's version (instead of models). This means that a given model can have multiple versions on different stages."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Queying model stages\n",
    "\n",
    "You can use the MLflow client to check all the possible stages a model can be:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.get_model_version_stages(model_name, version=\"latest\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see what model version is on each stage by getting the model from the registry:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.get_latest_versions(model_name, stages=[\"Staging\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that multiple versions can be in the same stage at the same time in Mlflow, however, this method returns the latest version (greater version) among all of them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Caution: Notice that stages are case sensitive."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transitioning models\n",
    "\n",
    "To transition a model to a particular stage, you can:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.transition_model_version_stage(model_name, version=2, stage=\"Staging\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, if there were an existing model version in that particular stage, it will remain there. Hence, it won't be replaced. Alternatively, you can indicate `archive_existing_versions=True` to tell MLflow to move the existing model's version to the stage `Archived`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.transition_model_version_stage(\n",
    "    model_name, version=3, stage=\"Staging\", archive_existing_versions=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading models from stages\n",
    "\n",
    "You can load a model in a particular stage directly from Python using the `load_model` function and the following `URI` format. Notice that for this method to success, you need to have all the libraries and dependencies already installed in the environment you are working at."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = mlflow.pyfunc.load_model(f\"models:/{model_name}/Staging\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Editing and deleting models\n",
    "\n",
    "Editing registered models is supported in both Mlflow and Azure ML, however, there are some differences between them that are important to notice:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Editing models\n",
    "\n",
    "You can edit model's description and tags from a model using Mlflow:\n",
    "\n",
    "> Renaming models is not supported in Azure ML as model objects are immmutable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.update_model_version(\n",
    "    model_name, version=1, description=\"A sample regressor model\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To edit tags, you have to use the method `set_model_version_tag` and `remove_model_version_tag`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.set_model_version_tag(model_name, version=\"1\", key=\"type\", value=\"regression\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Removing a tag:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.delete_model_version_tag(model_name, version=\"1\", key=\"type\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deleting a model version\n",
    "\n",
    "You can delete any model version in the registry using the MLflow client. However, Azure ML doesn't support deleting the entire model container. To achieve the same thing, you will need to delete all the model versions from a given model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "last_version = client.search_model_versions(filter_string=f\"name='{model_name}'\")[\n",
    "    0\n",
    "].version\n",
    "print(f\"Last model version is {last_version}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.delete_model_version(model_name, version=last_version)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10 - SDK V2",
   "language": "python",
   "name": "python310-sdkv2"
  },
  "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}