04-custom-train-retrieval.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5a7ef552-9756-457e-8734-fcb8899d035f",
   "metadata": {},
   "source": [
    "# Scaling Two-Tower training with Vertex AI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e167805c-1fd5-4b32-8233-b951636d888a",
   "metadata": {},
   "source": [
    "## Load env config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "03855d9e-afe2-4edc-8b78-6b1d44b109ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PREFIX = ndr-v1\n"
     ]
    }
   ],
   "source": [
    "# naming convention for all cloud resources\n",
    "VERSION        = \"v1\"                  # TODO\n",
    "PREFIX         = f'ndr-{VERSION}'      # TODO\n",
    "\n",
    "print(f\"PREFIX = {PREFIX}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7322bf4e-0bb0-486c-b041-6882ecee7239",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "PROJECT_ID               = \"hybrid-vertex\"\n",
      "PROJECT_NUM              = \"934903580331\"\n",
      "LOCATION                 = \"us-central1\"\n",
      "\n",
      "REGION                   = \"us-central1\"\n",
      "BQ_LOCATION              = \"US\"\n",
      "VPC_NETWORK_NAME         = \"ucaip-haystack-vpc-network\"\n",
      "\n",
      "VERTEX_SA                = \"934903580331-compute@developer.gserviceaccount.com\"\n",
      "\n",
      "PREFIX                   = \"ndr-v1\"\n",
      "VERSION                  = \"v1\"\n",
      "\n",
      "APP                      = \"sp\"\n",
      "MODEL_TYPE               = \"2tower\"\n",
      "FRAMEWORK                = \"tfrs\"\n",
      "DATA_VERSION             = \"v1\"\n",
      "TRACK_HISTORY            = \"5\"\n",
      "\n",
      "BUCKET_NAME              = \"ndr-v1-hybrid-vertex-bucket\"\n",
      "BUCKET_URI               = \"gs://ndr-v1-hybrid-vertex-bucket\"\n",
      "SOURCE_BUCKET            = \"spotify-million-playlist-dataset\"\n",
      "\n",
      "DATA_GCS_PREFIX          = \"data\"\n",
      "DATA_PATH                = \"gs://ndr-v1-hybrid-vertex-bucket/data\"\n",
      "VOCAB_SUBDIR             = \"vocabs\"\n",
      "VOCAB_FILENAME           = \"vocab_dict.pkl\"\n",
      "\n",
      "CANDIDATE_PREFIX         = \"candidates\"\n",
      "TRAIN_DIR_PREFIX         = \"train\"\n",
      "VALID_DIR_PREFIX         = \"valid\"\n",
      "\n",
      "VPC_NETWORK_FULL         = \"projects/934903580331/global/networks/ucaip-haystack-vpc-network\"\n",
      "\n",
      "BQ_DATASET               = \"spotify_e2e_test\"\n",
      "BQ_TABLE_TRAIN           = \"train_flatten_last_5\"\n",
      "BQ_TABLE_VALID           = \"train_flatten_valid_last_5\"\n",
      "BQ_TABLE_CANDIDATES      = \"candidates\"\n",
      "\n",
      "REPO_SRC                 = \"src\"\n",
      "PIPELINES_SUB_DIR        = \"feature_pipes\"\n",
      "\n",
      "REPOSITORY               = \"ndr-v1-spotify\"\n",
      "IMAGE_NAME               = \"train-v1\"\n",
      "REMOTE_IMAGE_NAME        = \"us-central1-docker.pkg.dev/hybrid-vertex/ndr-v1-spotify/train-v1\"\n",
      "DOCKERNAME               = \"tfrs\"\n",
      "\n",
      "SERVING_IMAGE_URI_CPU    = \"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest\"\n",
      "SERVING_IMAGE_URI_GPU    = \"us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-11:latest\"\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# staging GCS\n",
    "GCP_PROJECTS             = !gcloud config get-value project\n",
    "PROJECT_ID               = GCP_PROJECTS[0]\n",
    "\n",
    "# GCS bucket and paths\n",
    "BUCKET_NAME              = f'{PREFIX}-{PROJECT_ID}-bucket'\n",
    "BUCKET_URI               = f'gs://{BUCKET_NAME}'\n",
    "\n",
    "config = !gsutil cat {BUCKET_URI}/config/notebook_env.py\n",
    "print(config.n)\n",
    "exec(config.n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "319b0b73-d1bc-4a90-a4aa-864d2215d646",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from datetime import datetime\n",
    "from time import time\n",
    "import pandas as pd\n",
    "import logging\n",
    "import time\n",
    "from pprint import pprint\n",
    "import pickle as pkl\n",
    "\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' \n",
    "\n",
    "logging.disable(logging.WARNING)\n",
    "\n",
    "from google.cloud import aiplatform as vertex_ai\n",
    "from google.cloud import storage\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2e4b6da8-8246-410c-8718-6a0dea4ea269",
   "metadata": {},
   "outputs": [],
   "source": [
    "storage_client = storage.Client(project=PROJECT_ID)\n",
    "\n",
    "vertex_ai.init(project=PROJECT_ID,location=REGION)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c45c7a0-c35e-4d60-9b70-e020e82402ab",
   "metadata": {},
   "source": [
    "## Prepare Train Job Specs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3325fcee-3bd1-4c1c-8ff8-5fac4be8fc16",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'vocab_dict.pkl'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "VOCAB_FILENAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "03bbe4be-68ed-4084-b4cd-eee9a311da58",
   "metadata": {},
   "outputs": [],
   "source": [
    "filehandler = open(VOCAB_FILENAME, 'rb')\n",
    "vocab_dict = pkl.load(filehandler)\n",
    "filehandler.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d9dfaf03-9b67-438d-bcd6-131a75d9d51d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pl_name_src\n",
      "track_name_pl\n",
      "artist_name_pl\n",
      "album_name_pl\n",
      "artist_genres_pl\n",
      "tracks_playlist_titles_pl\n",
      "track_name_can\n",
      "artist_name_can\n",
      "album_name_can\n",
      "artist_genres_can\n",
      "track_pl_titles_can\n"
     ]
    }
   ],
   "source": [
    "for keys in vocab_dict:\n",
    "    print(keys)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa482a72-54a8-4ef0-9fbc-de576b18443e",
   "metadata": {},
   "source": [
    "### Training Accelerators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "03410e35-7e1c-481c-aef2-55c25a90bcc4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WORKER_MACHINE_TYPE            : n1-highmem-16\n",
      "REPLICA_COUNT                  : 1\n",
      "ACCELERATOR_TYPE               : NVIDIA_TESLA_T4\n",
      "PER_MACHINE_ACCELERATOR_COUNT  : 1\n",
      "DISTRIBUTE_STRATEGY            : single\n",
      "REDUCTION_SERVER_COUNT         : 0\n",
      "REDUCTION_SERVER_MACHINE_TYPE  : n1-highcpu-16\n"
     ]
    }
   ],
   "source": [
    "from src.two_tower_jt import train_utils\n",
    "\n",
    "gpu_dict = train_utils.get_accelerator_config(\n",
    "    key='t4', \n",
    "    worker_machine_type = 'n1-highmem-16', \n",
    "    reduction_n=0\n",
    ")\n",
    "\n",
    "WORKER_MACHINE_TYPE            = gpu_dict['WORKER_MACHINE_TYPE']\n",
    "REPLICA_COUNT                  = gpu_dict['REPLICA_COUNT']\n",
    "ACCELERATOR_TYPE               = gpu_dict['ACCELERATOR_TYPE']\n",
    "PER_MACHINE_ACCELERATOR_COUNT  = gpu_dict['PER_MACHINE_ACCELERATOR_COUNT']\n",
    "DISTRIBUTE_STRATEGY            = gpu_dict['DISTRIBUTE_STRATEGY']\n",
    "REDUCTION_SERVER_COUNT         = gpu_dict['REDUCTION_SERVER_COUNT']\n",
    "REDUCTION_SERVER_MACHINE_TYPE  = gpu_dict['REDUCTION_SERVER_MACHINE_TYPE'] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0d132c8-c958-4b3c-a4df-ae56f24ecefc",
   "metadata": {},
   "source": [
    "### Vertex Experiments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18fa718a-1d12-4921-b1ca-9dae7b7cb661",
   "metadata": {},
   "source": [
    "#### create an experiemnt and experiment run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b87a2c89-08c9-4b52-8f25-bfa60a92b82e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EXPERIMENT_NAME: scale-training-v1\n",
      "RUN_NAME: run-20230926-165148\n"
     ]
    }
   ],
   "source": [
    "EXPERIMENT_PREFIX = 'scale-training'                     # custom identifier for organizing experiments\n",
    "EXPERIMENT_NAME=f'{EXPERIMENT_PREFIX}-{VERSION}'\n",
    "RUN_NAME = f'run-{time.strftime(\"%Y%m%d-%H%M%S\")}'\n",
    "\n",
    "print(f\"EXPERIMENT_NAME: {EXPERIMENT_NAME}\")\n",
    "print(f\"RUN_NAME: {RUN_NAME}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e934b4be-dd2a-48ff-84de-be9525b741ac",
   "metadata": {},
   "source": [
    "#### create Managed TensorBoard instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e6952007-ed86-4e41-bffb-59b27107dcc6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TB_RESOURCE_NAME: projects/934903580331/locations/us-central1/tensorboards/6339229891896016896\n",
      "TB display name: scale-training-v1\n"
     ]
    }
   ],
   "source": [
    "# # create new TB instance\n",
    "TENSORBOARD_DISPLAY_NAME=f\"{EXPERIMENT_NAME}\"\n",
    "\n",
    "tensorboard = vertex_ai.Tensorboard.create(\n",
    "    display_name=TENSORBOARD_DISPLAY_NAME, \n",
    "    project=PROJECT_ID, \n",
    "    location=REGION\n",
    ")\n",
    "TB_RESOURCE_NAME = tensorboard.resource_name\n",
    "\n",
    "# use existing TB instance\n",
    "# TB_RESOURCE_NAME = 'projects/934903580331/locations/us-central1/tensorboards/XXXXXX'\n",
    "\n",
    "print(f\"TB_RESOURCE_NAME: {TB_RESOURCE_NAME}\")\n",
    "print(f\"TB display name: {tensorboard.display_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5a4cf3b-0e59-4604-96fa-8b658057bc04",
   "metadata": {},
   "source": [
    "### training config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3cb06e80-175e-45df-8f33-89f85987e766",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PROJECTION_DIM : 32\n"
     ]
    }
   ],
   "source": [
    "SEED = 1234\n",
    "\n",
    "# =================================================\n",
    "# trainconfig: GPU related\n",
    "# =================================================\n",
    "TF_GPU_THREAD_COUNT   = '8'      # '1' | '4' | '8'\n",
    "\n",
    "# =================================================\n",
    "# trainconfig: data input pipeline\n",
    "# =================================================\n",
    "BLOCK_LENGTH          = 64            # 1, 8, 16, 32, 64\n",
    "NUM_DATA_SHARDS       = 4          # 2, 4, 8, 16, 32, 64\n",
    "# TRAIN_PREFETCH=3\n",
    "\n",
    "# =================================================\n",
    "# trainconfig: training hparams\n",
    "# =================================================\n",
    "NUM_EPOCHS           = 15\n",
    "LEARNING_RATE        = 0.01\n",
    "BATCH_SIZE           = 4096           # 8192, 4096, 2048, 1024, 512 \n",
    "\n",
    "# dropout\n",
    "DROPOUT_RATE         = 0.33\n",
    "\n",
    "# model size\n",
    "EMBEDDING_DIM        = 128\n",
    "PROJECTION_DIM       = int(EMBEDDING_DIM / 4) # 50  \n",
    "LAYER_SIZES          = '[512,256,128]'\n",
    "MAX_TOKENS           = 20000     # vocab\n",
    "\n",
    "print(f\"PROJECTION_DIM : {PROJECTION_DIM}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9e0a3d78-0570-49fd-a2ae-018375f82ce1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VALID_FREQUENCY : 7\n",
      "VALID_STEPS     : 20\n",
      "EPOCH_STEPS     : 2003\n",
      "EMBED_FREQUENCY : 1\n",
      "HIST_FREQUENCY  : 0\n",
      "CHECKPOINT_FREQ : 500\n",
      "UPDATE_FREQ     : 500\n"
     ]
    }
   ],
   "source": [
    "# =================================================\n",
    "# trainconfig: train & valid steps\n",
    "# =================================================\n",
    "train_sample_cnt  = 8_205_265 # 8_205_265\n",
    "valid_samples_cnt = 82_959\n",
    "\n",
    "# validation & evaluation\n",
    "VALID_FREQUENCY   = NUM_EPOCHS // 2 # 20\n",
    "VALID_STEPS       = valid_samples_cnt // BATCH_SIZE # 100\n",
    "EPOCH_STEPS       = train_sample_cnt // BATCH_SIZE\n",
    "\n",
    "# =================================================\n",
    "# trainconfig: tensorboard\n",
    "# =================================================\n",
    "EMBED_FREQUENCY      = 1\n",
    "HIST_FREQUENCY       = 0\n",
    "CHECKPOINT_FREQ      = EPOCH_STEPS // 4 # 'epoch'\n",
    "UPDATE_FREQ          = EPOCH_STEPS // 4 # 'epoch'\n",
    "\n",
    "print(f\"VALID_FREQUENCY : {VALID_FREQUENCY}\")\n",
    "print(f\"VALID_STEPS     : {VALID_STEPS}\")\n",
    "print(f\"EPOCH_STEPS     : {EPOCH_STEPS}\")\n",
    "print(f\"EMBED_FREQUENCY : {EMBED_FREQUENCY}\")\n",
    "print(f\"HIST_FREQUENCY  : {HIST_FREQUENCY}\")\n",
    "print(f\"CHECKPOINT_FREQ : {CHECKPOINT_FREQ}\")\n",
    "print(f\"UPDATE_FREQ     : {UPDATE_FREQ}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c01dd24f-3005-4fa2-ac03-65d97cb3d001",
   "metadata": {},
   "source": [
    "### data sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f39edd30-6c16-4dc6-9381-9a2cdbf14c23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# =================================================\n",
    "# trainconfig: Data sources\n",
    "# =================================================\n",
    "TRAIN_DIR_PREFIX = f'data/{DATA_VERSION}/valid' # train\n",
    "VALID_DIR_PREFIX = f'data/{DATA_VERSION}/valid' \n",
    "CANDIDATE_PREFIX = f'data/{DATA_VERSION}/candidates'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2862f6a6-11f5-44be-ac64-a3cd0e8a6ed1",
   "metadata": {},
   "source": [
    "### training args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "651cc211-4660-406d-a586-5e60ebfeb805",
   "metadata": {},
   "outputs": [],
   "source": [
    "WORKER_CMD = [\"python\", \"-m\", \"src.two_tower_jt.task\"]\n",
    "# WORKER_CMD = [\"python\", \"./task.py\"]\n",
    "# WORKER_CMD = [\"python\", \"-m\", \"task\"]\n",
    "\n",
    "WORKER_ARGS = [\n",
    "    f'--project={PROJECT_ID}',\n",
    "    f'--train_output_gcs_bucket={BUCKET_NAME}',\n",
    "    f'--train_dir={BUCKET_NAME}',\n",
    "    f'--train_dir_prefix={TRAIN_DIR_PREFIX}',\n",
    "    f'--valid_dir={BUCKET_NAME}',\n",
    "    f'--valid_dir_prefix={VALID_DIR_PREFIX}',\n",
    "    f'--candidate_file_dir={BUCKET_NAME}',\n",
    "    f'--candidate_files_prefix={CANDIDATE_PREFIX}',\n",
    "    f'--experiment_name={EXPERIMENT_NAME}',\n",
    "    f'--experiment_run={RUN_NAME}',\n",
    "    f'--num_epochs={NUM_EPOCHS}',\n",
    "    f'--batch_size={BATCH_SIZE}',\n",
    "    f'--embedding_dim={EMBEDDING_DIM}',\n",
    "    f'--projection_dim={PROJECTION_DIM}',\n",
    "    f'--layer_sizes={LAYER_SIZES}',\n",
    "    f'--learning_rate={LEARNING_RATE}',\n",
    "    f'--valid_frequency={VALID_FREQUENCY}',\n",
    "    f'--valid_steps={VALID_STEPS}',\n",
    "    f'--epoch_steps={EPOCH_STEPS}',\n",
    "    f'--distribute={DISTRIBUTE_STRATEGY}',\n",
    "    f'--model_version={VERSION}',\n",
    "    f'--pipeline_version={VERSION}',\n",
    "    f'--seed={SEED}',\n",
    "    f'--max_tokens={MAX_TOKENS}',\n",
    "    f'--tb_resource_name={TB_RESOURCE_NAME}',\n",
    "    f'--embed_frequency={EMBED_FREQUENCY}',\n",
    "    f'--update_frequency={UPDATE_FREQ}',          # TODO - turn on\n",
    "    f'--hist_frequency={HIST_FREQUENCY}', \n",
    "    f'--tf_gpu_thread_count={TF_GPU_THREAD_COUNT}',\n",
    "    f'--block_length={BLOCK_LENGTH}',\n",
    "    f'--num_data_shards={NUM_DATA_SHARDS}',\n",
    "    f'--chkpt_freq={CHECKPOINT_FREQ}',\n",
    "    f'--dropout_rate={DROPOUT_RATE}',\n",
    "    # uncomment these to pass value of True (bool)\n",
    "    # f'--cache_train',                              # caches train_dataset\n",
    "    # f'--evaluate_model',                           # runs model.eval()\n",
    "    # f'--write_embeddings',                         # writes embeddings index in train job\n",
    "    f'--profiler',                                   # runs TB profiler\n",
    "    # f'--set_jit',                                  # enables XLA\n",
    "    f'--compute_batch_metrics',\n",
    "    f'--use_cross_layer',\n",
    "    f'--use_dropout',\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "354feb4e-43c7-4238-9d4c-8a37b6278a62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'container_spec': {'args': ['--project=hybrid-vertex',\n",
      "                              '--train_output_gcs_bucket=ndr-v1-hybrid-vertex-bucket',\n",
      "                              '--train_dir=ndr-v1-hybrid-vertex-bucket',\n",
      "                              '--train_dir_prefix=data/v1/valid',\n",
      "                              '--valid_dir=ndr-v1-hybrid-vertex-bucket',\n",
      "                              '--valid_dir_prefix=data/v1/valid',\n",
      "                              '--candidate_file_dir=ndr-v1-hybrid-vertex-bucket',\n",
      "                              '--candidate_files_prefix=data/v1/candidates',\n",
      "                              '--experiment_name=scale-training-v1',\n",
      "                              '--experiment_run=run-20230926-165148',\n",
      "                              '--num_epochs=15',\n",
      "                              '--batch_size=4096',\n",
      "                              '--embedding_dim=128',\n",
      "                              '--projection_dim=32',\n",
      "                              '--layer_sizes=[512,256,128]',\n",
      "                              '--learning_rate=0.01',\n",
      "                              '--valid_frequency=7',\n",
      "                              '--valid_steps=20',\n",
      "                              '--epoch_steps=2003',\n",
      "                              '--distribute=single',\n",
      "                              '--model_version=v1',\n",
      "                              '--pipeline_version=v1',\n",
      "                              '--seed=1234',\n",
      "                              '--max_tokens=20000',\n",
      "                              '--tb_resource_name=projects/934903580331/locations/us-central1/tensorboards/6339229891896016896',\n",
      "                              '--embed_frequency=1',\n",
      "                              '--update_frequency=500',\n",
      "                              '--hist_frequency=0',\n",
      "                              '--tf_gpu_thread_count=8',\n",
      "                              '--block_length=64',\n",
      "                              '--num_data_shards=4',\n",
      "                              '--chkpt_freq=500',\n",
      "                              '--dropout_rate=0.33',\n",
      "                              '--profiler',\n",
      "                              '--compute_batch_metrics',\n",
      "                              '--use_cross_layer',\n",
      "                              '--use_dropout'],\n",
      "                     'command': ['python', '-m', 'src.two_tower_jt.task'],\n",
      "                     'image_uri': 'us-central1-docker.pkg.dev/hybrid-vertex/ndr-v1-spotify/train-v1:latest'},\n",
      "  'machine_spec': {'accelerator_count': 1,\n",
      "                   'accelerator_type': 'NVIDIA_TESLA_T4',\n",
      "                   'machine_type': 'n1-highmem-16'},\n",
      "  'replica_count': 1}]\n"
     ]
    }
   ],
   "source": [
    "from util import workerpool_specs\n",
    "\n",
    "WORKER_POOL_SPECS = workerpool_specs.prepare_worker_pool_specs(\n",
    "    image_uri=f\"{REMOTE_IMAGE_NAME}:latest\",\n",
    "    args=WORKER_ARGS,\n",
    "    cmd=WORKER_CMD,\n",
    "    replica_count=REPLICA_COUNT,\n",
    "    machine_type=WORKER_MACHINE_TYPE,\n",
    "    accelerator_count=PER_MACHINE_ACCELERATOR_COUNT,\n",
    "    accelerator_type=ACCELERATOR_TYPE,\n",
    "    reduction_server_count=REDUCTION_SERVER_COUNT,\n",
    "    reduction_server_machine_type=REDUCTION_SERVER_MACHINE_TYPE,\n",
    ")\n",
    "\n",
    "from pprint import pprint\n",
    "pprint(WORKER_POOL_SPECS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d79957d8-b9f4-4c64-86df-908ad5a03bbd",
   "metadata": {},
   "source": [
    "### copy training package to GCS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5f034f37-89cc-41aa-ba3d-f648f924db07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " Copied training package and Dockerfile to gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230926-165148\n",
      "\n"
     ]
    }
   ],
   "source": [
    "BASE_OUTPUT_DIR = f'gs://{BUCKET_NAME}/{EXPERIMENT_NAME}/{RUN_NAME}'\n",
    "\n",
    "# copy training Dockerfile\n",
    "!gsutil -q cp $REPO_SRC/cloudbuild.yaml $BASE_OUTPUT_DIR/cloudbuild.yaml\n",
    "!gsutil -q cp $REPO_SRC/Dockerfile_tfrs $BASE_OUTPUT_DIR/Dockerfile_tfrs\n",
    "!gsutil -q cp vocab_dict.pkl $BASE_OUTPUT_DIR/vocab_dict.pkl\n",
    "\n",
    "# # # copy training application code\n",
    "! gsutil -q -m cp -r $REPO_SRC/two_tower_jt/* $BASE_OUTPUT_DIR/trainer\n",
    "\n",
    "print(f\"\\n Copied training package and Dockerfile to {BASE_OUTPUT_DIR}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "4026b578-8b8b-4d90-ac9c-c8ec5c1cfb81",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230926-165148/Dockerfile_tfrs\n",
      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230926-165148/cloudbuild.yaml\n",
      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230926-165148/vocab_dict.pkl\n",
      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230926-165148/trainer/\n"
     ]
    }
   ],
   "source": [
    "! gsutil ls $BASE_OUTPUT_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "556feb75-b1df-4a8d-a852-9e9be89f15fa",
   "metadata": {},
   "source": [
    "## submit training job to Vertex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "93b8a038-15b4-4a11-a458-e854f246d0e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JOB_NAME: train-v1-run-20230926-165148\n"
     ]
    }
   ],
   "source": [
    "vertex_ai.init(\n",
    "    project=PROJECT_ID,\n",
    "    location=REGION,\n",
    "    experiment=EXPERIMENT_NAME,\n",
    ")\n",
    "\n",
    "JOB_NAME = f'train-{VERSION}-{RUN_NAME}'\n",
    "print(f\"JOB_NAME: {JOB_NAME}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "36104f63-a11c-44f5-8f4c-6ff87b6784e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "job = vertex_ai.CustomJob(\n",
    "    display_name=JOB_NAME,\n",
    "    worker_pool_specs=WORKER_POOL_SPECS,\n",
    "    base_output_dir=BASE_OUTPUT_DIR,\n",
    "    staging_bucket=f\"{BASE_OUTPUT_DIR}/staging\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "1132d226-6862-4458-87a1-972c8b6deeb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "job.run(\n",
    "    tensorboard=TB_RESOURCE_NAME,\n",
    "    service_account=VERTEX_SA,\n",
    "    restart_job_on_worker_restart=False,\n",
    "    enable_web_access=True,\n",
    "    sync=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56675ff9-158e-486b-ad94-76a9c0fc0b0a",
   "metadata": {},
   "source": [
    "## TensorBoard Profiler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5fd5533-d364-45a6-b1de-2db034d6df6f",
   "metadata": {},
   "source": [
    "Once the profiler has uploaded trace logs to `BASE_OUTPUT_DIR/logs`, we can use the in-notebook tensoborad extension to view the profiler\n",
    "\n",
    "<img\n",
    "  src=\"img/tfrs-train-profiler-v1.png\"\n",
    "  alt=\"Alt text\"\n",
    "  title=\"train profiler\"\n",
    "  style=\"display: inline-block; margin: 0 auto; max-width: 1200px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "33a58184-e0c4-4403-ad80-da5fdb7dbaed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TB_LOGS_PATH: gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230925-203116/logs\n"
     ]
    }
   ],
   "source": [
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "TB_LOGS_PATH = f'{BASE_OUTPUT_DIR}/logs'\n",
    "\n",
    "print(f\"TB_LOGS_PATH: {TB_LOGS_PATH}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "d432d552-4416-4f4c-964f-bcca57b60a07",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load_ext tensorboard\n",
    "# %reload_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "171227db-d5f9-47e4-8b67-a528950233ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %tensorboard --logdir=$TB_LOGS_PATH"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8341545a-cadf-448a-b746-c5f74e33e58c",
   "metadata": {},
   "source": [
    "# Notes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "48e2fe35-7abe-4d76-9189-c58c4b5c1026",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m\n",
       "\u001b[0mjob\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mservice_account\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mnetwork\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mtimeout\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mrestart_job_on_worker_restart\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0menable_web_access\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mexperiment\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mForwardRef\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'aiplatform.Experiment'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mexperiment_run\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mForwardRef\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'aiplatform.ExperimentRun'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mtensorboard\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0msync\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mcreate_request_timeout\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "Run this configured CustomJob.\n",
       "\n",
       "Args:\n",
       "    service_account (str):\n",
       "        Optional. Specifies the service account for workload run-as account.\n",
       "        Users submitting jobs must have act-as permission on this run-as account.\n",
       "    network (str):\n",
       "        Optional. The full name of the Compute Engine network to which the job\n",
       "        should be peered. For example, projects/12345/global/networks/myVPC.\n",
       "        Private services access must already be configured for the network.\n",
       "        If left unspecified, the network set in aiplatform.init will be used.\n",
       "        Otherwise, the job is not peered with any network.\n",
       "    timeout (int):\n",
       "        The maximum job running time in seconds. The default is 7 days.\n",
       "    restart_job_on_worker_restart (bool):\n",
       "        Restarts the entire CustomJob if a worker\n",
       "        gets restarted. This feature can be used by\n",
       "        distributed training jobs that are not resilient\n",
       "        to workers leaving and joining a job.\n",
       "    enable_web_access (bool):\n",
       "        Whether you want Vertex AI to enable interactive shell access\n",
       "        to training containers.\n",
       "        https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell\n",
       "    experiment (Union[aiplatform.Experiment, str]):\n",
       "        Optional. The instance or name of an Experiment resource to which\n",
       "        this CustomJob will upload training parameters and metrics.\n",
       "\n",
       "        `service_account` is required with provided `experiment`.\n",
       "        For more information on configuring your service account please visit:\n",
       "        https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training\n",
       "    experiment_run (Union[aiplatform.ExperimentRun, str]):\n",
       "        Optional. The instance or name of an ExperimentRun resource to which\n",
       "        this CustomJob will upload training parameters and metrics.\n",
       "        This arg can only be set when `experiment` is set. If 'experiment'\n",
       "        is set but 'experiment_run` is not, an ExperimentRun resource\n",
       "        will still be auto-generated.\n",
       "    tensorboard (str):\n",
       "        Optional. The name of a Vertex AI\n",
       "        [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]\n",
       "        resource to which this CustomJob will upload Tensorboard\n",
       "        logs. Format:\n",
       "        ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``\n",
       "\n",
       "        The training script should write Tensorboard to following Vertex AI environment\n",
       "        variable:\n",
       "\n",
       "        AIP_TENSORBOARD_LOG_DIR\n",
       "\n",
       "        `service_account` is required with provided `tensorboard`.\n",
       "        For more information on configuring your service account please visit:\n",
       "        https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training\n",
       "    sync (bool):\n",
       "        Whether to execute this method synchronously. If False, this method\n",
       "        will unblock and it will be executed in a concurrent Future.\n",
       "    create_request_timeout (float):\n",
       "        Optional. The timeout for the create request in seconds.\n",
       "\u001b[0;31mFile:\u001b[0m      ~/.local/lib/python3.7/site-packages/google/cloud/aiplatform/jobs.py\n",
       "\u001b[0;31mType:\u001b[0m      method\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "?job.run"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "152126ee-261b-4b21-b594-8aeebeb667f0",
   "metadata": {},
   "source": [
    "**Finished**"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "tf2-gpu.2-10.m103",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-10:m103"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}