ml_project.ipynb

{
 "cells": [
  {
   "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",
    "KMS_KEY = \"projects/prj-d-kms-3i3k/locations/us-central1/keyRings/sample-keyring/cryptoKeys/prj-d-bu3machine-learning\"\n",
    "PIPELINE_RUNNER = \"pipeline-runner@prj-d-bu3machine-learning-ma6i.iam.gserviceaccount.com\"\n",
    "PEER_NETWORK=\"projects/316945073583/global/networks/vpc-d-shared-restricted\"\n",
    "HOST_NETWORK_PROJECT_ID=\"prj-d-shared-restricted-9enw\"\n",
    "SHARED_SUBNETWORK=\"https://www.googleapis.com/compute/v1/projects/prj-d-shared-restricted-9enw/regions/us-central1/subnetworks/sb-d-shared-restricted-us-central1\"\n",
    "# must have compute network user in shared vpc + roles/dataflow.admin\n",
    "DATAFLOW_RUNNER=\"notebook-runner@prj-d-bu3machine-learning-ma6i.iam.gserviceaccount.com\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "331e585f-bb53-495e-80d7-c963a2fb3bb3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "DATA_PATH = \"data\"\n",
    "KFP_COMPONENTS_PATH = \"components\"\n",
    "SRC = \"src\"\n",
    "BUILD = \"build\"\n",
    "Image = f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/pipelinerepo/pipeline_tutorial:latest\"\n",
    "GCLOUD_IMAGE=f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/pipelinerepo/gcloud\"\n",
    "\n",
    "#export GOOGLE_APPLICATION_CREDENTIALS = \n",
    "!mkdir -m 777 -p {SRC} {DATA_PATH} {KFP_COMPONENTS_PATH} {BUILD}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66521495",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile Dockerfile.python3.8slim\n",
    "FROM python:3.8-slim\n",
    "RUN pip install kfp==2.4.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "860e0348",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile Dockerfile.gcloud\n",
    "FROM google/cloud-sdk\n",
    "RUN pip install kfp==2.4.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bf4e9cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker build . -f ./Dockerfile.gcloud -t \"{GCLOUD_IMAGE}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c6a2413",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker build . -f ./Dockerfile.python3.8slim -t \"{REGION}-docker.pkg.dev/{PROJECT_ID}/pipelinerepo/python3.8-slim\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5b55c98",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker build . -f ./Dockerfile -t \"{Image}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb82b1d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!echo Y | gcloud auth configure-docker {REGION}-docker.pkg.dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dce47863",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker push \"{REGION}-docker.pkg.dev/{PROJECT_ID}/pipelinerepo/python3.8-slim\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd7bf3ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker push \"{GCLOUD_IMAGE}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c043e294",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker push \"{Image}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f60c6bc3-d97f-41c8-879c-1edb59a6cfb9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install protobuf==3.20.* tensorflow==2.8.0  tensorflow-hub==0.13.0 kfp==2.4.0"
   ]
  },
  {
   "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": "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": "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": "code",
   "execution_count": null,
   "id": "dd5ace09-afd5-4135-8e28-9f8b6c2b1c96",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mkdir -m 777 -p {KFP_COMPONENTS_PATH}/bq_dataset_component"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbbf2bee-cc42-4aff-8fff-3bb506bcc741",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Dataflow sources"
   ]
  },
  {
   "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=\"gs://vertexpipe/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": "code",
   "execution_count": null,
   "id": "0e613ca7-f7cd-46a7-bf35-1c4c2a67f0fb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile requirements.txt\n",
    "apache-beam[gcp]==2.52.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76881a2e-0ed8-4d0c-9e1f-512a93130312",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!gsutil cp -R {SRC} {BUCKET_URI}/{SRC}\n",
    "!gsutil cp requirements.txt {BUCKET_URI}/requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1052dd72-692d-4215-9ff5-e791f3fbc41a",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Dataflow arguments component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "167c96f7-6416-464e-9326-6350b1345cbb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@component(base_image=f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/pipelinerepo/python3.8-slim\")\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",
    "    subnetwork: 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",
    "        subnetwork,\n",
    "        \"--service_account_email\",\n",
    "        dataflow_sa,\n",
    "        \"--no_use_public_ips\",\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03de2873-3145-4a7f-abaf-3de1267ac708",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Costum training component"
   ]
  },
  {
   "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": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "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": "bbb09b15",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install google-cloud-pipeline-components==2.8.0"
   ]
  },
  {
   "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"
   ]
  },
  {
   "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": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "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",
    "    \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": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "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=\"compute default\"\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=\"compute default\"\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": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "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",
    "    import time\n",
    "    aiplatform.init(service_account=service_account)\n",
    "    list_monitors = aiplatform.ModelDeploymentMonitoringJob.list(filter=f'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",
    "        yaml = \"\"\"type: object\n",
    "        properties:\n",
    "        \"\"\"\n",
    "        schema = bq_table.schema\n",
    "        for feature in schema:\n",
    "            if feature.name == TARGET:\n",
    "                continue\n",
    "            if feature.field_type == \"STRING\":\n",
    "                f_type = \"string\"\n",
    "            else:\n",
    "                f_type = \"float\"\n",
    "            yaml += f\"\"\"  {feature.name}:\n",
    "            type: {f_type}\n",
    "        \"\"\"\n",
    "\n",
    "        yaml += \"\"\"required:\n",
    "        \"\"\"\n",
    "        for feature in schema:\n",
    "            if feature.name == \"income_bracket\":\n",
    "                continue\n",
    "            yaml += f\"\"\"- {feature.name}\n",
    "        \"\"\"\n",
    "        with open(\"monitoring_schema.yaml\", \"w\") as f:\n",
    "            f.write(yaml)\n",
    "        storage_client = storage.Client()\n",
    "        bucket = storage_client.bucket(bucket_name)\n",
    "        blob = bucket.blob(\"monitoring_schema.yaml\")\n",
    "        blob.upload_from_string(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\"{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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b9e2a15",
   "metadata": {},
   "outputs": [],
   "source": [
    "@component(\n",
    "        base_image = f\"{GCLOUD_IMAGE}\"\n",
    ")\n",
    "def arbitrary() -> str:\n",
    "    import subprocess\n",
    "    bash_command = \"gcloud config list\"\n",
    "    result = subprocess.run(bash_command, shell=True, capture_output=True, text=True)\n",
    "\n",
    "    print(\"Command Output:\")\n",
    "    print(result.stdout)\n",
    "\n",
    "    if result.stderr:\n",
    "        print(\"Error Output:\")\n",
    "        print(result.stderr)\n",
    "    bash_command = \"gcloud compute regions list\"\n",
    "    result = subprocess.run(bash_command, shell=True, capture_output=True, text=True)\n",
    "\n",
    "    print(\"Command Output:\")\n",
    "    print(result.stdout)\n",
    "\n",
    "    if result.stderr:\n",
    "        print(\"Error Output:\")\n",
    "        print(result.stderr)\n",
    "    return \"Finished\""
   ]
  },
  {
   "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",
    "    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",
    "    requirements_file_path: str=f'{BUCKET_URI}/requirements.txt',\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",
    "    subnetwork: str=SHARED_SUBNETWORK,\n",
    "    dataflow_sa: str=DATAFLOW_RUNNER,\n",
    "):\n",
    "    from google_cloud_pipeline_components.v1.bigquery import (\n",
    "        BigqueryCreateModelJobOp, BigqueryEvaluateModelJobOp,\n",
    "        BigqueryPredictModelJobOp, 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",
    "    #arbitrary container\n",
    "    arb_op = arbitrary()\n",
    "\n",
    "    # create the dataset\n",
    "    bq_dataset_op = BigqueryQueryJobOp(\n",
    "        query=create_bq_dataset_query,\n",
    "        project=project,\n",
    "        location=region,\n",
    "    ).after(arb_op)\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",
    "        subnetwork=subnetwork,\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",
    "        subnetwork=subnetwork,\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=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": "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",
    "service_account = PIPELINE_RUNNER\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",
    "                \"requirements_file_path\": f'{BUCKET_URI}/requirements.txt',\n",
    "                \"python_file_path\": f'{BUCKET_URI}/src/ingest_pipeline.py',\n",
    "                \"setup_file_uri\": f'{BUCKET_URI}/setup.py',\n",
    "                \"temp_location\": f'{BUCKET_URI}/temp_dataflow',\n",
    "                \"runner\": RUNNER,\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",
    "    \"encryption\": KMS_KEY,\n",
    "    \"service_account\": service_account\n",
    "}\n",
    "\n",
    "monitoring_config={\n",
    "    'email': 'ccolin@clsecteam.com', \n",
    "    'name': 'census_monitoring'  \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",
    "    encryption_spec_key_name=deployment_config.get(\"encryption\"),\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",
    "        \"requirements_file_path\": dataflow_config['requirements_file_path'],\n",
    "        \"python_file_path\": dataflow_config['python_file_path'],\n",
    "        \"dataflow_temp_location\": dataflow_config['temp_location'],\n",
    "        \"runner\": dataflow_config['runner'],\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",
    "        \"monitoring_name\": monitoring_config['name'],\n",
    "        \"monitoring_email\": monitoring_config['email'], \n",
    "        \n",
    "    },\n",
    "    enable_caching=False,\n",
    ")\n",
    "\n",
    "pipeline.submit(network=PEER_NETWORK)\n",
    "\n",
    "# might need to give service agent permission\n",
    "# $gcloud projects add-iam-policy-binding ARTIFACT-PROJECT --role=roles/artifactregistry.reader --member=serviceAccount:service-PRJ-NUMBER@gcp-sa-aiplatform-cc.iam.gserviceaccount.com"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bf3827a-b88c-4110-aa4e-96836472056a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model = tf.saved_model.load(f\"{BUCKET_URI}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73dc5b11-c9ac-4515-8446-07b30f2a94b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model.signatures"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m119",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m119"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}