{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9a6008b8",
   "metadata": {},
   "source": [
    "# MNIST Digits Recognition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4684c50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ######################################### E2E Demo example precondition ##############################################\n",
    "\n",
    "# # You need to generate secret to have ability read datasource files from S3 bucket by Spark application (Airflow DAG)\n",
    "# # Copy template \"object_store_secret.yaml.tpl\" from shared/aie-tutorials/current-release/Data-Analytics/Spark directory to e.g. user folder\n",
    "\n",
    "# # Users have to pull the following images manually and make available in local repo for clusters in airgap network to run the example.\n",
    "  # \"docker.io/kubeflowkatib/kubeflow-pipelines-launcher\", \"nikenano/launchernew:latest\", \"quay.io/aipipeline/kserve-component:v0.10.1\"\n",
    "\n",
    "import os\n",
    "\n",
    "namespace = os.getenv('NOTEBOOK_NAMESPACE')\n",
    "\n",
    "!sed -e \"s/\\$AUTH_TOKEN/$AUTH_TOKEN/\" /mnt/user/object_store_secret.yaml.tpl > /tmp/object_store_secret.yaml\n",
    "!kubectl apply -f /tmp/object_store_secret.yaml -n {namespace}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88cb0bbf-cf10-4b3a-97bf-60ff80878545",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# # Local airgap repo path prefix shall be added to images mentioned above in following yaml files.\n",
    "  # \"component/katib-launcher-component.yaml\", \"component/kubeflow-launcher-component.yaml\", \"component/kserve-component.yaml\"\n",
    "\n",
    "if [[ -n \"$AIRGAP_REGISTRY\" ]]; then\n",
    "    echo 'AIRGAP_REGISTRY is set - updating component YAML files'\n",
    "    for file in `ls component/*.yaml`; do\n",
    "        echo \"$file\";\n",
    "        sed \"\\%image: *$AIRGAP_REGISTRY%b; s%image: *%image: $AIRGAP_REGISTRY%g\" --in-place \"$file\";\n",
    "    done;\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89bb0c90-a09f-4197-b9a2-f7558efb31de",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# Specify proxy and no_proxy values for pipeline from notebook environment variables. Update manually if needed.\n",
    "http_proxy = os.environ[\"http_proxy\"] ## \"http://<proxy_host>:<port>\"\n",
    "# After default.svc append comma separated list of IP address, hostname, CIDR block, service name, etc.\n",
    "no_proxy = os.environ[\"no_proxy\"] ## \"localhost,.cluster.local,.svc,.default.svc, <host>:<port>, x.x.x.x/y,...\"\n",
    "# Specify airgap registry to use as a prefix for images\n",
    "airgap_registry = os.environ[\"AIRGAP_REGISTRY\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a8f39a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the name of the user mounted directory\n",
    "user_mounted_dir_name = \"user\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c1a88c5-b825-4ddc-ac35-dfa6dfd3c067",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import uuid\n",
    "import kfp\n",
    "import kfp.dsl as dsl\n",
    "from kfp import components\n",
    "from pathlib import Path\n",
    "\n",
    "######################################### Storage Parameters ##############################################\n",
    "\n",
    "# The path that is used inside the pods to mount training data\n",
    "mnt_path = \"/mnt/data/\"\n",
    "uuid = uuid.uuid4().hex[:4]\n",
    "\n",
    "# The initial training data is written to the user volume by the spark job in the Apache Parquet format.\n",
    "initial_training_data_dir = f\"{str(Path.home())}/{user_mounted_dir_name}/mnist-spark-data\"\n",
    "\n",
    "# The path is relative, final_training_data_dir should be in the same folder with the notebook\n",
    "final_training_data_dir=f\"training-{uuid}\"\n",
    "\n",
    "os.mkdir(\"/mnt/user/\" + final_training_data_dir, 0o777)\n",
    "\n",
    "######################################## KFP parameters ##################################################\n",
    "\n",
    "# Name of the Katib experiment\n",
    "name = f\"mnist-experiment-{uuid}\"\n",
    "print(f\"Katib experiment name: {name}\")\n",
    "\n",
    "# Number of epoch to train the model\n",
    "training_steps=\"100\"\n",
    "volume_str = f\"model-volume-{uuid}\"\n",
    "\n",
    "# Create volume to train and serve the model.\n",
    "!sed -e \"s/pvc-name/$volume_str/\" ./pvc.yaml > /tmp/pvc.yaml\n",
    "!kubectl apply -f /tmp/pvc.yaml\n",
    "\n",
    "namespace = os.getenv('NOTEBOOK_NAMESPACE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72049f0f-a0d2-4b37-94e2-57c433bcf0a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transformation of the training data from Apache Parquet format to the format that is required for the MNIST example\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "if not os.path.exists(final_training_data_dir):\n",
    "    os.makedirs(final_training_data_dir)\n",
    "    \n",
    "with open(final_training_data_dir + \"/train-images-idx3-ubyte.gz\", 'wb') as f1, \\\n",
    "     open(final_training_data_dir + \"/t10k-images-idx3-ubyte.gz\", 'wb') as f2, \\\n",
    "     open(final_training_data_dir + \"/train-labels-idx1-ubyte.gz\", 'wb') as f3, \\\n",
    "     open(final_training_data_dir + \"/t10k-labels-idx1-ubyte.gz\", 'wb') as f4:\n",
    "         mnist_parquet = pd.read_parquet(initial_training_data_dir)\n",
    "         x_train, x_test, y_train, y_test = mnist_parquet[\"content\"]\n",
    "         f1.write(x_train)\n",
    "         f2.write(x_test)\n",
    "         f3.write(y_train)\n",
    "         f4.write(y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90bf2234",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = []\n",
    "if http_proxy:\n",
    "    env.append({\"name\": \"HTTP_PROXY\", \"value\":http_proxy})\n",
    "    env.append({\"name\": \"HTTPS_PROXY\", \"value\":http_proxy})\n",
    "    env.append({\"name\": \"http_proxy\", \"value\":http_proxy})\n",
    "    env.append({\"name\": \"https_proxy\", \"value\":http_proxy})\n",
    "if no_proxy:\n",
    "    env.append({\"name\": \"no_proxy\", \"value\":no_proxy})\n",
    "    env.append({\"name\": \"NO_PROXY\", \"value\":no_proxy})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e726f371-df27-44ad-931a-3105da2496a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You should define 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 = {\n",
    "        \"objectiveMetricName\": \"loss\",\n",
    "        \"type\": \"minimize\",\n",
    "        \"goal\": 0.001,\n",
    "    }\n",
    "\n",
    "    # Algorithm specification.\n",
    "    algorithm = {\n",
    "        \"algorithmName\": \"random\",\n",
    "    }\n",
    "\n",
    "    # Experiment search space.\n",
    "    # In this example we tune learning rate and batch size.\n",
    "    parameters = [\n",
    "        {\n",
    "            \"name\": \"learning_rate\",\n",
    "            \"parameterType\": \"double\",\n",
    "            \"feasibleSpace\": {\n",
    "                \"min\": \"0.01\",\n",
    "                \"max\": \"0.05\",\n",
    "            },\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"batch_size\",\n",
    "            \"parameterType\": \"int\",\n",
    "            \"feasibleSpace\": {\n",
    "                \"min\": \"80\",\n",
    "                \"max\": \"100\",\n",
    "            },\n",
    "        },\n",
    "    ]\n",
    "\n",
    "    # Experiment Trial template.\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\": f\"{airgap_registry}gcr.io/mapr-252711/kubeflow/kfexamples/docker.io/liuhougangxa/tf-estimator-mnist:ezua-1.7\",\n",
    "                                    \"command\": [\n",
    "                                        \"python\",\n",
    "                                        \"/opt/model.py\",\n",
    "                                        f\"--tf-data-dir={mnt_path}{final_training_data_dir}/\",\n",
    "                                        f\"--tf-train-steps={str(training_steps)}\",\n",
    "                                        f\"--tf-export-dir={mnt_path}{final_training_data_dir}/\",\n",
    "                                        \"--tf-learning-rate=${trialParameters.learningRate}\",\n",
    "                                        \"--tf-batch-size=${trialParameters.batchSize}\"\n",
    "                                    ],\n",
    "                                    \"env\": env,\n",
    "                                    \"volumeMounts\": [\n",
    "                                        {\n",
    "                                            \"mountPath\": mnt_path,\n",
    "                                            \"name\": \"data-volume\"\n",
    "                                        }\n",
    "                                    ],\n",
    "                                    \"resources\":\n",
    "                                    {\n",
    "                                        \"limits\":{\n",
    "                                            \"cpu\":\"2\",\n",
    "                                            \"memory\":\"1Gi\"\n",
    "                                        },\n",
    "                                        \"requests\":{\n",
    "                                            \"cpu\":\"100m\",\n",
    "                                            \"memory\":\"50Mi\"\n",
    "                                        }\n",
    "                                    }      \n",
    "                                }\n",
    "                            ],\n",
    "                            \"imagePullSecrets\": [\n",
    "                                {\n",
    "                                    \"name\": \"hpe-imagepull-secrets\"\n",
    "                                }\n",
    "                            ],\n",
    "                            \"volumes\": [\n",
    "                                {\n",
    "                                    \"name\": \"data-volume\",\n",
    "                                    \"persistentVolumeClaim\": {\n",
    "                                        \"claimName\": \"user-pvc\"\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\": f\"{airgap_registry}gcr.io/mapr-252711/kubeflow/kfexamples/docker.io/liuhougangxa/tf-estimator-mnist:ezua-1.7\",\n",
    "                                    \"command\": [\n",
    "                                        \"python\",\n",
    "                                        \"/opt/model.py\",\n",
    "                                        f\"--tf-data-dir={mnt_path}{final_training_data_dir}/\",\n",
    "                                        f\"--tf-train-steps={str(training_steps)}\",\n",
    "                                        f\"--tf-export-dir={mnt_path}{final_training_data_dir}/\",\n",
    "                                        \"--tf-learning-rate=${trialParameters.learningRate}\",\n",
    "                                        \"--tf-batch-size=${trialParameters.batchSize}\"\n",
    "                                    ],\n",
    "                                    \"env\": env,\n",
    "                                    \"volumeMounts\": [\n",
    "                                        {\n",
    "                                            \"mountPath\": mnt_path,\n",
    "                                            \"name\": \"data-volume\"\n",
    "                                        }\n",
    "                                    ],\n",
    "                                    \"resources\":\n",
    "                                    {\n",
    "                                        \"limits\":{\n",
    "                                            \"cpu\":\"2\",\n",
    "                                            \"memory\":\"1Gi\"\n",
    "                                        },\n",
    "                                        \"requests\":{\n",
    "                                            \"cpu\":\"100m\",\n",
    "                                            \"memory\":\"50Mi\"\n",
    "                                        }\n",
    "                                    }\n",
    "                                }\n",
    "                            ],\n",
    "                            \"imagePullSecrets\": [\n",
    "                                {\n",
    "                                    \"name\": \"hpe-imagepull-secrets\"\n",
    "                                }\n",
    "                            ],\n",
    "                            \"volumes\": [\n",
    "                                {\n",
    "                                    \"name\": \"data-volume\",\n",
    "                                    \"persistentVolumeClaim\": {\n",
    "                                        \"claimName\": \"user-pvc\"\n",
    "                                    }\n",
    "                                }\n",
    "                            ]\n",
    "                        }\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "    # Configure parameters for the Trial template.\n",
    "    trial_template = {\n",
    "        \"primaryContainerName\": \"tensorflow\",\n",
    "        \"trialParameters\": [\n",
    "            {\n",
    "                \"name\": \"learningRate\",\n",
    "                \"description\": \"Learning rate for the training model\",\n",
    "                \"reference\": \"learning_rate\",\n",
    "            },\n",
    "            {\n",
    "                \"name\": \"batchSize\",\n",
    "                \"description\": \"Batch size for the model\",\n",
    "                \"reference\": \"batch_size\",\n",
    "            },\n",
    "        ],\n",
    "        \"trialSpec\": trial_spec,\n",
    "    }\n",
    "\n",
    "    # Create an Experiment from the above parameters.\n",
    "    experiment_spec = {\n",
    "        \"maxTrialCount\": max_trial_count,\n",
    "        \"maxFailedTrialCount\": max_failed_trial_count,\n",
    "        \"parallelTrialCount\": parallel_trial_count,\n",
    "        \"objective\": objective,\n",
    "        \"algorithm\": algorithm,\n",
    "        \"parameters\": parameters,\n",
    "        \"trialTemplate\": 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_file(\"component/katib-launcher-component.yaml\")\n",
    "\n",
    "    op = katib_experiment_launcher_op(\n",
    "        experiment_name=experiment_name,\n",
    "        experiment_namespace=experiment_namespace,\n",
    "        experiment_spec=experiment_spec,\n",
    "        experiment_timeout_minutes=60,\n",
    "        delete_finished_experiment=False)\n",
    "\n",
    "    return op"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d73c7c8-edfd-4945-85b1-a4c35817da05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function converts Katib Experiment HP results to args.\n",
    "\n",
    "from kfp import components\n",
    "@dsl.component(base_image = f\"{airgap_registry}python:3.9\")\n",
    "def convert_katib_results(katib_results: dict) -> str:\n",
    "    import json\n",
    "    import pprint\n",
    "    katib_results_json = 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(f\"Best Hyperparameters: {best_hps}\")\n",
    "    return \" \".join(best_hps)\n",
    "\n",
    "# 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, http_proxy, no_proxy):\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",
    "    best_hp_op = convert_katib_results(katib_results = katib_op.output)\n",
    "    best_hp_op.set_caching_options(False)\n",
    "    # Setting the proxy and no_proxy to allow download of the required packages inside the component.\n",
    "    if no_proxy:\n",
    "        no_proxy += ',.kubeflow'\n",
    "        best_hp_op.set_env_variable('no_proxy', no_proxy)\n",
    "        best_hp_op.set_env_variable('NO_PROXY', no_proxy)\n",
    "    if http_proxy : \n",
    "        best_hp_op.set_env_variable('http_proxy', http_proxy)\n",
    "        best_hp_op.set_env_variable('https_proxy', http_proxy)\n",
    "        best_hp_op.set_env_variable('HTTP_PROXY', http_proxy)\n",
    "        best_hp_op.set_env_variable('HTTPS_PROXY', http_proxy)\n",
    "    \n",
    "    best_hps = str(best_hp_op.output)\n",
    "\n",
    "    # Create the TFJob Chief and Worker specification with the best Hyperparameters.\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\": f\"{airgap_registry}gcr.io/mapr-252711/kubeflow/kfexamples/docker.io/liuhougangxa/tf-estimator-mnist:ezua-1.7\",\n",
    "                        \"command\": [\n",
    "                            \"sh\",\n",
    "                            \"-c\"\n",
    "                        ],\n",
    "                        \"args\": [\n",
    "                            f\"python /opt/model.py --tf-data-dir={mnt_path}{final_training_data_dir}/ --tf-export-dir=/mnt/export --tf-train-steps={training_steps} {best_hps}\"\n",
    "                        ],\n",
    "                        \"env\": env,\n",
    "                        \"volumeMounts\": [\n",
    "                            {\n",
    "                                \"mountPath\": \"/mnt/export\",\n",
    "                                \"name\": \"model-volume\"\n",
    "                            },\n",
    "                            {\n",
    "                                \"mountPath\": mnt_path,\n",
    "                                \"name\": \"data-volume\"\n",
    "                            }\n",
    "                        ],\n",
    "                        \"resources\":\n",
    "                        {\n",
    "                            \"limits\":{\n",
    "                                \"cpu\":\"2\",\n",
    "                                \"memory\": \"1Gi\"\n",
    "                            },\n",
    "                            \"requests\":{\n",
    "                                \"cpu\": \"100m\",\n",
    "                                \"memory\":\"50Mi\"\n",
    "                            }\n",
    "                        }\n",
    "                    }\n",
    "                ],\n",
    "                \"imagePullSecrets\": [\n",
    "                    {\n",
    "                        \"name\": \"hpe-imagepull-secrets\"\n",
    "                    }\n",
    "                ],\n",
    "                \"volumes\": [\n",
    "                    {\n",
    "                        \"name\": \"model-volume\",\n",
    "                        \"persistentVolumeClaim\": {\n",
    "                            \"claimName\": volume_str\n",
    "                        }\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"data-volume\",\n",
    "                        \"persistentVolumeClaim\": {\n",
    "                            \"claimName\": \"user-pvc\"\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\": f\"{airgap_registry}gcr.io/mapr-252711/kubeflow/kfexamples/docker.io/liuhougangxa/tf-estimator-mnist:ezua-1.7\",\n",
    "                        \"command\": [\n",
    "                            \"sh\",\n",
    "                            \"-c\",\n",
    "                        ],\n",
    "                        \"args\": [\n",
    "                          f\"python /opt/model.py --tf-data-dir={mnt_path}{final_training_data_dir}/ --tf-export-dir=/mnt/export --tf-train-steps={training_steps} {best_hps}\" \n",
    "                        ],\n",
    "                        \"env\": env,\n",
    "                        \"volumeMounts\": [\n",
    "                            {\n",
    "                                \"mountPath\": \"/mnt/export\",\n",
    "                                \"name\": \"model-volume\"\n",
    "                            },\n",
    "                            {\n",
    "                                \"mountPath\": mnt_path,\n",
    "                                \"name\": \"data-volume\"\n",
    "                            }\n",
    "                        ],\n",
    "                        \"resources\":\n",
    "                        {\n",
    "                            \"limits\":{\n",
    "                                \"cpu\":\"2\",\n",
    "                                \"memory\":\"1Gi\"\n",
    "                            },\n",
    "                            \"requests\":{\n",
    "                                \"cpu\": \"100m\",\n",
    "                                \"memory\":\"50Mi\"\n",
    "                            }\n",
    "                        }\n",
    "                    }\n",
    "                ],\n",
    "                \"imagePullSecrets\": [\n",
    "                    {\n",
    "                        \"name\": \"hpe-imagepull-secrets\"\n",
    "                    }\n",
    "                ],\n",
    "                \"volumes\": [\n",
    "                    {\n",
    "                        \"name\": \"model-volume\",\n",
    "                        \"persistentVolumeClaim\": {\n",
    "                            \"claimName\": volume_str\n",
    "                        }\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"data-volume\",\n",
    "                        \"persistentVolumeClaim\": {\n",
    "                            \"claimName\": \"user-pvc\"\n",
    "                        }\n",
    "                    }\n",
    "                ]\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "    # Create the KFP task for the TFJob.\n",
    "    tfjob_launcher_op = components.load_component_from_file(\"component/kubeflow-launcher-component.yaml\")\n",
    "    \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": "code",
   "execution_count": null,
   "id": "d1b910bd-f6b6-4b9b-a44d-f67bc0324e04",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_serving_task(model_name, model_namespace, tfjob_op, model_volume_op):\n",
    "    api_version = 'serving.kserve.io/v1beta1'\n",
    "    serving_component_url = 'component/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 = 'component/kfserving-component.yaml'\n",
    "\n",
    "    inference_service = f'''\n",
    "      apiVersion: \"{api_version}\"\n",
    "      kind: \"InferenceService\"\n",
    "      metadata:\n",
    "        name: {name}\n",
    "        namespace: {namespace}\n",
    "        annotations:\n",
    "          \"sidecar.istio.io/inject\": \"false\"\n",
    "      spec:\n",
    "        predictor:\n",
    "          tensorflow:\n",
    "            storageUri: \"pvc://{volume_str}/\"\n",
    "            resources:\n",
    "              limits:\n",
    "                cpu: '2'\n",
    "                memory: 1Gi\n",
    "              requests:\n",
    "                cpu: 100m\n",
    "                memory: 50Mi\n",
    "      '''\n",
    "\n",
    "    serving_launcher_op = components.load_component_from_file(serving_component_url)\n",
    "    serving_launcher_op(action=\"apply\", inferenceservice_yaml=inference_service).after(tfjob_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b3d30c7-e76f-4721-acbe-f0c300f63436",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dsl.pipeline(\n",
    "    name=name,\n",
    "    description=\"An end to end mnist example including hyperparameter tuning, train and inference\"\n",
    ")\n",
    "\n",
    "def mnist_pipeline(name: str, namespace: str):\n",
    "\n",
    "    # Run the hyperparameter tuning with Katib.\n",
    "    katib_op = create_katib_experiment_task(name, namespace, training_steps)\n",
    "\n",
    "    model_volume_op =  volume_str\n",
    "\n",
    "    # Run the distributive training with TFJob.\n",
    "    tfjob_op = create_tfjob_task(name, namespace, training_steps, katib_op, model_volume_op, http_proxy, no_proxy)\n",
    "\n",
    "    # Create the KServe inference.\n",
    "    create_serving_task(name, namespace, tfjob_op, model_volume_op)\n",
    "    print(\"Volume: \", volume_str)\n",
    "\n",
    "# Run the Kubeflow Pipeline in the user's namespace.\n",
    "kfp_client = kfp.Client()\n",
    "run_id = kfp_client.create_run_from_pipeline_func(pipeline_func=mnist_pipeline, namespace=namespace, \n",
    "                                                  arguments={\"name\": name,\n",
    "                                                             \"namespace\" : namespace}).run_id\n",
    "print(f\"Run ID: {run_id}\")\n",
    "\n",
    "kfp_client.wait_for_run_completion(run_id=run_id, timeout=36000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba476c30-2d3d-424b-8dfb-04277705739e",
   "metadata": {},
   "outputs": [],
   "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.state == \"SUCCEEDED\":\n",
    "    print(f\"Run {run_id} has been Succeeded\\n\")\n",
    "\n",
    "    # Specify the image URL here.\n",
    "    image = Image.open(\"image/9.bmp\")\n",
    "    data = np.array(image.convert('L').resize((28, 28))).astype(np.float64).reshape(-1, 28, 28, 1)\n",
    "    data_formatted = np.array2string(data, separator=\",\", formatter={\"float\": lambda x: \"%.1f\" % x})\n",
    "    json_request = f'{{ \"instances\" : {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 = f\"http://{name}-predictor.{namespace}.svc.cluster.local/v1/models/{name}:predict\"\n",
    "    response = requests.post(url, data=json_request)\n",
    "\n",
    "    print(\"Prediction for the image\")\n",
    "    display(image)\n",
    "    print(response.json())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "kubeflow_notebook": {
   "autosnapshot": false,
   "docker_image": "",
   "experiment": {
    "id": "",
    "name": ""
   },
   "experiment_name": "",
   "katib_metadata": {
    "algorithm": {
     "algorithmName": "grid"
    },
    "maxFailedTrialCount": 3,
    "maxTrialCount": 12,
    "objective": {
     "objectiveMetricName": "",
     "type": "minimize"
    },
    "parallelTrialCount": 3,
    "parameters": []
   },
   "katib_run": false,
   "pipeline_description": "",
   "pipeline_name": "",
   "snapshot_volumes": false,
   "steps_defaults": [],
   "volume_access_mode": "rwm",
   "volumes": []
  },
  "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": 5
}
