08-RaySGD-MLflow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "saving-california",
   "metadata": {},
   "source": [
    "# Bonus: RaySGD + MLflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "coated-receiver",
   "metadata": {},
   "source": [
    "## Distributed Deep Learning Made Simple: RaySGD\n",
    "\n",
    "Distributed deep learning -- or optimization in general using TensorFlow/PyTorch -- has slowly been getting easier, even if it's not 100% transparent (or turn-key) quite yet.\n",
    "\n",
    "Native solutions (official TF/PyTorch distributed code) are getting easier; Horovod has matured and is straightforward. But one of the simplest approaches of all comes via the Ray project.\n",
    "\n",
    "In particular, RaySGD provides a zero-ops and minimal API approach: https://docs.ray.io/en/master/raysgd/raysgd.html\n",
    "\n",
    "Let's take a look! \n",
    "\n",
    "First, we'll set up a toy example (adapted from https://docs.ray.io/en/master/raysgd/raysgd_tensorflow.html)\n",
    "\n",
    "Later, we'll try a real dataset and integrate with additional tools.\n",
    "\n",
    "Here, we create trivial dummy data and a Tensorflow Dataset loader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "christian-adrian",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def linear_dataset(size=100):\n",
    "    x = np.random.rand(size)\n",
    "    y = 2 * x\n",
    "\n",
    "    x = x.reshape((-1, 1))\n",
    "    y = y.reshape((-1, 1))\n",
    "\n",
    "    return x, y\n",
    "\n",
    "def simple_dataset(config):\n",
    "    batch_size = config[\"batch_size\"]\n",
    "    x_train, y_train = linear_dataset(size=NUM_TRAIN_SAMPLES)\n",
    "    x_test, y_test = linear_dataset(size=NUM_TEST_SAMPLES)\n",
    "\n",
    "    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
    "    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
    "    train_dataset = train_dataset.shuffle(NUM_TRAIN_SAMPLES).repeat().batch(\n",
    "        batch_size)\n",
    "    test_dataset = test_dataset.repeat().batch(batch_size)\n",
    "\n",
    "    return train_dataset, test_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "marked-chile",
   "metadata": {},
   "source": [
    "Next, we define our model using regular `tf.keras` components. We define a model-creation function, as RaySGD uses a factory pattern (basically a pattern that takes creator functions rather than object instances) for dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "transsexual-spokesman",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import Dense\n",
    "\n",
    "def simple_model(config):\n",
    "    model = Sequential([Dense(10, input_shape=(1, )), Dense(1)])\n",
    "\n",
    "    model.compile(\n",
    "        optimizer=\"sgd\",\n",
    "        loss=\"mean_squared_error\",\n",
    "        metrics=[\"mean_squared_error\"])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "after-mother",
   "metadata": {},
   "source": [
    "Now we get to the RaySGD code. After definining a minimal \"config\" object, we define a trainer function that accomplishes these steps:\n",
    "* instantiate a TFTrainer instance to wrap the actual distributed training\n",
    "* explicitly calculate starting model performance\n",
    "* train multiple epochs\n",
    "    * here we're explicitly calling `train` twice ... in a real example we would run more epochs with a control loop\n",
    "* calculate final stats, change in the loss, and a \"sanity check\" that the loss actually went down"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sealed-screen",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "from ray.util.sgd.tf.tf_trainer import TFTrainer, TFTrainable\n",
    "\n",
    "NUM_TRAIN_SAMPLES = 1000\n",
    "NUM_TEST_SAMPLES = 400\n",
    "\n",
    "def create_config(batch_size):\n",
    "\n",
    "    return {\n",
    "        \"batch_size\": batch_size,\n",
    "        \"fit_config\": {\n",
    "            \"steps_per_epoch\": NUM_TRAIN_SAMPLES // batch_size\n",
    "        },\n",
    "        \"evaluate_config\": {\n",
    "            \"steps\": NUM_TEST_SAMPLES // batch_size,\n",
    "        }\n",
    "    }\n",
    "\n",
    "def train_example(num_replicas=1, batch_size=128, use_gpu=False):\n",
    "    trainer = TFTrainer(\n",
    "        model_creator=simple_model,\n",
    "        data_creator=simple_dataset,\n",
    "        num_replicas=num_replicas,\n",
    "        use_gpu=use_gpu,\n",
    "        verbose=True,\n",
    "        config=create_config(batch_size))\n",
    "\n",
    "    # model baseline performance\n",
    "    start_stats = trainer.validate()\n",
    "    print(start_stats)\n",
    "\n",
    "    # train for 2 epochs\n",
    "    trainer.train()\n",
    "    trainer.train()\n",
    "\n",
    "    # model performance after training (should improve)\n",
    "    end_stats = trainer.validate()\n",
    "    print(end_stats)\n",
    "\n",
    "    # sanity check that training worked\n",
    "    dloss = end_stats[\"validation_loss\"] - start_stats[\"validation_loss\"]\n",
    "    dmse = (end_stats[\"validation_mean_squared_error\"] -\n",
    "            start_stats[\"validation_mean_squared_error\"])\n",
    "    print(f\"dLoss: {dloss}, dMSE: {dmse}\")\n",
    "\n",
    "    if dloss > 0 or dmse > 0:\n",
    "        print(\"training sanity check failed. loss increased!\")\n",
    "    else:\n",
    "        print(\"success!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "expired-perry",
   "metadata": {},
   "source": [
    "Ok, now that we're all set up, let's start Ray and run the training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "earned-nancy",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.init()\n",
    "\n",
    "train_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "controlled-bryan",
   "metadata": {},
   "source": [
    "Note that the Dashboard which defaults to `localhost:8265` is a key part of the Ray system, but may not be compatible with (and visible through) the binder container proxy.\n",
    "\n",
    "__How does this work?__\n",
    "\n",
    "In a nutshell,\n",
    "* `TFTrainer` wraps TensorFlow's `MultiWorkerMirroredStrategy` as described here: https://docs.ray.io/en/master/raysgd/raysgd_tensorflow.html\n",
    "* `MultiWorkerMirroredStrategy` is a synchronous distributed approach featuring multilateral reduce (e.g., AllReduce): https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/MultiWorkerMirroredStrategy\n",
    "\n",
    "But the \"Hello World\" isn't very impressive. Let's at least try a slightly more realistic, if not industrial strength, dataset and model.\n",
    "\n",
    "We'll train a shallow (1-layer) dense feed-forward network with ReLU activation on the R/ggplot2 diamonds data (https://ggplot2.tidyverse.org/reference/diamonds.html)\n",
    "\n",
    "Start with a data loader, model, and config builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "little-freeware",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "def diamonds_dataset(config):\n",
    "    batch_size = config[\"batch_size\"]\n",
    "    df = pd.read_csv('data/diamonds.csv')\n",
    "    df.drop(df.columns[0], axis=1, inplace=True)\n",
    "    df = pd.get_dummies(df, prefix=['cut_', 'color_', 'clarity_'])\n",
    "    y = df.price.to_numpy()\n",
    "    X = df.drop(columns=['price']).to_numpy()\n",
    "    train_size = 40_000\n",
    "    X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=train_size)\n",
    "    \n",
    "    train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))\n",
    "    test_dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test))\n",
    "    train_dataset = train_dataset.shuffle(len(X_train)).repeat().batch(\n",
    "        batch_size)\n",
    "    test_dataset = test_dataset.repeat().batch(batch_size)\n",
    "\n",
    "    return train_dataset, test_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "coordinated-server",
   "metadata": {},
   "outputs": [],
   "source": [
    "def diamonds_simple_model(config):\n",
    "    model = Sequential([Dense(30, input_shape=(26, ), activation='relu'), Dense(1)])\n",
    "\n",
    "    model.compile(\n",
    "        optimizer=\"adam\",\n",
    "        loss=\"mean_squared_error\",\n",
    "        metrics=[\"mean_squared_error\"])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "green-induction",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_diamonds_config(batch_size):\n",
    "    return {\n",
    "        \"batch_size\": batch_size,\n",
    "        \"fit_config\": {\n",
    "            \"steps_per_epoch\": 40000 // batch_size\n",
    "        },\n",
    "        \"evaluate_config\": {\n",
    "            \"steps\": 13940 // batch_size,\n",
    "        }\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "together-winning",
   "metadata": {},
   "source": [
    "For clarity, we'll define a bare-basics training function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "voluntary-regard",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_diamonds(num_replicas=1, batch_size=128, use_gpu=False):\n",
    "    trainer = TFTrainer(\n",
    "        model_creator=diamonds_simple_model,\n",
    "        data_creator=diamonds_dataset,\n",
    "        num_replicas=num_replicas,\n",
    "        use_gpu=use_gpu,\n",
    "        verbose=False,\n",
    "        config=create_diamonds_config(batch_size))\n",
    "\n",
    "    start_stats = trainer.validate()\n",
    "    print(start_stats)\n",
    "\n",
    "    for i in range(32):\n",
    "        trainer.train()\n",
    "\n",
    "    end_stats = trainer.validate()\n",
    "    print(end_stats)\n",
    "        \n",
    "train_diamonds()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "theoretical-district",
   "metadata": {},
   "source": [
    "We won't get famous for these results, but we definitely made progress.\n",
    "\n",
    "## MLflow\n",
    "\n",
    "One of the top open source frameworks for managing machine learning, from experiment to deployment, is MLflow (https://mlflow.org/)\n",
    "\n",
    "Created by Databricks and open-sourced under the Linux Foundation, MLflow has rapidly evolved to support a variety of key ML engineering tasks including\n",
    "* Experiment tracking\n",
    "    * Parameters, results, data, code, and model artefacts/assets\n",
    "* Tracking project environments for reproducibility\n",
    "* Deployment from a variety of model formats to various target (prediction/scoring) environments\n",
    "* Model registry\n",
    "    * Versioning, lineage/provenance\n",
    "\n",
    "More features are planned for the future; today, we'll look at just the original experiment tracking features.\n",
    "\n",
    "MLflow supports auto-instrumentation for a number of popular platforms (https://mlflow.org/docs/latest/tracking.html#automatic-logging) but these don't include RaySGD/Distributed TensorFlow yet... plus we want to see concretely how the pieces fit together.\n",
    "\n",
    "Before proceeding, start the MLflow UI server. On Binder/JupyterLab, open a new Terminal and type\n",
    "\n",
    "`mlflow ui`\n",
    "\n",
    "The UI is served on port 5000 by default, which should be accessible via the JupyterLab proxy (open a new tab and replace `/lab...` in your URL with `/proxy/5000/`)\n",
    "\n",
    "We'll start with the minimal code to save params and metrics to MLflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "russian-cliff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlflow import log_metric, log_param, end_run\n",
    "\n",
    "log_param(\"foo_count\", 42)\n",
    "log_param(\"bar_count\", 43)\n",
    "\n",
    "for i in range(10):\n",
    "    log_metric(\"score\", i)\n",
    "    \n",
    "end_run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "invisible-evaluation",
   "metadata": {},
   "source": [
    "This info should appear in the UI (though it may require a refresh)\n",
    "\n",
    "We can use a context manager for our runs and provide step indices to enable additional view."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "national-honduras",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "with mlflow.start_run():\n",
    "    for i in range(10):\n",
    "        log_metric(\"score\", i*i, step=i)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "monthly-music",
   "metadata": {},
   "source": [
    "Now let's look at a more realistic use of MLflow, including\n",
    "* Creating a named Experiment\n",
    "* Recording data for runs to this Experiment\n",
    "* Integrating with our RaySGD/TF model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "inappropriate-chorus",
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_id = mlflow.create_experiment(\"Diamonds RaySGD\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bored-edition",
   "metadata": {},
   "outputs": [],
   "source": [
    "from  mlflow.tracking import MlflowClient\n",
    "client = MlflowClient()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "negative-slope",
   "metadata": {},
   "outputs": [],
   "source": [
    "run = client.create_run(experiment_id) # returns mlflow.entities.Run\n",
    "client.log_param(run.info.run_id, \"hello\", \"world\")\n",
    "client.set_terminated(run.info.run_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "elementary-subject",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_diamonds_mlflow(num_replicas=1, batch_size=128, use_gpu=False):\n",
    "    trainer = TFTrainer(\n",
    "        model_creator=diamonds_simple_model,\n",
    "        data_creator=diamonds_dataset,\n",
    "        num_replicas=num_replicas,\n",
    "        use_gpu=use_gpu,\n",
    "        verbose=False,\n",
    "        config=create_diamonds_config(batch_size))\n",
    "\n",
    "    start_stats = trainer.validate()\n",
    "    print(start_stats)\n",
    "\n",
    "    ml_run = client.create_run(experiment_id)\n",
    "\n",
    "    for i in range(32):\n",
    "        train_stats = trainer.train()\n",
    "        if i % 2 == 0:\n",
    "            val_stats = trainer.validate()            \n",
    "            client.log_metric(ml_run.info.run_id, \"validation_loss\", val_stats[\"validation_loss\"], step=i)\n",
    "            client.log_metric(ml_run.info.run_id, \"training_loss\", train_stats[\"train_loss\"], step=i)\n",
    "        \n",
    "    client.set_terminated(ml_run.info.run_id)\n",
    "\n",
    "    end_stats = trainer.validate()\n",
    "    print(end_stats)\n",
    "        \n",
    "train_diamonds_mlflow()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "impressive-joseph",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "normal-enzyme",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}