{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiment Tracking with Bailo & MLFlow"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to core concepts covered in previous notebooks, Bailo also offers integrations that might be useful within the wider machine learning lifecycle. This example notebook will run through **experiment tracking** in particular, integrating with **MLFlow Tracking**. The following concepts will be covered:\n",
    "\n",
    "* Creating a new experiment using a Bailo model.\n",
    "* Conducting experiment runs and logging parameters/metrics.\n",
    "* Importing existing experiments from **MLFlow Tracking**.\n",
    "* Publishing results to the Bailo service.\n",
    "\n",
    "Prerequisites:\n",
    "\n",
    "* Completion of the basic notebooks, in particular **models_and_releases_demo_pytorch.ipynb**.\n",
    "* Python 3.9 or higher (including a notebook environment for this demo).\n",
    "* A local or remote Bailo service (see https://github.com/gchq/Bailo).\n",
    "* Dependencies for MLFlow.\n",
    "* A local or remote MLFlow Tracking server, if following the MLFlow steps.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "### Connecting with Bailo\n",
    "\n",
    "In order to create an `Experiment()` object, you will first need to have a Bailo `Model()` object, and thus a defined `Client()` too. We learned how to do this in a previous notebook, but this time we will create a new model with a custom schema which supports model metrics. *More on that later...*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install dependencies...\n",
    "! pip install bailo[mlflow]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Necessary import statements\n",
    "\n",
    "from bailo import Model, Client, Experiment, Schema, SchemaKind\n",
    "import mlflow\n",
    "import random\n",
    "\n",
    "# Instantiating the Bailo client\n",
    "\n",
    "client = Client(\"http://127.0.0.1:8080\") # <- INSERT BAILO URL (if not hosting locally)\n",
    "\n",
    "# Creating a demo model\n",
    "\n",
    "model = Model.create(client=client, name=\"YOLOv5\", description=\"YOLOv5 model for object detection.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up MLFlow Tracking\n",
    "\n",
    "In order to complete the integration element of this tutorial, we will need to set up a local instance of MLFlow Tracking, and create a sample experiment run. *This will not contain any actual model training and is only to demonstrate the functionality of Bailo.*\n",
    "\n",
    "Run `mlflow ui --port 5050` on the command line. This will run on **localhost:5050** and the UI can be accessed on a browser."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing a custom schema for tracking\n",
    "\n",
    "The Bailo UI is designed to display metrics in a particular way, therefore you will need to use a schema that supports this. *This is necessary to publish results*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defines the schema in an external script as it is quite large!\n",
    "%run -i set_schema.py\n",
    "\n",
    "# Assigns a random schema ID\n",
    "schema_id = random.randint(1, 1000000)\n",
    "\n",
    "# Creates the schema on Bailo\n",
    "schema = Schema.create(client=client, schema_id=str(schema_id), name=\"Experiment Tracking\", description=\"Demo tracking schema\", kind=SchemaKind.MODEL, json_schema=json_schema)\n",
    "\n",
    "# Model cards need to be instantiated with their mandatory fields before metrics can be published.\n",
    "model.card_from_schema(schema_id=str(schema_id))\n",
    "new_card = {\n",
    "  'overview': {\n",
    "    'tags': [],\n",
    "    'modelSummary': 'YOLOv5 model for object detection.',\n",
    "  }\n",
    "}\n",
    "model.update_model_card(model_card=new_card)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a new experiment\n",
    "\n",
    "Experiments with the Bailo client are created using the `Model.create_experiment()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment = model.create_experiment()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conducting experiment runs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running an experiment with the Bailo python client\n",
    "\n",
    "You can run experiments directly using the Bailo python client as follows.\n",
    "\n",
    "**NOTE**: This will only work for sequential experiment runs, so if you're running experiments in parallel then it would be better to use **MLFlow Tracking**. We will learn how to import completed experiments from MLFlow in the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Arbitrary params\n",
    "params = {\n",
    "    \"lr\": 0.01,\n",
    "    \"anchor_t\": 4.0,\n",
    "    \"scale\": 0.5,\n",
    "}\n",
    "\n",
    "# Arbitrary metrics\n",
    "metrics = {\n",
    "    \"accuracy\": 0.98,\n",
    "}\n",
    "\n",
    "for x in range(5):\n",
    "    experiment.start_run()\n",
    "    experiment.log_params(params)\n",
    "    experiment.log_metrics(metrics)\n",
    "    experiment.log_artifacts([\"weights.txt\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a dummy MLFlow experiment run\n",
    "\n",
    "This section conducts an arbitrary experiment run and logs the params/metrics to your local MLFlow server. *We need this for the next section*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting local tracking URI and experiment name\n",
    "mlflow.set_tracking_uri(uri=\"http://127.0.0.1:5050\")\n",
    "mlflow.set_experiment(\"Demonstrator\")\n",
    "\n",
    "# Logging the same metrics to the local MLFlow server\n",
    "with mlflow.start_run():\n",
    "    mlflow.log_params(params)\n",
    "    mlflow.log_metric(\"accuracy\", 0.86)\n",
    "    mlflow.log_artifact(\"weights.txt\")\n",
    "    mlflow.set_tag(\"Training Info\", \"YOLOv5 Demo Model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing existing experiments from MLFlow into Bailo\n",
    "\n",
    "You can import existing experiments into the `Experiment()` class by using the `Experiment.from_mlflow()` method. You must provide the **MLFlow tracking URI** and the experiment ID. To get the experiment ID, go to the link provided in the cell \"Creating a dummy MLFlow experiment run\". In the details section you will see \"Experiment ID\", copy this ID and add it to the argument **experiment_id**. ![image.png](img/experiment_id.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_mlflow = model.create_experiment()\n",
    "experiment_mlflow.from_mlflow(\n",
    "    tracking_uri=\"http://127.0.0.1:5050\", experiment_id=\"471968520867705095\"\n",
    ")  # <- INSERT MLFLOW EXPERIMENT ID. CAN BE FOUND ON THE UI."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Publishing results to Bailo\n",
    "\n",
    "Experiment runs can be published to the model card using the `Experiment.publish()` method, **one at a time**. This is because the intended use is only to publish the most successful run.\n",
    "Therefore, you must specify the **run_id** to publish, or specify an order so the client can select the best result. As well as this, you must specify the location of the metrics in your schema (in this case *performance.performanceMetrics* as per the schema we defined earlier).\n",
    "\n",
    "Examples for both scenarios can be seen below."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Publishing a specific run\n",
    "\n",
    "To publish a specific run, you must pass the **run_id** into the method. In this example, use the run ID created in one of our previous runs. We get this similarly to getting the experiment ID, by either by following the link provided in the terminal output when we create our experiment run OR by going to the mlflow uri, going to our experiment tab, which in this example is, \"Demonstrator\". Then we select our chosen run, go to the details table, find \"run ID\" and copy the ID. We paste this ID in the `publish()` function, as our **run_id**. ![image.png](img/run_id.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_mlflow.publish(mc_loc=\"performance.performanceMetrics\", run_id=\"c531286193a5499e803ae48cc898819b\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Publishing the best run\n",
    "\n",
    "To publish the best run, you must define what the best is for your use case. This can be done using the `select_by` parameter with a string e.g. `accuracy MIN|MAX`. Depending on the requirements, `accuracy` could be any metric you have defined in your experiment.\n",
    "\n",
    "In the below example, we will use `accuracy MAX` to publish the experiment run with the highest accuracy.\n",
    "\n",
    "If previous publish is run, this subsequent `accuracy MAX` call can not be made without error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_mlflow.publish(mc_loc=\"performance.performanceMetrics\", select_by=\"accuracy MAX\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If successful, our metrics should now be under the **Performance** tab of your model card on the UI! Additionally, our artifact will have been published as a new release (*this will have been done twice if you ran both the above steps*)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}