notebooks/official/training/hyperparameter_tuning_xgboost.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "title:generic,gcp"
      },
      "source": [
        "# Vertex AI Hyperparameter Tuning for XGBoost\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/hyperparameter_tuning_xgboost.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Ftraining%2Fhyperparameter_tuning_xgboost.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/training/hyperparameter_tuning_xgboost.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/hyperparameter_tuning_xgboost.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:mlops"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to Vertex AI hyperparameter tuning with XGBoost.\n",
        "\n",
        "Learn more about [Vertex AI hyperparameter tuning](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:mlops,stage2,get_started_vertex_training_xgboost"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to use **Vertex AI hyperparameter tuning** service for training an XGBoost model.\n",
        "\n",
        "This tutorial uses the following Vertex AI services:\n",
        "\n",
        "- **Vertex AI training**\n",
        "- **Vertex AI hyperparameter tuning** (uses **Vertex AI Vizier**)\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Train using a Python training application package.\n",
        "- Report accuracy during hyperparameter tuning.\n",
        "- Save the model artifacts to Cloud Storage using GCSFuse.\n",
        "- List the best model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:iris,lcn"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the [Iris dataset](https://www.tensorflow.org/datasets/catalog/iris) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). This dataset does not require any feature engineering. The version of the dataset in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the type of Iris flower species from a class of three species: setosa, virginica, or versicolor."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fc0ad661ebb"
      },
      "source": [
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0316df526f8"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_mlops"
      },
      "source": [
        "### Install Vertex AI SDK for Python\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b75757581291"
      },
      "outputs": [],
      "source": [
        "# install packages\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e9255e3b156f"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0c0b2427998a"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "435b8e413535"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92e68cfc3a90"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "46604f70e831"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ee0f8020b0c"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "294fe4e5a671"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5755d1a554f"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d2de92accb67"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aO4sKJfFox9R"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "306bf9ab7f7f"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5878e859ba4d"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
      },
      "source": [
        "#### Set hardware accelerators\n",
        "\n",
        "You can set hardware accelerators for training and prediction.\n",
        "\n",
        "Set the variables `TRAIN_GPU/TRAIN_NGPU` and `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa T4 GPUs allocated to each VM, you would specify:\n",
        "\n",
        "    (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
        "Learn more about [hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators).\n",
        "\n",
        "**Note**: TF releases before version 2.3 for GPU support are known to fail while loading the custom model in this tutorial. The issue is fixed in TF versions 2.3 and above. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cQUrG4Mbox9T"
      },
      "outputs": [],
      "source": [
        "TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
        "DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:training,prediction,xgboost"
      },
      "source": [
        "#### Set pre-built containers\n",
        "\n",
        "Set the pre-built Docker container image for training and prediction.\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers).\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XujRA5ueox9U"
      },
      "outputs": [],
      "source": [
        "TRAIN_VERSION = \"xgboost-cpu.1-1\"\n",
        "DEPLOY_VERSION = \"xgboost-cpu.1-1\"\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
        "    LOCATION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "machine:training"
      },
      "source": [
        "#### Set machine type\n",
        "\n",
        "Set the variable `TRAIN_COMPUTE` to configure the compute resources for VMs that you use for training. Learn more about the [machine types supported for training](https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UMPFgENkox9U"
      },
      "outputs": [],
      "source": [
        "TRAIN_COMPUTE = \"n1-standard-4\"\n",
        "print(\"Train machine type\", TRAIN_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "examine_training_package:xgboost"
      },
      "source": [
        "## Python training application package\n",
        "\n",
        "In this example, you use Vertex AI hyperparameter tuning service with a training job that executes a Python training application package.\n",
        "\n",
        "Learn more about [hyperparameter tuning in Vertex AI](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview).\n",
        "\n",
        "Take a look at how a Python package can be structured for running a custom training job in Vertex AI. The package contains the following directory structure:\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
        "- setup.cfg\n",
        "- setup.py\n",
        "- trainer\n",
        "  - \\_\\_init\\_\\_.py\n",
        "  - task.py\n",
        "\n",
        "The files *setup.cfg* and *setup.py* provide instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
        "The file *trainer/task.py* is the Python script that is executed when you run the custom training job. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "509c7082b6e4"
      },
      "source": [
        "### Create a folder structure as Python package\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f4wS4eISox9V"
      },
      "outputs": [],
      "source": [
        "# Make folder for Python training script\n",
        "! rm -rf custom\n",
        "! mkdir custom\n",
        "\n",
        "# Add package information\n",
        "! touch custom/README.md\n",
        "\n",
        "setup_cfg = \"[egg_info]\\n\\ntag_build =\\n\\ntag_date = 0\"\n",
        "! echo \"$setup_cfg\" > custom/setup.cfg\n",
        "\n",
        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'cloudml-hypertune',\\n\\n    ],\\n\\n    packages=setuptools.find_packages())\"\n",
        "! echo \"$setup_py\" > custom/setup.py\n",
        "\n",
        "pkg_info = \"Metadata-Version: 1.0\\n\\nName: Iris tabular classification\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: aferlitsch@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex\"\n",
        "! echo \"$pkg_info\" > custom/PKG-INFO\n",
        "\n",
        "# Make the training subfolder\n",
        "! mkdir custom/trainer\n",
        "! touch custom/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "taskpy_contents:iris,xgboost"
      },
      "source": [
        "### Create a training script\n",
        "\n",
        "Next, you create *task.py* script for your training job. Some noteable steps include:\n",
        "\n",
        "1. <u>Handling command-line arguments:</u>\n",
        "    - `model-dir`: The location to save the trained model. If no value is passed, the location to save the model is obtained from the environment variable `AIP_MODEL_DIR` that defaults to the staging bucket location.\n",
        "    - `dataset_data_url`: The location of the training data to download.\n",
        "    - `dataset_labels_url`: The location of the training labels to download.\n",
        "    - `boost-rounds`: Tunable hyperparameter\n",
        "1. <u>Data preprocessing (`get_data()`):</u>\n",
        "    - Download the dataset and split into training and test.\n",
        "1. <u>Training (`train_model()`):</u>\n",
        "    - Trains the model\n",
        "1. <u>Evaluation (`evaluate_model()`):</u>\n",
        "    - Evaluates the model.\n",
        "    - If hyperparameter tuning, reports the metric for accuracy.\n",
        "1. <u>Saving model artifacts:</u>\n",
        "    - Saves the model artifacts and evaluation metrics where the Cloud Storage location specified by `model-dir`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WiSnFuDoox9W"
      },
      "outputs": [],
      "source": [
        "%%writefile custom/trainer/task.py\n",
        "import datetime\n",
        "import os\n",
        "import subprocess\n",
        "import sys\n",
        "import pandas as pd\n",
        "import xgboost as xgb\n",
        "import hypertune\n",
        "import argparse\n",
        "import logging\n",
        "\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.metrics import accuracy_score\n",
        "\n",
        "parser = argparse.ArgumentParser()\n",
        "parser.add_argument('--model-dir', dest='model_dir',\n",
        "                    default=os.getenv('AIP_MODEL_DIR'), type=str, help='Model dir.')\n",
        "parser.add_argument(\"--dataset-data-url\", dest=\"dataset_data_url\",\n",
        "                    type=str, help=\"Download url for the training data.\")\n",
        "parser.add_argument(\"--dataset-labels-url\", dest=\"dataset_labels_url\",\n",
        "                    type=str, help=\"Download url for the training data labels.\")\n",
        "parser.add_argument(\"--boost-rounds\", dest=\"boost_rounds\",\n",
        "                    default=20, type=int, help=\"Number of boosted rounds\")\n",
        "args = parser.parse_args()\n",
        "\n",
        "logging.getLogger().setLevel(logging.INFO)\n",
        "\n",
        "def get_data():\n",
        "    logging.info(\"Downloading training data and labelsfrom: {}, {}\".format(args.dataset_data_url, args.dataset_labels_url))\n",
        "    # gsutil outputs everything to stderr. Hence, the need to divert it to stdout.\n",
        "    subprocess.check_call(['gsutil', 'cp', args.dataset_data_url, 'data.csv'], stderr=sys.stdout)\n",
        "    subprocess.check_call(['gsutil', 'cp', args.dataset_labels_url, 'labels.csv'], stderr=sys.stdout)\n",
        "\n",
        "\n",
        "    # Load data into pandas, then use `.values` to get NumPy arrays\n",
        "    data = pd.read_csv('data.csv').values\n",
        "    labels = pd.read_csv('labels.csv').values\n",
        "\n",
        "    # Convert one-column 2D array into 1D array for use with XGBoost\n",
        "    labels = labels.reshape((labels.size,))\n",
        "\n",
        "    train_data, test_data, train_labels, test_labels = train_test_split(data, labels, test_size=0.2, random_state=7)\n",
        "\n",
        "    # Load data into DMatrix object\n",
        "    dtrain = xgb.DMatrix(train_data, label=train_labels)\n",
        "    return dtrain, test_data, test_labels\n",
        "\n",
        "def train_model(dtrain):\n",
        "    logging.info(\"Start training ...\")\n",
        "    # Train XGBoost model\n",
        "    model = xgb.train({}, dtrain, num_boost_round=args.boost_rounds)\n",
        "    logging.info(\"Training completed\")\n",
        "    return model\n",
        "\n",
        "def evaluate_model(model, test_data, test_labels):\n",
        "    dtest = xgb.DMatrix(test_data)\n",
        "    pred = model.predict(dtest)\n",
        "    predictions = [round(value) for value in pred]\n",
        "    # evaluate predictions\n",
        "    accuracy = accuracy_score(test_labels, predictions)\n",
        "    logging.info(f\"Evaluation completed with model accuracy: {accuracy}\")\n",
        "\n",
        "    # report metric for hyperparameter tuning\n",
        "    hpt = hypertune.HyperTune()\n",
        "    hpt.report_hyperparameter_tuning_metric(\n",
        "        hyperparameter_metric_tag='accuracy',\n",
        "        metric_value=accuracy\n",
        "    )\n",
        "    return accuracy\n",
        "\n",
        "\n",
        "dtrain, test_data, test_labels = get_data()\n",
        "model = train_model(dtrain)\n",
        "accuracy = evaluate_model(model, test_data, test_labels)\n",
        "\n",
        "# GCSFuse conversion\n",
        "gs_prefix = 'gs://'\n",
        "gcsfuse_prefix = '/gcs/'\n",
        "if args.model_dir.startswith(gs_prefix):\n",
        "    args.model_dir = args.model_dir.replace(gs_prefix, gcsfuse_prefix)\n",
        "    dirpath = os.path.split(args.model_dir)[0]\n",
        "    if not os.path.isdir(dirpath):\n",
        "        os.makedirs(dirpath)\n",
        "\n",
        "# Export the classifier to a file\n",
        "gcs_model_path = os.path.join(args.model_dir, 'model.bst')\n",
        "logging.info(\"Saving model artifacts to {}\". format(gcs_model_path))\n",
        "model.save_model(gcs_model_path)\n",
        "\n",
        "logging.info(\"Saving metrics to {}/metrics.json\". format(args.model_dir))\n",
        "gcs_metrics_path = os.path.join(args.model_dir, 'metrics.json')\n",
        "with open(gcs_metrics_path, \"w\") as f:\n",
        "    f.write(f\"{'accuracy: {accuracy}'}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tarball_training_script"
      },
      "source": [
        "### Store training script on Cloud Storage bucket\n",
        "\n",
        "Compress the whole training folder as a tar ball and then store it in a Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dnmdycf6ox9X"
      },
      "outputs": [],
      "source": [
        "! rm -f custom.tar custom.tar.gz\n",
        "! tar cvf custom.tar custom\n",
        "! gzip custom.tar\n",
        "! gsutil cp custom.tar.gz $BUCKET_URI/trainer_iris.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_machine_specification"
      },
      "source": [
        "## Define machine configuration\n",
        "\n",
        "Define the machine configuration for your custom hyperparameter tuning job. This tells Vertex AI what type of machine instance to provision when the job gets started. \n",
        "\n",
        "You can specify the following parameters: \n",
        "  - `machine_type`: The type of GCP instance to provision -- e.g., n1-standard-8.\n",
        "  - `accelerator_type`: The type, if any, of hardware accelerator. In this tutorial if you previously set the variable `TRAIN_GPU`, you're using a GPU. Otherwise you're using a CPU.\n",
        "  - `accelerator_count`: Number of accelerators."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "train_custom_job_machine_specification"
      },
      "outputs": [],
      "source": [
        "if TRAIN_GPU:\n",
        "    machine_spec = {\n",
        "        \"machine_type\": TRAIN_COMPUTE,\n",
        "        \"accelerator_type\": TRAIN_GPU,\n",
        "        \"accelerator_count\": TRAIN_NGPU,\n",
        "    }\n",
        "else:\n",
        "    machine_spec = {\"machine_type\": TRAIN_COMPUTE, \"accelerator_count\": 0}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_disk_specification"
      },
      "source": [
        "## Define disk configuration (optional)\n",
        "\n",
        "Optionally, define the disk configuration for your custom hyperparameter tuning job. This tells Vertex AI what type and size of disk to provision in each machine instance for the hyperparameter tuning.\n",
        "\n",
        "You can specify the following parameters:\n",
        "  - `boot_disk_type`: Either SSD or Standard. SSD is faster, and Standard is less expensive. Defaults to SSD.\n",
        "  - `boot_disk_size_gb`: Size of disk in GB."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "train_custom_job_disk_specification"
      },
      "outputs": [],
      "source": [
        "DISK_TYPE = \"pd-ssd\"  # [ pd-ssd, pd-standard]\n",
        "DISK_SIZE = 100  # GB\n",
        "\n",
        "disk_spec = {\"boot_disk_type\": DISK_TYPE, \"boot_disk_size_gb\": DISK_SIZE}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "prepare_custom_cmdargs:iris,xgboost"
      },
      "source": [
        "## Set worker pool specs\n",
        "\n",
        "Specify the following worker pool specs for your custom training container:\n",
        "\n",
        "- `args`: The command-line arguments to pass to the executable that is set as the entry point into the container.\n",
        "  - `--model-dir`: Specifies where to store the model artifacts in the Cloud Storage bucket.\n",
        "  - `--dataset-data-url`: The location of the training data to download.\n",
        "  - `--dataset-labels-url`: The location of the training labels to download.\n",
        "  - `--boost-rounds`: Sets the tunable hyperparameter `num_boost_round` while training the XGBoost model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AoUfpBqVox9Y"
      },
      "outputs": [],
      "source": [
        "# Set path to save model\n",
        "MODEL_DIR = \"{}/aiplatform-custom-job\".format(BUCKET_URI)\n",
        "# Set the source path to the dataset\n",
        "DATASET_DIR = \"gs://cloud-samples-data/ai-platform/iris\"\n",
        "\n",
        "# Set the command-line arguments\n",
        "CMDARGS = [\n",
        "    \"--dataset-data-url=\" + DATASET_DIR + \"/iris_data.csv\",\n",
        "    \"--dataset-labels-url=\" + DATASET_DIR + \"/iris_target.csv\",\n",
        "]\n",
        "\n",
        "# Set the worker pool specs\n",
        "worker_pool_spec = [\n",
        "    {\n",
        "        \"replica_count\": 1,\n",
        "        \"machine_spec\": machine_spec,\n",
        "        \"disk_spec\": disk_spec,\n",
        "        \"python_package_spec\": {\n",
        "            \"executor_image_uri\": TRAIN_IMAGE,\n",
        "            \"package_uris\": [BUCKET_URI + \"/trainer_iris.tar.gz\"],\n",
        "            \"python_module\": \"trainer.task\",\n",
        "            \"args\": CMDARGS,\n",
        "        },\n",
        "    }\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_custom_job:mbsdk"
      },
      "source": [
        "## Create a custom training job\n",
        "\n",
        "Use the `CustomJob` class to create a custom training job with the following parameters:\n",
        "\n",
        "- `display_name`: A human readable name for the custom job.\n",
        "- `worker_pool_specs`: The specification for the corresponding VM instances.\n",
        "- `base_output_dir`: The Cloud Storage location for storing the model artifacts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_custom_job:mbsdk"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.CustomJob(\n",
        "    display_name=\"iris\",\n",
        "    worker_pool_specs=worker_pool_spec,\n",
        "    base_output_dir=MODEL_DIR,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_hpt_job:mbsdk"
      },
      "source": [
        "## Create a hyperparameter tuning job\n",
        "\n",
        "Use the `HyperparameterTuningJob` class to create a hyperparameter tuning job with the following parameters:\n",
        "\n",
        "- `display_name`: A human readable name for the custom job.\n",
        "- `custom_job`: The CustomJob object created for training. The worker pool spec from this custom job applies to the CustomJobs created in all the trials.\n",
        "- `metrics_spec`: The metrics to optimize. The dictionary key is the metric_id, which is reported by your training job, and the dictionary value is the optimization goal of the metric('minimize' or 'maximize').\n",
        "- `parameter_spec`: The parameters to optimize. The dictionary key is the metric_id, which is passed into your training job as a command line key word argument, and the dictionary value is the parameter specification of the metric.\n",
        "- `search_algorithm`: The search algorithm to use. Takes `grid`, `random` and `None` as values. Hyperparameter tuning for custom training uses [Vertex AI Vizier](https://cloud.google.com/vertex-ai/docs/vizier/overview) for training jobs. \n",
        "- `max_trial_count`: The maximum number of trials to perform."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_hpt_job:mbsdk"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
        "\n",
        "hpt_job = aiplatform.HyperparameterTuningJob(\n",
        "    display_name=\"iris\",\n",
        "    custom_job=job,\n",
        "    metric_spec={\n",
        "        \"accuracy\": \"maximize\",\n",
        "    },\n",
        "    parameter_spec={\n",
        "        \"boost-rounds\": hpt.IntegerParameterSpec(min=10, max=100, scale=\"linear\"),\n",
        "    },\n",
        "    search_algorithm=None,\n",
        "    max_trial_count=6,\n",
        "    parallel_trial_count=1,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "run_hpt_job:mbsdk"
      },
      "source": [
        "## Run the hyperparameter tuning job\n",
        "\n",
        "Use the `run()` method to execute the hyperparameter tuning job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "run_hpt_job:mbsdk"
      },
      "outputs": [],
      "source": [
        "hpt_job.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpt_trial_results:mbsdk"
      },
      "source": [
        "### Display the hyperparameter tuning job trial results\n",
        "\n",
        "Once the hyperparameter tuning job successfully finishes, you can access the results from each trial using the `trials` property."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hpt_trial_results:mbsdk"
      },
      "outputs": [],
      "source": [
        "print(hpt_job.trials)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "best_trial:mbsdk"
      },
      "source": [
        "## Fetch the best trial\n",
        "\n",
        "Identify the best trial and print the details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "best_trial:mbsdk"
      },
      "outputs": [],
      "source": [
        "# Initialize a tuple to identify the best configuration\n",
        "best = (None, None, None, 0.0)\n",
        "# Iterate through the trails and update the best configuration\n",
        "for trial in hpt_job.trials:\n",
        "    # Keep track of the best outcome\n",
        "    if float(trial.final_measurement.metrics[0].value) > best[3]:\n",
        "        try:\n",
        "            best = (\n",
        "                trial.id,\n",
        "                float(trial.parameters[0].value),\n",
        "                float(trial.parameters[1].value),\n",
        "                float(trial.final_measurement.metrics[0].value),\n",
        "            )\n",
        "        except:\n",
        "            best = (\n",
        "                trial.id,\n",
        "                float(trial.parameters[0].value),\n",
        "                None,\n",
        "                float(trial.final_measurement.metrics[0].value),\n",
        "            )\n",
        "\n",
        "# print details of the best configuration\n",
        "print(best)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "get_best_model"
      },
      "source": [
        "## List the best model\n",
        "\n",
        "The model artifacts for the best model are saved at: \n",
        "\n",
        "    MODEL_DIR/<best_trial_id>/model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "get_best_model"
      },
      "outputs": [],
      "source": [
        "# Fetch the best model\n",
        "BEST_MODEL_DIR = MODEL_DIR + \"/\" + best[0] + \"/model\"\n",
        "\n",
        "! gsutil ls {BEST_MODEL_DIR}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup:mbsdk"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
        "\n",
        "- HyperparameterTuning Job \n",
        "- Cloud Storage bucket\n",
        "- Locally generated files"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JyWy23gDox9a"
      },
      "outputs": [],
      "source": [
        "# Delete the hyperparameter tuning job\n",
        "hpt_job.delete()\n",
        "\n",
        "# Delete the Cloud Storage bucket\n",
        "delete_bucket = False  # Set True to delete the bucket\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI\n",
        "\n",
        "# Delete the locally generated files\n",
        "! rm -rf custom/\n",
        "! rm custom.tar.gz"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "hyperparameter_tuning_xgboost.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}