notebooks/official/pipelines/get_started_with_machine_management.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": "06cb539bac9c"
      },
      "source": [
        "# Get started with machine management for Vertex AI Pipelines\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/pipelines/get_started_with_machine_management.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/pipelines/get_started_with_machine_management.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\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/pipelines/get_started_with_machine_management.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                     \n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c16c6df3c48d"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to manage machine resources when training as a component in `Vertex AI Pipelines`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8b10d2fb975d"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to convert a self-contained custom training component into a `Vertex AI CustomJob`, whereby:\n",
        "\n",
        "    - The training job and artifacts are trackable.\n",
        "    - Set machine resources, such as machine-type, cpu/gpu, memory, disk, etc.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- `Vertex AI Pipelines`\n",
        "\n",
        "The steps performed in this tutorial include:\n",
        "\n",
        "- Create a custom component with a self-contained training job.\n",
        "- Execute pipeline using component-level settings for machine resources\n",
        "- Convert the self-contained training component into a `Vertex AI CustomJob`.\n",
        "- Execute pipeline using customjob-level settings for machine resources "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "39c8466c1f07"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset is the MNIST dataset. The dataset consists of 28x28 grayscale images of the digits 0 .. 9."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "35bee437737d"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_mlops"
      },
      "source": [
        "## Installations\n",
        "\n",
        "Install the packages required for executing this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "install_mlops"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
        "                         google-cloud-pipeline-components --quiet\n",
        "! pip3 install --upgrade kfp --quiet\n",
        "! pip3 install --upgrade tensorflow==2.7 --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-ZBOjErv5mM"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "before_you_begin"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "before_you_begin:nogpu"
      },
      "source": [
        "### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "set_project_id"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "region"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gcp_authenticate"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FvQeFm3Gv5mR"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ad1138a125ea"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ce6043da7b33"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0367eac06a10"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "21ad4dbb4a61"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c13224697bfb"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bucket:mbsdk"
      },
      "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": "bucket"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**Only 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": "create_bucket"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account"
      },
      "source": [
        "#### Service Account\n",
        "\n",
        "**If you don't know your service account**, try to get your service account using `gcloud` command by executing the second cell below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "set_service_account"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "autoset_service_account"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "if (\n",
        "    SERVICE_ACCOUNT == \"\"\n",
        "    or SERVICE_ACCOUNT is None\n",
        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
        "):\n",
        "    # Get your service account from gcloud\n",
        "    if not IS_COLAB:\n",
        "        shell_output = !gcloud auth list 2>/dev/null\n",
        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
        "\n",
        "    if IS_COLAB:\n",
        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
        "        # print(\"shell_output=\", shell_output)\n",
        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "\n",
        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account:pipelines"
      },
      "source": [
        "#### Set service account access for Vertex AI Pipelines\n",
        "\n",
        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step -- you only need to run these once per service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "set_service_account:pipelines"
      },
      "outputs": [],
      "source": [
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
        "\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "setup_vars"
      },
      "source": [
        "### Set up variables\n",
        "\n",
        "Next, set up some variables used throughout the tutorial.\n",
        "\n",
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7650fa22e03d"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "import tensorflow as tf\n",
        "from google.cloud import aiplatform\n",
        "from google_cloud_pipeline_components.v1.custom_job import \\\n",
        "    create_custom_training_job_from_component\n",
        "from kfp import compiler, dsl\n",
        "from kfp.dsl import component"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6183d4018b5b"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dcf75deec9a4"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accelerators:training,prediction,ngpu,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 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 will fail to load the custom model in this tutorial. It is 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": "accelerators:training,prediction,ngpu,mbsdk"
      },
      "outputs": [],
      "source": [
        "TRAIN_GPU, TRAIN_NGPU = (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80, 1)\n",
        "\n",
        "DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:training,prediction"
      },
      "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": "container:training,prediction"
      },
      "outputs": [],
      "source": [
        "TF = \"2.5\".replace(\".\", \"-\")\n",
        "TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
        "DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
        "\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
        "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "machine:training,prediction"
      },
      "source": [
        "#### Set machine type\n",
        "\n",
        "Next, set the machine type to use for training and prediction.\n",
        "\n",
        "- Set the variables `TRAIN_COMPUTE` and `DEPLOY_COMPUTE` to configure  the compute resources for the VMs you will use for for training and prediction.\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": "machine:training,prediction"
      },
      "outputs": [],
      "source": [
        "TRAIN_COMPUTE = \"n1-standard-4\"\n",
        "print(\"Train machine type\", TRAIN_COMPUTE)\n",
        "\n",
        "DEPLOY_COMPUTE = \"n1-standard-4\"\n",
        "print(\"Deploy machine type\", DEPLOY_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d9a7b0163bef"
      },
      "source": [
        "## Create a self-contained custom training component\n",
        "\n",
        "First, you create a component that self-contains the entire training step. This component trains a simple MNIST model using TensorFlow framework. The training is wholly self-contained in the component:\n",
        "\n",
        "    - Get and preprocess the data.\n",
        "    - Get/build the model.\n",
        "    - Train the model.\n",
        "    - Save the model.\n",
        "    \n",
        "The component takes the following parameters:\n",
        "\n",
        "- `model_dir`: The Cloud Storage location to save the trained model artifacts.\n",
        "- `epochs`: The number of epochs to train the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7492cdd9eb6d"
      },
      "outputs": [],
      "source": [
        "@component(\n",
        "    base_image=TRAIN_IMAGE,\n",
        "    packages_to_install=[\"tensorflow\"],\n",
        ")\n",
        "def self_contained_training_component(\n",
        "    model_dir: str,\n",
        "    epochs: int,\n",
        ") -> str:\n",
        "    import numpy as np\n",
        "\n",
        "    def get_data():\n",
        "        from tensorflow.keras.datasets import mnist\n",
        "\n",
        "        (x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
        "        x_train = (x_train / 255.0).astype(np.float32)\n",
        "        x_test = (x_test / 255.0).astype(np.float32)\n",
        "\n",
        "        return (x_train, y_train, x_test, y_test)\n",
        "\n",
        "    def get_model():\n",
        "        from tensorflow.keras import Sequential\n",
        "        from tensorflow.keras.layers import Dense, Flatten\n",
        "\n",
        "        model = Sequential(\n",
        "            [\n",
        "                Flatten(input_shape=(28, 28, 1)),\n",
        "                Dense(128, activation=\"relu\"),\n",
        "                Dense(256, activation=\"relu\"),\n",
        "                Dense(128, activation=\"relu\"),\n",
        "                Dense(10, activation=\"softmax\"),\n",
        "            ]\n",
        "        )\n",
        "\n",
        "        model.compile(\n",
        "            optimizer=\"Adam\", loss=\"sparse_categorical_crossentropy\", metrics=[\"acc\"]\n",
        "        )\n",
        "\n",
        "        return model\n",
        "\n",
        "    def train_model(x_train, y_train, model, epochs):\n",
        "        history = model.fit(x_train, y_train, epochs=epochs)\n",
        "        return history\n",
        "\n",
        "    (x_train, y_train, _, _) = get_data()\n",
        "    model = get_model()\n",
        "    train_model(x_train, y_train, model, epochs)\n",
        "\n",
        "    model.save(model_dir)\n",
        "    return model_dir\n",
        "\n",
        "\n",
        "compiler.Compiler().compile(self_contained_training_component, \"demo_componet.yaml\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "408dfb113b5e"
      },
      "source": [
        "## Create the self-contained-training pipeline\n",
        "\n",
        "Next, you create the pipeline for training this component, consisting of the following steps:\n",
        "\n",
        "- *Train the model*. For this component, you set the following component level resources:\n",
        "    - `cpu_limit`: The number of CPUs for the container's VM instance.\n",
        "    - `memory_limit`: The amount of memory for the container's VM instance.\n",
        "    - `node_selector_constraint` The type of GPU for the container's VM instance.\n",
        "    - `gpu_limit`: The number of GPUs for the container's VM instance.\n",
        "- *Import model artifacts into a Model Container artifact*.\n",
        "- *Upload the Container artifact into a `Vertex AI Model` resource*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a5ea18c76110"
      },
      "outputs": [],
      "source": [
        "PIPELINE_ROOT = \"{}/pipeline_root/machine_settings\".format(BUCKET_URI)\n",
        "\n",
        "CPU_LIMIT = \"8\"  # vCPUs\n",
        "MEMORY_LIMIT = \"8G\"\n",
        "\n",
        "\n",
        "@dsl.pipeline(\n",
        "    name=\"component-level-set-resources\",\n",
        "    description=\"A simple pipeline that requests component-level machine resource\",\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        ")\n",
        "def pipeline(epochs: int, model_dir: str, project: str = PROJECT_ID):\n",
        "    from google_cloud_pipeline_components.types import artifact_types\n",
        "    from google_cloud_pipeline_components.v1.model import ModelUploadOp\n",
        "    from kfp.dsl import importer_node\n",
        "\n",
        "    training_job_task = (\n",
        "        self_contained_training_component(epochs=epochs, model_dir=model_dir)\n",
        "        .set_display_name(\"self-contained-training\")\n",
        "        .set_cpu_limit(CPU_LIMIT)\n",
        "        .set_memory_limit(MEMORY_LIMIT)\n",
        "        .add_node_selector_constraint(\"NVIDIA_TESLA_K80\")\n",
        "        .set_gpu_limit(TRAIN_NGPU)\n",
        "    )\n",
        "\n",
        "    import_unmanaged_model_task = importer_node.importer(\n",
        "        artifact_uri=training_job_task.output,\n",
        "        artifact_class=artifact_types.UnmanagedContainerModel,\n",
        "        metadata={\n",
        "            \"containerSpec\": {\n",
        "                \"imageUri\": DEPLOY_IMAGE,\n",
        "            },\n",
        "        },\n",
        "    ).after(training_job_task)\n",
        "\n",
        "    _ = ModelUploadOp(\n",
        "        project=project,\n",
        "        display_name=\"mnist_model\",\n",
        "        unmanaged_container_model=import_unmanaged_model_task.outputs[\"artifact\"],\n",
        "    ).after(import_unmanaged_model_task)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4da63151e961"
      },
      "source": [
        "### Compile and execute the pipeline\n",
        "\n",
        "Next, you compile the pipeline and then execute it. The pipeline takes the following parameters, which are passed as the dictionary `parameter_values`:\n",
        "\n",
        "- `model_dir`: The Cloud Storage location to save the model artifacts.\n",
        "- `epochs`: The number of epochs to train the model.\n",
        "- `project`: Your project ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "371dab32ab03"
      },
      "outputs": [],
      "source": [
        "compiler.Compiler().compile(\n",
        "    pipeline_func=pipeline,\n",
        "    package_path=\"component_level_settings.yaml\",\n",
        ")\n",
        "\n",
        "pipeline = aiplatform.PipelineJob(\n",
        "    display_name=\"component-level-settings\",\n",
        "    template_path=\"component_level_settings.yaml\",\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        "    parameter_values={\"model_dir\": BUCKET_URI, \"epochs\": 20, \"project\": PROJECT_ID},\n",
        "    enable_caching=False,\n",
        ")\n",
        "\n",
        "pipeline.run()\n",
        "\n",
        "! rm -rf component_level_settings.yaml"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view_pipleline_results:bqml"
      },
      "source": [
        "### View the pipeline results\n",
        "\n",
        "Once the pipeline has completed, you can view the artifact outputs for each component step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "291d81126fef"
      },
      "outputs": [],
      "source": [
        "PROJECT_NUMBER = pipeline.gca_resource.name.split(\"/\")[1]\n",
        "print(PROJECT_NUMBER)\n",
        "\n",
        "\n",
        "def print_pipeline_output(job, output_task_name):\n",
        "    JOB_ID = job.name\n",
        "    print(JOB_ID)\n",
        "    for _ in range(len(job.gca_resource.job_detail.task_details)):\n",
        "        TASK_ID = job.gca_resource.job_detail.task_details[_].task_id\n",
        "        EXECUTE_OUTPUT = (\n",
        "            PIPELINE_ROOT\n",
        "            + \"/\"\n",
        "            + PROJECT_NUMBER\n",
        "            + \"/\"\n",
        "            + JOB_ID\n",
        "            + \"/\"\n",
        "            + output_task_name\n",
        "            + \"_\"\n",
        "            + str(TASK_ID)\n",
        "            + \"/executor_output.json\"\n",
        "        )\n",
        "        GCP_RESOURCES = (\n",
        "            PIPELINE_ROOT\n",
        "            + \"/\"\n",
        "            + PROJECT_NUMBER\n",
        "            + \"/\"\n",
        "            + JOB_ID\n",
        "            + \"/\"\n",
        "            + output_task_name\n",
        "            + \"_\"\n",
        "            + str(TASK_ID)\n",
        "            + \"/gcp_resources\"\n",
        "        )\n",
        "        EVAL_METRICS = (\n",
        "            PIPELINE_ROOT\n",
        "            + \"/\"\n",
        "            + PROJECT_NUMBER\n",
        "            + \"/\"\n",
        "            + JOB_ID\n",
        "            + \"/\"\n",
        "            + output_task_name\n",
        "            + \"_\"\n",
        "            + str(TASK_ID)\n",
        "            + \"/evaluation_metrics\"\n",
        "        )\n",
        "        if tf.io.gfile.exists(EXECUTE_OUTPUT):\n",
        "            ! gsutil cat $EXECUTE_OUTPUT\n",
        "            return EXECUTE_OUTPUT\n",
        "        elif tf.io.gfile.exists(GCP_RESOURCES):\n",
        "            ! gsutil cat $GCP_RESOURCES\n",
        "            return GCP_RESOURCES\n",
        "        elif tf.io.gfile.exists(EVAL_METRICS):\n",
        "            ! gsutil cat $EVAL_METRICS\n",
        "            return EVAL_METRICS\n",
        "\n",
        "    return None\n",
        "\n",
        "\n",
        "print(\"self-contained-training\")\n",
        "artifacts = print_pipeline_output(pipeline, \"self-contained-training\")\n",
        "print(\"\\n\\n\")\n",
        "print(\"importer\")\n",
        "artifacts = print_pipeline_output(pipeline, \"importer\")\n",
        "print(\"\\n\\n\")\n",
        "print(\"model-upload\")\n",
        "artifacts = print_pipeline_output(pipeline, \"model-upload\")\n",
        "output = !gsutil cat $artifacts\n",
        "output = json.loads(output[0])\n",
        "model_id = output[\"artifacts\"][\"model\"][\"artifacts\"][0][\"metadata\"][\"resourceName\"]\n",
        "print(\"\\n\")\n",
        "print(\"MODEL ID\", model_id)\n",
        "print(\"\\n\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "delete_pipeline"
      },
      "source": [
        "### Delete a pipeline job\n",
        "\n",
        "After a pipeline job is completed, you can delete the pipeline job with the method `delete()`.  Prior to completion, a pipeline job can be canceled with the method `cancel()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1UTEiNi9J39s"
      },
      "outputs": [],
      "source": [
        "pipeline.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2311daeb87ee"
      },
      "source": [
        "### Delete the model\n",
        "\n",
        "You can delete the `Model` resource generated by your pipeline with the `delete()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0806f57602fd"
      },
      "outputs": [],
      "source": [
        "model = aiplatform.Model(model_id)\n",
        "model.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2b8f4f5cde8d"
      },
      "source": [
        "## Convert self-contained training component to a `Vertex AI CustomJob`.\n",
        "\n",
        "Next, you use the utility `create_custom_training_job_from_component()` into a `Vertex AI CustomJob`. This provides the benefits of:\n",
        "\n",
        "- Adds additional ML Metadata tracking as a custom job.\n",
        "- Can set resource controls specific to the custom job.\n",
        "    - `machine_type`: The machine (VM) instance for the `CustomJob`.\n",
        "    - `accelerator_type`: The type (if any) of GPU or TPU.\n",
        "    - `accerlator_count`: The number of HW acclerators (GPU/TPU) or zero.\n",
        "    - `replica_count`: The number of VM instances for the job (Default is 1).\n",
        "    - `boot_disk_type`: Type of the boot disk (default is \"pd-ssd\"). \n",
        "    - `boot_disk_size_gb`:  Size in GB of the boot disk (default is 100GB)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8268b4ebe07d"
      },
      "outputs": [],
      "source": [
        "custom_job_op = create_custom_training_job_from_component(\n",
        "    self_contained_training_component,\n",
        "    display_name=\"test-component\",\n",
        "    machine_type=TRAIN_COMPUTE,\n",
        "    accelerator_type=TRAIN_GPU.name,\n",
        "    accelerator_count=TRAIN_NGPU,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1903cc4ad80f"
      },
      "source": [
        "### Create the CustomJob pipeline\n",
        "\n",
        "Next, you create the pipeline for training this component, consisting of the following steps:\n",
        "\n",
        "- *Train the model*. For this component, you set the following custom-job level resources:\n",
        "    - `machine_type`: The machine (VM) instance.\n",
        "    - `accelerator_type`: The type of GPU for the container's VM instance.\n",
        "    - `accelerator_count`: The number of GPUs for the container's VM instance.\n",
        "    - `replica_count`: The number of machine (VM) instances.\n",
        "- *Import model artifacts into a Model Container artifact*.\n",
        "- *Upload the Container artifact into a `Vertex AI Model` resource*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vRMAa6YEqAhq"
      },
      "outputs": [],
      "source": [
        "@dsl.pipeline(\n",
        "    name=\"customjob-set-resources\",\n",
        "    description=\"A simple pipeline that requests customjob-level machine resource\",\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        ")\n",
        "def pipeline(\n",
        "    epochs: int, model_dir: str, project: str = PROJECT_ID, region: str = REGION\n",
        "):\n",
        "    from google_cloud_pipeline_components.types import artifact_types\n",
        "    from google_cloud_pipeline_components.v1.model import ModelUploadOp\n",
        "    from kfp.dsl import importer_node\n",
        "\n",
        "    training_job_task = custom_job_op(\n",
        "        epochs=epochs, model_dir=model_dir, project=project, location=region\n",
        "    )\n",
        "\n",
        "    import_unmanaged_model_task = importer_node.importer(\n",
        "        artifact_uri=training_job_task.outputs[\"Output\"],\n",
        "        artifact_class=artifact_types.UnmanagedContainerModel,\n",
        "        metadata={\n",
        "            \"containerSpec\": {\n",
        "                \"imageUri\": DEPLOY_IMAGE,\n",
        "            },\n",
        "        },\n",
        "    ).after(training_job_task)\n",
        "\n",
        "    _ = ModelUploadOp(\n",
        "        project=project,\n",
        "        display_name=\"mnist_model\",\n",
        "        unmanaged_container_model=import_unmanaged_model_task.outputs[\"artifact\"],\n",
        "    ).after(import_unmanaged_model_task)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cf17b10aba45"
      },
      "source": [
        "### Compile and execute the pipeline\n",
        "\n",
        "Next, you compile the pipeline and then execute it. The pipeline takes the following parameters, which are passed as the dictionary `parameter_values`:\n",
        "\n",
        "- `model_dir`: The Cloud Storage location to save the model artifacts.\n",
        "- `epochs`: The number of epochs to train the model.\n",
        "- `project`: Your project ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5af9ac9acf12"
      },
      "outputs": [],
      "source": [
        "compiler.Compiler().compile(\n",
        "    pipeline_func=pipeline,\n",
        "    package_path=\"customjob_level_settings.yaml\",\n",
        ")\n",
        "\n",
        "pipeline = aiplatform.PipelineJob(\n",
        "    display_name=\"customjob-level-settings\",\n",
        "    template_path=\"customjob_level_settings.yaml\",\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        "    parameter_values={\"model_dir\": BUCKET_URI, \"epochs\": 20, \"project\": PROJECT_ID},\n",
        "    enable_caching=False,\n",
        ")\n",
        "\n",
        "pipeline.run()\n",
        "\n",
        "! rm -rf customjob_level_settings.yaml"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view_pipleline_results:bqml"
      },
      "source": [
        "### View the pipeline results\n",
        "\n",
        "Once the pipeline has completed, you can view the artifact outputs for each component step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1fcf1627d99c"
      },
      "outputs": [],
      "source": [
        "print(\"self-contained-training-component\")\n",
        "artifacts = print_pipeline_output(pipeline, \"self-contained-training-component\")\n",
        "print(\"\\n\\n\")\n",
        "print(\"importer\")\n",
        "artifacts = print_pipeline_output(pipeline, \"importer\")\n",
        "print(\"\\n\\n\")\n",
        "print(\"model-upload\")\n",
        "artifacts = print_pipeline_output(pipeline, \"model-upload\")\n",
        "output = !gsutil cat $artifacts\n",
        "output = json.loads(output[0])\n",
        "model_id = output[\"artifacts\"][\"model\"][\"artifacts\"][0][\"metadata\"][\"resourceName\"]\n",
        "print(\"\\n\")\n",
        "print(\"MODEL ID\", model_id)\n",
        "print(\"\\n\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "delete_pipeline"
      },
      "source": [
        "### Delete a pipeline job\n",
        "\n",
        "After a pipeline job is completed, you can delete the pipeline job with the method `delete()`.  Prior to completion, a pipeline job can be canceled with the method `cancel()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1UTEiNi9J39s"
      },
      "outputs": [],
      "source": [
        "pipeline.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c8199abb1687"
      },
      "source": [
        "### Delete the model\n",
        "\n",
        "You can delete the `Model` resource generated by your pipeline with the `delete()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "96056ca9f28c"
      },
      "outputs": [],
      "source": [
        "model = aiplatform.Model(model_id)\n",
        "model.delete()"
      ]
    },
    {
      "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:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dab6f335219e"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Set this to true only if you'd like to delete your bucket\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "get_started_with_machine_management.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}