census_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "00a9d81e-1da7-4ced-bddc-fd56107ebd1a",
   "metadata": {},
   "source": [
    "The purpose of this notebook is to build our pipeline components, test each component and eventually build a pipeline out of these components and test that. Once the pipeline is built, we can move the code to a compile_pipeline.py file to be compiled as a CI/CD step. The code for running the pipeline will exist in a runpipeline.py file which will be run as part of the CI/CD as well. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37821695-8796-44e7-8ae2-775c54555e7d",
   "metadata": {
    "tags": []
   },
   "source": [
    "Set the values for PROJECT_ID, REGION, and BUCKET_URI according to your project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9030884f-2daf-4bf5-abbc-c9366b82d649",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "PROJECT_ID = \"prj-d-bu3machine-learning-ma6i\"\n",
    "REGION = \"us-central1\"\n",
    "BUCKET_URI = \"gs://bkt-d-vertexbucket\"\n",
    "DATAFLOW_SUBNET=\"https://www.googleapis.com/compute/v1/projects/prj-d-shared-restricted-9enw/regions/us-central1/subnetworks/sb-d-shared-restricted-us-central1\"\n",
    "KMS_KEY = \"projects/prj-d-kms-3i3k/locations/us-central1/keyRings/sample-keyring/cryptoKeys/prj-d-bu3machine-learning\"\n",
    "COMPUTE_ENGINE_SA = \"401570045548-compute@developer.gserviceaccount.com\"\n",
    "DATAFLOW_RUNNER = \"notebook-runner@prj-d-bu3machine-learning-ma6i.iam.gserviceaccount.com\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "816e581c-a463-410e-9708-fb27916927af",
   "metadata": {},
   "source": [
    "You can set these variables to your desired values or leave them as is"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8231198a-8389-4cc0-a698-c2b1c24b1d1b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "KFP_COMPONENTS_PATH = \"components\"\n",
    "SRC = \"src\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "331e585f-bb53-495e-80d7-c963a2fb3bb3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# create a src directory for dataflow source code\n",
    "!mkdir -m 777 -p {SRC}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7feca03d-c55a-44ae-bae7-814d1a93a236",
   "metadata": {},
   "source": [
    "This is the image we will use to run pipeline components. Replace the name of the artifact project with that of yours, e.g.:\n",
    "##### \"us-central1-docker.pkg.dev/{prj-c-bu3artifacts-####}/c-publish-artifacts/vertexpipeline:v2\"\n",
    "As part of the project inflation pipelines, the image from the Dockerfile in this repository is built and pushed to project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33be8e30-0c49-43ca-ae3e-2eb644012f87",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "Image = \"us-central1-docker.pkg.dev/prj-c-bu3artifacts-zfe5/c-publish-artifacts/vertexpipeline:v2\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b32b50cb-20ad-4ab2-9a3e-a1e330596944",
   "metadata": {},
   "source": [
    "while working in the dev environment's notebook, we install the required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f60c6bc3-d97f-41c8-879c-1edb59a6cfb9",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade tensorflow==2.8.0  tensorflow-hub==0.13.0 kfp\n",
    "!pip install --upgrade google-cloud-pipeline-components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ddb16f5-a901-4a41-981e-5ec33d482715",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install --quiet protobuf==3.20.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56c3e887-506d-4d97-b417-4f0c2fd65c8c",
   "metadata": {},
   "source": [
    "copy data to bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3040dc95-9cf7-4815-8091-2bcfb03f32d3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!gsutil cp -r data {BUCKET_URI}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70f54986-8c76-47a1-914b-07be39d2f948",
   "metadata": {},
   "source": [
    "import the required libraries to run experiments and build components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02c4447b-b375-4595-82c3-706bd18ee837",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path as path\n",
    "from urllib.parse import urlparse\n",
    "import os\n",
    "from six.moves import urllib\n",
    "import tempfile\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import tensorflow_hub as hub\n",
    "from google.cloud import aiplatform\n",
    "from google.cloud import bigquery\n",
    "from google.api_core.exceptions import GoogleAPIError\n",
    "from kfp import compiler, dsl\n",
    "from kfp.dsl import component\n",
    "from kfp.dsl import Input, Output, Model, Metrics, OutputPath\n",
    "from typing import NamedTuple\n",
    "\n",
    "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e18331e6-eaa0-40ea-a0e8-a4b2e164e1a4",
   "metadata": {
    "tags": []
   },
   "source": [
    "set up global variables for running experiments and pipeline components (you can leave them as is to avoid any errors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4223f59-9012-4064-bee9-67c7aa96cbbf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "DATA_URL = f'{BUCKET_URI}/data'\n",
    "TRAINING_FILE = 'adult.data.csv'\n",
    "EVAL_FILE = 'adult.test.csv'\n",
    "TRAINING_URL = '%s/%s' % (DATA_URL, TRAINING_FILE)\n",
    "EVAL_URL = '%s/%s' % (DATA_URL, EVAL_FILE)\n",
    "DATASET_ID = 'census_dataset'\n",
    "TRAINING_TABLE_ID = 'census_train_table'\n",
    "EVAL_TABLE_ID = 'census_eval_table'\n",
    "RUNNER = \"DataflowRunner\"\n",
    "REGION=\"us-central1\"\n",
    "JOB_NAME=\"census-ingest\"\n",
    "CSV_SCHEMA = [\n",
    "      bigquery.SchemaField(\"age\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"workclass\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"fnlwgt\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"education\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"education_num\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"marital_status\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"occupation\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"relationship\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"race\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"gender\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"capital_gain\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"capital_loss\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"hours_per_week\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"native_country\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"income_bracket\", \"STRING\"),\n",
    "  ]\n",
    "\n",
    "UNUSED_COLUMNS = [\"fnlwgt\", \"education_num\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "976a9327-b077-477b-b765-9ff9480c9fd5",
   "metadata": {
    "tags": []
   },
   "source": [
    "# BigQuery dataset creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd5ace09-afd5-4135-8e28-9f8b6c2b1c96",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# create a directory to save the component\n",
    "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/bq_dataset_component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b38ab3-f86c-4905-badc-7daaf2f61dbb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "create_bq_dataset_query = f\"\"\"\n",
    "CREATE SCHEMA IF NOT EXISTS {DATASET_ID}\n",
    "\"\"\"\n",
    "\n",
    "with open(\n",
    "    f\"{KFP_COMPONENTS_PATH}/bq_dataset_component/create_bq_dataset.sql\", \"w\"\n",
    ") as q:\n",
    "    q.write(create_bq_dataset_query)\n",
    "q.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbbf2bee-cc42-4aff-8fff-3bb506bcc741",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Dataflow sources\n",
    "The below cells generate the dataflow source code for ingesting data from our gcs bucket to the biquery dataset create above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "459aac56-76d0-4e0d-bc39-b48e180f78fa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!touch {SRC}/__init__.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28a08119-b131-4100-a2f2-944088d21a3f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile src/ingest_pipeline.py\n",
    "from __future__ import absolute_import\n",
    "import logging\n",
    "import argparse\n",
    "import apache_beam as beam\n",
    "from apache_beam.io import ReadFromText, ReadAllFromText\n",
    "from apache_beam.dataframe.io import read_csv\n",
    "from apache_beam.io.gcp.bigquery import WriteToBigQuery\n",
    "from apache_beam.options.pipeline_options import PipelineOptions\n",
    "from apache_beam.io.gcp.internal.clients import bigquery\n",
    "from apache_beam.options.pipeline_options import SetupOptions\n",
    "\n",
    "def get_bigquery_schema():\n",
    "    \"\"\"\n",
    "    A function to get the BigQuery schema.\n",
    "    Returns:\n",
    "        A list of BigQuery schema.\n",
    "    \"\"\"\n",
    "\n",
    "    table_schema = bigquery.TableSchema()\n",
    "    columns = (('age', 'FLOAT64', 'nullable'),\n",
    "               ('workclass', 'STRING', 'nullable'),\n",
    "               ('fnlwgt', 'FLOAT64', 'nullable'),\n",
    "               ('education', 'STRING', 'nullable'),\n",
    "               ('education_num', 'FLOAT64', 'nullable'),\n",
    "               ('marital_status', 'STRING', 'nullable'),\n",
    "               ('occupation', 'STRING', 'nullable'),\n",
    "               (\"relationship\", \"STRING\", 'nullable'),\n",
    "               (\"race\", \"STRING\", 'nullable'),\n",
    "               (\"gender\", \"STRING\", 'nullable'),\n",
    "               (\"capital_gain\", \"FLOAT64\", 'nullable'),\n",
    "               (\"capital_loss\", \"FLOAT64\", 'nullable'),\n",
    "               (\"hours_per_week\", \"FLOAT64\", 'nullable'),\n",
    "               (\"native_country\", \"STRING\", 'nullable'),\n",
    "               (\"income_bracket\", \"STRING\", 'nullable')\n",
    "              )\n",
    "\n",
    "    for column in columns:\n",
    "        column_schema = bigquery.TableFieldSchema()\n",
    "        column_schema.name = column[0]\n",
    "        column_schema.type = column[1]\n",
    "        column_schema.mode = column[2]\n",
    "        table_schema.fields.append(column_schema)\n",
    "\n",
    "    return table_schema\n",
    "\n",
    "def get_args():\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--url', dest='url', default=\"BUCKET_URI/data/adult.data.csv\",\n",
    "                        help='url of the data to be downloaded')\n",
    "    parser.add_argument('--bq-dataset', dest='dataset_id', required=False,\n",
    "                        default='census_dataset', help='Dataset name used in BigQuery.')\n",
    "    parser.add_argument('--bq-table', dest='table_id', required=False,\n",
    "                        default='census_train_table', help='Table name used in BigQuery.')\n",
    "    parser.add_argument('--bq-project', dest='project_id', required=False,\n",
    "                        default='majid-test-407120', help='project id')\n",
    "    args, pipeline_args = parser.parse_known_args()\n",
    "    return args, pipeline_args\n",
    "\n",
    "def transform(line):\n",
    "    values = line.split(\",\")\n",
    "    d = {}\n",
    "    fields = [\"age\",\"workclass\",\"fnlwgt\",\"education\",\"education_num\",\n",
    "              \"marital_status\",\"occupation\",\"relationship\",\"race\",\"gender\",\n",
    "              \"capital_gain\",\"capital_loss\",\"hours_per_week\",\"native_country\",\"income_bracket\"]\n",
    "    for i in range(len(fields)):\n",
    "        d[fields[i]] = values[i].strip()\n",
    "    return d\n",
    "\n",
    "def load_data_into_bigquery(args, pipeline_args):\n",
    "    options = PipelineOptions(pipeline_args)\n",
    "    options.view_as(SetupOptions).save_main_session = True\n",
    "    p = beam.Pipeline(options=options)\n",
    " \n",
    "    (p \n",
    "     | 'Create PCollection' >> beam.Create([args.url])\n",
    "     | 'ReadFromText' >> ReadAllFromText(skip_header_lines=1)\n",
    "     | 'string to bq row' >> beam.Map(lambda s: transform(s))\n",
    "     | 'WriteToBigQuery' >> WriteToBigQuery(\n",
    "        table=args.table_id,\n",
    "        dataset=args.dataset_id,\n",
    "        project=args.project_id,\n",
    "        schema=get_bigquery_schema(),\n",
    "        create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,\n",
    "        write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,\n",
    "    )\n",
    "    )\n",
    "\n",
    "    job = p.run()\n",
    "    if options.get_all_options()['runner'] == 'DirectRunner':\n",
    "        job.wait_until_finish()\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    args, pipeline_args = get_args()\n",
    "    logging.getLogger().setLevel(logging.INFO)\n",
    "    load_data_into_bigquery(args, pipeline_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab1f80b6-047f-4f33-9297-94d589395303",
   "metadata": {},
   "source": [
    "Now upload the dataflow source code to the bucket to be accessible by Dataflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76881a2e-0ed8-4d0c-9e1f-512a93130312",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!gsutil cp -R {SRC} {BUCKET_URI}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1052dd72-692d-4215-9ff5-e791f3fbc41a",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Dataflow arguments component\n",
    "This components sole purpose is to prepare the command line arguments for Dataflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "167c96f7-6416-464e-9326-6350b1345cbb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@component(base_image=Image)\n",
    "def build_dataflow_args(\n",
    "    bq_dataset: str,\n",
    "    url: str,\n",
    "    bq_table: str,\n",
    "    job_name: str,\n",
    "    runner: str,\n",
    "    bq_project: str,\n",
    "    subnet: str,\n",
    "    dataflow_sa: str,\n",
    ") -> list:\n",
    "    return [\n",
    "        \"--job_name\",\n",
    "        job_name,\n",
    "        \"--runner\",\n",
    "        runner,\n",
    "        \"--url\",\n",
    "        url,\n",
    "        \"--bq-dataset\",\n",
    "        bq_dataset,\n",
    "        \"--bq-table\",\n",
    "        bq_table,\n",
    "        \"--bq-project\",\n",
    "        bq_project,\n",
    "        \"--subnetwork\",\n",
    "        subnet,\n",
    "        \"--no_use_public_ips\",\n",
    "        \"--worker_zone\",\n",
    "        \"us-central1-c\",\n",
    "        \"--service_account_email\",\n",
    "        dataflow_sa,\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03de2873-3145-4a7f-abaf-3de1267ac708",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Costum training component\n",
    "The model training involves 3 simple steps:\n",
    "1. Reading the data from bigquery using the tensorflow-io library (make sure you set the default kms key for your project above in advance)\n",
    "2. Creating a keras model\n",
    "3. Training the model on the data and optionally storing logs in a tensorboard directory (/tblogs) to enable tensorboard usage in future\n",
    "\n",
    "After we test the training function, it can be wrapped in a vertex pipeline coponent using the @componet decorator as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b026b981-cbc0-43bd-af56-36a9f1a2f514",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/custom_training_component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "309ed10c-2434-4dd5-9f92-4352ff9e4b78",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@component(\n",
    "    base_image=Image,\n",
    "    output_component_file=f\"{KFP_COMPONENTS_PATH}/custom_training_component/training.yaml\"\n",
    ")\n",
    "def custom_train_model(\n",
    "    project: str,\n",
    "    table: str,\n",
    "    dataset: str,\n",
    "    tb_log_dir: str,\n",
    "    model: Output[Model],\n",
    "    epochs: int = 5,\n",
    "    batch_size: int = 32,\n",
    "    lr: float = 0.01, # not used here but can be passed to an optimizer\n",
    "):\n",
    "    \n",
    "    from tensorflow.python.framework import ops\n",
    "    from tensorflow.python.framework import dtypes\n",
    "    from tensorflow_io.bigquery import BigQueryClient\n",
    "    from tensorflow_io.bigquery import BigQueryReadSession\n",
    "    from tensorflow import feature_column\n",
    "    from google.cloud import bigquery\n",
    "    \n",
    "    import tensorflow as tf\n",
    "    CSV_SCHEMA = [\n",
    "      bigquery.SchemaField(\"age\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"workclass\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"fnlwgt\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"education\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"education_num\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"marital_status\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"occupation\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"relationship\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"race\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"gender\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"capital_gain\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"capital_loss\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"hours_per_week\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"native_country\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"income_bracket\", \"STRING\"),\n",
    "  ]\n",
    "\n",
    "    UNUSED_COLUMNS = [\"fnlwgt\", \"education_num\"]\n",
    "    def transform_row(row_dict):\n",
    "        # Trim all string tensors\n",
    "        trimmed_dict = { column:\n",
    "                      (tf.strings.strip(tensor) if tensor.dtype == 'string' else tensor) \n",
    "                      for (column,tensor) in row_dict.items()\n",
    "                      }\n",
    "        # Extract feature column\n",
    "        income_bracket = trimmed_dict.pop('income_bracket')\n",
    "        # Convert feature column to 0.0/1.0\n",
    "        income_bracket_float = tf.cond(tf.equal(tf.strings.strip(income_bracket), '>50K'), \n",
    "                     lambda: tf.constant(1.0), \n",
    "                     lambda: tf.constant(0.0))\n",
    "        return (trimmed_dict, income_bracket_float)\n",
    "\n",
    "    def read_bigquery(table_name, dataset=dataset):\n",
    "        tensorflow_io_bigquery_client = BigQueryClient()\n",
    "        read_session = tensorflow_io_bigquery_client.read_session(\n",
    "          \"projects/\" + project,\n",
    "          project, table, dataset,\n",
    "          list(field.name for field in CSV_SCHEMA \n",
    "               if not field.name in UNUSED_COLUMNS),\n",
    "          list(dtypes.double if field.field_type == 'FLOAT64' \n",
    "               else dtypes.string for field in CSV_SCHEMA\n",
    "               if not field.name in UNUSED_COLUMNS),\n",
    "          requested_streams=2)\n",
    "\n",
    "        dataset = read_session.parallel_read_rows()\n",
    "        transformed_ds = dataset.map(transform_row)\n",
    "        return transformed_ds\n",
    "\n",
    "    training_ds = read_bigquery(table).shuffle(10000).batch(batch_size)\n",
    "\n",
    "\n",
    "\n",
    "    feature_columns = []\n",
    "    def get_categorical_feature_values(column):\n",
    "        query = 'SELECT DISTINCT TRIM({}) FROM `{}`.{}.{}'.format(column, project, dataset, table)\n",
    "        client = bigquery.Client(project=project)\n",
    "        dataset_ref = client.dataset(dataset)\n",
    "        job_config = bigquery.QueryJobConfig()\n",
    "        query_job = client.query(query, job_config=job_config)\n",
    "        result = query_job.to_dataframe()\n",
    "        return result.values[:,0]\n",
    "\n",
    "    # numeric cols\n",
    "    for header in ['capital_gain', 'capital_loss', 'hours_per_week']:\n",
    "        feature_columns.append(feature_column.numeric_column(header))\n",
    "\n",
    "    # categorical cols\n",
    "    for header in ['workclass', 'marital_status', 'occupation', 'relationship',\n",
    "                   'race', 'native_country', 'education']:\n",
    "        categorical_feature = feature_column.categorical_column_with_vocabulary_list(\n",
    "            header, get_categorical_feature_values(header))\n",
    "        categorical_feature_one_hot = feature_column.indicator_column(categorical_feature)\n",
    "        feature_columns.append(categorical_feature_one_hot)\n",
    "\n",
    "    # bucketized cols\n",
    "    age = feature_column.numeric_column('age')\n",
    "    age_buckets = feature_column.bucketized_column(age, boundaries=[18, 25, 30, 35, 40, 45, 50, 55, 60, 65])\n",
    "    feature_columns.append(age_buckets)\n",
    "\n",
    "    feature_layer = tf.keras.layers.DenseFeatures(feature_columns)\n",
    "\n",
    "\n",
    "    Dense = tf.keras.layers.Dense\n",
    "    keras_model = tf.keras.Sequential(\n",
    "      [\n",
    "        feature_layer,\n",
    "          Dense(100, activation=tf.nn.relu, kernel_initializer='uniform'),\n",
    "          Dense(75, activation=tf.nn.relu),\n",
    "          Dense(50, activation=tf.nn.relu),\n",
    "          Dense(25, activation=tf.nn.relu),\n",
    "          Dense(1, activation=tf.nn.sigmoid)\n",
    "      ])\n",
    "\n",
    "    tensorboard = tf.keras.callbacks.TensorBoard(log_dir=tb_log_dir)\n",
    "    # Compile Keras model\n",
    "    keras_model.compile(loss='binary_crossentropy', metrics=['accuracy'])\n",
    "    keras_model.fit(training_ds, epochs=epochs, callbacks=[tensorboard])\n",
    "    keras_model.save(model.path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ac20209-57ca-4b3e-9740-e2fc531591a8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from google_cloud_pipeline_components.v1.custom_job import utils\n",
    "custom_job_distributed_training_op = utils.create_custom_training_job_op_from_component(\n",
    "    custom_train_model, replica_count=1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "044b4e52-4596-41d9-ac6c-9dcfb7ab4bea",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Custom Evaluation component\n",
    "This step is very similar to the training component. At the end we flag the model for deployment only if its' accuracy is over 80 percent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3daaae2-229c-4deb-b264-6b86a80364b6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/custom_eval_component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b8f2b3a-0562-465f-987c-a68782b48c0c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# evaluation component\n",
    "@component(\n",
    "    base_image=Image,\n",
    "    output_component_file=f\"{KFP_COMPONENTS_PATH}/custom_eval_component/eval.yaml\"\n",
    ")\n",
    "def custom_eval_model(\n",
    "    model_dir: str,\n",
    "    project: str,\n",
    "    table: str,\n",
    "    dataset: str,\n",
    "    tb_log_dir: str,\n",
    "    model: Input[Model],\n",
    "    metrics: Output[Metrics],\n",
    "    batch_size: int = 32,\n",
    ")-> NamedTuple(\"Outputs\", [(\"dep_decision\", str)]):\n",
    "    from tensorflow.python.framework import ops\n",
    "    from tensorflow.python.framework import dtypes\n",
    "    from tensorflow_io.bigquery import BigQueryClient\n",
    "    from tensorflow_io.bigquery import BigQueryReadSession\n",
    "    from tensorflow import feature_column\n",
    "    from google.cloud import bigquery\n",
    "    \n",
    "    \n",
    "    import tensorflow as tf\n",
    "    CSV_SCHEMA = [\n",
    "      bigquery.SchemaField(\"age\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"workclass\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"fnlwgt\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"education\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"education_num\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"marital_status\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"occupation\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"relationship\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"race\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"gender\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"capital_gain\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"capital_loss\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"hours_per_week\", \"FLOAT64\"),\n",
    "      bigquery.SchemaField(\"native_country\", \"STRING\"),\n",
    "      bigquery.SchemaField(\"income_bracket\", \"STRING\"),\n",
    "  ]\n",
    "\n",
    "    UNUSED_COLUMNS = [\"fnlwgt\", \"education_num\"]\n",
    "    def transform_row(row_dict):\n",
    "        # Trim all string tensors\n",
    "        trimmed_dict = { column:\n",
    "                      (tf.strings.strip(tensor) if tensor.dtype == 'string' else tensor) \n",
    "                      for (column,tensor) in row_dict.items()\n",
    "                      }\n",
    "        # Extract feature column\n",
    "        income_bracket = trimmed_dict.pop('income_bracket')\n",
    "        # Convert feature column to 0.0/1.0\n",
    "        income_bracket_float = tf.cond(tf.equal(tf.strings.strip(income_bracket), '>50K'), \n",
    "                     lambda: tf.constant(1.0), \n",
    "                     lambda: tf.constant(0.0))\n",
    "        return (trimmed_dict, income_bracket_float)\n",
    "\n",
    "    def read_bigquery(table_name, dataset=dataset):\n",
    "        tensorflow_io_bigquery_client = BigQueryClient()\n",
    "        read_session = tensorflow_io_bigquery_client.read_session(\n",
    "          \"projects/\" + project,\n",
    "          project, table, dataset,\n",
    "          list(field.name for field in CSV_SCHEMA \n",
    "               if not field.name in UNUSED_COLUMNS),\n",
    "          list(dtypes.double if field.field_type == 'FLOAT64' \n",
    "               else dtypes.string for field in CSV_SCHEMA\n",
    "               if not field.name in UNUSED_COLUMNS),\n",
    "          requested_streams=2)\n",
    "\n",
    "        dataset = read_session.parallel_read_rows()\n",
    "        transformed_ds = dataset.map(transform_row)\n",
    "        return transformed_ds\n",
    "\n",
    "    eval_ds = read_bigquery(table).batch(batch_size)\n",
    "    keras_model = tf.keras.models.load_model(model.path)\n",
    "    tensorboard = tf.keras.callbacks.TensorBoard(log_dir=tb_log_dir)\n",
    "    loss, accuracy = keras_model.evaluate(eval_ds, callbacks=[tensorboard])\n",
    "    metrics.log_metric(\"accuracy\", accuracy)\n",
    "    # Deploy the model only if its accuracy is higher than 80 percent\n",
    "    if accuracy > 0.8:\n",
    "        dep_decision = \"true\"\n",
    "        keras_model.save(model_dir)\n",
    "    else:\n",
    "        dep_decision = \"false\"\n",
    "    return (dep_decision,)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dc0d316-9094-4eee-b7e7-6e07c09b83d0",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Deployment component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ab4b7c5-f952-4710-8d4e-8105f3a1cb4e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/deployment_component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2797433c-4d68-4ee4-8b64-c27520a6e503",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@component(\n",
    "    base_image=Image,\n",
    "    output_component_file=f\"{KFP_COMPONENTS_PATH}/deployment_component/deploy.yaml\"\n",
    ")\n",
    "def deploy_model(\n",
    "        serving_container_image_uri: str,\n",
    "        model_name: str,\n",
    "        model_dir: str,\n",
    "        endpoint_name: str,\n",
    "        project_id: str,\n",
    "        region: str,\n",
    "        split: int,\n",
    "        min_nodes: int,\n",
    "        max_nodes: int,\n",
    "        encryption: str,\n",
    "        service_account: str,\n",
    "        model: Input[Model],\n",
    "        vertex_model: Output[Model],\n",
    "        vertex_endpoint: Output[Model]\n",
    "):\n",
    "    from google.cloud import aiplatform    \n",
    "    aiplatform.init(service_account=service_account)\n",
    "    def create_endpoint():\n",
    "        endpoints = aiplatform.Endpoint.list(\n",
    "        filter=f'display_name=\"{endpoint_name}\"',\n",
    "        order_by='create_time desc',\n",
    "        project=project_id,\n",
    "        location=region,\n",
    "        )\n",
    "        if len(endpoints) > 0:\n",
    "            endpoint = endpoints[0] # most recently created\n",
    "        else:\n",
    "            endpoint = aiplatform.Endpoint.create(\n",
    "                display_name=endpoint_name,\n",
    "                project=project_id,\n",
    "                location=region,\n",
    "                encryption_spec_key_name=encryption\n",
    "        )\n",
    "        return endpoint\n",
    "\n",
    "    endpoint = create_endpoint()\n",
    "    \n",
    "\n",
    "    def upload_model():\n",
    "        listed_model = aiplatform.Model.list(\n",
    "        filter=f'display_name=\"{model_name}\"',\n",
    "        project=project_id,\n",
    "        location=region,\n",
    "        )\n",
    "        if len(listed_model) > 0:\n",
    "            model_version = listed_model[0]\n",
    "            model_upload = aiplatform.Model.upload(\n",
    "                    display_name=model_name,\n",
    "                    parent_model=model_version.resource_name,\n",
    "                    artifact_uri=model_dir,\n",
    "                    serving_container_image_uri=serving_container_image_uri,\n",
    "                    location=region,\n",
    "                    project=project_id,\n",
    "                    encryption_spec_key_name=encryption\n",
    "            )\n",
    "        else:\n",
    "            model_upload = aiplatform.Model.upload(\n",
    "                    display_name=model_name,\n",
    "                    artifact_uri=model_dir,\n",
    "                    serving_container_image_uri=serving_container_image_uri,\n",
    "                    location=region,\n",
    "                    project=project_id,\n",
    "                    encryption_spec_key_name=encryption,\n",
    "                \n",
    "            )\n",
    "        return model_upload\n",
    "    \n",
    "    uploaded_model = upload_model()\n",
    "    \n",
    "    # Save data to the output params\n",
    "    vertex_model.uri = uploaded_model.resource_name\n",
    "    def deploy_to_endpoint(model, endpoint):\n",
    "        deployed_models = endpoint.list_models()\n",
    "        if len(deployed_models) > 0:\n",
    "            latest_model_id = deployed_models[-1].id\n",
    "            print(\"your objects properties:\", deployed_models[0].create_time.__dir__())\n",
    "            model_deploy = uploaded_model.deploy(\n",
    "                # machine_type=\"n1-standard-4\",\n",
    "                endpoint=endpoint,\n",
    "                traffic_split={\"0\": 25, latest_model_id: 75},\n",
    "                deployed_model_display_name=model_name,\n",
    "                min_replica_count=min_nodes,\n",
    "                max_replica_count=max_nodes,\n",
    "                encryption_spec_key_name=encryption,\n",
    "                service_account=service_account\n",
    "            )\n",
    "        else:\n",
    "            model_deploy = uploaded_model.deploy(\n",
    "            # machine_type=\"n1-standard-4\",\n",
    "            endpoint=endpoint,\n",
    "            traffic_split={\"0\": 100},\n",
    "            min_replica_count=min_nodes,\n",
    "            max_replica_count=max_nodes,\n",
    "            deployed_model_display_name=model_name,\n",
    "            encryption_spec_key_name=encryption,\n",
    "            service_account=service_account\n",
    "        )\n",
    "        return model_deploy.resource_name\n",
    "\n",
    "    vertex_endpoint.uri = deploy_to_endpoint(vertex_model, endpoint)\n",
    "    vertex_endpoint.metadata['resourceName']=endpoint.resource_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb3ab72c-475c-4688-9f8f-b9799fe0bc4d",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Model monitoring component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d191bd2b-41bc-4367-b292-a366b63561da",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/monitoring_component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57a49a83-cc9d-4c7d-8260-85451063968f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@component(\n",
    "    base_image=Image,\n",
    "    output_component_file=f\"{KFP_COMPONENTS_PATH}/monitoring_component/monitoring.yaml\"\n",
    ")\n",
    "def create_monitoring(\n",
    "    monitoring_name: str,\n",
    "    project_id: str,\n",
    "    region: str,\n",
    "    endpoint: Input[Model],\n",
    "    bq_data_uri: str,\n",
    "    bucket_name: str,\n",
    "    email: str,\n",
    "    encryption: str,\n",
    "    service_account: str,\n",
    "):\n",
    "    from google.cloud.aiplatform import model_monitoring\n",
    "    from google.cloud import aiplatform\n",
    "    from google.cloud import bigquery\n",
    "    from google.cloud import storage\n",
    "    from collections import OrderedDict\n",
    "    import time\n",
    "    import yaml\n",
    "    def ordered_dict_representer(self, value):  # can be a lambda if that's what you prefer\n",
    "        return self.represent_mapping('tag:yaml.org,2002:map', value.items())\n",
    "    yaml.add_representer(OrderedDict, ordered_dict_representer)\n",
    "    \n",
    "    aiplatform.init(service_account=service_account)\n",
    "    list_monitors = aiplatform.ModelDeploymentMonitoringJob.list(filter=f'(state=\"JOB_STATE_SUCCEEDED\" OR state=\"JOB_STATE_RUNNING\") AND display_name=\"{monitoring_name}\"', project=project_id)\n",
    "    if len(list_monitors) == 0:\n",
    "        alerting_config = model_monitoring.EmailAlertConfig(\n",
    "            user_emails=[email], enable_logging=True\n",
    "        )\n",
    "        # schedule config\n",
    "        MONITOR_INTERVAL = 1\n",
    "        schedule_config = model_monitoring.ScheduleConfig(monitor_interval=MONITOR_INTERVAL)\n",
    "        # sampling strategy\n",
    "        SAMPLE_RATE = 0.5 \n",
    "        logging_sampling_strategy = model_monitoring.RandomSampleConfig(sample_rate=SAMPLE_RATE)\n",
    "        # drift config\n",
    "        DRIFT_THRESHOLD_VALUE = 0.05\n",
    "        DRIFT_THRESHOLDS = {\n",
    "            \"capital_gain\": DRIFT_THRESHOLD_VALUE,\n",
    "            \"capital_loss\": DRIFT_THRESHOLD_VALUE,\n",
    "        }\n",
    "        drift_config = model_monitoring.DriftDetectionConfig(drift_thresholds=DRIFT_THRESHOLDS)\n",
    "        # Skew config\n",
    "        DATASET_BQ_URI = bq_data_uri\n",
    "        TARGET = \"income_bracket\"\n",
    "        SKEW_THRESHOLD_VALUE = 0.5\n",
    "        SKEW_THRESHOLDS = {\n",
    "            \"capital_gain\": SKEW_THRESHOLD_VALUE,\n",
    "            \"capital_loss\": SKEW_THRESHOLD_VALUE,\n",
    "        }\n",
    "        skew_config = model_monitoring.SkewDetectionConfig(\n",
    "            data_source=DATASET_BQ_URI, skew_thresholds=SKEW_THRESHOLDS, target_field=TARGET\n",
    "        )\n",
    "        # objective config out of skew and drift configs\n",
    "        objective_config = model_monitoring.ObjectiveConfig(\n",
    "            skew_detection_config=skew_config,\n",
    "            drift_detection_config=drift_config,\n",
    "            explanation_config=None,\n",
    "        )\n",
    "\n",
    "        bqclient = bigquery.Client()\n",
    "        table = bigquery.TableReference.from_string(DATASET_BQ_URI[5:])\n",
    "        bq_table = bqclient.get_table(table)\n",
    "        schema = bq_table.schema\n",
    "        schemayaml = OrderedDict({\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {},\n",
    "            \"required\": []\n",
    "        })\n",
    "        for feature in schema:\n",
    "            if feature.name in [\"income_bracket\"]:\n",
    "                continue\n",
    "            if feature.field_type == \"STRING\":\n",
    "                f_type = \"string\"\n",
    "            else:\n",
    "                f_type = \"number\"\n",
    "            schemayaml['properties'][feature.name] = {\"type\": f_type}\n",
    "            if feature.name not in [\"fnlwgt\", \"education_num\"]:\n",
    "                schemayaml['required'].append(feature.name)\n",
    "            \n",
    "        with open(\"monitoring_schema.yaml\", \"w\") as yaml_file:\n",
    "            yaml.dump(schemayaml, yaml_file, default_flow_style=False)\n",
    "        storage_client = storage.Client()\n",
    "        bucket = storage_client.bucket(bucket_name)\n",
    "        blob = bucket.blob(\"monitoring_schema.yaml\")\n",
    "        blob.upload_from_filename(\"monitoring_schema.yaml\")\n",
    "\n",
    "        monitoring_job = aiplatform.ModelDeploymentMonitoringJob.create(\n",
    "            display_name=monitoring_name,\n",
    "            project=project_id,\n",
    "            location=region,\n",
    "            endpoint=endpoint.metadata['resourceName'],\n",
    "            logging_sampling_strategy=logging_sampling_strategy,\n",
    "            schedule_config=schedule_config,\n",
    "            alert_config=alerting_config,\n",
    "            objective_configs=objective_config,\n",
    "            analysis_instance_schema_uri=f\"gs://{bucket_name}/monitoring_schema.yaml\",\n",
    "            encryption_spec_key_name=encryption,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60faba58-e64a-444f-a80a-819af74fa488",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Pipeline build and compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bfb7d67-49b9-416b-90c1-568648631f7d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@dsl.pipeline(name=\"census-income-pipeline\")\n",
    "def pipeline(\n",
    "    create_bq_dataset_query: str,\n",
    "    project: str,\n",
    "    deployment_project: str,\n",
    "    region: str,\n",
    "    model_dir: str,\n",
    "    bucket_name: str,\n",
    "    monitoring_name: str,\n",
    "    monitoring_email: str,\n",
    "    encryption: str,\n",
    "    service_account: str,\n",
    "    prod_service_account: str,\n",
    "    dataflow_subnet: str,\n",
    "    train_data_url: str=TRAINING_URL,\n",
    "    eval_data_url: str=EVAL_URL,\n",
    "    bq_dataset: str=DATASET_ID,\n",
    "    bq_train_table: str=TRAINING_TABLE_ID,\n",
    "    bq_eval_table: str=EVAL_TABLE_ID,\n",
    "    job_name: str=JOB_NAME,\n",
    "    python_file_path: str=f'{BUCKET_URI}/src/ingest_pipeline.py',\n",
    "    dataflow_temp_location: str=f'{BUCKET_URI}/temp_dataflow',\n",
    "    runner: str=RUNNER,\n",
    "    lr: float=0.01, \n",
    "    epochs: int=5,\n",
    "    batch_size: int=32,\n",
    "    base_train_dir: str=f'{BUCKET_URI}/training', \n",
    "    tb_log_dir: str=f'{BUCKET_URI}/tblogs',\n",
    "    deployment_image: str=\"us-docker.pkg.dev/cloud-aiplatform/prediction/tf2-cpu.2-8:latest\",\n",
    "    deployed_model_name: str='income_bracket_predictor',\n",
    "    endpoint_name: str='census_endpoint',\n",
    "    min_nodes: int=2,\n",
    "    max_nodes: int=4,\n",
    "    traffic_split: int=25,\n",
    "    dataflow_sa: str=DATAFLOW_RUNNER,\n",
    "):\n",
    "    from google_cloud_pipeline_components.v1.bigquery import (\n",
    "        BigqueryQueryJobOp)\n",
    "    from google_cloud_pipeline_components.v1.dataflow import \\\n",
    "        DataflowPythonJobOp\n",
    "    from google_cloud_pipeline_components.v1.wait_gcp_resources import \\\n",
    "        WaitGcpResourcesOp\n",
    "    \n",
    "    from google_cloud_pipeline_components.types import artifact_types\n",
    "    from google_cloud_pipeline_components.v1.batch_predict_job import \\\n",
    "        ModelBatchPredictOp\n",
    "    from google_cloud_pipeline_components.v1.model import ModelUploadOp\n",
    "    from kfp.dsl import importer_node\n",
    "    from google_cloud_pipeline_components.v1.endpoint import EndpointCreateOp, ModelDeployOp\n",
    "    \n",
    "    # create the dataset\n",
    "    bq_dataset_op = BigqueryQueryJobOp(\n",
    "        query=create_bq_dataset_query,\n",
    "        project=project,\n",
    "        location=region,\n",
    "    )\n",
    "\n",
    "    # instantiate dataflow args\n",
    "    dataflow_args_train = build_dataflow_args(\n",
    "        job_name=f\"{job_name}train\",\n",
    "        url=train_data_url,\n",
    "        bq_dataset=bq_dataset,\n",
    "        bq_table=bq_train_table,\n",
    "        runner=runner,\n",
    "        bq_project=project,\n",
    "        subnet=dataflow_subnet,\n",
    "        dataflow_sa=dataflow_sa,\n",
    "    ).after(bq_dataset_op)\n",
    "    dataflow_args_eval = build_dataflow_args(\n",
    "        job_name=f\"{job_name}eval\",\n",
    "        url=eval_data_url,\n",
    "        bq_dataset=bq_dataset,\n",
    "        bq_table=bq_eval_table,\n",
    "        runner=runner,\n",
    "        bq_project=project,\n",
    "        subnet=dataflow_subnet,\n",
    "        dataflow_sa=dataflow_sa,\n",
    "    ).after(bq_dataset_op)\n",
    "\n",
    "    # run dataflow job\n",
    "    dataflow_python_train_op = DataflowPythonJobOp(\n",
    "        python_module_path=python_file_path,\n",
    "        args=dataflow_args_train.output,\n",
    "        project=project,\n",
    "        location=region,\n",
    "        temp_location=f\"{dataflow_temp_location}/train\",\n",
    "    ).after(dataflow_args_train)\n",
    "    dataflow_python_eval_op = DataflowPythonJobOp(\n",
    "        python_module_path=python_file_path,\n",
    "        args=dataflow_args_eval.output,\n",
    "        project=project,\n",
    "        location=region,\n",
    "        temp_location=f\"{dataflow_temp_location}/eval\",\n",
    "    ).after(dataflow_args_eval)\n",
    "\n",
    "    dataflow_wait_train_op = WaitGcpResourcesOp(\n",
    "        gcp_resources=dataflow_python_train_op.outputs[\"gcp_resources\"]\n",
    "    ).after(dataflow_python_train_op)\n",
    "    dataflow_wait_eval_op = WaitGcpResourcesOp(\n",
    "        gcp_resources=dataflow_python_eval_op.outputs[\"gcp_resources\"]\n",
    "    ).after(dataflow_python_eval_op)\n",
    "        \n",
    "    # create and train model\n",
    "    custom_training_task = custom_job_distributed_training_op(\n",
    "        lr=lr,\n",
    "        epochs=epochs,\n",
    "        project=project,\n",
    "        table=bq_train_table,\n",
    "        dataset=bq_dataset,\n",
    "        location=region,\n",
    "        base_output_directory=base_train_dir,\n",
    "        tb_log_dir=tb_log_dir,\n",
    "        batch_size=batch_size\n",
    "    ).after(dataflow_wait_train_op)\n",
    "    \n",
    "    custom_eval_task = custom_eval_model(\n",
    "        model_dir=model_dir,\n",
    "        project=project,\n",
    "        table=bq_eval_table,\n",
    "        dataset=bq_dataset,\n",
    "        tb_log_dir=tb_log_dir,\n",
    "        model=custom_training_task.outputs[\"model\"],\n",
    "        batch_size=batch_size,\n",
    "    )\n",
    "    custom_eval_task.after(custom_training_task)\n",
    "    custom_eval_task.after(dataflow_wait_eval_op)\n",
    "    with dsl.If(\n",
    "        custom_eval_task.outputs[\"dep_decision\"] == \"true\",\n",
    "        name=\"deploy_decision\",\n",
    "    ):\n",
    "        model_deploy_op = deploy_model(\n",
    "            serving_container_image_uri=deployment_image,\n",
    "            model_name=deployed_model_name,\n",
    "            endpoint_name=endpoint_name,\n",
    "            project_id=deployment_project,\n",
    "            region=region,\n",
    "            split=traffic_split,\n",
    "            model=custom_training_task.outputs['model'],\n",
    "            model_dir=model_dir,\n",
    "            min_nodes=min_nodes,\n",
    "            max_nodes=max_nodes,\n",
    "            encryption=encryption,\n",
    "            service_account=prod_service_account\n",
    "        ).after(custom_eval_task)\n",
    "   \n",
    "        monitroing_job = create_monitoring(\n",
    "            monitoring_name=monitoring_name,\n",
    "            project_id=deployment_project,\n",
    "            region=region,\n",
    "            endpoint=model_deploy_op.outputs['vertex_endpoint'],\n",
    "            bq_data_uri=f\"bq://{project}.{bq_dataset}.{bq_train_table}\",\n",
    "            bucket_name=bucket_name,\n",
    "            email=monitoring_email,\n",
    "            encryption=encryption,\n",
    "            service_account=service_account\n",
    "        ).after(model_deploy_op)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "777fbce5-34fa-43e8-b1b8-0e95d6f20e3c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "compiler.Compiler().compile(pipeline_func=pipeline, package_path=\"./common/vertex-ai-pipeline/pipeline_package.yaml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2899022-8116-4f8d-bff3-07550040f411",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Pipeline run"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44012e12-9746-48a1-9161-ff9f843a7d90",
   "metadata": {
    "tags": []
   },
   "source": [
    "Finally to test the pipeline end to end, set the input arguments for each component accordingly.\n",
    "Note that there are two service accounts supplied. One for the current project and one for the prod environment. The reason behind it the CI/CD design that runs the pipeline in one environment (non-prod) and deploys the model to prod."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78d133e2-a353-45bb-bd5d-f741f626e777",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "timestamp = datetime.now().strftime(\"%d_%H_%M_%S\")\n",
    "pipelineroot = f'{BUCKET_URI}/pipelineroot'\n",
    "# In dev, these two service accounts are the same as the deployment environment is the same as where pipeline runs\n",
    "service_account = COMPUTE_ENGINE_SA\n",
    "prod_service_account = COMPUTE_ENGINE_SA\n",
    "\n",
    "data_config={\n",
    " \"train_data_url\": TRAINING_URL,\n",
    " \"eval_data_url\": EVAL_URL,\n",
    " \"bq_dataset\": DATASET_ID,\n",
    " \"bq_train_table\": TRAINING_TABLE_ID,\n",
    " \"bq_eval_table\": EVAL_TABLE_ID,\n",
    "}\n",
    "\n",
    "dataflow_config={\n",
    "                \"job_name\": JOB_NAME,\n",
    "                \"python_file_path\": f'{BUCKET_URI}/src/ingest_pipeline.py',\n",
    "                \"temp_location\": f'{BUCKET_URI}/temp_dataflow',\n",
    "                \"runner\": RUNNER,\n",
    "                \"subnet\": DATAFLOW_SUBNET,\n",
    "}\n",
    "\n",
    "train_config={\n",
    "             'lr': 0.01, \n",
    "             'epochs': 5, \n",
    "             'base_train_dir': f'{BUCKET_URI}/training', \n",
    "             'tb_log_dir': f'{BUCKET_URI}/tblogs',\n",
    "}\n",
    "\n",
    "deployment_config={\n",
    "    'image': \"us-docker.pkg.dev/cloud-aiplatform/prediction/tf2-cpu.2-8:latest\",\n",
    "    'model_name': \"income_bracket_predictor\",\n",
    "    'endpoint_name': \"census_income_endpoint\",\n",
    "    'min_nodes': 2,\n",
    "    'max_nodes': 4,\n",
    "    'deployment_project': PROJECT_ID,\n",
    "    # important to replace the envryption key here with the key in your own dev environment.\n",
    "    # format would be: projects/prj-d-kms-####/locations/us-central1/keyRings/sample-keyring/cryptoKeys/prj-d-bu3machine-learning\n",
    "    \"encryption\": KMS_KEY,\n",
    "    \"service_account\": service_account,\n",
    "    \"prod_service_account\": prod_service_account,\n",
    "}\n",
    "\n",
    "monitoring_config={\n",
    "    'email': 'ccolin@clsecteam.com', \n",
    "    'name': 'census_monitoring'  \n",
    "}\n",
    "\n",
    "\n",
    "\n",
    "pipeline = aiplatform.PipelineJob(\n",
    "    display_name=f\"census_income_{timestamp}\",\n",
    "    template_path='./common/vertex-ai-pipeline/pipeline_package.yaml',\n",
    "    pipeline_root=pipelineroot,\n",
    "    # important to replace the envryption key here with the key in your own dev environment.\n",
    "    # format would be: projects/prj-d-kms-####/locations/us-central1/keyRings/sample-keyring/cryptoKeys/prj-d-bu3machine-learning\n",
    "    encryption_spec_key_name=KMS_KEY,\n",
    "    parameter_values={\n",
    "        \"create_bq_dataset_query\": create_bq_dataset_query,\n",
    "        \"bq_dataset\": data_config['bq_dataset'],\n",
    "        \"bq_train_table\": data_config['bq_train_table'],\n",
    "        \"bq_eval_table\": data_config['bq_eval_table'],\n",
    "        \"job_name\": dataflow_config['job_name'],\n",
    "        \"train_data_url\": data_config['train_data_url'],\n",
    "        \"eval_data_url\": data_config['eval_data_url'],\n",
    "        \"python_file_path\": dataflow_config['python_file_path'],\n",
    "        \"dataflow_temp_location\": dataflow_config['temp_location'],\n",
    "        \"runner\": dataflow_config['runner'],\n",
    "        \"dataflow_subnet\": dataflow_config['subnet'],\n",
    "        \"project\": PROJECT_ID,\n",
    "        \"region\": REGION,\n",
    "        \"model_dir\": f\"{BUCKET_URI}\",\n",
    "        \"bucket_name\": BUCKET_URI[5:],\n",
    "        \"epochs\": train_config['epochs'],\n",
    "        \"lr\": train_config['lr'],\n",
    "        \"base_train_dir\": train_config['base_train_dir'],\n",
    "        \"tb_log_dir\": train_config['tb_log_dir'],\n",
    "        \"deployment_image\": deployment_config['image'],\n",
    "        \"deployed_model_name\": deployment_config[\"model_name\"],\n",
    "        \"endpoint_name\": deployment_config[\"endpoint_name\"],\n",
    "        \"min_nodes\": deployment_config[\"min_nodes\"],\n",
    "        \"max_nodes\": deployment_config[\"max_nodes\"],\n",
    "        \"deployment_project\": deployment_config[\"deployment_project\"],\n",
    "        \"encryption\": deployment_config.get(\"encryption\"),\n",
    "        \"service_account\": deployment_config[\"service_account\"],\n",
    "        \"prod_service_account\": deployment_config[\"prod_service_account\"],\n",
    "        \"monitoring_name\": monitoring_config['name'],\n",
    "        \"monitoring_email\": monitoring_config['email'], \n",
    "        \n",
    "    },\n",
    "    enable_caching=False,\n",
    ")\n",
    "\n",
    "\n",
    "pipeline.run(service_account=service_account)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3b5cc0e-5577-4c9a-8d91-6cda3cc23049",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m118",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m118"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}