2-1_Experiment_Tracking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 2.1 - Experiment Tracking with MLflow / W&B\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/zenml-io/zenbytes/blob/main/2-1_Experiment_Tracking.ipynb)\n",
    "\n",
    "***Key Concepts:*** *Experiment Tracking, MLflow, Weights & Biases*\n",
    "\n",
    "When training models, you often need to run hundreds of experiments with different types of models and different hyperparameters to find what works best. Keeping track of every experiment and how each design decision affected the model performance is hardly feasible without additional tools. That is why experiment trackers like [TensorBoard](https://www.tensorflow.org/tensorboard), [Weights & Biases](https://wandb.ai/site), or [MLflow](https://mlflow.org/) are often one of the first touchpoints ML practitioners have with MLOps as they progress through the ML journey. In addition, these tools are invaluable for larger ML teams, as they allow them to share experiment results and collaborate during experimentation.\n",
    "\n",
    "Since there are many excellent experiment tracking tools, we should aim to prevent vendor lock-in by writing modular ML code that allows us to switch between different tools easily. That is precisely what ZenML will do for us.\n",
    "\n",
    "This lesson will focus on how to effectively track experiments with [MLflow](https://mlflow.org/), which is one of the most popular open-source MLOps tools and used by many ML teams in practice.\n",
    "\n",
    "For research-heavy ML teams, we have also included a short bonus section at the end on how to use [Weights & Biases](https://wandb.ai/site) instead of MLflow in your ZenML pipelines.\n",
    "\n",
    "Please run the following commands to install both tools with their respective dependencies. This will also restart the kernel of your notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install \"zenml[server]\"\n",
    "!zenml integration install sklearn mlflow wandb -y\n",
    "!rm -rf .zen\n",
    "!zenml init\n",
    "%pip install pyparsing==2.4.2  # required for Colab\n",
    "\n",
    "import IPython\n",
    "\n",
    "# automatically restart kernel\n",
    "IPython.Application.instance().kernel.do_shutdown(restart=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Colab Note:** On Colab, you need an [ngrok account](https://dashboard.ngrok.com/signup) to view some of the visualizations later. Please set up an account, then set your user token below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NGROK_TOKEN = \"\"  # TODO: set your ngrok token if you are working on Colab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from zenml.environment import Environment\n",
    "\n",
    "if Environment.in_google_colab():  # Colab only setup\n",
    "\n",
    "    IN_COLAB = True\n",
    "\n",
    "    # clone zenbytes repo to get source code of previous lessons\n",
    "    !git clone https://github.com/zenml-io/zenbytes.git  # noqa\n",
    "    !mv zenbytes/steps .\n",
    "    !mv zenbytes/pipelines .\n",
    "\n",
    "    # install ngrok and expose port 8080\n",
    "    !pip install pyngrok\n",
    "    !ngrok authtoken {NGROK_TOKEN}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's import our pipeline definition and some of the pipeline steps that we built in the previous lessons:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pipelines.digits_pipeline import digits_pipeline\n",
    "from steps.evaluator import evaluator\n",
    "from steps.importer import importer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Register an MLFlow Experiment Tracker\n",
    "\n",
    "[MLflow](https://mlflow.org/) is an amazing open-source MLOps platform that provides powerful tools to handle various ML lifecycle steps, such as experiment tracking, code packaging, model deployment, and more. We will focus on the [MLflow Tracking](https://mlflow.org/docs/latest/tracking.html) component in this lesson, but we will learn about other MLflow features in later lessons.\n",
    "\n",
    "To use MLFlow in a ZenML pipeline, we first need to add MLflow into our ZenML MLOps stack. \n",
    "We do this by registering a new experiment tracker with ZenML that we then add into our current stack:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Register the MLflow experiment tracker\n",
    "!zenml experiment-tracker register mlflow_tracker --flavor=mlflow\n",
    "\n",
    "# Create a new stack that includes an MLflow experiment\n",
    "!zenml stack register mlflow_exp_tracker_stack -a default -o default -e mlflow_tracker\n",
    "\n",
    "# Set the new stack as active\n",
    "!zenml stack set mlflow_exp_tracker_stack\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also turn off MLflow warnings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from absl import logging as absl_logging\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "absl_logging.set_verbosity(-10000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use MLFlow in a ZenML Pipeline\n",
    "\n",
    "To integrate the MLFlow experiment tracker into our previously defined ZenML pipeline, we only need to adjust the `svc_trainer` step. Let us define a new `svc_trainer_mlflow` step in which we use MLflow's [`mlflow.sklearn.autolog()`](https://www.mlflow.org/docs/latest/python_api/mlflow.sklearn.html#mlflow.sklearn.autolog) feature to automatically log all relevant attributes and metrics of our model to MLflow. \n",
    "\n",
    "By adding an `experiment_tracker=mlflow_tracker` parameter in the `@step` decorator, ZenML automatically takes care of initializing MLflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "import numpy as np\n",
    "from sklearn.base import ClassifierMixin\n",
    "from sklearn.svm import SVC\n",
    "from zenml.steps import step\n",
    "\n",
    "@step(enable_cache=False, experiment_tracker=\"mlflow_tracker\")\n",
    "def svc_trainer_mlflow(\n",
    "    X_train: np.ndarray,\n",
    "    y_train: np.ndarray,\n",
    ") -> ClassifierMixin:\n",
    "    \"\"\"Train an sklearn SVC classifier and log to MLflow.\"\"\"\n",
    "    mlflow.sklearn.autolog()  # log all model hparams and metrics to MLflow\n",
    "    model = SVC(gamma=1e-3)\n",
    "    model.fit(X_train, y_train)\n",
    "    return model\n",
    "\n",
    "svc_mlflow_pipeline = digits_pipeline(\n",
    "    importer=importer(),\n",
    "    trainer=svc_trainer_mlflow(),\n",
    "    evaluator=evaluator(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's do the same for our decision tree trainer step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.tree import DecisionTreeClassifier\n",
    "\n",
    "@step(enable_cache=False, experiment_tracker=\"mlflow_tracker\")\n",
    "def tree_trainer_with_mlflow(\n",
    "    X_train: np.ndarray,\n",
    "    y_train: np.ndarray,\n",
    ") -> ClassifierMixin:\n",
    "    \"\"\"Train an sklearn decision tree classifier and log to MLflow.\"\"\"\n",
    "    mlflow.sklearn.autolog()  # log all model hparams and metrics to MLflow\n",
    "    model = DecisionTreeClassifier()\n",
    "    model.fit(X_train, y_train)\n",
    "    return model\n",
    "\n",
    "tree_mlflow_pipeline = digits_pipeline(\n",
    "    importer=importer(),\n",
    "    trainer=tree_trainer_with_mlflow(),\n",
    "    evaluator=evaluator(),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that's it, we're all set up! Now all `pipeline.run()` calls will automatically log all hyperparameters and metrics to MLflow. Let's try it out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "svc_mlflow_pipeline.run(unlisted=True)\n",
    "tree_mlflow_pipeline.run(unlisted=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compare our runs within the MLflow UI, run the following cell, then open http://127.0.0.1:4997/ in your browser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This will start a serving process for mlflow\n",
    "#  - if you want to continue in the notebook you need to manually\n",
    "#  interrupt the kernel\n",
    "\n",
    "from zenml.environment import Environment\n",
    "from zenml.integrations.mlflow.mlflow_utils import get_tracking_uri\n",
    "\n",
    "\n",
    "def open_mlflow_ui(port=4997):\n",
    "\n",
    "    if Environment.in_google_colab():  # Colab only setup\n",
    "        from pyngrok import ngrok\n",
    "\n",
    "        public_url = ngrok.connect(port)\n",
    "        print(f\"\\x1b[31mIn Colab, use this URL instead: {public_url}!\\x1b[0m\")\n",
    "\n",
    "    !mlflow ui --backend-store-uri=\"{get_tracking_uri()}\" --port={port}\n",
    "\n",
    "\n",
    "open_mlflow_ui()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first thing you will see is an overview of all your runs, as shown below:\n",
    "\n",
    "![MLflow UI](_assets/2-1/mlflow_ui.png)\n",
    "\n",
    "Click on the `Parameters >` tab on top of the table to see *all* hyperparameters of your model. Now you can see at a glance which model performed best and which hyperparameters changed between different runs. In our case, we can see that the SVC model with `gamma=0.001` achieved the best test accuracy of `0.969`.\n",
    "\n",
    "If we click on one of the links in the `Start Time` column, we can see additional details of the respective run. In particular, we can find a `model.pkl` file under the `Artifacts` tab, which we could now use to deploy our model in an inference/production environment. In the next lesson, `2-2_Local_Deployment.ipynb`, we will learn how to do this automatically as part of our pipelines with the [MLflow Models](https://mlflow.org/docs/latest/models.html)  component.\n",
    "\n",
    "If you would like to inspect the MLflow logs of your runs manually, you can find\n",
    "the logging location using the `experiment_tracker_url` metadata field of the \n",
    "trainer step of your pipeline run, e.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from zenml.post_execution import get_unlisted_runs\n",
    "\n",
    "pipeline_run = get_unlisted_runs()[0]\n",
    "step = pipeline_run.get_step(\"trainer\")\n",
    "experiment_tracker_url = step.metadata[\"experiment_tracker_url\"].value\n",
    "\n",
    "print(experiment_tracker_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Alternative Tool: Weights & Biases\n",
    "\n",
    "Of course, MLflow is not the only tool you can use for experiment tracking. The following example will show how we could achieve the same with another experiment tracking tool: [Weights & Biases](https://wandb.ai/site).\n",
    "This example assumes you already have some familiarity with W&B. In particular, you need a Weights & Biases account, which you can set up for free [here](https://wandb.ai/login?signup=true).\n",
    "\n",
    "You then need to define the three variables below to authorize yourself in W&B and to tell ZenML which entity/project you want to log to:\n",
    "- `WANDB_API_KEY`: your API key, which you can retrieve at [https://wandb.ai/authorize](https://wandb.ai/authorize). Make sure never to share this key (in particular, make sure to remove the key before pushing this notebook to any public Git repositories!).\n",
    "- `WANDB_ENTITY`: the entity (team or user) that owns the project you want to log to. If you are using W&B alone, just use your username here.\n",
    "- `WANDB_PROJECT`: the name of the W&B project you want to log to. If you have never used W&B before or want to start a new project, simply type the new project name here, e.g., \"zenbytes\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WANDB_API_KEY = None  # TODO: replace this with your W&B API key\n",
    "WANDB_ENTITY = None  # TODO: replace this with your W&B entity name\n",
    "WANDB_PROJECT = \"zenbytes\"  # TODO: replace this with your W&B project name (if you want to log to a specific project)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Register the W&B experiment tracker\n",
    "!zenml experiment-tracker register wandb_tracker --flavor=wandb --api_key={WANDB_API_KEY} --entity={WANDB_ENTITY} --project_name={WANDB_PROJECT}\n",
    "\n",
    "# Create a new stack that includes a W&B experiment tracker\n",
    "!zenml stack register wandb_stack -a default -o default -e wandb_tracker\n",
    "\n",
    "# Set the new stack as active\n",
    "!zenml stack set wandb_stack"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we need to add wandb to our `svc_trainer` step and use that to initialize our ZenML pipeline. The overall setup is the same as for the MLflow example above: we simply add an `@enable_wandb` decorator to our step, and then we can use `wandb` functionality within the step as we wish.\n",
    "\n",
    "The main difference to the MLflow example before is that W&B has no sklearn autolog functionality. Instead, we need to call `wandb.log(...)` for each value we want to log to Weights & Biases.\n",
    "\n",
    "Since we also want to log our test score, we need to adjust our `evaluator` step accordingly as well.\n",
    "\n",
    "Note that, despite wandb being used in different steps within a pipeline, ZenML handles initializing wandb and ensures that the experiment name is the same as the pipeline name and that the experiment run name is the same as the pipeline run name. This establishes a lineage between pipelines in ZenML and experiments in wandb."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.base import ClassifierMixin\n",
    "from sklearn.svm import SVC\n",
    "import wandb\n",
    "from zenml.steps import step\n",
    "\n",
    "from pipelines.digits_pipeline import digits_pipeline\n",
    "from steps.importer import importer\n",
    "\n",
    "\n",
    "@step(enable_cache=False, experiment_tracker=\"wandb_tracker\")\n",
    "def svc_trainer_wandb(\n",
    "    X_train: np.ndarray,\n",
    "    y_train: np.ndarray,\n",
    ") -> ClassifierMixin:\n",
    "    \"\"\"Train an sklearn SVC classifier and log to W&B.\"\"\"\n",
    "    gamma = 1e-3\n",
    "    wandb.log({\"gamma\": gamma})  # log gamma hparam to wandb\n",
    "    model = SVC(gamma=gamma)\n",
    "    model.fit(X_train, y_train)\n",
    "    return model\n",
    "\n",
    "@step(enable_cache=False, experiment_tracker=\"wandb_tracker\")\n",
    "def evaluator_wandb(\n",
    "    X_test: np.ndarray,\n",
    "    y_test: np.ndarray,\n",
    "    model: ClassifierMixin,\n",
    ") -> float:\n",
    "    \"\"\"Calculate the accuracy on the test set and log to W&B.\"\"\"\n",
    "    test_acc = model.score(X_test, y_test)\n",
    "    wandb.log({\"test acc\": test_acc})  # log test_acc to wandb\n",
    "    print(f\"Test accuracy: {test_acc}\")\n",
    "    return test_acc\n",
    "\n",
    "svc_wandb_pipeline = digits_pipeline(\n",
    "    importer=importer(),\n",
    "    trainer=svc_trainer_wandb(),\n",
    "    evaluator=evaluator_wandb(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, execute the cell below to run your pipeline with different gamma values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "svc_wandb_pipeline.run(unlisted=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the cell below and follow the link to see the run in your Weights & Biases \n",
    "project:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_step = get_unlisted_runs()[0].get_step(\"trainer\")\n",
    "experiment_tracker_url = trainer_step.metadata[\"experiment_tracker_url\"].value\n",
    "print(experiment_tracker_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a more detailed example of how to use Weights & Biases experiment tracking in your ZenML pipeline, see the [ZenML wandb_tracking example](https://github.com/zenml-io/zenml/tree/main/examples/wandb_tracking)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the [next lesson](2-2_Local_Deployment.ipynb), we will learn how to deploy models locally with MLflow."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "zenml310",
   "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.10.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "569b3361e3ec4d7692543ddda480ca8173a6c158bb706498f2e35ca1687a80ea"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}