07-prediction-serving.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4004af63",
   "metadata": {},
   "source": [
    "# 07 - Prediction Serving\n",
    "\n",
    "The purpose of the notebook is to show how to use the deployed model for online and batch prediction.\n",
    "The notebook covers the following tasks:\n",
    "1. Test the endpoints for online prediction.\n",
    "2. Use the uploaded custom model for batch prediction.\n",
    "3. Run a the batch prediction pipeline using `Vertex Pipelines`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dad1f75",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d02a9d5",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7f3ce81",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from datetime import datetime\n",
    "import tensorflow as tf\n",
    "\n",
    "from google.cloud import aiplatform as vertex_ai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e73bc25",
   "metadata": {},
   "source": [
    "### Setup Google Cloud project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29ea9b0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT = '[your-project-id]' # Change to your project id.\n",
    "REGION = 'us-central1' # Change to your region.\n",
    "BUCKET = '[your-bucket-name]' # Change to your bucket name.\n",
    "\n",
    "if PROJECT == \"\" or PROJECT is None or PROJECT == \"[your-project-id]\":\n",
    "    # Get your GCP project id from gcloud\n",
    "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
    "    PROJECT = shell_output[0]\n",
    "    \n",
    "if BUCKET == \"\" or BUCKET is None or BUCKET == \"[your-bucket-name]\":\n",
    "    # Get your bucket name to GCP project id\n",
    "    BUCKET = PROJECT\n",
    "    # Try to create the bucket if it doesn't exists\n",
    "    ! gsutil mb -l $REGION gs://$BUCKET\n",
    "    print(\"\")\n",
    "    \n",
    "print(\"Project ID:\", PROJECT)\n",
    "print(\"Region:\", REGION)\n",
    "print(\"Bucket name:\", BUCKET)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecba79b0",
   "metadata": {},
   "source": [
    "### Set configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "537732be",
   "metadata": {},
   "outputs": [],
   "source": [
    "VERSION = 'v01'\n",
    "DATASET_DISPLAY_NAME = 'chicago-taxi-tips'\n",
    "MODEL_DISPLAY_NAME = f'{DATASET_DISPLAY_NAME}-classifier-{VERSION}'\n",
    "ENDPOINT_DISPLAY_NAME = f'{DATASET_DISPLAY_NAME}-classifier'\n",
    "\n",
    "SERVE_BQ_DATASET_NAME = 'playground_us' # Change to your serving BigQuery dataset name.\n",
    "SERVE_BQ_TABLE_NAME = 'chicago_taxitrips_prep' # Change to your serving BigQuery table name."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e508dd0",
   "metadata": {},
   "source": [
    "## 1. Making Online Predicitons\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38be76f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertex_ai.init(\n",
    "    project=PROJECT,\n",
    "    location=REGION,\n",
    "    staging_bucket=BUCKET\n",
    ")\n",
    "\n",
    "endpoint_name = vertex_ai.Endpoint.list(\n",
    "    filter=f'display_name={ENDPOINT_DISPLAY_NAME}', \n",
    "    order_by=\"update_time\")[-1].gca_resource.name\n",
    "\n",
    "endpoint = vertex_ai.Endpoint(endpoint_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6b8053d",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_instances = [  \n",
    "    {\n",
    "        \"dropoff_grid\": [\"POINT(-87.6 41.9)\"],\n",
    "        \"euclidean\": [2064.2696],\n",
    "        \"loc_cross\": [\"\"],\n",
    "        \"payment_type\": [\"Credit Card\"],\n",
    "        \"pickup_grid\": [\"POINT(-87.6 41.9)\"],\n",
    "        \"trip_miles\": [1.37],\n",
    "        \"trip_day\": [12],\n",
    "        \"trip_hour\": [16],\n",
    "        \"trip_month\": [2],\n",
    "        \"trip_day_of_week\": [4],\n",
    "        \"trip_seconds\": [555]\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7cb447e",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = endpoint.predict(test_instances).predictions\n",
    "\n",
    "for prediction in predictions:\n",
    "    print(prediction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "330d9dfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "explanations = endpoint.explain(test_instances).explanations\n",
    "\n",
    "for explanation in explanations:\n",
    "    print(explanation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddc90ffa",
   "metadata": {},
   "source": [
    "## 2. Batch Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "046757e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "WORKSPACE = f\"gs://{BUCKET}/{DATASET_DISPLAY_NAME}/\"\n",
    "SERVING_DATA_DIR = os.path.join(WORKSPACE, 'serving_data')\n",
    "SERVING_INPUT_DATA_DIR = os.path.join(SERVING_DATA_DIR, 'input_data')\n",
    "SERVING_OUTPUT_DATA_DIR = os.path.join(SERVING_DATA_DIR, 'output_predictions')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e8fbc4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "if tf.io.gfile.exists(SERVING_DATA_DIR):\n",
    "    print(\"Removing previous serving data...\")\n",
    "    tf.io.gfile.rmtree(SERVING_DATA_DIR)\n",
    "    \n",
    "print(\"Creating serving data directory...\")\n",
    "tf.io.gfile.mkdir(SERVING_DATA_DIR)\n",
    "print(\"Serving data directory is ready.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f7b60fa",
   "metadata": {},
   "source": [
    "### Extract serving data to Cloud Storage as JSONL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04bb69ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.common import datasource_utils\n",
    "from src.preprocessing import etl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfd4cf91",
   "metadata": {},
   "outputs": [],
   "source": [
    "LIMIT = 10000\n",
    "\n",
    "sql_query = datasource_utils.get_serving_source_query(\n",
    "    bq_dataset_name=SERVE_BQ_DATASET_NAME, \n",
    "    bq_table_name=SERVE_BQ_TABLE_NAME,\n",
    "    limit=LIMIT\n",
    ")\n",
    "\n",
    "print(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f5afb73",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = f\"extract-{DATASET_DISPLAY_NAME}-serving-{datetime.now().strftime('%Y%m%d%H%M%S')}\"\n",
    "\n",
    "args = {\n",
    "    'job_name': job_name,\n",
    "    #'runner': 'DataflowRunner',\n",
    "    'sql_query': sql_query,\n",
    "    'exported_data_prefix': os.path.join(SERVING_INPUT_DATA_DIR, \"data-\"),\n",
    "    'temporary_dir': os.path.join(WORKSPACE, 'tmp'),\n",
    "    'gcs_location': os.path.join(WORKSPACE, 'bq_tmp'),\n",
    "    'project': PROJECT,\n",
    "    'region': REGION,\n",
    "    'setup_file': './setup.py'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "588e1949",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.get_logger().setLevel('ERROR')\n",
    "\n",
    "print(\"Data extraction started...\")\n",
    "etl.run_extract_pipeline(args)\n",
    "print(\"Data extraction completed.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a036944a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gsutil ls {SERVING_INPUT_DATA_DIR}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ff0d72b",
   "metadata": {},
   "source": [
    "### Submit the batch prediction job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb72b16e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name =  vertex_ai.Model.list(\n",
    "    filter=f'display_name={MODEL_DISPLAY_NAME}',\n",
    "    order_by=\"update_time\")[-1].gca_resource.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dac58bf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_resources =  {\n",
    "    \"machine_type\": 'n1-standard-2',\n",
    "    #'accelerator_count': 1,\n",
    "    #'accelerator_type': 'NVIDIA_TESLA_T4'\n",
    "    \"starting_replica_count\": 1,\n",
    "    \"max_replica_count\": 10,\n",
    "}\n",
    "\n",
    "job_display_name = f\"{MODEL_DISPLAY_NAME}-prediction-job-{datetime.now().strftime('%Y%m%d%H%M%S')}\"\n",
    "\n",
    "vertex_ai.BatchPredictionJob.create(\n",
    "    job_display_name=job_display_name,\n",
    "    model_name=model_name,\n",
    "    gcs_source=SERVING_INPUT_DATA_DIR + '/*.jsonl',\n",
    "    gcs_destination_prefix=SERVING_OUTPUT_DATA_DIR,\n",
    "    instances_format='jsonl',\n",
    "    predictions_format='jsonl',\n",
    "    sync=True,\n",
    "    **job_resources,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f281a9",
   "metadata": {},
   "source": [
    "## 3. Run the batch prediction pipeline using Vertex Pipelines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "809ba028",
   "metadata": {},
   "outputs": [],
   "source": [
    "WORKSPACE = f\"gs://{BUCKET}/{DATASET_DISPLAY_NAME}/\"\n",
    "ARTIFACT_STORE = os.path.join(WORKSPACE, 'tfx_artifacts')\n",
    "PIPELINE_NAME = f'{MODEL_DISPLAY_NAME}-predict-pipeline'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "769a1d9e",
   "metadata": {},
   "source": [
    "### Set the pipeline configurations for the Vertex AI run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5add19d",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"PROJECT\"] = PROJECT\n",
    "os.environ[\"REGION\"] = REGION\n",
    "os.environ[\"GCS_LOCATION\"] = f\"gs://{BUCKET}/{DATASET_DISPLAY_NAME}\"\n",
    "os.environ[\"MODEL_DISPLAY_NAME\"] = MODEL_DISPLAY_NAME\n",
    "os.environ[\"PIPELINE_NAME\"] = PIPELINE_NAME\n",
    "os.environ[\"ARTIFACT_STORE_URI\"] = ARTIFACT_STORE\n",
    "os.environ[\"BATCH_PREDICTION_BQ_DATASET_NAME\"] = SERVE_BQ_DATASET_NAME\n",
    "os.environ[\"BATCH_PREDICTION_BQ_TABLE_NAME\"] = SERVE_BQ_TABLE_NAME\n",
    "os.environ[\"SERVE_LIMIT\"] = \"1000\"\n",
    "os.environ[\"BEAM_RUNNER\"] = \"DirectRunner\"\n",
    "os.environ[\"TFX_IMAGE_URI\"] = f\"gcr.io/{PROJECT}/{DATASET_DISPLAY_NAME}:{VERSION}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6d0e2ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "from src.tfx_pipelines import config\n",
    "importlib.reload(config)\n",
    "\n",
    "for key, value in config.__dict__.items():\n",
    "    if key.isupper(): print(f'{key}: {value}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f128b46e",
   "metadata": {},
   "source": [
    "### (Optional) Build the ML container image\n",
    "\n",
    "This is the `TFX` runtime environment for the training pipeline steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f24fa5e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!echo $TFX_IMAGE_URI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3949cc7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud builds submit --tag $TFX_IMAGE_URI . --timeout=15m --machine-type=e2-highcpu-8"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98a9890d",
   "metadata": {},
   "source": [
    "### Compile pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09c8a3a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.tfx_pipelines import runner\n",
    "\n",
    "pipeline_definition_file = f'{config.PIPELINE_NAME}.json'\n",
    "pipeline_definition = runner.compile_prediction_pipeline(pipeline_definition_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bc2792a",
   "metadata": {},
   "source": [
    "### Submit run to Vertex Pipelines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37dcc92d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kfp.v2.google.client import AIPlatformClient\n",
    "\n",
    "pipeline_client = AIPlatformClient(\n",
    "    project_id=PROJECT, region=REGION)\n",
    "                 \n",
    "pipeline_client.create_run_from_job_spec(\n",
    "    job_spec_path=pipeline_definition_file\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e0d5bef",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "name": "common-cpu.m79",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/base-cpu:m79"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}