notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MLOps using Kubeflow and Feature Store\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Before start\n",
    "- Google Cloud Storage Bucket / Pipelines / Feature Store must be created on same region {suggested: europe-west3(Frankfurt)}\n",
    "- Pipelines considerations when failing if jobs is \"in progress\", needs to improve adding code until previous import is done\n",
    "- Some issues with $tfx 2.7 so needed to downgrade to 2.5\n",
    "- Align on aiplatform sdk\n",
    "- Have good mlops use case for using feature store serving (costs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "#Create new virtual environment\n",
    "conda create -n myenv python=3.7\n",
    "conda activate myenv\n",
    "\n",
    "#Launch Jupyter notebook using teh virtual environment\n",
    "pip install --user ipykernel\n",
    "python -m ipykernel install --user --name=myenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# The Google Cloud Notebook product has specific requirements\n",
    "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
    "\n",
    "# Google Cloud Notebook\n",
    "USER_FLAG = \"\"\n",
    "if IS_GOOGLE_CLOUD_NOTEBOOK:\n",
    "    USER_FLAG = \"--user\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "from google.cloud import aiplatform\n",
    "from google_cloud_pipeline_components import aiplatform as gcc_aip\n",
    "from kfp.v2 import compiler, dsl\n",
    "from kfp.v2.google.client import AIPlatformClient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports for environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "from google_cloud_pipeline_components import aiplatform as gcc_aip\n",
    "\n",
    "from kfp.v2 import compiler, dsl\n",
    "from kfp.v2.google.client import AIPlatformClient\n",
    "from google.cloud.aiplatform_v1 import (\n",
    "    FeaturestoreOnlineServingServiceClient,\n",
    "    FeaturestoreServiceClient,\n",
    ")\n",
    "from google.cloud.aiplatform_v1.types import FeatureSelector, IdMatcher\n",
    "from google.cloud.aiplatform_v1.types import entity_type as entity_type_pb2\n",
    "from google.cloud.aiplatform_v1.types import feature as feature_pb2\n",
    "from google.cloud.aiplatform_v1.types import featurestore as featurestore_pb2\n",
    "from google.cloud.aiplatform_v1.types import (\n",
    "    featurestore_online_service as featurestore_online_service_pb2,\n",
    ")\n",
    "from google.cloud.aiplatform_v1.types import (\n",
    "    featurestore_service as featurestore_service_pb2,\n",
    ")\n",
    "from google.cloud.aiplatform_v1.types import io as io_pb2\n",
    "from google.protobuf.duration_pb2 import Duration\n",
    "from google.cloud import aiplatform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "\n",
    "from infrastructure.src.main.resources import load_component\n",
    "from infrastructure.src.main.resources import feature_store_helper\n",
    "\n",
    "from infrastructure.src.main.resources.generator import generator_component\n",
    "from infrastructure.src.main.resources.ingester import ingester_component\n",
    "from infrastructure.src.main.resources.trainer import trainer_component"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "PROJECT_ID = \"\"\n",
    "\n",
    "# Get your Google Cloud project ID from gcloud\n",
    "if not os.getenv(\"IS_TESTING\"):\n",
    "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
    "    PROJECT_ID = shell_output[0]\n",
    "    print(\"Project ID: \", PROJECT_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bucket, BigQuery and Feature Store Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BUCKET_NAME = \"gs://\" + PROJECT_ID + \"movielens\"\n",
    "REGION = \"europe-west4\"\n",
    "\n",
    "API_ENDPOINT = \"europe-west4-aiplatform.googleapis.com\"\n",
    "INPUT_CSV_FILE = \"\"\n",
    "FEATURESTORE_ID = \"movie_prediction\"\n",
    "FEATURE_STORE_REGION = REGION\n",
    "ENTITY_TYPE_ID=\"movie_entity\"\n",
    "ENTITY_ID_FIELD=\"user_id\"\n",
    "\n",
    "# BigQuery parameters\n",
    "BIGQUERY_DATASET_ID = f\"{PROJECT_ID}.movielens_dataset\"\n",
    "BIGQUERY_LOCATION = \"EU\"\n",
    "BIGQUERY_TABLE_ID = f\"{BIGQUERY_DATASET_ID}.training_dataset\"\n",
    "BIGQUERY_RAW_TABLE_ID = f\"{BIGQUERY_DATASET_ID}.raw_dataset\"\n",
    "BIGQUERY_INPUT_URI=f\"bq://{BIGQUERY_RAW_TABLE_ID}\"\n",
    "\n",
    "\n",
    "# https://www.kaggle.com/prajitdatta/movielens-100k-dataset\n",
    "# Dataset parameters\n",
    "RAW_DATA_PATH = BUCKET_NAME+\"/raw_data/u.data\"\n",
    "\n",
    "# u.data -- The full u data set, 100000 ratings by 943 users on 1682 items.\n",
    "# Each user has rated at least 20 movies. Users and items are numbered consecutively from 1. The data is randomly ordered. \n",
    "# This is a tab separated list of user id | item id | rating | timestamp.\n",
    "# The time stamps are unix seconds since 1/1/1970 UTC\n",
    "\n",
    "# Pipeline parameters\n",
    "PIPELINE_NAME = \"movie-prediction\"\n",
    "ENABLE_CACHING = False\n",
    "PIPELINE_ROOT = f\"{BUCKET_NAME}/pipeline\"\n",
    "PIPELINE_SPEC_PATH = \"./pipeline/metadata_pipeline.json\"\n",
    "OUTPUT_COMPONENT_SPEC = \"output-component.yaml\"\n",
    "\n",
    "BIGQUERY_TMP_FILE = (\n",
    "    \"tmp.json\" \n",
    ")\n",
    "BIGQUERY_MAX_ROWS = 5 \n",
    "\n",
    "TFRECORD_FILE = (\n",
    "    f\"{BUCKET_NAME}/trainer_input_path/*\"  \n",
    ")\n",
    "\n",
    "LOGGER_PUBSUB_TOPIC = \"logger-pubsub-topic\"\n",
    "LOGGER_CLOUD_FUNCTION = \"logger-cloud-function\"\n",
    "\n",
    "# Trainer parameters\n",
    "TRAINING_ARTIFACTS_DIR = f\"{BUCKET_NAME}/artifacts\"\n",
    "TRAINING_REPLICA_COUNT = 1\n",
    "TRAINING_MACHINE_TYPE = \"n1-standard-4\"\n",
    "TRAINING_ACCELERATOR_TYPE = \"ACCELERATOR_TYPE_UNSPECIFIED\"\n",
    "TRAINING_ACCELERATOR_COUNT = 0\n",
    "\n",
    "# Deployer parameters\n",
    "TRAINED_POLICY_DISPLAY_NAME = \"movielens-trained-policy\"\n",
    "ENDPOINT_DISPLAY_NAME = \"movielens-endpoint\"\n",
    "ENDPOINT_MACHINE_TYPE = \"n1-standard-4\"\n",
    "ENDPOINT_REPLICA_COUNT = 1  \n",
    "ENDPOINT_ACCELERATOR_TYPE = \"ACCELERATOR_TYPE_UNSPECIFIED\"  \n",
    "ENDPOINT_ACCELERATOR_COUNT = 0 \n",
    "\n",
    "# Prediction container parameters\n",
    "PREDICTION_CONTAINER = \"prediction-container\"\n",
    "PREDICTION_CONTAINER_DIR = \"./infrastructure/src/main/resources/prediction_container\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bucket Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! gsutil mb -l $REGION $BUCKET_NAME\n",
    "\n",
    "# Download the sample data into your RAW_DATA_PATH\n",
    "! gsutil cp \"raw_data/u.data\" $RAW_DATA_PATH"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Featurestore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_store_helper.create_featurestore(\n",
    "    project=PROJECT_ID,\n",
    "    featurestore_id=FEATURESTORE_ID,\n",
    "    location=FEATURE_STORE_REGION,\n",
    "    api_endpoint=API_ENDPOINT,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Listing Featurestore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get featurestore list to verify, can use this code later to cleanup\n",
    "featurestore_list = feature_store_helper.list_featurestore(\n",
    "    project=PROJECT_ID, location=FEATURE_STORE_REGION,\n",
    "    api_endpoint=API_ENDPOINT\n",
    ")\n",
    "print(featurestore_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Entity Featurestore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_store_helper.create_entity_type(project=PROJECT_ID,\n",
    "                                        location=FEATURE_STORE_REGION,\n",
    "                                        api_endpoint=API_ENDPOINT,\n",
    "                                        featurestore_id=FEATURESTORE_ID,\n",
    "                                        entity_type_id=ENTITY_TYPE_ID,\n",
    "                                        description=\"movies\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Features for Featurestore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_store_helper.create_feature(project=PROJECT_ID,\n",
    "                                    location=FEATURE_STORE_REGION,\n",
    "                                    api_endpoint=API_ENDPOINT,\n",
    "                                    featurestore_id=FEATURESTORE_ID,\n",
    "                                    entity_type_id=ENTITY_TYPE_ID,\n",
    "                                    feature_id=\"user_id\",\n",
    "                                    value_type=aiplatform.gapic.Feature.ValueType.STRING)\n",
    "feature_store_helper.create_feature(project=PROJECT_ID,\n",
    "                                    location=FEATURE_STORE_REGION,\n",
    "                                    api_endpoint=API_ENDPOINT,\n",
    "                                    featurestore_id=FEATURESTORE_ID,\n",
    "                                    entity_type_id=ENTITY_TYPE_ID,\n",
    "                                    feature_id=\"item_id\",\n",
    "                                    value_type=aiplatform.gapic.Feature.ValueType.STRING)\n",
    "feature_store_helper.create_feature(project=PROJECT_ID,\n",
    "                                    location=FEATURE_STORE_REGION,\n",
    "                                    api_endpoint=API_ENDPOINT,\n",
    "                                    featurestore_id=FEATURESTORE_ID,\n",
    "                                    entity_type_id=ENTITY_TYPE_ID,\n",
    "                                    feature_id=\"rating\",\n",
    "                                    value_type=aiplatform.gapic.Feature.ValueType.STRING)\n",
    "feature_store_helper.create_feature(project=PROJECT_ID,\n",
    "                                    location=FEATURE_STORE_REGION,\n",
    "                                    api_endpoint=API_ENDPOINT,\n",
    "                                    featurestore_id=FEATURESTORE_ID,\n",
    "                                    entity_type_id=ENTITY_TYPE_ID,\n",
    "                                    feature_id=\"timestamp\",\n",
    "                                    value_type=aiplatform.gapic.Feature.ValueType.STRING)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Execute Cloudbuild file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cloudbuild_yaml = \"\"\"steps:\n",
    "- name: \"gcr.io/kaniko-project/executor:latest\"\n",
    "  args: [\"--destination=gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest\",\n",
    "         \"--cache=false\",\n",
    "         \"--cache-ttl=99h\"]\n",
    "  env: [\"AIP_STORAGE_URI={ARTIFACTS_DIR}\",\n",
    "        \"PROJECT_ID={PROJECT_ID}\",\n",
    "        \"LOGGER_PUBSUB_TOPIC={LOGGER_PUBSUB_TOPIC}\"]\n",
    "options:\n",
    "  machineType: \"E2_HIGHCPU_8\"\n",
    "\"\"\".format(\n",
    "    PROJECT_ID=PROJECT_ID,\n",
    "    PREDICTION_CONTAINER=PREDICTION_CONTAINER,\n",
    "    ARTIFACTS_DIR=TRAINING_ARTIFACTS_DIR,\n",
    "    LOGGER_PUBSUB_TOPIC=LOGGER_PUBSUB_TOPIC,\n",
    ")\n",
    "\n",
    "with open(f\"{PREDICTION_CONTAINER_DIR}/cloudbuild.yaml\", \"w\") as fp:\n",
    "    fp.write(cloudbuild_yaml)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! gcloud builds submit --config $PREDICTION_CONTAINER_DIR/cloudbuild.yaml $PREDICTION_CONTAINER_DIR\n",
    "#gcloud builds submit --config src/prediction_container/cloudbuild.yaml  src/prediction_container\n",
    "#TESTER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RUN code locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "importlib.reload(load_component)\n",
    "load_component.load_raw_data_to_bigquery(PROJECT_ID,RAW_DATA_PATH,BIGQUERY_DATASET_ID, BIGQUERY_LOCATION, BIGQUERY_RAW_TABLE_ID)\n",
    "\n",
    "importlib.reload(generator_component)\n",
    "generator_component.generate_movielens_dataset_for_bigquery(PROJECT_ID,RAW_DATA_PATH,8,20,20,3, BIGQUERY_TMP_FILE, BIGQUERY_DATASET_ID, BIGQUERY_LOCATION,BIGQUERY_TABLE_ID)\n",
    "\n",
    "importlib.reload(ingester_component)\n",
    "ingester_component.ingest_bigquery_dataset_into_tfrecord(PROJECT_ID, BIGQUERY_TABLE_ID,TFRECORD_FILE, 10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# code for manually verifying your tfrecord dataset\n",
    "import tensorflow as tf \n",
    "raw_dataset = tf.data.TFRecordDataset(TFRECORD_FILE)\n",
    "\n",
    "for raw_record in raw_dataset.take(1):\n",
    "    example = tf.train.Example()\n",
    "    example.ParseFromString(raw_record.numpy())\n",
    "    print(example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "importlib.reload(trainer_component)\n",
    "trainer_component.training_op"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from infrastructure.src.main.resources import feature_store_helper\n",
    "importlib.reload(feature_store_helper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_store_helper.import_feature_values(project=PROJECT_ID,\n",
    "                                           location=FEATURE_STORE_REGION,\n",
    "                                    api_endpoint=API_ENDPOINT,\n",
    "                                    featurestore_id=FEATURESTORE_ID,\n",
    "                                    entity_type_id=ENTITY_TYPE_ID,\n",
    "                                    bigquery_uri=BIGQUERY_TABLE_ID,\n",
    "                                    entity_id_field=ENTITY_ID_FIELD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from infrastructure.src.main.resources import bigquery_to_featurestore\n",
    "importlib.reload(bigquery_to_featurestore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bigquery_to_featurestore.import_feature_values(project=PROJECT_ID,\n",
    "                                               location=FEATURE_STORE_REGION,\n",
    "                                    api_endpoint=API_ENDPOINT,\n",
    "                                    featurestore_id=FEATURESTORE_ID,\n",
    "                                    entity_type_id=ENTITY_TYPE_ID,\n",
    "                                    bigquery_uri=BIGQUERY_TABLE_ID,\n",
    "                                    entity_id_field=ENTITY_ID_FIELD)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kubeflow Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from kfp.components import load_component_from_url\n",
    "from kfp.components import create_component_from_func\n",
    "\n",
    "from infrastructure.src.main.resources.trainer import trainer_component\n",
    "from infrastructure.src.main.resources.generator import generator_component\n",
    "from infrastructure.src.main.resources.ingester import ingester_component\n",
    "from infrastructure.src.main.resources import load_component\n",
    "from infrastructure.src.main.resources import bigquery_to_featurestore\n",
    "\n",
    "# so jupyter kernel reloads the modules when we change them\n",
    "importlib.reload(load_component)\n",
    "importlib.reload(generator_component)\n",
    "importlib.reload(ingester_component)\n",
    "importlib.reload(bigquery_to_featurestore)\n",
    "\n",
    "\n",
    "@dsl.pipeline(pipeline_root=PIPELINE_ROOT, name=f\"{PIPELINE_NAME}-startup\")\n",
    "def pipeline(\n",
    "    # Pipeline configs\n",
    "    project_id: str,\n",
    "    raw_data_path: str,\n",
    "    training_artifacts_dir: str,\n",
    "    featurestore_id: str,\n",
    "    entity_type_id: str,\n",
    "    bigquery_uri: str,\n",
    "    entity_id_field: str,\n",
    "    # BigQuery configs\n",
    "    bigquery_dataset_id: str,\n",
    "    bigquery_location: str,\n",
    "    bigquery_table_id: str,\n",
    "    bigquery_raw_table_id: str,\n",
    "    bigquery_max_rows: int = 10000,\n",
    "    # TF-Agents RL configs\n",
    "    batch_size: int = 8,\n",
    "    rank_k: int = 20,\n",
    "    num_actions: int = 20,\n",
    "    driver_steps: int = 3,\n",
    "    num_epochs: int = 5,\n",
    "    tikhonov_weight: float = 0.01,\n",
    "    agent_alpha: float = 10,\n",
    ") -> None:\n",
    "    \n",
    "    load_op = create_component_from_func(\n",
    "    func=load_component.load_raw_data_to_bigquery,\n",
    "    output_component_file=\"./outputs/load-output-component.yaml\",\n",
    "    packages_to_install=[\n",
    "      \"google-cloud-bigquery==2.20.0\",\n",
    "    ],\n",
    "  )\n",
    "    load_task = load_op(\n",
    "        project_id=project_id,\n",
    "        raw_data_path=raw_data_path,\n",
    "        bigquery_dataset_id=bigquery_dataset_id,\n",
    "        bigquery_location=bigquery_location,\n",
    "        bigquery_table_id=bigquery_raw_table_id,\n",
    "    )\n",
    "    \n",
    "    preprocess_op = create_component_from_func(\n",
    "    func=bigquery_to_featurestore.import_feature_values,\n",
    "    output_component_file=\"./outputs/preprocess-output-component.yaml\",\n",
    "    packages_to_install=[\n",
    "      \"google-cloud-aiplatform\",\n",
    "    ],\n",
    "    )\n",
    "    preprocess_task = preprocess_op(\n",
    "        project=project_id,\n",
    "        featurestore_id=featurestore_id,\n",
    "        entity_type_id=entity_type_id,\n",
    "        bigquery_uri=bigquery_uri,\n",
    "        bigquery_table_id=load_task.outputs[\"bigquery_table_id\"],\n",
    "        entity_id_field=entity_id_field, \n",
    "    )\n",
    "    \n",
    "    generate_op = create_component_from_func(\n",
    "    func=generator_component.generate_movielens_dataset_for_bigquery,\n",
    "    base_image=\"tensorflow/tensorflow:2.5.0\",\n",
    "    output_component_file=\"./outputs/generate-output-component.yaml\",\n",
    "    packages_to_install=[\n",
    "      \"google-cloud-bigquery==2.20.0\",\n",
    "      \"tensorflow==2.5.0\",\n",
    "      \"Image\",\n",
    "      \"tf-agents==0.8.0\",\n",
    "    ],\n",
    "  )\n",
    "\n",
    "    # Run the Generator component.\n",
    "    generate_task = generate_op(\n",
    "        project_id=project_id,\n",
    "        raw_data_path=raw_data_path,\n",
    "        batch_size=batch_size,\n",
    "        rank_k=rank_k,\n",
    "        num_actions=num_actions,\n",
    "        driver_steps=driver_steps,\n",
    "        bigquery_tmp_file=BIGQUERY_TMP_FILE,\n",
    "        bigquery_dataset_id=bigquery_dataset_id,\n",
    "        bigquery_location=bigquery_location,\n",
    "        bigquery_table_id=bigquery_table_id,\n",
    "        feature_id=preprocess_task.outputs[\"featurestore_id\"],\n",
    "    )\n",
    "    \n",
    "    ingest_op = create_component_from_func(\n",
    "    func=ingester_component.ingest_bigquery_dataset_into_tfrecord,\n",
    "    base_image=\"tensorflow/tensorflow:2.5.0\",\n",
    "    output_component_file=f\"ingest-{OUTPUT_COMPONENT_SPEC}\",\n",
    "    packages_to_install=[\n",
    "      \"google-cloud-bigquery==2.20.0\",\n",
    "      \"tensorflow==2.5.0\",\n",
    "    ],\n",
    "  )\n",
    "\n",
    "    # Run the Ingester component.\n",
    "    ingest_task = ingest_op(\n",
    "        project_id=project_id,\n",
    "        bigquery_table_id=generate_task.outputs[\"bigquery_table_id\"],\n",
    "        bigquery_max_rows=bigquery_max_rows,\n",
    "        tfrecord_file=TFRECORD_FILE,\n",
    "    )\n",
    "\n",
    "    # Run the Trainer component and submit custom job to Vertex AI.\n",
    "    train_op = create_component_from_func(\n",
    "      func=trainer_component.training_op,\n",
    "      output_component_file=f\"trainer-{OUTPUT_COMPONENT_SPEC}\",\n",
    "      packages_to_install=[\n",
    "          \"tensorflow==2.5.0\",\n",
    "          \"tf-agents==0.8.0\",\n",
    "      ])\n",
    "\n",
    "    train_task = train_op(\n",
    "      training_artifacts_dir=training_artifacts_dir,\n",
    "      # tfrecord_file=ingest_task.outputs[\"tfrecord_file\"],\n",
    "      tfrecord_file=\"gs://mlops-vertex-capgemini/trainer_input_path/tf\",\n",
    "      num_epochs=num_epochs,\n",
    "      rank_k=rank_k,\n",
    "      num_actions=num_actions,\n",
    "      tikhonov_weight=tikhonov_weight,\n",
    "      agent_alpha=agent_alpha)\n",
    "\n",
    "    worker_pool_specs = [\n",
    "        {\n",
    "            \"containerSpec\": {\n",
    "                \"imageUri\": train_task.container.image,\n",
    "            },\n",
    "            \"replicaCount\": TRAINING_REPLICA_COUNT,\n",
    "            \"machineSpec\": {\n",
    "                \"machineType\": TRAINING_MACHINE_TYPE,\n",
    "                \"acceleratorType\": TRAINING_ACCELERATOR_TYPE,\n",
    "                \"acceleratorCount\": TRAINING_ACCELERATOR_COUNT,\n",
    "            },\n",
    "        },\n",
    "    ]\n",
    "    train_task.custom_job_spec = {\n",
    "        \"displayName\": train_task.name,\n",
    "        \"jobSpec\": {\n",
    "            \"workerPoolSpecs\": worker_pool_specs,\n",
    "        },\n",
    "    }\n",
    "\n",
    "    # # Run the Deployer components.\n",
    "    # # Upload the trained policy as a model.\n",
    "    model_upload_op = gcc_aip.ModelUploadOp(\n",
    "      project=project_id,\n",
    "      location=REGION,\n",
    "      display_name=TRAINED_POLICY_DISPLAY_NAME,\n",
    "      artifact_uri=training_artifacts_dir,\n",
    "      serving_container_image_uri=f\"gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest\",\n",
    "    )\n",
    "    train_task.after(ingest_task)\n",
    "    # # Model uploading has to occur after training completes.\n",
    "    model_upload_op.after(train_task)\n",
    "    # # Create a Vertex AI endpoint. (This operation can occur in parallel with\n",
    "    # # the Generator, Ingester, Trainer components.)\n",
    "    endpoint_create_op = gcc_aip.EndpointCreateOp(\n",
    "      project=project_id,\n",
    "      location=REGION,\n",
    "      display_name=ENDPOINT_DISPLAY_NAME)\n",
    "    # Deploy the uploaded, trained policy to the created endpoint. (This operation\n",
    "    # has to occur after both model uploading and endpoint creation complete.)\n",
    "    model_deploy_op = gcc_aip.ModelDeployOp(\n",
    "      project=project_id,\n",
    "      endpoint=endpoint_create_op.outputs[\"endpoint\"],\n",
    "      model=model_upload_op.outputs[\"model\"],\n",
    "      #endpoint=\"495264017615421440\",\n",
    "      #model=\"944845526120005632\",\n",
    "      deployed_model_display_name=TRAINED_POLICY_DISPLAY_NAME,\n",
    "      machine_type=ENDPOINT_MACHINE_TYPE)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile the authored pipeline.\n",
    "compiler.Compiler().compile(pipeline_func=pipeline, package_path=PIPELINE_SPEC_PATH)\n",
    "\n",
    "# Createa Vertex AI client.\n",
    "api_client = AIPlatformClient(project_id=PROJECT_ID, region=REGION)\n",
    "\n",
    "# Create a pipeline run job.\n",
    "response = api_client.create_run_from_job_spec(\n",
    "    job_spec_path=PIPELINE_SPEC_PATH,\n",
    "    parameter_values={\n",
    "        # Pipeline configs\n",
    "        \"project_id\": PROJECT_ID,\n",
    "        \"raw_data_path\": RAW_DATA_PATH,\n",
    "        \"training_artifacts_dir\": TRAINING_ARTIFACTS_DIR,\n",
    "        \"featurestore_id\": FEATURESTORE_ID,\n",
    "        \"entity_type_id\": ENTITY_TYPE_ID,\n",
    "        \"bigquery_uri\": BIGQUERY_INPUT_URI,\n",
    "        \"entity_id_field\": ENTITY_ID_FIELD,\n",
    "        # BigQuery configs\n",
    "        \"bigquery_dataset_id\": BIGQUERY_DATASET_ID,\n",
    "        \"bigquery_location\": BIGQUERY_LOCATION,\n",
    "        \"bigquery_table_id\": BIGQUERY_TABLE_ID,\n",
    "        \"bigquery_raw_table_id\": BIGQUERY_RAW_TABLE_ID,\n",
    "    },\n",
    "    enable_caching=ENABLE_CACHING,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run locally for debugging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_id=PROJECT_ID\n",
    "raw_data_path= RAW_DATA_PATH\n",
    "training_artifacts_dir= TRAINING_ARTIFACTS_DIR\n",
    "# BigQuery configs\n",
    "bigquery_dataset_id= BIGQUERY_DATASET_ID\n",
    "bigquery_location= BIGQUERY_LOCATION\n",
    "bigquery_table_id= BIGQUERY_TABLE_ID\n",
    "tfrecord_file=\"gs://mlops-vertex-capgemini/trainer_input_path/tf\"\n",
    "bigquery_max_rows= 10000\n",
    "# TF-Agents RL configs\n",
    "batch_size=  8\n",
    "rank_k= 20\n",
    "num_actions= 20\n",
    "driver_steps= 3\n",
    "num_epochs= 5\n",
    "tikhonov_weight: float = 0.01\n",
    "agent_alpha: float = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"start\")\n",
    "\"\"\"The Trainer component for training a policy on TFRecord files.\"\"\"\n",
    "# Import for the function return value type.\n",
    "from typing import NamedTuple  # pylint: disable=unused-import\n",
    "\n",
    "from kfp import components\n",
    "\n",
    "import collections\n",
    "from typing import Dict, List, NamedTuple  # pylint: disable=redefined-outer-name,reimported\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "from tf_agents import agents\n",
    "from tf_agents import policies\n",
    "from tf_agents import trajectories\n",
    "from tf_agents.bandits.agents import lin_ucb_agent\n",
    "from tf_agents.policies import policy_saver\n",
    "from tf_agents.specs import tensor_spec\n",
    "\n",
    "import logging\n",
    "\n",
    "per_arm = False  # Using the non-per-arm version of the movie environment.\n",
    "\n",
    "# Mapping from feature name to serialized value\n",
    "feature_description = {\n",
    "    \"step_type\": tf.io.FixedLenFeature((), tf.string),\n",
    "    \"observation\": tf.io.FixedLenFeature((), tf.string),\n",
    "    \"action\": tf.io.FixedLenFeature((), tf.string),\n",
    "    \"policy_info\": tf.io.FixedLenFeature((), tf.string),\n",
    "    \"next_step_type\": tf.io.FixedLenFeature((), tf.string),\n",
    "    \"reward\": tf.io.FixedLenFeature((), tf.string),\n",
    "    \"discount\": tf.io.FixedLenFeature((), tf.string),\n",
    "}\n",
    "\n",
    "def _parse_record(raw_record: tf.Tensor) -> Dict[str, tf.Tensor]:\n",
    "    \"\"\"Parses a serialized `tf.train.Example` proto.\n",
    "    Args:\n",
    "    raw_record: A serialized data record of a `tf.train.Example` proto.\n",
    "    Returns:\n",
    "    A dict mapping feature names to values as `tf.Tensor` objects of type\n",
    "    string containing serialized protos, following `feature_description`.\n",
    "    \"\"\"\n",
    "    return tf.io.parse_single_example(raw_record, feature_description)\n",
    "\n",
    "def build_trajectory(\n",
    "    parsed_record: Dict[str, tf.Tensor],\n",
    "    policy_info: policies.utils.PolicyInfo) -> trajectories.Trajectory:\n",
    "    \"\"\"Builds a `trajectories.Trajectory` object from `parsed_record`.\n",
    "    Args:\n",
    "    parsed_record: A dict mapping feature names to values as `tf.Tensor`\n",
    "        objects of type string containing serialized protos.\n",
    "    policy_info: Policy information specification.\n",
    "    Returns:\n",
    "    A `trajectories.Trajectory` object that contains values as de-serialized\n",
    "    `tf.Tensor` objects from `parsed_record`.\n",
    "    \"\"\"\n",
    "    return trajectories.Trajectory(\n",
    "        step_type=tf.expand_dims(\n",
    "            tf.io.parse_tensor(parsed_record[\"step_type\"], out_type=tf.int32),\n",
    "            axis=1),\n",
    "        observation=tf.expand_dims(\n",
    "            tf.io.parse_tensor(\n",
    "                parsed_record[\"observation\"], out_type=tf.float32),\n",
    "            axis=1),\n",
    "        action=tf.expand_dims(\n",
    "            tf.io.parse_tensor(parsed_record[\"action\"], out_type=tf.int32),\n",
    "            axis=1),\n",
    "        policy_info=policy_info,\n",
    "        next_step_type=tf.expand_dims(\n",
    "            tf.io.parse_tensor(\n",
    "                parsed_record[\"next_step_type\"], out_type=tf.int32),\n",
    "            axis=1),\n",
    "        reward=tf.expand_dims(\n",
    "            tf.io.parse_tensor(parsed_record[\"reward\"], out_type=tf.float32),\n",
    "            axis=1),\n",
    "        discount=tf.expand_dims(\n",
    "            tf.io.parse_tensor(parsed_record[\"discount\"], out_type=tf.float32),\n",
    "            axis=1))\n",
    "\n",
    "def train_policy_on_trajectory(\n",
    "    agent: agents.TFAgent,\n",
    "    tfrecord_file: str,\n",
    "    num_epochs: int\n",
    ") -> NamedTuple(\"TrainOutputs\", [\n",
    "    (\"policy\", policies.TFPolicy),\n",
    "    (\"train_loss\", Dict[str, List[float]]),\n",
    "]):\n",
    "    \"\"\"Trains the policy in `agent` on the dataset of `tfrecord_file`.\n",
    "    Parses `tfrecord_file` as `tf.train.Example` objects, packages them into\n",
    "    `trajectories.Trajectory` objects, and trains the agent's policy on these\n",
    "    trajectory objects.\n",
    "    Args:\n",
    "    agent: A TF-Agents agent that carries the policy to train.\n",
    "    tfrecord_file: Path to the TFRecord file containing the training dataset.\n",
    "    num_epochs: Number of epochs to train the policy.\n",
    "    Returns:\n",
    "    A NamedTuple of (a trained TF-Agents policy, a dict mapping from\n",
    "    \"epoch<i>\" to lists of loss values produced at each training step).\n",
    "    \"\"\"\n",
    "    raw_dataset = tf.data.TFRecordDataset([tfrecord_file])\n",
    "    parsed_dataset = raw_dataset.map(_parse_record)\n",
    "\n",
    "    train_loss = collections.defaultdict(list)\n",
    "    for epoch in range(num_epochs):\n",
    "        for parsed_record in parsed_dataset:\n",
    "            trajectory = build_trajectory(parsed_record, agent.policy.info_spec)\n",
    "            loss, _ = agent.train(trajectory)\n",
    "            train_loss[f\"epoch{epoch + 1}\"].append(loss.numpy())\n",
    "\n",
    "    train_outputs = collections.namedtuple(\n",
    "        \"TrainOutputs\",\n",
    "        [\"policy\", \"train_loss\"])\n",
    "    return train_outputs(agent.policy, train_loss)\n",
    "\n",
    "def execute_training_and_save_policy(\n",
    "    training_artifacts_dir: str,\n",
    "    tfrecord_file: str,\n",
    "    num_epochs: int,\n",
    "    rank_k: int,\n",
    "    num_actions: int,\n",
    "    tikhonov_weight: float,\n",
    "    agent_alpha: float) -> None:\n",
    "    \"\"\"Executes training for the policy and saves the policy.\n",
    "    Args:\n",
    "    training_artifacts_dir: Path to store the Trainer artifacts (trained\n",
    "        policy).\n",
    "    tfrecord_file: Path to file to write the ingestion result TFRecords.\n",
    "    num_epochs: Number of training epochs.\n",
    "    rank_k: Rank for matrix factorization in the movie environment; also\n",
    "        the observation dimension.\n",
    "    num_actions: Number of actions (movie items) to choose from.\n",
    "    tikhonov_weight: LinUCB Tikhonov regularization weight of the Trainer.\n",
    "    agent_alpha: LinUCB exploration parameter that multiplies the confidence\n",
    "        intervals of the Trainer.\n",
    "    \"\"\"\n",
    "    # Define time step and action specs for one batch.\n",
    "    time_step_spec = trajectories.TimeStep(\n",
    "        step_type=tensor_spec.TensorSpec(\n",
    "            shape=(), dtype=tf.int32, name=\"step_type\"),\n",
    "        reward=tensor_spec.TensorSpec(\n",
    "            shape=(), dtype=tf.float32, name=\"reward\"),\n",
    "        discount=tensor_spec.BoundedTensorSpec(\n",
    "            shape=(), dtype=tf.float32, name=\"discount\", minimum=0.,\n",
    "            maximum=1.),\n",
    "        observation=tensor_spec.TensorSpec(\n",
    "            shape=(rank_k,), dtype=tf.float32,\n",
    "            name=\"observation\"))\n",
    "\n",
    "    action_spec = tensor_spec.BoundedTensorSpec(\n",
    "        shape=(),\n",
    "        dtype=tf.int32,\n",
    "        name=\"action\",\n",
    "        minimum=0,\n",
    "        maximum=num_actions - 1)\n",
    "\n",
    "    # Define RL agent/algorithm.\n",
    "    agent = lin_ucb_agent.LinearUCBAgent(\n",
    "        time_step_spec=time_step_spec,\n",
    "        action_spec=action_spec,\n",
    "        tikhonov_weight=tikhonov_weight,\n",
    "        alpha=agent_alpha,\n",
    "        dtype=tf.float32,\n",
    "        accepts_per_arm_features=per_arm)\n",
    "    agent.initialize()\n",
    "    logging.info(\"TimeStep Spec (for each batch):\\n%s\\n\", agent.time_step_spec)\n",
    "    logging.info(\"Action Spec (for each batch):\\n%s\\n\", agent.action_spec)\n",
    "\n",
    "    # Perform off-policy training.\n",
    "    policy, _ = train_policy_on_trajectory(\n",
    "        agent=agent,\n",
    "        tfrecord_file=tfrecord_file,\n",
    "        num_epochs=num_epochs)\n",
    "\n",
    "    # Save trained policy.\n",
    "    logging.info(\"saving policy\")\n",
    "    saver = policy_saver.PolicySaver(policy)\n",
    "    saver.save(training_artifacts_dir)\n",
    "\n",
    "execute_training_and_save_policy(\n",
    "    training_artifacts_dir=training_artifacts_dir,\n",
    "    tfrecord_file=tfrecord_file,\n",
    "    num_epochs=num_epochs,\n",
    "    rank_k=rank_k,\n",
    "    num_actions=num_actions,\n",
    "    tikhonov_weight=tikhonov_weight,\n",
    "    agent_alpha=agent_alpha)\n",
    "\n",
    "outputs = collections.namedtuple(\n",
    "    \"Outputs\",\n",
    "    [\"training_artifacts_dir\"])\n",
    "\n",
    "print(outputs(training_artifacts_dir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from typing import Dict, List\n",
    "\n",
    "import fastapi\n",
    "\n",
    "from google.cloud import pubsub_v1\n",
    "\n",
    "import tensorflow as tf\n",
    "import tf_agents\n",
    "from tf_agents import policies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AIP_STORAGE_URI=\"gs://mlops-vertex-capgemini/artifacts\"\n",
    "tf.saved_model.load(AIP_STORAGE_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulation of retraining pipeline on streaming data using pubsub and cloud functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulator parameters\n",
    "SIMULATOR_PUBSUB_TOPIC = (\n",
    "    \"simulator-pubsub-topic\" \n",
    ")\n",
    "SIMULATOR_CLOUD_FUNCTION = (\n",
    "    \"simulator-cloud-function\"  \n",
    ")\n",
    "SIMULATOR_SCHEDULER_JOB = (\n",
    "    \"simulator-scheduler-job\"  \n",
    ")\n",
    "SIMULATOR_SCHEDULE = \"*/5 * * * *\"  \n",
    "SIMULATOR_SCHEDULER_MESSAGE = (\n",
    "    \"simulator-message\"  \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! gcloud pubsub topics create $SIMULATOR_PUBSUB_TOPIC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scheduler_job_args = \" \".join(\n",
    "    [\n",
    "        SIMULATOR_SCHEDULER_JOB,\n",
    "        f\"--schedule='{SIMULATOR_SCHEDULE}'\",\n",
    "        f\"--topic={SIMULATOR_PUBSUB_TOPIC}\",\n",
    "        f\"--message-body={SIMULATOR_SCHEDULER_MESSAGE}\",\n",
    "    ]\n",
    ")\n",
    "\n",
    "! echo $scheduler_job_args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! gcloud scheduler jobs create pubsub $scheduler_job_args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoints = ! gcloud beta ai endpoints list \\\n",
    "    --region=$REGION \\\n",
    "    --filter=display_name=$ENDPOINT_DISPLAY_NAME\n",
    "print(\"\\n\".join(endpoints), \"\\n\")\n",
    "\n",
    "ENDPOINT_ID = endpoints[2].split(\" \")[0]\n",
    "print(f\"ENDPOINT_ID={ENDPOINT_ID}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### More ENV variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE=8\n",
    "RANK_K=20\n",
    "NUM_ACTIONS=20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENV_VARS = \",\".join(\n",
    "    [\n",
    "        f\"PROJECT_ID={PROJECT_ID}\",\n",
    "        f\"REGION={REGION}\",\n",
    "        f\"ENDPOINT_ID={ENDPOINT_ID}\",\n",
    "        f\"RAW_DATA_PATH={RAW_DATA_PATH}\",\n",
    "        f\"BATCH_SIZE={BATCH_SIZE}\",\n",
    "        f\"RANK_K={RANK_K}\",\n",
    "        f\"NUM_ACTIONS={NUM_ACTIONS}\",\n",
    "    ]\n",
    ")\n",
    "\n",
    "! echo $ENV_VARS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use of simluator code\n",
    "\n",
    "! gcloud functions deploy $SIMULATOR_CLOUD_FUNCTION \\\n",
    "    --region=$REGION \\\n",
    "    --trigger-topic=$SIMULATOR_PUBSUB_TOPIC \\\n",
    "    --runtime=python37 \\\n",
    "    --memory=512MB \\\n",
    "    --timeout=200s \\\n",
    "    --source=infrastructure/src/main/resources/simulator \\\n",
    "    --entry-point=simulate \\\n",
    "    --stage-bucket=$BUCKET_NAME \\\n",
    "    --update-env-vars=$ENV_VARS\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import aiplatform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! gcloud pubsub topics create $LOGGER_PUBSUB_TOPIC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENV_VARS = \",\".join(\n",
    "    [\n",
    "        f\"PROJECT_ID={PROJECT_ID}\",\n",
    "        f\"RAW_DATA_PATH={RAW_DATA_PATH}\",\n",
    "        f\"BATCH_SIZE={BATCH_SIZE}\",\n",
    "        f\"RANK_K={RANK_K}\",\n",
    "        f\"NUM_ACTIONS={NUM_ACTIONS}\",\n",
    "        f\"BIGQUERY_TMP_FILE={BIGQUERY_TMP_FILE}\",\n",
    "        f\"BIGQUERY_DATASET_ID={BIGQUERY_DATASET_ID}\",\n",
    "        f\"BIGQUERY_LOCATION={BIGQUERY_LOCATION}\",\n",
    "        f\"BIGQUERY_TABLE_ID={BIGQUERY_TABLE_ID}\",\n",
    "    ]\n",
    ")\n",
    "\n",
    "! echo $ENV_VARS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! gcloud functions deploy $LOGGER_CLOUD_FUNCTION \\\n",
    "    --region=$REGION \\\n",
    "    --trigger-topic=$LOGGER_PUBSUB_TOPIC \\\n",
    "    --runtime=python37 \\\n",
    "    --memory=512MB \\\n",
    "    --timeout=200s \\\n",
    "    --source=infrastructure/src/main/resources/logger \\\n",
    "    --entry-point=log_prediction_to_bigquery \\\n",
    "    --stage-bucket=$BUCKET_NAME \\\n",
    "    --update-env-vars=$ENV_VARS\n",
    "\n",
    "TRIGGER_SCHEDULE = \"*/30 * * * *\"  # Schedule to trigger the pipeline. Eg. \"*/30 * * * *\" means every 30 mins.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dsl.pipeline(\n",
    "    pipeline_root=PIPELINE_ROOT,\n",
    "    name=f\"{PIPELINE_NAME}-retraining\")\n",
    "def pipeline(\n",
    "    # Pipeline configs\n",
    "    project_id: str,\n",
    "    training_artifacts_dir: str,\n",
    "\n",
    "    # BigQuery configs\n",
    "    bigquery_table_id: str,\n",
    "    bigquery_max_rows: int = 10000,\n",
    "\n",
    "    # TF-Agents RL configs\n",
    "    rank_k: int = 20,\n",
    "    num_actions: int = 20,\n",
    "    num_epochs: int = 5,\n",
    "    tikhonov_weight: float = 0.01,\n",
    "    agent_alpha: float = 10) -> None:\n",
    "\n",
    "  # Run the Ingester component.\n",
    "  ingest_op = create_component_from_func(\n",
    "      func=ingester_component.ingest_bigquery_dataset_into_tfrecord,\n",
    "      output_component_file=f\"ingester-{OUTPUT_COMPONENT_SPEC}\",\n",
    "      packages_to_install=[\n",
    "          \"google-cloud-bigquery==2.20.0\",\n",
    "          \"tensorflow==2.5.0\",\n",
    "      ])\n",
    "  ingest_task = ingest_op(\n",
    "      project_id=project_id,\n",
    "      bigquery_table_id=bigquery_table_id,\n",
    "      bigquery_max_rows=bigquery_max_rows,\n",
    "      tfrecord_file=TFRECORD_FILE)\n",
    "\n",
    "  # Run the Trainer component and submit custom job to Vertex AI.\n",
    "  train_op = create_component_from_func(\n",
    "      func=trainer_component.training_op,\n",
    "      output_component_file=f\"trainer-{OUTPUT_COMPONENT_SPEC}\",\n",
    "      packages_to_install=[\n",
    "          \"tensorflow==2.5.0\",\n",
    "          \"tf-agents==0.8.0\",\n",
    "      ])\n",
    "  train_task = train_op(\n",
    "      training_artifacts_dir=training_artifacts_dir,\n",
    "      tfrecord_file=ingest_task.outputs[\"tfrecord_file\"],\n",
    "      num_epochs=num_epochs,\n",
    "      rank_k=rank_k,\n",
    "      num_actions=num_actions,\n",
    "      tikhonov_weight=tikhonov_weight,\n",
    "      agent_alpha=agent_alpha)\n",
    "\n",
    "  worker_pool_specs = [\n",
    "      {\n",
    "          \"containerSpec\": {\n",
    "              \"imageUri\":train_task.container.image,\n",
    "          },\n",
    "          \"replicaCount\": TRAINING_REPLICA_COUNT,\n",
    "          \"machineSpec\": {\n",
    "              \"machineType\": TRAINING_MACHINE_TYPE,\n",
    "              \"acceleratorType\": TRAINING_ACCELERATOR_TYPE,\n",
    "              \"acceleratorCount\": TRAINING_ACCELERATOR_COUNT,\n",
    "          },\n",
    "      },\n",
    "  ]\n",
    "  train_task.custom_job_spec = {\n",
    "      \"displayName\": train_task.name,\n",
    "      \"jobSpec\": {\n",
    "          \"workerPoolSpecs\": worker_pool_specs,\n",
    "      }\n",
    "  }\n",
    "\n",
    "  # Run the Deployer components.\n",
    "  # Upload the trained policy as a model.\n",
    "  model_upload_op = gcc_aip.ModelUploadOp(\n",
    "      project=project_id,\n",
    "      display_name=TRAINED_POLICY_DISPLAY_NAME,\n",
    "      artifact_uri=training_artifacts_dir,\n",
    "      serving_container_image_uri=f\"gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest\",\n",
    "  )\n",
    "  # Model uploading has to occur after training completes.\n",
    "  model_upload_op.after(train_task)\n",
    "  # Create a Vertex AI endpoint. (This operation can occur in parallel with\n",
    "  # the Generator, Ingester, Trainer components.)\n",
    "  endpoint_create_op = gcc_aip.EndpointCreateOp(\n",
    "      project=project_id,\n",
    "      display_name=ENDPOINT_DISPLAY_NAME)\n",
    "  # Deploy the uploaded, trained policy to the created endpoint. (This operation\n",
    "  # has to occur after both model uploading and endpoint creation complete.)\n",
    "  model_deploy_op = gcc_aip.ModelDeployOp(\n",
    "      project=project_id,\n",
    "      endpoint=endpoint_create_op.outputs[\"endpoint\"],\n",
    "      model=model_upload_op.outputs[\"model\"],\n",
    "      deployed_model_display_name=TRAINED_POLICY_DISPLAY_NAME,\n",
    "      machine_type=ENDPOINT_MACHINE_TYPE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile the authored pipeline.\n",
    "compiler.Compiler().compile(pipeline_func=pipeline,                                                     \n",
    "                            package_path=PIPELINE_SPEC_PATH)\n",
    "\n",
    "# Createa Vertex AI client.\n",
    "api_client = AIPlatformClient(\n",
    "    project_id=PROJECT_ID,\n",
    "    region=REGION)\n",
    "\n",
    "# Schedule a recurring pipeline.\n",
    "response = api_client.create_schedule_from_job_spec(\n",
    "    job_spec_path=PIPELINE_SPEC_PATH,\n",
    "    schedule=TRIGGER_SCHEDULE,\n",
    "    parameter_values={\n",
    "        # Pipeline configs\n",
    "        \"project_id\": PROJECT_ID,\n",
    "        \"training_artifacts_dir\": TRAINING_ARTIFACTS_DIR,\n",
    "\n",
    "        # BigQuery config\n",
    "        \"bigquery_table_id\": BIGQUERY_TABLE_ID,\n",
    "    })\n",
    "response[\"name\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cleanup environment and delete featurestore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# feature_store_helper.cleanup_featurestore(PROJECT_ID, FEATURESTORE_ID)\n",
    "\n",
    "# # # Delete endpoint resource.\n",
    "# ! gcloud ai endpoints delete $ENDPOINT_ID --quiet --region $REGION\n",
    "\n",
    "# # # Delete Pub/Sub topics.\n",
    "# ! gcloud pubsub topics delete $SIMULATOR_PUBSUB_TOPIC --quiet\n",
    "# ! gcloud pubsub topics delete $LOGGER_PUBSUB_TOPIC --quiet\n",
    "\n",
    "# # Delete Cloud Functions.\n",
    "# ! gcloud functions delete $SIMULATOR_CLOUD_FUNCTION --quiet\n",
    "# ! gcloud functions delete $LOGGER_CLOUD_FUNCTION --quiet\n",
    "\n",
    "# # Delete Scheduler job.\n",
    "# ! gcloud scheduler jobs delete $SIMULATOR_SCHEDULER_JOB --quiet\n",
    "\n",
    "# Delete Cloud Storage objects that were created.\n",
    "# ! gsutil -m rm -r $PIPELINE_ROOT\n",
    "# ! gsutil -m rm -r $TRAINING_ARTIFACTS_DIR"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "77e81ca774b11d7c56e2cca515643be85a0de53eafc7203808e7c65d1c7cb7e5"
  },
  "kernelspec": {
   "display_name": "Python 3.8.0 64-bit ('3.8.0')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.8.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}