19_scale_deploy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "absent-toner",
   "metadata": {},
   "source": [
    "# Training and Deploying TensorFlow Models at Scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "graduate-relation",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No GPU was detected.\n"
     ]
    }
   ],
   "source": [
    "# FIXME: meke autocompletion working again\n",
    "%config Completer.use_jedi = False\n",
    "\n",
    "import os\n",
    "\n",
    "# OpenAI gym\n",
    "import gym\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "# Get smooth animations\n",
    "mpl.rc('animation', html='jshtml')\n",
    "\n",
    "physical_devices = tf.config.list_physical_devices('GPU')\n",
    "\n",
    "if not physical_devices:\n",
    "    print(\"No GPU was detected.\")\n",
    "else:\n",
    "    # https://stackoverflow.com/a/60699372\n",
    "    tf.config.experimental.set_memory_growth(physical_devices[0], True)\n",
    "    \n",
    "from tensorflow import keras"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "instant-breeding",
   "metadata": {},
   "source": [
    "## Deploying TensorFlow models to TensorFlow Serving\n",
    "*TensorFlow Serving (TFS)* provides simple REST and gRPC APIs and handle model versioning and graceful updates and rollbacks (blue-green or stop-the-world).\n",
    "\n",
    "*Note: It's generally good idea to include and compile preprocessing logic in the model so that clients can send raw data and not duplicate these preprocessing steps. This, however, implies that all the preprocessing etc. must be done with TF-only functions. Otherwise it won't be compiled and saved in the computation graph.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "important-smoke",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch MNIST dataset\n",
    "(X_train_full, y_train_full), (X_test, y_test) = keras.datasets.mnist.load_data()\n",
    "\n",
    "# Scale the data\n",
    "max_value = 255.\n",
    "X_train_full = X_train_full[..., np.newaxis].astype(np.float32) / max_value\n",
    "X_test = X_test[..., np.newaxis].astype(np.float32) / max_value\n",
    "\n",
    "# Split the raw training data to training and validation sets\n",
    "valid_split = 5000\n",
    "X_valid, X_train = X_train_full[:valid_split], X_train_full[valid_split:]\n",
    "y_valid, y_train = y_train_full[:valid_split], y_train_full[valid_split:]\n",
    "\n",
    "# Use first couple of test instances for predictions\n",
    "X_new = X_test[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tough-education",
   "metadata": {},
   "source": [
    "### Save/Load a SavedModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "superior-separate",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "1719/1719 [==============================] - 3s 2ms/step - loss: 1.1140 - accuracy: 0.7066 - val_loss: 0.3715 - val_accuracy: 0.9024\n",
      "Epoch 2/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.3695 - accuracy: 0.8981 - val_loss: 0.2990 - val_accuracy: 0.9144\n",
      "Epoch 3/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.3154 - accuracy: 0.9100 - val_loss: 0.2651 - val_accuracy: 0.9272\n",
      "Epoch 4/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.2765 - accuracy: 0.9222 - val_loss: 0.2436 - val_accuracy: 0.9334\n",
      "Epoch 5/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.2556 - accuracy: 0.9276 - val_loss: 0.2257 - val_accuracy: 0.9364\n",
      "Epoch 6/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.2367 - accuracy: 0.9321 - val_loss: 0.2121 - val_accuracy: 0.9396\n",
      "Epoch 7/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.2198 - accuracy: 0.9391 - val_loss: 0.1970 - val_accuracy: 0.9454\n",
      "Epoch 8/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.2057 - accuracy: 0.9425 - val_loss: 0.1880 - val_accuracy: 0.9476\n",
      "Epoch 9/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.1940 - accuracy: 0.9458 - val_loss: 0.1778 - val_accuracy: 0.9524\n",
      "Epoch 10/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.1798 - accuracy: 0.9481 - val_loss: 0.1685 - val_accuracy: 0.9542\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7f87a5434400>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set RNG state\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Build model v1\n",
    "model = keras.models.Sequential([\n",
    "    keras.layers.Flatten(input_shape=[28, 28, 1]),\n",
    "    keras.layers.Dense(100, activation=\"relu\"),\n",
    "    keras.layers.Dense(10, activation=\"softmax\"),\n",
    "])\n",
    "\n",
    "model.compile(\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    optimizer=keras.optimizers.SGD(lr=1e-2),\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "model.fit(X_train, y_train, epochs=10, validation_data=(X_valid, y_valid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "victorian-israeli",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 1.  , 0.  , 0.  ],\n",
       "       [0.  , 0.  , 0.99, 0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],\n",
       "       [0.  , 0.97, 0.01, 0.  , 0.  , 0.  , 0.  , 0.01, 0.  , 0.  ]],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test model's predictions\n",
    "np.round(model.predict(X_new), 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "continued-magnet",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'data/my_mnist_model/0001'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_version = \"0001\"\n",
    "model_name = \"my_mnist_model\"\n",
    "model_path = os.path.join(\"data\", model_name, model_version)\n",
    "model_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "affiliated-bullet",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf data/{model_name}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "framed-bruce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: data/my_mnist_model/0001/assets\n"
     ]
    }
   ],
   "source": [
    "# Save the model in `SavedModel` format/structure\n",
    "tf.saved_model.save(model, model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "advance-amsterdam",
   "metadata": {},
   "source": [
    "Let's print the structure of `SavedModel`:\n",
    "* `saved_model.pb` is a protobuf-serialized computaiton graph of the model\n",
    "* `variables/` contains all the weights, possibly split into multiple files\n",
    "* `assets/` contains additional data such as examples, dictionary tables, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "opposed-thesis",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34mdata/my_mnist_model\u001b[00m\n",
      "└── \u001b[01;34m0001\u001b[00m\n",
      "    ├── \u001b[01;34massets\u001b[00m\n",
      "    ├── saved_model.pb\n",
      "    └── \u001b[01;34mvariables\u001b[00m\n",
      "        ├── variables.data-00000-of-00001\n",
      "        └── variables.index\n",
      "\n",
      "3 directories, 3 files\n"
     ]
    }
   ],
   "source": [
    "!tree data/{model_name}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "competent-newfoundland",
   "metadata": {},
   "source": [
    "The `saved_model_cli` tool can be handy to describe the saved model as well as running predictions (for debugging)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "stretch-office",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-03-05 08:32:36.977697: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2021-03-05 08:32:36.977722: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n",
      "The given SavedModel contains the following tag-sets:\n",
      "'serve'\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli show --dir {model_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "comfortable-desire",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-03-05 08:32:39.920243: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2021-03-05 08:32:39.920270: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n",
      "The given SavedModel MetaGraphDef contains SignatureDefs with the following keys:\n",
      "SignatureDef key: \"__saved_model_init_op\"\n",
      "SignatureDef key: \"serving_default\"\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli show --dir {model_path} --tag_set serve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "intensive-berlin",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-03-05 08:32:42.201065: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2021-03-05 08:32:42.201091: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n",
      "The given SavedModel SignatureDef contains the following input(s):\n",
      "  inputs['flatten_input'] tensor_info:\n",
      "      dtype: DT_FLOAT\n",
      "      shape: (-1, 28, 28, 1)\n",
      "      name: serving_default_flatten_input:0\n",
      "The given SavedModel SignatureDef contains the following output(s):\n",
      "  outputs['dense_1'] tensor_info:\n",
      "      dtype: DT_FLOAT\n",
      "      shape: (-1, 10)\n",
      "      name: StatefulPartitionedCall:0\n",
      "Method name is: tensorflow/serving/predict\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli show --dir {model_path} --tag_set serve --signature_def serving_default"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "located-weather",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-03-05 08:32:45.003892: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2021-03-05 08:32:45.003918: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n",
      "\n",
      "MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:\n",
      "\n",
      "signature_def['__saved_model_init_op']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['__saved_model_init_op'] tensor_info:\n",
      "        dtype: DT_INVALID\n",
      "        shape: unknown_rank\n",
      "        name: NoOp\n",
      "  Method name is: \n",
      "\n",
      "signature_def['serving_default']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "    inputs['flatten_input'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: (-1, 28, 28, 1)\n",
      "        name: serving_default_flatten_input:0\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['dense_1'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: (-1, 10)\n",
      "        name: StatefulPartitionedCall:0\n",
      "  Method name is: tensorflow/serving/predict\n",
      "\n",
      "Defined Functions:\n",
      "  Function Name: '__call__'\n",
      "    Option #1\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          flatten_input: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='flatten_input')\n",
      "        Argument #2\n",
      "          DType: bool\n",
      "          Value: True\n",
      "        Argument #3\n",
      "          DType: NoneType\n",
      "          Value: None\n",
      "    Option #2\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          inputs: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='inputs')\n",
      "        Argument #2\n",
      "          DType: bool\n",
      "          Value: True\n",
      "        Argument #3\n",
      "          DType: NoneType\n",
      "          Value: None\n",
      "    Option #3\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          inputs: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='inputs')\n",
      "        Argument #2\n",
      "          DType: bool\n",
      "          Value: False\n",
      "        Argument #3\n",
      "          DType: NoneType\n",
      "          Value: None\n",
      "    Option #4\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          flatten_input: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='flatten_input')\n",
      "        Argument #2\n",
      "          DType: bool\n",
      "          Value: False\n",
      "        Argument #3\n",
      "          DType: NoneType\n",
      "          Value: None\n",
      "\n",
      "  Function Name: '_default_save_signature'\n",
      "    Option #1\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          flatten_input: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='flatten_input')\n",
      "\n",
      "  Function Name: 'call_and_return_all_conditional_losses'\n",
      "    Option #1\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          inputs: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='inputs')\n",
      "        Argument #2\n",
      "          DType: bool\n",
      "          Value: False\n",
      "        Argument #3\n",
      "          DType: NoneType\n",
      "          Value: None\n",
      "    Option #2\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          inputs: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='inputs')\n",
      "        Argument #2\n",
      "          DType: bool\n",
      "          Value: True\n",
      "        Argument #3\n",
      "          DType: NoneType\n",
      "          Value: None\n",
      "    Option #3\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          flatten_input: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='flatten_input')\n",
      "        Argument #2\n",
      "          DType: bool\n",
      "          Value: True\n",
      "        Argument #3\n",
      "          DType: NoneType\n",
      "          Value: None\n",
      "    Option #4\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          flatten_input: TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32, name='flatten_input')\n",
      "        Argument #2\n",
      "          DType: bool\n",
      "          Value: False\n",
      "        Argument #3\n",
      "          DType: NoneType\n",
      "          Value: None\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli show --dir {model_path} --all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bottom-blond",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the input instances to a numpy file\n",
    "inputs_path = os.path.join(\"data\", \"my_mnist_tests.npy\")\n",
    "np.save(inputs_path, X_new)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "dimensional-confusion",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'flatten_input'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the name of the input layer in the model\n",
    "input_name = model.input_names[0]\n",
    "input_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "developed-needle",
   "metadata": {},
   "source": [
    "Use the CLI to make predictions on these test instances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ongoing-stack",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-03-05 08:32:47.632133: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2021-03-05 08:32:47.632163: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n",
      "2021-03-05 08:32:49.435242: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set\n",
      "2021-03-05 08:32:49.435398: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory\n",
      "2021-03-05 08:32:49.435411: W tensorflow/stream_executor/cuda/cuda_driver.cc:326] failed call to cuInit: UNKNOWN ERROR (303)\n",
      "2021-03-05 08:32:49.435429: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (mpc-xps): /proc/driver/nvidia/version does not exist\n",
      "2021-03-05 08:32:49.435679: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2021-03-05 08:32:49.435891: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set\n",
      "WARNING:tensorflow:From /home/matyama/.cache/pypoetry/virtualenvs/homl-lPHrmr2i-py3.8/lib/python3.8/site-packages/tensorflow/python/tools/saved_model_cli.py:445: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.\n",
      "INFO:tensorflow:Restoring parameters from data/my_mnist_model/0001/variables/variables\n",
      "2021-03-05 08:32:49.470117: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:196] None of the MLIR optimization passes are enabled (registered 0 passes)\n",
      "2021-03-05 08:32:49.488440: I tensorflow/core/platform/profile_utils/cpu_utils.cc:112] CPU Frequency: 2400000000 Hz\n",
      "Result for output key dense_1:\n",
      "[[1.1432634e-04 1.5143777e-07 9.8062563e-04 2.7728046e-03 3.7552679e-06\n",
      "  7.6340642e-05 3.9137930e-08 9.9556726e-01 5.3547199e-05 4.3099097e-04]\n",
      " [8.1703230e-04 3.5410001e-05 9.8826653e-01 7.0426976e-03 1.2937582e-07\n",
      "  2.3344452e-04 2.5721604e-03 9.6440456e-10 1.0325429e-03 8.7997570e-08]\n",
      " [4.4410775e-05 9.7030008e-01 9.0604592e-03 2.2620754e-03 4.8724114e-04\n",
      "  2.8746312e-03 2.2719912e-03 8.3600031e-03 4.0411805e-03 2.9791440e-04]]\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli run --dir {model_path} --tag_set serve \\\n",
    "                     --signature_def serving_default    \\\n",
    "                     --inputs {input_name}={inputs_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "alone-makeup",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 1.  , 0.  , 0.  ],\n",
       "       [0.  , 0.  , 0.99, 0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],\n",
       "       [0.  , 0.97, 0.01, 0.  , 0.  , 0.  , 0.  , 0.01, 0.  , 0.  ]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.round(\n",
    "    [[1.1347984e-04, 1.5187356e-07, 9.7032893e-04, 2.7640699e-03, 3.7826971e-06, 7.6876910e-05, 3.9140293e-08, 9.9559116e-01, 5.3502394e-05, 4.2665208e-04],\n",
    "    [8.2443521e-04, 3.5493889e-05, 9.8826385e-01, 7.0466995e-03, 1.2957400e-07, 2.3389691e-04, 2.5639210e-03, 9.5886099e-10, 1.0314899e-03, 8.7952529e-08],\n",
    "    [4.4693781e-05, 9.7028232e-01, 9.0526715e-03, 2.2641101e-03, 4.8766597e-04, 2.8800720e-03, 2.2714981e-03, 8.3753867e-03, 4.0439744e-03, 2.9759688e-04]],\n",
    "    2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "purple-dancing",
   "metadata": {},
   "source": [
    "### TensorFlow Serving\n",
    "Serving the model with TF Serving is as simple as running following docker container\n",
    "```bash\n",
    "docker run \\\n",
    "   -it \\\n",
    "   --rm \\\n",
    "   -p 8500:8500 \\\n",
    "   -p 8501:8501 \\\n",
    "   -v \"$PWD/data/my_mnist_model:/models/my_mnist_model\" \\\n",
    "   -e MODEL_NAME=my_mnist_model \\\n",
    "   tensorflow/serving\n",
    "```\n",
    "Let's note that\n",
    "* Port 8500 will expose the gRPC API\n",
    "* Port 8501 will expose the REST API\n",
    "* We mount the SavedModel directory to `/models` which is the default location TF Serving looks at\n",
    "\n",
    "Some features of TF Serving:\n",
    "* Automatic detection of new models and their versions\n",
    "* Graceful model updates\n",
    "* Simple rollback to previous version (deleting version directory)\n",
    "* Batching incomming requests for better GPU utilization on larger input batches and followup redistribution of predictions to appropriate clients"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acoustic-purple",
   "metadata": {},
   "source": [
    "#### REST API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "floppy-logistics",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\'{\"signature_name\": \"serving_default\", \"instances\": [[[[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]], [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]], [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]], [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]], [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]], [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]], [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]], [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.32941177487373...'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "# Build a request JSON\n",
    "input_data_json = json.dumps({\n",
    "    \"signature_name\": \"serving_default\",\n",
    "    \"instances\": X_new.tolist(),\n",
    "})\n",
    "\n",
    "# Print first portion of the JSON data\n",
    "repr(input_data_json)[:1500] + \"...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "cloudy-token",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['predictions'])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import requests\n",
    "\n",
    "# Send a POST request with the input JSON payload\n",
    "response = requests.post(\n",
    "    f\"http://localhost:8501/v1/models/{model_name}:predict\",\n",
    "    data=input_data_json,\n",
    ")\n",
    "\n",
    "# Raise an exception in case of error\n",
    "response.raise_for_status()\n",
    "\n",
    "# Parse the response JSON\n",
    "response = response.json()\n",
    "\n",
    "# Show the response fields\n",
    "response.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "driven-andrews",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 1.  , 0.  , 0.  ],\n",
       "       [0.  , 0.  , 0.99, 0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],\n",
       "       [0.  , 0.97, 0.01, 0.  , 0.  , 0.  , 0.  , 0.01, 0.  , 0.  ]])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Convert predictions to a numpy array and display them\n",
    "y_proba = np.array(response[\"predictions\"])\n",
    "y_proba.round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "solved-builder",
   "metadata": {},
   "source": [
    "#### gRPC API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "generic-titanium",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "outputs {\n",
       "  key: \"dense_1\"\n",
       "  value {\n",
       "    dtype: DT_FLOAT\n",
       "    tensor_shape {\n",
       "      dim {\n",
       "        size: 3\n",
       "      }\n",
       "      dim {\n",
       "        size: 10\n",
       "      }\n",
       "    }\n",
       "    float_val: 0.00011432634346419945\n",
       "    float_val: 1.5143777432058414e-07\n",
       "    float_val: 0.0009806256275624037\n",
       "    float_val: 0.002772804582491517\n",
       "    float_val: 3.755267925953376e-06\n",
       "    float_val: 7.634064240846783e-05\n",
       "    float_val: 3.913793023002654e-08\n",
       "    float_val: 0.995567262172699\n",
       "    float_val: 5.354719905881211e-05\n",
       "    float_val: 0.0004309909709263593\n",
       "    float_val: 0.0008170322980731726\n",
       "    float_val: 3.541000114637427e-05\n",
       "    float_val: 0.9882665276527405\n",
       "    float_val: 0.0070426976308226585\n",
       "    float_val: 1.2937582027916505e-07\n",
       "    float_val: 0.00023344451619777828\n",
       "    float_val: 0.0025721604470163584\n",
       "    float_val: 9.64404556214049e-10\n",
       "    float_val: 0.0010325429029762745\n",
       "    float_val: 8.799757011956899e-08\n",
       "    float_val: 4.441077544470318e-05\n",
       "    float_val: 0.9703000783920288\n",
       "    float_val: 0.009060459211468697\n",
       "    float_val: 0.0022620754316449165\n",
       "    float_val: 0.00048724113730713725\n",
       "    float_val: 0.0028746312018483877\n",
       "    float_val: 0.0022719912230968475\n",
       "    float_val: 0.008360003121197224\n",
       "    float_val: 0.004041180480271578\n",
       "    float_val: 0.0002979144046548754\n",
       "  }\n",
       "}\n",
       "model_spec {\n",
       "  name: \"my_mnist_model\"\n",
       "  version {\n",
       "    value: 1\n",
       "  }\n",
       "  signature_name: \"serving_default\"\n",
       "}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import grpc\n",
    "\n",
    "from tensorflow_serving.apis import prediction_service_pb2_grpc\n",
    "from tensorflow_serving.apis.predict_pb2 import PredictRequest\n",
    "\n",
    "# Build the prediction protobuf request\n",
    "request = PredictRequest()\n",
    "request.model_spec.name = model_name\n",
    "request.model_spec.signature_name = \"serving_default\"\n",
    "request.inputs[input_name].CopyFrom(tf.make_tensor_proto(X_new))\n",
    "\n",
    "# Create a gRPC channel and service stub\n",
    "channel = grpc.insecure_channel('localhost:8500')\n",
    "predict_service = prediction_service_pb2_grpc.PredictionServiceStub(channel)\n",
    "\n",
    "# Request the predictions and show the response\n",
    "response = predict_service.Predict(request, timeout=10.0)\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "personal-sheep",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 1.  , 0.  , 0.  ],\n",
       "       [0.  , 0.  , 0.99, 0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],\n",
       "       [0.  , 0.97, 0.01, 0.  , 0.  , 0.  , 0.  , 0.01, 0.  , 0.  ]],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the name of the output layer from the model\n",
    "output_name = model.output_names[0]\n",
    "\n",
    "# Get the protobuf outputs from the response\n",
    "outputs_proto = response.outputs[output_name]\n",
    "\n",
    "# Parse and display the predictions\n",
    "y_proba = tf.make_ndarray(outputs_proto)\n",
    "y_proba.round(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "undefined-second",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 1.  , 0.  , 0.  ],\n",
       "       [0.  , 0.  , 0.99, 0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],\n",
       "       [0.  , 0.97, 0.01, 0.  , 0.  , 0.  , 0.  , 0.01, 0.  , 0.  ]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Alternatively one can parse the predictions directly to numpy (if TF is not available)\n",
    "shape = [dim.size for dim in outputs_proto.tensor_shape.dim]\n",
    "y_proba = np.array(outputs_proto.float_val).reshape(shape)\n",
    "y_proba.round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "opponent-humanity",
   "metadata": {},
   "source": [
    "### Deploying a new model version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "strange-mason",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "1719/1719 [==============================] - 3s 2ms/step - loss: 1.1567 - accuracy: 0.6691 - val_loss: 0.3418 - val_accuracy: 0.9042\n",
      "Epoch 2/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.3376 - accuracy: 0.9032 - val_loss: 0.2674 - val_accuracy: 0.9242\n",
      "Epoch 3/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.2779 - accuracy: 0.9187 - val_loss: 0.2227 - val_accuracy: 0.9366\n",
      "Epoch 4/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.2362 - accuracy: 0.9318 - val_loss: 0.2032 - val_accuracy: 0.9432\n",
      "Epoch 5/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.2109 - accuracy: 0.9388 - val_loss: 0.1831 - val_accuracy: 0.9478\n",
      "Epoch 6/10\n",
      "1719/1719 [==============================] - 3s 2ms/step - loss: 0.1951 - accuracy: 0.9430 - val_loss: 0.1740 - val_accuracy: 0.9498\n",
      "Epoch 7/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.1798 - accuracy: 0.9475 - val_loss: 0.1604 - val_accuracy: 0.9540\n",
      "Epoch 8/10\n",
      "1719/1719 [==============================] - 3s 2ms/step - loss: 0.1654 - accuracy: 0.9518 - val_loss: 0.1542 - val_accuracy: 0.9562\n",
      "Epoch 9/10\n",
      "1719/1719 [==============================] - 2s 1ms/step - loss: 0.1569 - accuracy: 0.9555 - val_loss: 0.1460 - val_accuracy: 0.9572\n",
      "Epoch 10/10\n",
      "1719/1719 [==============================] - 3s 1ms/step - loss: 0.1419 - accuracy: 0.9581 - val_loss: 0.1358 - val_accuracy: 0.9616\n"
     ]
    }
   ],
   "source": [
    "# Set RNG state\n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)\n",
    "\n",
    "# Build model v2\n",
    "model = keras.models.Sequential([\n",
    "    keras.layers.Flatten(input_shape=[28, 28, 1]),\n",
    "    keras.layers.Dense(50, activation=\"relu\"),\n",
    "    keras.layers.Dense(50, activation=\"relu\"),\n",
    "    keras.layers.Dense(10, activation=\"softmax\")\n",
    "])\n",
    "\n",
    "model.compile(\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    optimizer=keras.optimizers.SGD(lr=1e-2),\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "history = model.fit(X_train, y_train, epochs=10, validation_data=(X_valid, y_valid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "decreased-granny",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: data/my_mnist_model/0002/assets\n"
     ]
    }
   ],
   "source": [
    "# Make new model version and path\n",
    "model_version = \"0002\"\n",
    "model_path = os.path.join(\"data\", model_name, model_version)\n",
    "\n",
    "# Save the model as a SavedModel\n",
    "tf.saved_model.save(model, model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "touched-conflict",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34mdata/my_mnist_model\u001b[00m\n",
      "├── \u001b[01;34m0001\u001b[00m\n",
      "│   ├── \u001b[01;34massets\u001b[00m\n",
      "│   ├── saved_model.pb\n",
      "│   └── \u001b[01;34mvariables\u001b[00m\n",
      "│       ├── variables.data-00000-of-00001\n",
      "│       └── variables.index\n",
      "└── \u001b[01;34m0002\u001b[00m\n",
      "    ├── \u001b[01;34massets\u001b[00m\n",
      "    ├── saved_model.pb\n",
      "    └── \u001b[01;34mvariables\u001b[00m\n",
      "        ├── variables.data-00000-of-00001\n",
      "        └── variables.index\n",
      "\n",
      "6 directories, 6 files\n"
     ]
    }
   ],
   "source": [
    "!tree data/{model_name}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "treated-december",
   "metadata": {},
   "source": [
    "TF Serving will automatically detect new model version and gracefully deploy it. As in a *blue-green deployment*, it will process all pending requests with the old model while collecting new ones which will be then processed by the new model. This can be reconfigured to immediately swap active model and process all unprocessed requests with the new version. However, in the latter configuration there is a short period when then the service is unavailable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "touched-graham",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 1.  , 0.  , 0.  ],\n",
       "       [0.  , 0.  , 0.99, 0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],\n",
       "       [0.  , 0.99, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ]])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Request predictions once again\n",
    "response = requests.post(\n",
    "    f\"http://localhost:8501/v1/models/{model_name}:predict\",\n",
    "    data=input_data_json,\n",
    ")\n",
    "response.raise_for_status()\n",
    "\n",
    "# Parse the response\n",
    "response = response.json()\n",
    "\n",
    "# Extract and show predictions\n",
    "y_proba = np.array(response[\"predictions\"])\n",
    "y_proba.round(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sunset-screw",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}