{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6df76a1f",
   "metadata": {},
   "source": [
    "# Using MLflow with Tune\n",
    "\n",
    "(tune-mlflow-ref)=\n",
    "\n",
    ":::{warning}\n",
    "If you are using these MLflow integrations with {ref}`ray-client`, it is recommended that you setup a\n",
    "remote Mlflow tracking server instead of one that is backed by the local filesystem.\n",
    ":::\n",
    "\n",
    "[MLflow](https://mlflow.org/) is an open source platform to manage the ML lifecycle, including experimentation,\n",
    "reproducibility, deployment, and a central model registry. It currently offers four components, including\n",
    "MLflow Tracking to record and query experiments, including code, data, config, and results.\n",
    "\n",
    "```{image} /images/mlflow.png\n",
    ":align: center\n",
    ":alt: MLflow\n",
    ":height: 80px\n",
    ":target: https://www.mlflow.org/\n",
    "```\n",
    "\n",
    "Ray Tune currently offers two lightweight integrations for MLflow Tracking.\n",
    "One is the {ref}`MLflowLoggerCallback <tune-mlflow-logger>`, which automatically logs\n",
    "metrics reported to Tune to the MLflow Tracking API.\n",
    "\n",
    "The other one is the {ref}`@mlflow_mixin <tune-mlflow-mixin>` decorator, which can be\n",
    "used with the function API. It automatically\n",
    "initializes the MLflow API with Tune's training information and creates a run for each Tune trial.\n",
    "Then within your training function, you can just use the\n",
    "MLflow like you would normally do, e.g. using `mlflow.log_metrics()` or even `mlflow.autolog()`\n",
    "to log to your training process.\n",
    "\n",
    "```{contents}\n",
    ":backlinks: none\n",
    ":local: true\n",
    "```\n",
    "\n",
    "## Running an MLflow Example\n",
    "\n",
    "In the following example we're going to use both of the above methods, namely the `MLflowLoggerCallback` and\n",
    "the `mlflow_mixin` decorator to log metrics.\n",
    "Let's start with a few crucial imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0e47339",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "import time\n",
    "\n",
    "import mlflow\n",
    "\n",
    "from ray import tune\n",
    "from ray.tune.integration.mlflow import MLflowLoggerCallback, mlflow_mixin"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, let's define an easy objective function (a Tune `Trainable`) that iteratively computes steps and evaluates\n",
    "intermediate scores that we report to Tune."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def evaluation_fn(step, width, height):\n",
    "    return (0.1 + width * step / 100) ** (-1) + height * 0.1\n",
    "\n",
    "\n",
    "def easy_objective(config):\n",
    "    width, height = config[\"width\"], config[\"height\"]\n",
    "\n",
    "    for step in range(config.get(\"steps\", 100)):\n",
    "        # Iterative training function - can be any arbitrary training procedure\n",
    "        intermediate_score = evaluation_fn(step, width, height)\n",
    "        # Feed the score back to Tune.\n",
    "        tune.report(iterations=step, mean_loss=intermediate_score)\n",
    "        time.sleep(0.1)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Given an MLFlow tracking URI, you can now simply use the `MLflowLoggerCallback` as a `callback` argument to\n",
    "your `tune.run()` call:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def tune_function(mlflow_tracking_uri, finish_fast=False):\n",
    "    tune.run(\n",
    "        easy_objective,\n",
    "        name=\"mlflow\",\n",
    "        num_samples=5,\n",
    "        callbacks=[\n",
    "            MLflowLoggerCallback(\n",
    "                tracking_uri=mlflow_tracking_uri,\n",
    "                experiment_name=\"example\",\n",
    "                save_artifact=True,\n",
    "            )\n",
    "        ],\n",
    "        config={\n",
    "            \"width\": tune.randint(10, 100),\n",
    "            \"height\": tune.randint(0, 100),\n",
    "            \"steps\": 5 if finish_fast else 100,\n",
    "        },\n",
    "    )"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "To use the `mlflow_mixin` decorator, you can simply decorate the objective function from earlier.\n",
    "Note that we also use `mlflow.log_metrics(...)` to log metrics to MLflow.\n",
    "Otherwise, the decorated version of our objective is identical to its original."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "@mlflow_mixin\n",
    "def decorated_easy_objective(config):\n",
    "    # Hyperparameters\n",
    "    width, height = config[\"width\"], config[\"height\"]\n",
    "\n",
    "    for step in range(config.get(\"steps\", 100)):\n",
    "        # Iterative training function - can be any arbitrary training procedure\n",
    "        intermediate_score = evaluation_fn(step, width, height)\n",
    "        # Log the metrics to mlflow\n",
    "        mlflow.log_metrics(dict(mean_loss=intermediate_score), step=step)\n",
    "        # Feed the score back to Tune.\n",
    "        tune.report(iterations=step, mean_loss=intermediate_score)\n",
    "        time.sleep(0.1)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "With this new objective function ready, you can now create a Tune run with it as follows:"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def tune_decorated(mlflow_tracking_uri, finish_fast=False):\n",
    "    # Set the experiment, or create a new one if does not exist yet.\n",
    "    mlflow.set_tracking_uri(mlflow_tracking_uri)\n",
    "    mlflow.set_experiment(experiment_name=\"mixin_example\")\n",
    "    tune.run(\n",
    "        decorated_easy_objective,\n",
    "        name=\"mlflow\",\n",
    "        num_samples=5,\n",
    "        config={\n",
    "            \"width\": tune.randint(10, 100),\n",
    "            \"height\": tune.randint(0, 100),\n",
    "            \"steps\": 5 if finish_fast else 100,\n",
    "            \"mlflow\": {\n",
    "                \"experiment_name\": \"mixin_example\",\n",
    "                \"tracking_uri\": mlflow.get_tracking_uri(),\n",
    "            },\n",
    "        },\n",
    "    )"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "If you hapen to have an MLFlow tracking URI, you can set it below in the `mlflow_tracking_uri` variable and set\n",
    "`smoke_test=False`.\n",
    "Otherwise, you can just run a quick test of the `tune_function` and `tune_decorated` functions without using MLflow."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "smoke_test = True\n",
    "\n",
    "if smoke_test:\n",
    "    mlflow_tracking_uri = os.path.join(tempfile.gettempdir(), \"mlruns\")\n",
    "else:\n",
    "    mlflow_tracking_uri = \"<MLFLOW_TRACKING_URI>\"\n",
    "\n",
    "tune_function(mlflow_tracking_uri, finish_fast=smoke_test)\n",
    "if not smoke_test:\n",
    "    df = mlflow.search_runs(\n",
    "        [mlflow.get_experiment_by_name(\"example\").experiment_id]\n",
    "    )\n",
    "    print(df)\n",
    "\n",
    "tune_decorated(mlflow_tracking_uri, finish_fast=smoke_test)\n",
    "if not smoke_test:\n",
    "    df = mlflow.search_runs(\n",
    "        [mlflow.get_experiment_by_name(\"mixin_example\").experiment_id]\n",
    "    )\n",
    "    print(df)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "id": "f0df0817",
   "metadata": {},
   "source": [
    "This completes our Tune and MLflow walk-through.\n",
    "In the following sections you can find more details on the API of the Tune-MLflow integration.\n",
    "\n",
    "## MLflow AutoLogging\n",
    "\n",
    "You can also check out {doc}`here </tune/examples/includes/mlflow_ptl_example>` for an example on how you can\n",
    "leverage MLflow auto-logging, in this case with Pytorch Lightning\n",
    "\n",
    "## MLflow Logger API\n",
    "\n",
    "(tune-mlflow-logger)=\n",
    "\n",
    "```{eval-rst}\n",
    ".. autoclass:: ray.tune.integration.mlflow.MLflowLoggerCallback\n",
    "   :noindex:\n",
    "```\n",
    "\n",
    "## MLflow Mixin API\n",
    "\n",
    "(tune-mlflow-mixin)=\n",
    "\n",
    "```{eval-rst}\n",
    ".. autofunction:: ray.tune.integration.mlflow.mlflow_mixin\n",
    "   :noindex:\n",
    "```\n",
    "\n",
    "## More MLflow Examples\n",
    "\n",
    "- {doc}`/tune/examples/includes/mlflow_ptl_example`: Example for using [MLflow](https://github.com/mlflow/mlflow/)\n",
    "  and [Pytorch Lightning](https://github.com/PyTorchLightning/pytorch-lightning) with Ray Tune."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}