{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Kubeflow Pipelines e2e mnist example\n",
    "\n",
    "In this notebook you will create e2e mnist Kubeflow Pipeline to perform:\n",
    "- Hyperparameter tuning using Katib\n",
    "- Distributive training with the best hyperparameters using TFJob\n",
    "- Serve the trained model using KServe\n",
    "\n",
    "Reference documentation:\n",
    "\n",
    "- https://www.kubeflow.org/docs/components/training/tftraining/\n",
    "- https://www.kubeflow.org/docs/components/katib/\n",
    "- https://www.kubeflow.org/docs/external-add-ons/kserve/\n",
    "\n",
    "**Note**: This Pipeline runs in the multi-user mode. Follow [this guide](https://www.kubeflow.org/docs/components/pipelines/sdk/connect-api/#multi-user-mode) to give your Notebook access to Kubeflow Pipelines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: kfp==1.8.4 in /opt/conda/lib/python3.8/site-packages (1.8.4)\n",
      "Requirement already satisfied: jsonschema<4,>=3.0.1 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (3.2.0)\n",
      "Requirement already satisfied: PyYAML<6,>=5.3 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (5.4.1)\n",
      "Requirement already satisfied: kfp-server-api<2.0.0,>=1.1.2 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (1.6.0)\n",
      "Requirement already satisfied: click<8,>=7.1.1 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (7.1.2)\n",
      "Requirement already satisfied: kubernetes<19,>=8.0.0 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (12.0.1)\n",
      "Requirement already satisfied: uritemplate<4,>=3.0.1 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (3.0.1)\n",
      "Requirement already satisfied: google-cloud-storage<2,>=1.20.0 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (1.41.1)\n",
      "Requirement already satisfied: protobuf<4,>=3.13.0 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (3.17.3)\n",
      "Requirement already satisfied: typing-extensions<4,>=3.10.0.2 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (3.10.0.2)\n",
      "Requirement already satisfied: google-api-python-client<2,>=1.7.8 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (1.12.10)\n",
      "Requirement already satisfied: cloudpickle<2,>=1.3.0 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (1.6.0)\n",
      "Requirement already satisfied: kfp-pipeline-spec<0.2.0,>=0.1.10 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (0.1.13)\n",
      "Requirement already satisfied: google-auth<2,>=1.6.1 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (1.34.0)\n",
      "Requirement already satisfied: strip-hints<1,>=0.1.8 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (0.1.10)\n",
      "Requirement already satisfied: docstring-parser<1,>=0.7.3 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (0.13)\n",
      "Requirement already satisfied: pydantic<2,>=1.8.2 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (1.9.0)\n",
      "Requirement already satisfied: fire<1,>=0.3.1 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (0.4.0)\n",
      "Requirement already satisfied: absl-py<=0.11,>=0.9 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (0.11.0)\n",
      "Requirement already satisfied: requests-toolbelt<1,>=0.8.0 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (0.9.1)\n",
      "Requirement already satisfied: Deprecated<2,>=1.2.7 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (1.2.13)\n",
      "Requirement already satisfied: tabulate<1,>=0.8.6 in /opt/conda/lib/python3.8/site-packages (from kfp==1.8.4) (0.8.9)\n",
      "Requirement already satisfied: six in /opt/conda/lib/python3.8/site-packages (from absl-py<=0.11,>=0.9->kfp==1.8.4) (1.16.0)\n",
      "Requirement already satisfied: wrapt<2,>=1.10 in /opt/conda/lib/python3.8/site-packages (from Deprecated<2,>=1.2.7->kfp==1.8.4) (1.13.3)\n",
      "Requirement already satisfied: termcolor in /opt/conda/lib/python3.8/site-packages (from fire<1,>=0.3.1->kfp==1.8.4) (1.1.0)\n",
      "Requirement already satisfied: google-auth-httplib2>=0.0.3 in /opt/conda/lib/python3.8/site-packages (from google-api-python-client<2,>=1.7.8->kfp==1.8.4) (0.1.0)\n",
      "Requirement already satisfied: httplib2<1dev,>=0.15.0 in /opt/conda/lib/python3.8/site-packages (from google-api-python-client<2,>=1.7.8->kfp==1.8.4) (0.20.4)\n",
      "Requirement already satisfied: google-api-core<3dev,>=1.21.0 in /opt/conda/lib/python3.8/site-packages (from google-api-python-client<2,>=1.7.8->kfp==1.8.4) (1.29.0)\n",
      "Requirement already satisfied: requests<3.0.0dev,>=2.18.0 in /opt/conda/lib/python3.8/site-packages (from google-api-core<3dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp==1.8.4) (2.27.1)\n",
      "Requirement already satisfied: setuptools>=40.3.0 in /opt/conda/lib/python3.8/site-packages (from google-api-core<3dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp==1.8.4) (49.6.0.post20210108)\n",
      "Requirement already satisfied: pytz in /opt/conda/lib/python3.8/site-packages (from google-api-core<3dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp==1.8.4) (2021.1)\n",
      "Requirement already satisfied: googleapis-common-protos<2.0dev,>=1.6.0 in /opt/conda/lib/python3.8/site-packages (from google-api-core<3dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp==1.8.4) (1.53.0)\n",
      "Requirement already satisfied: packaging>=14.3 in /opt/conda/lib/python3.8/site-packages (from google-api-core<3dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp==1.8.4) (20.9)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /opt/conda/lib/python3.8/site-packages (from google-auth<2,>=1.6.1->kfp==1.8.4) (0.2.8)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /opt/conda/lib/python3.8/site-packages (from google-auth<2,>=1.6.1->kfp==1.8.4) (4.8)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /opt/conda/lib/python3.8/site-packages (from google-auth<2,>=1.6.1->kfp==1.8.4) (4.2.2)\n",
      "Requirement already satisfied: google-resumable-media<3.0dev,>=1.3.0 in /opt/conda/lib/python3.8/site-packages (from google-cloud-storage<2,>=1.20.0->kfp==1.8.4) (2.2.1)\n",
      "Requirement already satisfied: google-cloud-core<3.0dev,>=1.6.0 in /opt/conda/lib/python3.8/site-packages (from google-cloud-storage<2,>=1.20.0->kfp==1.8.4) (2.2.2)\n",
      "Requirement already satisfied: google-crc32c<2.0dev,>=1.0 in /opt/conda/lib/python3.8/site-packages (from google-resumable-media<3.0dev,>=1.3.0->google-cloud-storage<2,>=1.20.0->kfp==1.8.4) (1.3.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.0,!=3.0.1,!=3.0.2,!=3.0.3,<4,>=2.4.2 in /opt/conda/lib/python3.8/site-packages (from httplib2<1dev,>=0.15.0->google-api-python-client<2,>=1.7.8->kfp==1.8.4) (2.4.7)\n",
      "Requirement already satisfied: pyrsistent>=0.14.0 in /opt/conda/lib/python3.8/site-packages (from jsonschema<4,>=3.0.1->kfp==1.8.4) (0.17.3)\n",
      "Requirement already satisfied: attrs>=17.4.0 in /opt/conda/lib/python3.8/site-packages (from jsonschema<4,>=3.0.1->kfp==1.8.4) (21.2.0)\n",
      "Requirement already satisfied: certifi in /opt/conda/lib/python3.8/site-packages (from kfp-server-api<2.0.0,>=1.1.2->kfp==1.8.4) (2021.5.30)\n",
      "Requirement already satisfied: python-dateutil in /opt/conda/lib/python3.8/site-packages (from kfp-server-api<2.0.0,>=1.1.2->kfp==1.8.4) (2.8.1)\n",
      "Requirement already satisfied: urllib3>=1.15 in /opt/conda/lib/python3.8/site-packages (from kfp-server-api<2.0.0,>=1.1.2->kfp==1.8.4) (1.26.5)\n",
      "Requirement already satisfied: websocket-client!=0.40.0,!=0.41.*,!=0.42.*,>=0.32.0 in /opt/conda/lib/python3.8/site-packages (from kubernetes<19,>=8.0.0->kfp==1.8.4) (1.0.1)\n",
      "Requirement already satisfied: requests-oauthlib in /opt/conda/lib/python3.8/site-packages (from kubernetes<19,>=8.0.0->kfp==1.8.4) (1.3.1)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /opt/conda/lib/python3.8/site-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.1->kfp==1.8.4) (0.4.8)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /opt/conda/lib/python3.8/site-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<3dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp==1.8.4) (2.0.12)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.8/site-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<3dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp==1.8.4) (3.2)\n",
      "Requirement already satisfied: wheel in /opt/conda/lib/python3.8/site-packages (from strip-hints<1,>=0.1.8->kfp==1.8.4) (0.36.2)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /opt/conda/lib/python3.8/site-packages (from requests-oauthlib->kubernetes<19,>=8.0.0->kfp==1.8.4) (3.2.0)\n",
      "Requirement already satisfied: kubeflow-katib==0.12.0 in /opt/conda/lib/python3.8/site-packages (0.12.0)\n",
      "Requirement already satisfied: six>=1.10 in /opt/conda/lib/python3.8/site-packages (from kubeflow-katib==0.12.0) (1.16.0)\n",
      "Requirement already satisfied: setuptools>=21.0.0 in /opt/conda/lib/python3.8/site-packages (from kubeflow-katib==0.12.0) (49.6.0.post20210108)\n",
      "Requirement already satisfied: urllib3>=1.15.1 in /opt/conda/lib/python3.8/site-packages (from kubeflow-katib==0.12.0) (1.26.5)\n",
      "Requirement already satisfied: certifi>=14.05.14 in /opt/conda/lib/python3.8/site-packages (from kubeflow-katib==0.12.0) (2021.5.30)\n",
      "Requirement already satisfied: kubernetes>=12.0.0 in /opt/conda/lib/python3.8/site-packages (from kubeflow-katib==0.12.0) (12.0.1)\n",
      "Requirement already satisfied: pyyaml>=3.12 in /opt/conda/lib/python3.8/site-packages (from kubernetes>=12.0.0->kubeflow-katib==0.12.0) (5.4.1)\n",
      "Requirement already satisfied: python-dateutil>=2.5.3 in /opt/conda/lib/python3.8/site-packages (from kubernetes>=12.0.0->kubeflow-katib==0.12.0) (2.8.1)\n",
      "Requirement already satisfied: requests-oauthlib in /opt/conda/lib/python3.8/site-packages (from kubernetes>=12.0.0->kubeflow-katib==0.12.0) (1.3.1)\n",
      "Requirement already satisfied: websocket-client!=0.40.0,!=0.41.*,!=0.42.*,>=0.32.0 in /opt/conda/lib/python3.8/site-packages (from kubernetes>=12.0.0->kubeflow-katib==0.12.0) (1.0.1)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.8/site-packages (from kubernetes>=12.0.0->kubeflow-katib==0.12.0) (2.27.1)\n",
      "Requirement already satisfied: google-auth>=1.0.1 in /opt/conda/lib/python3.8/site-packages (from kubernetes>=12.0.0->kubeflow-katib==0.12.0) (1.34.0)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /opt/conda/lib/python3.8/site-packages (from google-auth>=1.0.1->kubernetes>=12.0.0->kubeflow-katib==0.12.0) (4.8)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /opt/conda/lib/python3.8/site-packages (from google-auth>=1.0.1->kubernetes>=12.0.0->kubeflow-katib==0.12.0) (4.2.2)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /opt/conda/lib/python3.8/site-packages (from google-auth>=1.0.1->kubernetes>=12.0.0->kubeflow-katib==0.12.0) (0.2.8)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /opt/conda/lib/python3.8/site-packages (from pyasn1-modules>=0.2.1->google-auth>=1.0.1->kubernetes>=12.0.0->kubeflow-katib==0.12.0) (0.4.8)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /opt/conda/lib/python3.8/site-packages (from requests->kubernetes>=12.0.0->kubeflow-katib==0.12.0) (2.0.12)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.8/site-packages (from requests->kubernetes>=12.0.0->kubeflow-katib==0.12.0) (3.2)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /opt/conda/lib/python3.8/site-packages (from requests-oauthlib->kubernetes>=12.0.0->kubeflow-katib==0.12.0) (3.2.0)\n"
     ]
    }
   ],
   "source": [
    "# Install required packages (Kubeflow Pipelines and Katib SDK).\n",
    "!pip install kfp==1.8.4\n",
    "!pip install kubeflow-katib==0.12.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kfp\n",
    "import kfp.dsl as dsl\n",
    "from kfp import components\n",
    "\n",
    "from kubeflow.katib import ApiClient\n",
    "from kubeflow.katib import V1beta1ExperimentSpec\n",
    "from kubeflow.katib import V1beta1AlgorithmSpec\n",
    "from kubeflow.katib import V1beta1ObjectiveSpec\n",
    "from kubeflow.katib import V1beta1ParameterSpec\n",
    "from kubeflow.katib import V1beta1FeasibleSpace\n",
    "from kubeflow.katib import V1beta1TrialTemplate\n",
    "from kubeflow.katib import V1beta1TrialParameterSpec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the Pipelines tasks\n",
    "\n",
    "To run this Pipeline, you should define:\n",
    "1. Katib hyperparameter tuning\n",
    "2. TFJob training\n",
    "3. KServe inference\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1. Katib hyperparameter tuning task\n",
    "\n",
    "Create the Kubeflow Pipelines task for the Katib hyperparameter tuning. This Experiment uses \"random\" algorithm and TFJob for the Trial's worker.\n",
    "\n",
    "The Katib Experiment is similar to this example: https://github.com/kubeflow/katib/blob/master/examples/v1beta1/kubeflow-training-operator/tfjob-mnist-with-summaries.yaml."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You should define the Experiment name, namespace and number of training steps in the arguments.\n",
    "def create_katib_experiment_task(experiment_name, experiment_namespace, training_steps):\n",
    "    # Trial count specification.\n",
    "    max_trial_count = 5\n",
    "    max_failed_trial_count = 3\n",
    "    parallel_trial_count = 2\n",
    "\n",
    "    # Objective specification.\n",
    "    objective = V1beta1ObjectiveSpec(\n",
    "        type=\"minimize\",\n",
    "        goal=0.001,\n",
    "        objective_metric_name=\"loss\"\n",
    "    )\n",
    "\n",
    "    # Algorithm specification.\n",
    "    algorithm = V1beta1AlgorithmSpec(\n",
    "        algorithm_name=\"random\",\n",
    "    )\n",
    "\n",
    "    # Experiment search space.\n",
    "    # In this example we tune learning rate and batch size.\n",
    "    parameters = [\n",
    "        V1beta1ParameterSpec(\n",
    "            name=\"learning_rate\",\n",
    "            parameter_type=\"double\",\n",
    "            feasible_space=V1beta1FeasibleSpace(\n",
    "                min=\"0.01\",\n",
    "                max=\"0.05\"\n",
    "            ),\n",
    "        ),\n",
    "        V1beta1ParameterSpec(\n",
    "            name=\"batch_size\",\n",
    "            parameter_type=\"int\",\n",
    "            feasible_space=V1beta1FeasibleSpace(\n",
    "                min=\"80\",\n",
    "                max=\"100\"\n",
    "            ),\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    # Experiment Trial template.\n",
    "    # TODO (andreyvelich): Use community image for the mnist example.\n",
    "    trial_spec = {\n",
    "        \"apiVersion\": \"kubeflow.org/v1\",\n",
    "        \"kind\": \"TFJob\",\n",
    "        \"spec\": {\n",
    "            \"tfReplicaSpecs\": {\n",
    "                \"Chief\": {\n",
    "                    \"replicas\": 1,\n",
    "                    \"restartPolicy\": \"OnFailure\",\n",
    "                    \"template\": {\n",
    "                        \"metadata\": {\n",
    "                            \"annotations\": {\n",
    "                                \"sidecar.istio.io/inject\": \"false\"\n",
    "                            }\n",
    "                        },\n",
    "                        \"spec\": {\n",
    "                            \"containers\": [\n",
    "                                {\n",
    "                                    \"name\": \"tensorflow\",\n",
    "                                    \"image\": \"docker.io/liuhougangxa/tf-estimator-mnist\",\n",
    "                                    \"command\": [\n",
    "                                        \"python\",\n",
    "                                        \"/opt/model.py\",\n",
    "                                        \"--tf-train-steps=\" + str(training_steps),\n",
    "                                        \"--tf-learning-rate=${trialParameters.learningRate}\",\n",
    "                                        \"--tf-batch-size=${trialParameters.batchSize}\"\n",
    "                                    ]\n",
    "                                }\n",
    "                            ]\n",
    "                        }\n",
    "                    }\n",
    "                },\n",
    "                \"Worker\": {\n",
    "                    \"replicas\": 1,\n",
    "                    \"restartPolicy\": \"OnFailure\",\n",
    "                    \"template\": {\n",
    "                        \"metadata\": {\n",
    "                            \"annotations\": {\n",
    "                                \"sidecar.istio.io/inject\": \"false\"\n",
    "                            }\n",
    "                        },\n",
    "                        \"spec\": {\n",
    "                            \"containers\": [\n",
    "                                {\n",
    "                                    \"name\": \"tensorflow\",\n",
    "                                    \"image\": \"docker.io/liuhougangxa/tf-estimator-mnist\",\n",
    "                                    \"command\": [\n",
    "                                        \"python\",\n",
    "                                        \"/opt/model.py\",\n",
    "                                        \"--tf-train-steps=\" + str(training_steps),\n",
    "                                        \"--tf-learning-rate=${trialParameters.learningRate}\",\n",
    "                                        \"--tf-batch-size=${trialParameters.batchSize}\"\n",
    "                                    ]\n",
    "                                }\n",
    "                            ]\n",
    "                        }\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "    # Configure parameters for the Trial template.\n",
    "    trial_template = V1beta1TrialTemplate(\n",
    "        primary_container_name=\"tensorflow\",\n",
    "        trial_parameters=[\n",
    "            V1beta1TrialParameterSpec(\n",
    "                name=\"learningRate\",\n",
    "                description=\"Learning rate for the training model\",\n",
    "                reference=\"learning_rate\"\n",
    "            ),\n",
    "            V1beta1TrialParameterSpec(\n",
    "                name=\"batchSize\",\n",
    "                description=\"Batch size for the model\",\n",
    "                reference=\"batch_size\"\n",
    "            ),\n",
    "        ],\n",
    "        trial_spec=trial_spec\n",
    "    )\n",
    "\n",
    "    # Create an Experiment from the above parameters.\n",
    "    experiment_spec = V1beta1ExperimentSpec(\n",
    "        max_trial_count=max_trial_count,\n",
    "        max_failed_trial_count=max_failed_trial_count,\n",
    "        parallel_trial_count=parallel_trial_count,\n",
    "        objective=objective,\n",
    "        algorithm=algorithm,\n",
    "        parameters=parameters,\n",
    "        trial_template=trial_template\n",
    "    )\n",
    "\n",
    "    # Create the KFP task for the Katib Experiment.\n",
    "    # Experiment Spec should be serialized to a valid Kubernetes object.\n",
    "    katib_experiment_launcher_op = components.load_component_from_url(\n",
    "        \"https://raw.githubusercontent.com/kubeflow/pipelines/master/components/kubeflow/katib-launcher/component.yaml\")\n",
    "    op = katib_experiment_launcher_op(\n",
    "        experiment_name=experiment_name,\n",
    "        experiment_namespace=experiment_namespace,\n",
    "        experiment_spec=ApiClient().sanitize_for_serialization(experiment_spec),\n",
    "        experiment_timeout_minutes=60,\n",
    "        delete_finished_experiment=False)\n",
    "\n",
    "    return op"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2. TFJob training task\n",
    "\n",
    "Create the Kubeflow Pipelines task for the TFJob training. In this example TFJob runs the Chief and Worker with 1 replica.\n",
    "\n",
    "Learn more about TFJob replica specifications in the Kubeflow docs: https://www.kubeflow.org/docs/components/training/tftraining/#what-is-tfjob."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function converts Katib Experiment HP results to args.\n",
    "def convert_katib_results(katib_results) -> str:\n",
    "    import json\n",
    "    import pprint\n",
    "    katib_results_json = json.loads(katib_results)\n",
    "    print(\"Katib results:\")\n",
    "    pprint.pprint(katib_results_json)\n",
    "    best_hps = []\n",
    "    for pa in katib_results_json[\"currentOptimalTrial\"][\"parameterAssignments\"]:\n",
    "        if pa[\"name\"] == \"learning_rate\":\n",
    "            best_hps.append(\"--tf-learning-rate=\" + pa[\"value\"])\n",
    "        elif pa[\"name\"] == \"batch_size\":\n",
    "            best_hps.append(\"--tf-batch-size=\" + pa[\"value\"])\n",
    "    print(\"Best Hyperparameters: {}\".format(best_hps))\n",
    "    return \" \".join(best_hps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You should define the TFJob name, namespace, number of training steps, output of Katib and model volume tasks in the arguments.\n",
    "def create_tfjob_task(tfjob_name, tfjob_namespace, training_steps, katib_op, model_volume_op):\n",
    "    import json\n",
    "    # Get parameters from the Katib Experiment.\n",
    "    # Parameters are in the format \"--tf-learning-rate=0.01 --tf-batch-size=100\"\n",
    "    convert_katib_results_op = components.func_to_container_op(convert_katib_results)\n",
    "    best_hp_op = convert_katib_results_op(katib_op.output)\n",
    "    best_hps = str(best_hp_op.output)\n",
    "\n",
    "    # Create the TFJob Chief and Worker specification with the best Hyperparameters.\n",
    "    # TODO (andreyvelich): Use community image for the mnist example.\n",
    "    tfjob_chief_spec = {\n",
    "        \"replicas\": 1,\n",
    "        \"restartPolicy\": \"OnFailure\",\n",
    "        \"template\": {\n",
    "            \"metadata\": {\n",
    "                \"annotations\": {\n",
    "                    \"sidecar.istio.io/inject\": \"false\"\n",
    "                }\n",
    "            },\n",
    "            \"spec\": {\n",
    "                \"containers\": [\n",
    "                    {\n",
    "                        \"name\": \"tensorflow\",\n",
    "                        \"image\": \"docker.io/liuhougangxa/tf-estimator-mnist\",\n",
    "                        \"command\": [\n",
    "                            \"sh\",\n",
    "                            \"-c\"\n",
    "                        ],\n",
    "                        \"args\": [\n",
    "                            \"python /opt/model.py --tf-export-dir=/mnt/export --tf-train-steps={} {}\".format(training_steps, best_hps)\n",
    "                        ],\n",
    "                        \"volumeMounts\": [\n",
    "                            {\n",
    "                                \"mountPath\": \"/mnt/export\",\n",
    "                                \"name\": \"model-volume\"\n",
    "                            }\n",
    "                        ]\n",
    "                    }\n",
    "                ],\n",
    "                \"volumes\": [\n",
    "                    {\n",
    "                        \"name\": \"model-volume\",\n",
    "                        \"persistentVolumeClaim\": {\n",
    "                            \"claimName\": str(model_volume_op.outputs[\"name\"])\n",
    "                        }\n",
    "                    }\n",
    "                ]\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "    tfjob_worker_spec = {\n",
    "        \"replicas\": 1,\n",
    "        \"restartPolicy\": \"OnFailure\",\n",
    "        \"template\": {\n",
    "            \"metadata\": {\n",
    "                \"annotations\": {\n",
    "                    \"sidecar.istio.io/inject\": \"false\"\n",
    "                }\n",
    "            },\n",
    "            \"spec\": {\n",
    "                \"containers\": [\n",
    "                    {\n",
    "                        \"name\": \"tensorflow\",\n",
    "                        \"image\": \"docker.io/liuhougangxa/tf-estimator-mnist\",\n",
    "                        \"command\": [\n",
    "                            \"sh\",\n",
    "                            \"-c\",\n",
    "                        ],\n",
    "                        \"args\": [\n",
    "                          \"python /opt/model.py --tf-export-dir=/mnt/export --tf-train-steps={} {}\".format(training_steps, best_hps) \n",
    "                        ],\n",
    "                    }\n",
    "                ],\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "    # Create the KFP task for the TFJob.\n",
    "    tfjob_launcher_op = components.load_component_from_url(\n",
    "        \"https://raw.githubusercontent.com/kubeflow/pipelines/master/components/kubeflow/launcher/component.yaml\")\n",
    "    op = tfjob_launcher_op(\n",
    "        name=tfjob_name,\n",
    "        namespace=tfjob_namespace,\n",
    "        chief_spec=json.dumps(tfjob_chief_spec),\n",
    "        worker_spec=json.dumps(tfjob_worker_spec),\n",
    "        tfjob_timeout_minutes=60,\n",
    "        delete_finished_tfjob=False)\n",
    "    return op"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3. KServe inference\n",
    "\n",
    "Create the Kubeflow Pipelines task for the KServe inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_serving_task(model_name, model_namespace, tfjob_op, model_volume_op):\n",
    "\n",
    "    api_version = 'serving.kserve.io/v1beta1'\n",
    "    serving_component_url = 'https://raw.githubusercontent.com/kubeflow/pipelines/master/components/kserve/component.yaml'\n",
    "\n",
    "    # Uncomment the following two lines if you are using KFServing v0.6.x or v0.5.x\n",
    "    # api_version = 'serving.kubeflow.org/v1beta1'\n",
    "    # serving_component_url = 'https://raw.githubusercontent.com/kubeflow/pipelines/master/components/kubeflow/kfserving/component.yaml'\n",
    "\n",
    "    inference_service = '''\n",
    "apiVersion: \"{}\"\n",
    "kind: \"InferenceService\"\n",
    "metadata:\n",
    "  name: {}\n",
    "  namespace: {}\n",
    "  annotations:\n",
    "    \"sidecar.istio.io/inject\": \"false\"\n",
    "spec:\n",
    "  predictor:\n",
    "    tensorflow:\n",
    "      storageUri: \"pvc://{}/\"\n",
    "'''.format(api_version, model_name, model_namespace, str(model_volume_op.outputs[\"name\"]))\n",
    "\n",
    "    serving_launcher_op = components.load_component_from_url(serving_component_url)\n",
    "    serving_launcher_op(action=\"apply\", inferenceservice_yaml=inference_service).after(tfjob_op)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the Kubeflow Pipeline\n",
    "\n",
    "You should create the Kubeflow Pipeline from the above tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'size': {{pipelineparam:op=model-volume;name=size}}, 'name': {{pipelineparam:op=model-volume;name=name}}, 'manifest': {{pipelineparam:op=model-volume;name=manifest}}}\n",
      "{{pipelineparam:op=model-volume;name=name}}\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<a href=\"/pipeline/#/experiments/details/a4d379f6-310d-4d9c-af53-6bd2f6118036\" target=\"_blank\" >Experiment details</a>."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<a href=\"/pipeline/#/runs/details/9519f884-8baf-4768-a728-29de8ef5b4e6\" target=\"_blank\" >Run details</a>."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run ID:  9519f884-8baf-4768-a728-29de8ef5b4e6\n"
     ]
    }
   ],
   "source": [
    "name=\"mnist-e2e\"\n",
    "namespace=\"kubeflow-user-example-com\"\n",
    "training_steps=\"200\"\n",
    "\n",
    "@dsl.pipeline(\n",
    "    name=\"End to End Pipeline\",\n",
    "    description=\"An end to end mnist example including hyperparameter tuning, train and inference\"\n",
    ")\n",
    "def mnist_pipeline(name=name, namespace=namespace, training_steps=training_steps):\n",
    "    # Run the hyperparameter tuning with Katib.\n",
    "    katib_op = create_katib_experiment_task(name, namespace, training_steps)\n",
    "\n",
    "    # Create volume to train and serve the model.\n",
    "    model_volume_op = dsl.VolumeOp(\n",
    "        name=\"model-volume\",\n",
    "        resource_name=\"model-volume\",\n",
    "        size=\"1Gi\",\n",
    "        modes=dsl.VOLUME_MODE_RWO\n",
    "    )\n",
    "\n",
    "    # Run the distributive training with TFJob.\n",
    "    tfjob_op = create_tfjob_task(name, namespace, training_steps, katib_op, model_volume_op)\n",
    "\n",
    "    # Create the KServe inference.\n",
    "    create_serving_task(name, namespace, tfjob_op, model_volume_op)\n",
    "# Run the Kubeflow Pipeline in the user's namespace.\n",
    "\n",
    "kfp_client=kfp.Client()\n",
    "run_id = kfp_client.create_run_from_pipeline_func(mnist_pipeline, namespace=namespace, arguments={}).run_id\n",
    "print(\"Run ID: \", run_id)"
   ]
  },
  {
   "attachments": {
    "f947c4a5-dc78-4ba4-8e47-ae73d8f0ecea.png": {
     "image/png": "[encoded data, 112716 characters]"
    }
   },
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The finished Pipeline should look as follows. ![Kubeflow Pipeline.png](attachment:f947c4a5-dc78-4ba4-8e47-ae73d8f0ecea.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict from the trained model\n",
    "\n",
    "Once Kubeflow Pipeline is finished, you are able to call the API endpoint with [mnist image](https://raw.githubusercontent.com/kubeflow/katib/master/examples/v1beta1/kubeflow-pipelines/images/9.bmp) to predict from the trained model.\n",
    "\n",
    "**Note**: If you are using Kubeflow + Dex setup and runing this Notebook outside of your Kubernetes cluster, follow [this guide](https://github.com/kserve/kserve/tree/master/docs/samples/istio-dex#authentication) to get Session ID for the API requests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run 9519f884-8baf-4768-a728-29de8ef5b4e6 has been Succeeded\n",
      "\n",
      "Prediction for the image\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-8-a8037f6b811f>:13: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  data = np.array(image.convert('L').resize((28, 28))).astype(np.float).reshape(-1, 28, 28, 1)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAABwAAAAcCAAAAABXZoBIAAAA1ElEQVR4nN3QPwtBYRQG8EMU0e0uZLIw+QKXRZlMGC0GX8CglE0pk0VxPwQmE5YrJYPVIjYMlImSwXNiMOi97319AM/6O6fzh+g/Y5hr5mrRNByseAZba4D7EnlSN8wy3uAYXJOwDEw0ohKwD9mtxehqRLQBCnZr8GPkJ/Ll79y0m37GiIjiK2AQsGMYiIbryyvjmZO20U9gAIcjTg43GhfethOROToO+En6xRUlZhnSjd+I6BY7xVIRY79w4XapR9IOSTWWYSWUqE0xlH771R7UrULefm5U2pxVCt0AAAAASUVORK5CYII=",
      "text/plain": [
       "<PIL.BmpImagePlugin.BmpImageFile image mode=L size=28x28 at 0x7F1A671899A0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'predictions': [{'predictions': [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], 'classes': 9}]}\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "import requests\n",
    "\n",
    "# Pipeline Run should be succeeded.\n",
    "kfp_run = kfp_client.get_run(run_id=run_id)\n",
    "if kfp_run.run.status == \"Succeeded\":\n",
    "    print(\"Run {} has been Succeeded\\n\".format(run_id))\n",
    "\n",
    "    # Specify the image URL here.\n",
    "    image_url = \"https://raw.githubusercontent.com/kubeflow/katib/master/examples/v1beta1/kubeflow-pipelines/images/9.bmp\"\n",
    "    image = Image.open(requests.get(image_url, stream=True).raw)\n",
    "    data = np.array(image.convert('L').resize((28, 28))).astype(np.float).reshape(-1, 28, 28, 1)\n",
    "    data_formatted = np.array2string(data, separator=\",\", formatter={\"float\": lambda x: \"%.1f\" % x})\n",
    "    json_request = '{{ \"instances\" : {} }}'.format(data_formatted)\n",
    "\n",
    "    # Specify the prediction URL. If you are runing this notebook outside of Kubernetes cluster, you should set the Cluster IP.\n",
    "    url = \"http://{}-predictor-default.{}.svc.cluster.local/v1/models/{}:predict\".format(name, namespace, name)\n",
    "    response = requests.post(url, data=json_request)\n",
    "\n",
    "    print(\"Prediction for the image\")\n",
    "    display(image)\n",
    "    print(response.json())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
