notebooks/official/custom/get_started_vertex_training_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": [
        "# Get started with Vertex AI Training 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/custom/get_started_vertex_training_xgboost.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Run 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%2Fcustom%2Fget_started_vertex_training_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/custom/get_started_vertex_training_xgboost.ipynb\">\n",
        "        <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
        "      Open in Vertex AI 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/custom/get_started_vertex_training_xgboost.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  \n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:mlops"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use Vertex AI Training for XGBoost models.\n",
        "\n",
        "Learn more about [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
      "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 Training for training a XGBoost custom model.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- Vertex AI Training\n",
        "- Vertex AI model resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Training using a Python package.\n",
        "- Report accuracy when hyperparameter tuning.\n",
        "- Save the model artifacts to Cloud Storage using GCSFuse.\n",
        "- Create a Vertex AI model resource."
      ]
    },
    {
      "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": "install_mlops"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a2c2cb2109a0"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ncRJ_Dfdox9L"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade google-cloud-aiplatform  --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ff555b32bab8"
      },
      "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": "f09b4dff629a"
      },
      "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": "ee775571c2b5"
      },
      "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": "c066e7925949"
      },
      "source": [
        "### Set Google Cloud project information and 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). 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": "be175254a715"
      },
      "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 $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a56633b047ee"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0e24370a7912"
      },
      "outputs": [],
      "source": [
        "import google.cloud.aiplatform as aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "setup_vars"
      },
      "source": [
        "### Set up variables\n",
        "\n",
        "Next, set up some variables used throughout the tutorial."
      ]
    },
    {
      "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 Tesla K80 GPUs allocated to each VM, you would specify:\n",
        "\n",
        "    (aiplatform.gapic..AcceleratorType.NVIDIA_TESLA_K80, 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 2.3 for GPU support is expected to fail to load the custom model in this tutorial. It's a known issue and fixed in TF 2.3. 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",
        "\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",
        "Next, set the machine type to use for training.\n",
        "\n",
        "- Set the variable `TRAIN_COMPUTE` to configure  the compute resources for the VMs used for training.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
        "*Note: The following is not supported for training:*\n",
        "\n",
        " - `standard`: 2 vCPUs\n",
        " - `highcpu`: 2, 4 and 8 vCPUs\n",
        "\n",
        "*Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs*."
      ]
    },
    {
      "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": "xgboost_intro"
      },
      "source": [
        "## Introduction to XGBoost training\n",
        "\n",
        "Once you have trained a XGBoost model, save it at a Cloud Storage location, so it can subsequently be uploaded to a Vertex AI model resource.\n",
        "The XGBoost package does not have support to save the model to a Cloud Storage location. Instead, do the following steps to save to a Cloud Storage location.\n",
        "\n",
        "1. Save the in-memory model to the local filesystem (e.g., model.bst).\n",
        "2. Use gsutil to copy the local copy to the specified Cloud Storage location.\n",
        "\n",
        "*Note*: You can do hyperparameter tuning with a XGBoost model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "examine_training_package:xgboost"
      },
      "source": [
        "### Examine the training package"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06692103c9e5"
      },
      "source": [
        "#### Package layout\n",
        "\n",
        "Before you start the training, look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\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` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when `trainer/task.py` is referred to in the worker pool specification, the directory slash is replaced with a dot and the file suffix (`.py`) suffix is dropped: (`trainer.task`).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b4b1b6caa966"
      },
      "source": [
        "#### Package Assembly"
      ]
    },
    {
      "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 the task script for the Python training package\n",
        "\n",
        "Next, you create the `task.py` script for driving the training package. Some noteable steps include:\n",
        "\n",
        "- Command-line arguments:\n",
        "    - `model-dir`: The location to save the trained model. When using Vertex AI custom training, the location is specified in the environment variable: `AIP_MODEL_DIR`,\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",
        "- Data preprocessing (`get_data()`):\n",
        "    - Download the dataset and split into training and test.\n",
        "- Training (`train_model()`):\n",
        "    - Trains the model\n",
        "- Evaluation (`evaluate_model()`):\n",
        "    - Evaluates the model.\n",
        "    - If hyperparameter tuning, reports the metric for accuracy.\n",
        "- Model artifact saving\n",
        "    - Saves the model artifacts and evaluation metrics where the Cloud Storage location specified by `model-dir`.\n",
        "    \n",
        "*Note:* The training script uses GCSFuse which mounts the Cloud Storage bucket as a network filesystem, allowing the script to perform filesystem operations (e.g., read and write) within a Cloud Storage bucket."
      ]
    },
    {
      "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",
        "import numpy as np\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 so we need to divert it to stdout.\n",
        "    subprocess.check_call(['gsutil', 'cp', args.dataset_data_url, 'data.csv'], stderr=sys.stdout)\n",
        "    # gsutil outputs everything to stderr so we need to divert it to 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",
        "    params = {\n",
        "        'objective': 'multi:softprob',\n",
        "        'num_class': 3\n",
        "    }\n",
        "    model = xgb.train(params, 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 = [np.around(value) for value in pred]\n",
        "    # evaluate predictions\n",
        "    try:\n",
        "        accuracy = accuracy_score(test_labels, predictions)\n",
        "    except:\n",
        "        accuracy = 0.0\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 your Cloud Storage bucket\n",
        "\n",
        "Next, you package the training folder into a compressed tar ball, and then store it in your 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": "create_custom_pp_training_job:mbsdk"
      },
      "source": [
        "### Create and run custom training job\n",
        "\n",
        "\n",
        "To train a custom model, you perform two steps: 1) create a custom training job, and 2) run the job.\n",
        "\n",
        "#### Create custom training job\n",
        "\n",
        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the custom training job.\n",
        "- `container_uri`: The training container image.\n",
        "\n",
        "- `python_package_gcs_uri`: The location of the Python training package as a tarball.\n",
        "- `python_module_name`: The relative path to the training script in the Python package.\n",
        "- `model_serving_container_uri`: The container image for deploying the model.\n",
        "\n",
        "*Note:* There is no requirements parameter. You specify any requirements in the `setup.py` script in your Python package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rVEMz1xqox9X"
      },
      "outputs": [],
      "source": [
        "DISPLAY_NAME = \"iris\"\n",
        "\n",
        "job = aiplatform.CustomPythonPackageTrainingJob(\n",
        "    display_name=DISPLAY_NAME,\n",
        "    python_package_gcs_uri=f\"{BUCKET_URI}/trainer_iris.tar.gz\",\n",
        "    python_module_name=\"trainer.task\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
        "    model_serving_container_image_uri=DEPLOY_IMAGE,\n",
        "    project=PROJECT_ID,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "prepare_custom_cmdargs:iris,xgboost"
      },
      "source": [
        "### Prepare your command-line arguments\n",
        "\n",
        "Now define the command-line arguments 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` : For our demonstrations, we use this command-line argument to specify where to store the model artifacts.\n",
        "      - direct: You pass the Cloud Storage location as a command line argument to your training script (set variable `DIRECT = True`), or\n",
        "      - indirect: The service passes the Cloud Storage location as the environment variable `AIP_MODEL_DIR` to your training script (set variable `DIRECT = False`). In this case, you tell the service the model artifact location in the job specification.\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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AoUfpBqVox9Y"
      },
      "outputs": [],
      "source": [
        "MODEL_DIR = \"{}/{}\".format(BUCKET_URI, \"model\")\n",
        "DATASET_DIR = \"gs://cloud-samples-data/ai-platform/iris\"\n",
        "\n",
        "ROUNDS = 20\n",
        "\n",
        "DIRECT = False\n",
        "if DIRECT:\n",
        "    CMDARGS = [\n",
        "        \"--dataset-data-url=\" + DATASET_DIR + \"/iris_data.csv\",\n",
        "        \"--dataset-labels-url=\" + DATASET_DIR + \"/iris_target.csv\",\n",
        "        \"--boost-rounds=\" + str(ROUNDS),\n",
        "        \"--model_dir=\" + MODEL_DIR,\n",
        "    ]\n",
        "else:\n",
        "    CMDARGS = [\n",
        "        \"--dataset-data-url=\" + DATASET_DIR + \"/iris_data.csv\",\n",
        "        \"--dataset-labels-url=\" + DATASET_DIR + \"/iris_target.csv\",\n",
        "        \"--boost-rounds=\" + str(ROUNDS),\n",
        "    ]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "run_custom_job:mbsdk"
      },
      "source": [
        "#### Run the custom training job\n",
        "\n",
        "Next, you run the custom job to start the training job by invoking the method `run`, with the following parameters:\n",
        "\n",
        "- `model_display_name`: The human readable name for the `Model` resource.\n",
        "- `args`: The command-line arguments to pass to the training script.\n",
        "- `replica_count`: The number of compute instances for training (replica_count = 1 is single node training).\n",
        "- `machine_type`: The machine type for the compute instances.\n",
        "- `accelerator_type`: The hardware accelerator type.\n",
        "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
        "- `base_output_dir`: The Cloud Storage location to write the model artifacts to.\n",
        "- `sync`: Whether to block until completion of the job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JCruQq1aox9Y"
      },
      "outputs": [],
      "source": [
        "if TRAIN_GPU:\n",
        "    model = job.run(\n",
        "        model_display_name=\"iris\",\n",
        "        args=CMDARGS,\n",
        "        replica_count=1,\n",
        "        machine_type=TRAIN_COMPUTE,\n",
        "        accelerator_type=TRAIN_GPU.name,\n",
        "        accelerator_count=TRAIN_NGPU,\n",
        "        base_output_dir=MODEL_DIR,\n",
        "        sync=False,\n",
        "    )\n",
        "else:\n",
        "    model = job.run(\n",
        "        model_display_name=\"iris\",\n",
        "        args=CMDARGS,\n",
        "        replica_count=1,\n",
        "        machine_type=TRAIN_COMPUTE,\n",
        "        base_output_dir=MODEL_DIR,\n",
        "        sync=False,\n",
        "    )\n",
        "\n",
        "model_path_to_deploy = MODEL_DIR"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "list_job"
      },
      "source": [
        "### List a custom training job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KBM_KLMSox9Y"
      },
      "outputs": [],
      "source": [
        "_job = job.list(filter=f\"display_name={DISPLAY_NAME}\")\n",
        "print(_job)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "custom_job_wait:mbsdk"
      },
      "source": [
        "### Wait for completion of custom training job\n",
        "\n",
        "Next, wait for the custom training job to complete. Alternatively, one can set the parameter `sync` to `True` in the `run()` method to block until the custom training job is completed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lHPMHbSyox9Z"
      },
      "outputs": [],
      "source": [
        "model.wait()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "delete_job"
      },
      "source": [
        "### Delete a custom training job\n",
        "\n",
        "After a training job is completed, you can delete the training job with the method `delete()`.  Prior to completion, a training job can be canceled with the method `cancel()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tlYg7Sp-ox9Z"
      },
      "outputs": [],
      "source": [
        "job.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup:mbsdk"
      },
      "source": [
        "## Clean up\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",
        "- Custom Job (Custom Training job is removed in the previous step)\n",
        "- Model\n",
        "- Cloud Storage Bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JyWy23gDox9a"
      },
      "outputs": [],
      "source": [
        "delete_bucket = False\n",
        "\n",
        "model.delete()\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_vertex_training_xgboost.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}