notebooks/community/model_garden/model_garden_gemma_fine_tuning_batch_deployment_on_rov.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Get started with Gemma on Ray on Vertex AI\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/community/model_garden/model_garden_gemma_fine_tuning_batch_deployment_on_rov.ipynb\">\n",
        "      <img width=\"32px\" 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://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_gemma_fine_tuning_batch_deployment_on_rov.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/blob/main/notebooks/community/model_garden/model_garden_gemma_fine_tuning_batch_deployment_on_rov.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> 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/community/model_garden/model_garden_gemma_fine_tuning_batch_deployment_on_rov.ipynb\">\n",
        "      <img width=\"32px\" 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": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use Ray on Vertex AI for fine-tuning and serving Gemma on Vertex AI.\n",
        "\n",
        "Learn more about [Ray on Vertex AI](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/overview)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d975e698c9a4"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you'll learn how to distribute Gemma Supervised tuning on Ray on Vertex AI. Furthermore, you'll learn how to deploy the trained model seamlessly for offline predictions using Ray Data on Ray on Vertex AI.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Ray on Vertex AI\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create a Ray cluster on Vertex AI\n",
        "- Tune Gemma with Ray Train on Ray on Vertex AI\n",
        "- Serving Gemma with Ray Data for offline predictions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08d289fa873f"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The [Extreme Summarization (XSum) dataset](https://huggingface.co/datasets/EdinburghNLP/xsum) is a dataset about abstractive single-document summarization systems."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aed92deeb4a0"
      },
      "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),\n",
        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing),\n",
        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cUKGInbrh0Cw"
      },
      "source": [
        "<b>Note</b>: This tutorial uses the Ray Jobs API via public Ray Dashboard. The Ray dashboard address is accessible from outside the VPC, including the public internet. To learn more about  private versus public connectivity, see the [Private and public connectivity](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/create-cluster#private_and_public_connectivity) section in the [Create a Ray cluster on Vertex AI](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/create-cluster) documentation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager).\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,artifactregistry.googleapis.com,cloudbuild.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Installation\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "# Install the packages\n",
        "import os\n",
        "\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    USER = \"--user\"\n",
        "else:\n",
        "    USER = \"\"\n",
        "\n",
        "! pip3 install {USER} google-cloud-aiplatform[ray]==1.48.0 -q --no-warn-conflicts\n",
        "! pip3 install {USER} google-cloud-aiplatform[tensorboard]==1.48.0 -q --no-warn-conflicts\n",
        "! pip3 install {USER} torch==2.2.1 datasets==2.17.0 transformers==4.38.1 evaluate==0.4.1 rouge-score==0.1.2 nltk==3.8.1 bitsandbytes==0.42.0 peft==0.8.2 accelerate==0.27.1 -q --no-warn-conflicts\n",
        "! pip3 install {USER} tensorflow==2.15.0 -q --no-warn-conflicts\n",
        "! pip3 install {USER} etils==1.5.0 fsspec==2023.10.0 gcsfs==2023.10.0 -q --no-warn-conflicts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "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": "XRvKdaPDTznN"
      },
      "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": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CylYbUxrx3W-"
      },
      "source": [
        "### Set Google Cloud project information\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Project ID\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "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": "fPrDj6HE9_EU"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "timestamp"
      },
      "source": [
        "#### Timestamp\n",
        "\n",
        "You create a timestamp to make resources you create unique in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W6Le1schAziq"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "#### Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = f\"your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}\n",
        "\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account"
      },
      "source": [
        "#### Service Account\n",
        "\n",
        "Set service account and grant the service account access to Vertex AI TensorBoard."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a_C_BMVpzhug"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lR2GDIkzp4af"
      },
      "outputs": [],
      "source": [
        "! gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "   --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "   --role=\"roles/storage.admin\"\n",
        "\n",
        "! gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "   --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "   --role=\"roles/aiplatform.user\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ek1-iTbPjzdJ"
      },
      "source": [
        "### Set tutorial folder\n",
        "\n",
        "Set up the folder to use in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BbfKRabXj3la"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path as path\n",
        "\n",
        "root_path = path.cwd()\n",
        "tutorial_path = root_path / \"tutorial\"\n",
        "data_path = tutorial_path / \"data\"\n",
        "src_path = tutorial_path / \"src\"\n",
        "experiments_path = tutorial_path / \"experiments\"\n",
        "models_path = tutorial_path / \"models\"\n",
        "build_path = tutorial_path / \"build\"\n",
        "tests_path = tutorial_path / \"tests\"\n",
        "\n",
        "data_path.mkdir(parents=True, exist_ok=True)\n",
        "src_path.mkdir(parents=True, exist_ok=True)\n",
        "experiments_path.mkdir(parents=True, exist_ok=True)\n",
        "models_path.mkdir(parents=True, exist_ok=True)\n",
        "build_path.mkdir(parents=True, exist_ok=True)\n",
        "tests_path.mkdir(parents=True, exist_ok=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k9ryiScCEapt"
      },
      "source": [
        "### Set a Ray cluster on Vertex AI\n",
        "\n",
        "Before running the code below, make sure to [set up](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/set-up) Ray on Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_mlq_1NLEonF"
      },
      "outputs": [],
      "source": [
        "import vertex_ray\n",
        "from google.cloud import aiplatform as vertex_ai\n",
        "from vertex_ray import NodeImages, Resources"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dorIZFvjnGKL"
      },
      "source": [
        "#### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VOOgvRJoQ6Xj"
      },
      "outputs": [],
      "source": [
        "vertex_ai.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15PYjIyOcx1d"
      },
      "source": [
        "#### Build the custom cluster image\n",
        "\n",
        " It's necessary to utilize Ray Custom cluster image support since certain dependencies are required.\n",
        "\n",
        " To use a custom cluster image, the first step is to build the image. Below there are the steps to cover:\n",
        "\n",
        "*  Prepare the requirements file\n",
        "*  Create the Dockerfile for the custom image\n",
        "*  Create the Docker image repository\n",
        "*  Build the Ray cluster custom image\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F5vW6XpnkFeR"
      },
      "source": [
        "##### Prepare the requirements file\n",
        "\n",
        "Prepare the `requirements` file that includes the dependencies your Ray application needs to run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "83KUQQylbrJR"
      },
      "outputs": [],
      "source": [
        "requirements = \"\"\"\n",
        "ipython==8.22.2\n",
        "torch==2.2.1\n",
        "ray==2.10.0\n",
        "ray[data]==2.10.0\n",
        "ray[train]==2.10.0\n",
        "ray[tune]==2.10.0\n",
        "datasets==2.17.0\n",
        "transformers==4.38.1\n",
        "evaluate==0.4.1\n",
        "rouge-score==0.1.2\n",
        "nltk==3.8.1\n",
        "accelerate==0.27.1\n",
        "bitsandbytes==0.42.0\n",
        "peft==0.8.2\n",
        "trl==0.7.10\n",
        "# flash-attn==2.5.5\n",
        "pyarrow==15.0.2\n",
        "fsspec==2023.10.0\n",
        "gcsfs==2023.10.0\n",
        "etils==1.7.0\n",
        "importlib-resources==6.1.2\n",
        "\"\"\"\n",
        "\n",
        "with open(build_path / \"requirements.txt\", \"w\") as rfile:\n",
        "    rfile.write(requirements)\n",
        "rfile.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5VHxHjCyKCyi"
      },
      "source": [
        "##### Create the Dockerfile\n",
        "\n",
        "Create the Dockerfile for the custom image by leveraging one of the prebuilt Ray on Vertex AI base images.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iu7bgcdeXZIS"
      },
      "outputs": [],
      "source": [
        "CUSTOM_BASE_IMAGE = \"us-docker.pkg.dev/vertex-ai/training/ray-gpu.2-9.py310:latest\"  # @param [\"us-docker.pkg.dev/vertex-ai/training/ray-cpu.2-4.py310:latest\", \"us-docker.pkg.dev/vertex-ai/training/ray-cpu.2-9.py310:latest\", \"us-docker.pkg.dev/vertex-ai/training/ray-gpu.2-4.py310:latest\", \"us-docker.pkg.dev/vertex-ai/training/ray-gpu.2-9.py310:latest\", \"europe-docker.pkg.dev/vertex-ai/training/ray-cpu.2-4.py310:latest\", \"europe-docker.pkg.dev/vertex-ai/training/ray-cpu.2-9.py310:latest\", \"europe-docker.pkg.dev/vertex-ai/training/ray-gpu.2-4.py310:latest\", \"europe-docker.pkg.dev/vertex-ai/training/ray-gpu.2-9.py310:latest\", \"asia-docker.pkg.dev/vertex-ai/training/ray-cpu.2-4.py310:latest\", \"asia-docker.pkg.dev/vertex-ai/training/ray-cpu.2-9.py310:latest\", \"asia-docker.pkg.dev/vertex-ai/training/ray-gpu.2-4.py310:latest\", \"asia-docker.pkg.dev/vertex-ai/training/ray-gpu.2-9.py310:latest\"] {allow-input: true}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sRSzXRpxKB_Q"
      },
      "outputs": [],
      "source": [
        "dockerfile = f\"\"\"\n",
        "FROM {CUSTOM_BASE_IMAGE}\n",
        "\n",
        "# Install training libraries.\n",
        "ENV PIP_ROOT_USER_ACTION=ignore\n",
        "COPY requirements.txt .\n",
        "RUN pip install -r requirements.txt\n",
        "\"\"\"\n",
        "\n",
        "with open(build_path / \"Dockerfile\", \"w\") as image_file:\n",
        "    image_file.write(dockerfile)\n",
        "image_file.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JTU6Ai5_RT1f"
      },
      "source": [
        "##### Create the Docker image repository\n",
        "\n",
        "To store the custom cluster image, create a Docker repository in the Artifact Registry."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DUoNxKGIRiGb"
      },
      "outputs": [],
      "source": [
        "REPO_NAME = f\"your-repo-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zh0iT2qVKBDi"
      },
      "outputs": [],
      "source": [
        "! gcloud artifacts repositories create {REPO_NAME} --repository-format=docker \\\n",
        "    --location={REGION} --description=\"Tutorial repository\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o9lYOI15SFgh"
      },
      "source": [
        "##### Build the Ray cluster custom image\n",
        "\n",
        "Finally, build the Ray cluster custom image using Cloud Build."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FR6BP7cgRLuD"
      },
      "outputs": [],
      "source": [
        "NODE_TRAIN_IMAGE = f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPO_NAME}/train\"\n",
        "BUILD_MACHINE_TYPE = \"E2_HIGHCPU_32\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9VQKg4rMTLTa"
      },
      "outputs": [],
      "source": [
        "! gcloud builds submit --region={REGION} --tag={NODE_TRAIN_IMAGE} \\\n",
        "    --machine-type={BUILD_MACHINE_TYPE} --timeout=3600 {build_path}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_O1xUMt7Z6r0"
      },
      "source": [
        "#### Create the Ray cluster\n",
        "\n",
        "With the custom image, create the Ray cluster using the custom image via Ray on Vertex AI SDK for Python."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eZkHOH3v2i1p"
      },
      "outputs": [],
      "source": [
        "CLUSTER_NAME = f\"your-cluster-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dd2TmEdyVLo1"
      },
      "source": [
        "###### Set the Ray cluster configuration\n",
        "\n",
        "Use the Vertex AI Python SDK for Ray on Vertex AI to set the cluster configuration.\n",
        "\n",
        "To know more about the cluster configuration, see the [documentation](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/create-cluster#ray-on-vertex-ai-sdk)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m0BBynZrFn7z"
      },
      "outputs": [],
      "source": [
        "HEAD_NODE_MACHINE_TYPE = \"n1-standard-16\"  # @param {type:\"string\"}\n",
        "HEAD_NODE_COUNT = 1  # @param {type:\"integer\"}\n",
        "\n",
        "WORKER_NODE_MACHINE_TYPE = \"a2-highgpu-1g\"  # @param {type:\"string\"}\n",
        "WORKER_NODE_COUNT = 1  # @param {type:\"integer\"}\n",
        "WORKER_ACCELERATION_TYPE = \"NVIDIA_TESLA_A100\"  # @param {type:\"string\"}\n",
        "WORKER_ACCELERATION_COUNT = 1  # @param {type:\"integer\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ppNbbhgm3_GE"
      },
      "outputs": [],
      "source": [
        "HEAD_NODE_TYPE = Resources(\n",
        "    machine_type=HEAD_NODE_MACHINE_TYPE,\n",
        "    node_count=HEAD_NODE_COUNT,\n",
        ")\n",
        "\n",
        "WORKER_NODE_TYPES = [\n",
        "    Resources(\n",
        "        machine_type=WORKER_NODE_MACHINE_TYPE,\n",
        "        node_count=WORKER_NODE_COUNT,\n",
        "        accelerator_type=WORKER_ACCELERATION_TYPE,\n",
        "        accelerator_count=WORKER_ACCELERATION_COUNT,\n",
        "    )\n",
        "]\n",
        "\n",
        "CUSTOM_IMAGES = NodeImages(\n",
        "    head=NODE_TRAIN_IMAGE,\n",
        "    worker=NODE_TRAIN_IMAGE,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BfkpGqRjksaW"
      },
      "source": [
        "##### Create the Ray cluster\n",
        "\n",
        "Create the Ray cluster with the predefined custom configuration. Creating a cluster can take several minutes, depending on its configuration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J-g6kLwqUj5n"
      },
      "outputs": [],
      "source": [
        "ray_cluster_name = vertex_ray.create_ray_cluster(\n",
        "    head_node_type=HEAD_NODE_TYPE,\n",
        "    worker_node_types=WORKER_NODE_TYPES,\n",
        "    custom_images=CUSTOM_IMAGES,\n",
        "    cluster_name=CLUSTER_NAME,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xmBlsbHAc2uO"
      },
      "source": [
        "##### Get the Ray cluster\n",
        "\n",
        "Use the Ray on Vertex AI SDK for Python to get the Ray cluster."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9UzG2WyXbZJi"
      },
      "outputs": [],
      "source": [
        "ray_clusters = vertex_ray.list_ray_clusters()\n",
        "ray_cluster_resource_name = ray_clusters[-1].cluster_resource_name\n",
        "ray_cluster = vertex_ray.get_ray_cluster(ray_cluster_resource_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g7ZKdv5-GCWr"
      },
      "outputs": [],
      "source": [
        "print(\"Ray cluster on Vertex AI:\", ray_cluster_resource_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "Import required libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "# General\n",
        "import io\n",
        "import logging\n",
        "import os\n",
        "import random\n",
        "import shutil\n",
        "import string\n",
        "import time\n",
        "\n",
        "import datasets\n",
        "import evaluate\n",
        "import pandas as pd\n",
        "# Ray - Training\n",
        "import ray\n",
        "import torch\n",
        "import transformers\n",
        "from etils import epath\n",
        "from google.cloud import storage\n",
        "from huggingface_hub import login\n",
        "from peft import PeftModel\n",
        "from ray.job_submission import JobStatus, JobSubmissionClient\n",
        "# Ray - Batch Serving\n",
        "from ray.tune import ExperimentAnalysis\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h9_UttTcNGYN"
      },
      "outputs": [],
      "source": [
        "print(\"Ray version: \", ray.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xFgb-sZbBi8i"
      },
      "source": [
        "### Set variables\n",
        "\n",
        "Initiate some tutorial variables."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zykxFjqUB9jt"
      },
      "outputs": [],
      "source": [
        "# Training\n",
        "HF_TOKEN = \"[your-hugging-face-token]\"  # @param {type:\"string\"}\n",
        "EXPERIMENTS_FOLDER_URI = epath.Path(BUCKET_URI) / \"experiments\"\n",
        "TENSORBOARD_NAME = f\"rov-xsum-gemma-tb-{TIMESTAMP}\"\n",
        "\n",
        "# Serving\n",
        "MODELS_PATH = epath.Path(BUCKET_URI) / \"models\"\n",
        "PREDICTIONS_FOLDER_URI = epath.Path(BUCKET_URI) / \"predictions\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mYQNboaOBk6B"
      },
      "source": [
        "### Define helpers\n",
        "\n",
        "Define an helper function to monitor the status of Ray job using Ray Dashboard API in your notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MrvG6VDIm9hG"
      },
      "outputs": [],
      "source": [
        "def monitor_job(client, job_id):\n",
        "    \"\"\"Monitors the status of Ray job using Ray Dashboard API\"\"\"\n",
        "\n",
        "    logging.basicConfig(\n",
        "        level=logging.INFO,\n",
        "        format=f\"%(asctime)s.%(msecs)03d %(levelname)s {job_id} -- %(message)s\",\n",
        "        datefmt=\"%Y-%m-%d %H:%M:%S\",\n",
        "        force=True,\n",
        "    )\n",
        "\n",
        "    while True:\n",
        "        job_status = client.get_job_status(job_id)\n",
        "\n",
        "        if job_status == JobStatus.SUCCEEDED:\n",
        "            logging.info(\"Job succeeded!\")\n",
        "            break\n",
        "\n",
        "        elif job_status == JobStatus.FAILED:\n",
        "            logging.info(\"Job failed!\")\n",
        "            break\n",
        "\n",
        "        else:\n",
        "            logging.info(\"Job is running...\")\n",
        "            time.sleep(60)\n",
        "\n",
        "    return job_status\n",
        "\n",
        "\n",
        "def read_json_files(bucket_name, prefix=None):\n",
        "    \"\"\"Reads JSON files from a cloud storage bucket and returns a Pandas DataFrame\"\"\"\n",
        "\n",
        "    # Set up storage client\n",
        "    storage_client = storage.Client()\n",
        "    bucket = storage_client.bucket(bucket_name)\n",
        "    blobs = bucket.list_blobs(prefix=prefix)\n",
        "\n",
        "    dfs = []\n",
        "\n",
        "    for blob in blobs:\n",
        "        if blob.name.endswith(\".json\"):\n",
        "            file_bytes = blob.download_as_bytes()\n",
        "            file_string = file_bytes.decode(\"utf-8\")\n",
        "            with io.StringIO(file_string) as json_file:\n",
        "                df = pd.read_json(json_file, lines=True)\n",
        "            dfs.append(df)\n",
        "\n",
        "    return pd.concat(dfs, ignore_index=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZkAwZBj71ipc"
      },
      "source": [
        "### Libraries settings\n",
        "\n",
        "Initiate some libraries settings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B3hT5TUM1lGc"
      },
      "outputs": [],
      "source": [
        "login(token=HF_TOKEN)\n",
        "datasets.disable_progress_bar()\n",
        "transformers.set_seed(8)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "assh71yp4G_O"
      },
      "source": [
        "### Create a Vertex AI TensorBoard instance\n",
        "\n",
        "Create a Vertex AI TensorBoard instance for tracking and monitoring your tuning jobs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mKyJ1vREZRIY"
      },
      "outputs": [],
      "source": [
        "tensorboard = vertex_ai.Tensorboard.create(\n",
        "    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=REGION\n",
        ")\n",
        "\n",
        "vertex_ai.init(\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        "    staging_bucket=BUCKET_URI,\n",
        "    experiment_tensorboard=tensorboard,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8BH0TlP3PtTB"
      },
      "source": [
        "## Fine-Tune Gemma with Ray Train\n",
        "\n",
        "In this tutorial, you fine-tune Gemma 2B (`gemma-2b-it`) for summarizing newspaper articles using HuggingFace Transformer on Ray on Vertex AI. In an effort to make this notebook easily reproducible, you write a simple Python `trainer.py` script and submit it to the Ray cluster on Vertex AI using the Ray Jobs API via the public Ray Dashboard.\n",
        "\n",
        "As mentioned at the beginning, **consider this option for experimentation only.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Du1IAdWos0AF"
      },
      "source": [
        "### Initialize the Ray package\n",
        "\n",
        "Create an `__init__.py` file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AWrX2WLUs4um"
      },
      "outputs": [],
      "source": [
        "with open(src_path / \"__init__.py\", \"a\") as init_file:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mVL71KFDstVR"
      },
      "source": [
        "### Prepare the train script\n",
        "\n",
        "Create the `src/train.py` file which is the Python script for initializing Gemma fine-tuning using HuggingFace TRL library."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FNQXWOHwtIDi"
      },
      "outputs": [],
      "source": [
        "train_script = '''\n",
        "# training libraries\n",
        "import os\n",
        "import numpy as np\n",
        "import torch\n",
        "from huggingface_hub import login\n",
        "import datasets\n",
        "import transformers\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, Seq2SeqTrainingArguments\n",
        "from peft import LoraConfig\n",
        "from trl import SFTTrainer\n",
        "import evaluate\n",
        "import ray\n",
        "import ray.train.huggingface.transformers\n",
        "\n",
        "def train_func(config):\n",
        "    # Helpers\n",
        "    def formatting_func(example):\n",
        "        \"\"\"Helper function for formatting data for instruction tuning according to Gemma documentation.\"\"\"\n",
        "        output_texts = []\n",
        "        for i in range(len(example)):\n",
        "          messages = [\n",
        "            {\"role\": \"user\",\n",
        "             \"content\": f\"Summarize the following ARTICLE in one sentence.\\\\n###ARTICLE: {example['document'][i]}\"},\n",
        "            {\"role\": \"assistant\",\n",
        "             \"content\": f\"{example['summary'][i]}<eos>\"} # Make minor gemma fixes #2029\n",
        "             ]\n",
        "          output_texts.append(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False))\n",
        "        return output_texts\n",
        "\n",
        "    def compute_metrics(eval_preds):\n",
        "        \"\"\"Helper function for computing metrics\"\"\"\n",
        "        preds, labels = eval_preds\n",
        "        preds = preds[0]\n",
        "\n",
        "        preds = np.where(preds != -100, preds, tokenizer.pad_token_id)\n",
        "        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
        "\n",
        "        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)\n",
        "        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
        "\n",
        "        metrics = rouge.compute(predictions=decoded_preds,\n",
        "                                references=decoded_labels,\n",
        "                                rouge_types=['rouge1', 'rouge2', 'rougeL', 'rougeLsum'],\n",
        "                                use_aggregator=True, use_stemmer=True)\n",
        "        metrics = {k: round(v * 100, 4) for k, v in metrics.items()}\n",
        "        return metrics\n",
        "\n",
        "    def preprocess_logits_for_metrics(logits, labels):\n",
        "        \"\"\"Helper function for logits preprocessing for metrics\"\"\"\n",
        "        preds = torch.argmax(logits, dim=-1)\n",
        "        return preds, labels\n",
        "\n",
        "    # Setting training\n",
        "    login(token=os.environ['HF_TOKEN'], add_to_git_credential=True)\n",
        "    transformers.set_seed(8)\n",
        "\n",
        "    # Load dataset\n",
        "    dataset_id = \"xsum\"\n",
        "    dataset = datasets.load_dataset(dataset_id, trust_remote_code=True)\n",
        "    train_dataset = dataset[\"train\"]\n",
        "    eval_dataset = dataset[\"test\"]\n",
        "\n",
        "    # Preprocess dataset\n",
        "    model_id = \"google/gemma-2b-it\"\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "    tokenizer.padding_side = 'right'\n",
        "\n",
        "    # Prepare model\n",
        "    bnb_config = BitsAndBytesConfig(\n",
        "        load_in_4bit=True,\n",
        "        bnb_4bit_use_double_quant=True,\n",
        "        bnb_4bit_quant_type=\"nf4\",\n",
        "        bnb_4bit_compute_dtype=torch.bfloat16,\n",
        "    )\n",
        "    model = AutoModelForCausalLM.from_pretrained(model_id,\n",
        "                                                 quantization_config=bnb_config,\n",
        "                                                 device_map={'': torch.cuda.current_device()},\n",
        "                                                 torch_dtype=torch.bfloat16,\n",
        "                                                 # attn_implementation=\"flash_attention_2\"\n",
        "                                                 )\n",
        "    lora_config = LoraConfig(\n",
        "        r=32,\n",
        "        lora_alpha=32,\n",
        "        lora_dropout=0.05,\n",
        "        target_modules=\"all-linear\",\n",
        "        bias=\"none\",\n",
        "        task_type=\"CAUSAL_LM\"\n",
        "    )\n",
        "\n",
        "    # model.gradient_checkpointing_enable()\n",
        "    rouge = evaluate.load(\"rouge\")\n",
        "\n",
        "    training_args = Seq2SeqTrainingArguments(\n",
        "        output_dir=\"checkpoints\",\n",
        "        per_device_train_batch_size=config.get(\"per_device_train_batch_size\"),\n",
        "        per_device_eval_batch_size=config.get(\"per_device_eval_batch_size\"),\n",
        "        gradient_accumulation_steps=config.get(\"gradient_accumulation_steps\"),\n",
        "        logging_strategy=\"steps\",\n",
        "        save_strategy=\"steps\",\n",
        "        evaluation_strategy=\"steps\",\n",
        "        max_steps=config.get(\"max_steps\"),\n",
        "        save_steps=config.get(\"save_steps\"),\n",
        "        logging_steps=config.get(\"logging_steps\"),\n",
        "        learning_rate=config.get(\"learning_rate\"),\n",
        "        optim=\"paged_adamw_8bit\",\n",
        "        bf16=False,\n",
        "        fp16=True,\n",
        "        report_to=\"none\",\n",
        "        predict_with_generate=True,\n",
        "        ddp_find_unused_parameters=False,\n",
        "        gradient_checkpointing=True,\n",
        "        push_to_hub=False,\n",
        "        disable_tqdm=False,\n",
        "        load_best_model_at_end=False\n",
        "    )\n",
        "\n",
        "    max_seq_length = 512\n",
        "    trainer = SFTTrainer(\n",
        "        model=model,\n",
        "        tokenizer=tokenizer,\n",
        "        args=training_args,\n",
        "        train_dataset=train_dataset,\n",
        "        eval_dataset=eval_dataset,\n",
        "        max_seq_length=max_seq_length,\n",
        "        compute_metrics=compute_metrics,\n",
        "        preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n",
        "        peft_config=lora_config,\n",
        "        formatting_func=formatting_func\n",
        "    )\n",
        "    # model.config.use_cache = False\n",
        "\n",
        "    callback = ray.train.huggingface.transformers.RayTrainReportCallback()\n",
        "    trainer.add_callback(callback)\n",
        "    trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)\n",
        "    trainer.train()\n",
        "'''\n",
        "\n",
        "with open(src_path / \"train.py\", \"w\") as f:\n",
        "    f.write(train_script)\n",
        "f.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pznqJKCpEcL3"
      },
      "source": [
        "### Prepare the distributed training script\n",
        "\n",
        "Create `src/trainer.py` file which is the Python script for executing the Ray distributed training job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rCFfhM7RP4RI"
      },
      "outputs": [],
      "source": [
        "trainer_script = \"\"\"\n",
        "# libraries\n",
        "import argparse\n",
        "\n",
        "# training libraries\n",
        "from train import train_func\n",
        "\n",
        "# ray libraries\n",
        "import ray\n",
        "import ray.train.huggingface.transformers\n",
        "from ray.train import ScalingConfig, RunConfig, CheckpointConfig\n",
        "from ray.train.torch import TorchTrainer\n",
        "\n",
        "\n",
        "# helpers\n",
        "def get_args():\n",
        "    parser = argparse.ArgumentParser(description='Supervised tuning Gemma on Ray on Vertex AI')\n",
        "\n",
        "    # some gemma parameters\n",
        "    parser.add_argument(\"--train_batch_size\", type=int, default=1, help=\"train batch size\")\n",
        "    parser.add_argument(\"--eval_batch_size\", type=int, default=1, help=\"eval batch size\")\n",
        "    parser.add_argument(\"--gradient_accumulation_steps\", type=int, default=4, help=\"gradient accumulation steps\")\n",
        "    parser.add_argument(\"--learning_rate\", type=float, default=2e-4, help=\"learning rate\")\n",
        "    parser.add_argument(\"--max_steps\", type=int, default=100, help=\"max steps\")\n",
        "    parser.add_argument(\"--save_steps\", type=int, default=10, help=\"save steps\")\n",
        "    parser.add_argument(\"--logging_steps\", type=int, default=10, help=\"logging steps\")\n",
        "\n",
        "    # ray parameters\n",
        "    parser.add_argument('--num-workers', dest='num_workers', type=int, default=1, help='Number of workers')\n",
        "    parser.add_argument('--use-gpu', dest='use_gpu', action='store_true', default=False, help='Use GPU')\n",
        "    parser.add_argument('--experiment-name', dest='experiment_name', type=str, default='gemma-on-rov', help='Experiment name')\n",
        "    parser.add_argument('--logging-dir', dest='logging_dir', type=str, help='Logging directory')\n",
        "    args = parser.parse_args()\n",
        "    return args\n",
        "\n",
        "\n",
        "def main():\n",
        "\n",
        "    args = get_args()\n",
        "    config = vars(args)\n",
        "\n",
        "    # initialize ray session\n",
        "    ray.shutdown()\n",
        "    ray.init()\n",
        "\n",
        "    # training config\n",
        "    train_loop_config = {\n",
        "        \"per_device_train_batch_size\": config['train_batch_size'],\n",
        "        \"per_device_eval_batch_size\": config['eval_batch_size'],\n",
        "        \"gradient_accumulation_steps\": config['gradient_accumulation_steps'],\n",
        "        \"learning_rate\": config['learning_rate'],\n",
        "        \"max_steps\": config['max_steps'],\n",
        "        \"save_steps\": config['save_steps'],\n",
        "        \"logging_steps\": config['logging_steps'],\n",
        "    }\n",
        "    scaling_config = ScalingConfig(num_workers=config['num_workers'], use_gpu=config['use_gpu'])\n",
        "    run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=5,\n",
        "                          checkpoint_score_attribute=\"loss\",\n",
        "                          checkpoint_score_order=\"min\"),\n",
        "                           storage_path=config['logging_dir'],\n",
        "                           name=config['experiment_name'])\n",
        "    trainer = TorchTrainer(\n",
        "        train_loop_per_worker=train_func,\n",
        "        train_loop_config=train_loop_config,\n",
        "        run_config=run_config,\n",
        "        scaling_config=scaling_config\n",
        "    )\n",
        "    # train\n",
        "    result = trainer.fit()\n",
        "\n",
        "    ray.shutdown()\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()\n",
        "\"\"\"\n",
        "\n",
        "with open(src_path / \"trainer.py\", \"w\") as f:\n",
        "    f.write(trainer_script)\n",
        "f.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PYcmfEvZ3C1i"
      },
      "source": [
        "### Submit a Ray job using the Ray Jobs API\n",
        "\n",
        "Submit the script to the Ray cluster on Vertex AI using the Ray Jobs API with  the public Ray dashboard address.\n",
        "\n",
        "Initiate the client to submit the job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FljDHRQ63EP4"
      },
      "outputs": [],
      "source": [
        "client = JobSubmissionClient(\n",
        "    address=\"vertex_ray://{}\".format(ray_cluster.dashboard_address)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uDHIlGlQJ2oi"
      },
      "source": [
        "Set some job configuration including experiment name, job id, training entrypoint and more."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dxkqJ9vntz7C"
      },
      "outputs": [],
      "source": [
        "train_id = \"\".join(random.choices(string.ascii_lowercase   string.digits, k=3))\n",
        "train_experiment_name = f\"rov-dialog-gemma-tune-{train_id}\"\n",
        "train_submission_id = f\"ray-job-{train_id}\"\n",
        "train_entrypoint = f\"python3 trainer.py --experiment-name={train_experiment_name} --logging-dir={EXPERIMENTS_FOLDER_URI} --num-workers={WORKER_NODE_COUNT} --use-gpu\"\n",
        "train_experiment_uri = EXPERIMENTS_FOLDER_URI / train_experiment_name\n",
        "train_runtime_env = {\n",
        "    \"working_dir\": str(src_path),\n",
        "    \"env_vars\": {\"HF_TOKEN\": HF_TOKEN, \"TORCH_NCCL_ASYNC_ERROR_HANDLING\": \"3\"},\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yhwkH-NzgEHS"
      },
      "source": [
        "Submit the job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TUPB6YKpur4f"
      },
      "outputs": [],
      "source": [
        "train_job_id = client.submit_job(\n",
        "    submission_id=train_submission_id,\n",
        "    entrypoint=train_entrypoint,\n",
        "    runtime_env=train_runtime_env,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8T0MaYZegF_1"
      },
      "source": [
        "Check the status of the job while is running using the `monitor_job` function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0uCeasefuso8"
      },
      "outputs": [],
      "source": [
        "train_job_status = monitor_job(client, train_job_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67FI31SqKFdF"
      },
      "source": [
        "### Check training artifacts\n",
        "\n",
        "After the Ray training job has completed, see the model artifacts in the Cloud Storage location.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q_JOYhXQKK8U"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -l {train_experiment_uri}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dAOAsmGYZ2l9"
      },
      "source": [
        "### Log metrics in Vertex AI TensorBoard\n",
        "\n",
        "Use Vertex AI TensorBoard for validating your training job by logging resulting metrics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TSMCKEp9Z8jx"
      },
      "outputs": [],
      "source": [
        "vertex_ai.upload_tb_log(\n",
        "    tensorboard_id=tensorboard.name,\n",
        "    tensorboard_experiment_name=train_experiment_name,\n",
        "    logdir=str(train_experiment_uri),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u2RfyLpW0XO7"
      },
      "source": [
        "## Serving tuned Gemma model with Ray Data for offline predictions\n",
        "\n",
        "Using Ray on Vertex AI for developing AI/ML applications offers various benefits. In this scenario, you can use Cloud storage to conveniently store model checkpoints, metrics and more. This allows you to quickly consume the model for AI/ML downstreaming tasks including generating batch predictions using Ray Data.  \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "700NkAsLQHJI"
      },
      "source": [
        "### Generate predictions (locally)\n",
        "\n",
        "Generate predictions locally to validate the tuned model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fN5YQ60PXbYq"
      },
      "source": [
        "#### Download Ray training checkpoints\n",
        "\n",
        "Download all resulting checkpoints from Ray job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uBf0_agLhrLV"
      },
      "outputs": [],
      "source": [
        "! gsutil -q cp -r {train_experiment_uri}/* {experiments_path}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VavpxaY1hlSN"
      },
      "source": [
        "#### Get the best checkpoint\n",
        "\n",
        "Use the `ExperimentAnalysis` method to retrieve the the best checkpoint according to relevant metrics and mode."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nUBEKpK1gVx9"
      },
      "outputs": [],
      "source": [
        "experiment_analysis = ExperimentAnalysis(experiments_path)\n",
        "log_path = experiment_analysis.get_best_trial(metric=\"eval_rougeLsum\", mode=\"max\")\n",
        "best_checkpoint = experiment_analysis.get_best_checkpoint(\n",
        "    log_path, metric=\"eval_rougeLsum\", mode=\"max\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "icFuBCwJXSub"
      },
      "source": [
        "#### Load the model after training\n",
        "\n",
        "After training the model, load the model as described in the Hugging Face [documentation](https://huggingface.co/docs/trl/use_model#use-adapters-peft).\n",
        "\n",
        "Set the model and adapters path. Also set the path to store the resulting tuned model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "370W-GHE7vM2"
      },
      "outputs": [],
      "source": [
        "base_model_path = \"google/gemma-2b-it\"\n",
        "peft_model_path = epath.Path(best_checkpoint.path) / \"checkpoint\"\n",
        "tuned_model_path = models_path / \"xsum-tuned-gemma-it\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9UCc4hvGjES8"
      },
      "source": [
        "Initiate the associated Gemma tokenizer and base model. Also initiate the resulting adapters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zbctXuXQAhLP"
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(base_model_path)\n",
        "tokenizer.padding_side = \"right\"\n",
        "\n",
        "base_model = AutoModelForCausalLM.from_pretrained(\n",
        "    base_model_path, device_map=\"auto\", torch_dtype=torch.float16\n",
        ")\n",
        "peft_model = PeftModel.from_pretrained(\n",
        "    base_model,\n",
        "    peft_model_path,\n",
        "    device_map=\"auto\",\n",
        "    torch_dtype=torch.bfloat16,\n",
        "    is_trainable=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Jytit8inMGy"
      },
      "source": [
        "Merge the base model and adapters to save the tuned model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aRGvb7hRnK-a"
      },
      "outputs": [],
      "source": [
        "tuned_model = peft_model.merge_and_unload()\n",
        "tuned_model.save_pretrained(tuned_model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t4WRLix49saQ"
      },
      "source": [
        "#### Generate summaries\n",
        "\n",
        "Generate summaries with the tuned model. Load the validation set of the tutorial dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PrsR5HjNunzW"
      },
      "outputs": [],
      "source": [
        "dataset = datasets.load_dataset(\n",
        "    \"xsum\", split=\"validation\", cache_dir=data_path, trust_remote_code=True\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dm5eFJNxoAHk"
      },
      "source": [
        "Sample one article to summarize."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fZjT8zAOPivt"
      },
      "outputs": [],
      "source": [
        "sample = dataset.select([random.randint(0, len(dataset) - 1)])\n",
        "document = sample[\"document\"][0]\n",
        "reference_summary = sample[\"summary\"][0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "16GiJy_EoEgi"
      },
      "source": [
        "Prepare the associated prompt following the [Gemma documentation](https://ai.google.dev/gemma/docs/formatting)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GFhA8R3wXPKj"
      },
      "outputs": [],
      "source": [
        "messages = [\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": f\"Summarize the following ARTICLE in one sentence.\\\\n###ARTICLE: {document}\",\n",
        "    },\n",
        "]\n",
        "\n",
        "prompt = tokenizer.apply_chat_template(\n",
        "    messages, tokenize=False, add_generation_prompt=True\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u7Lvss80objd"
      },
      "source": [
        "Initiate the text-generation pipeline for generating summaries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wq6m6G-b_w7b"
      },
      "outputs": [],
      "source": [
        "tuned_gemma_pipeline = pipeline(\n",
        "    \"text-generation\", model=tuned_model, tokenizer=tokenizer, max_new_tokens=50\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yb5GNgv1ovpc"
      },
      "source": [
        "Generate the associated summary."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EuF2CMGx6AGi"
      },
      "outputs": [],
      "source": [
        "generated_tuned_gemma_summary = tuned_gemma_pipeline(\n",
        "    prompt, do_sample=True, temperature=0.1, add_special_tokens=True\n",
        ")[0][\"generated_text\"][len(prompt) :]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OKjU_Mhgo8yO"
      },
      "source": [
        "Print the generated summary."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Oq6HLy4Pu8oC"
      },
      "outputs": [],
      "source": [
        "print(f\"Reference summary: {reference_summary}\")\n",
        "print(\"-\" * 100)\n",
        "print(f\"Tuned generated summary: {generated_tuned_gemma_summary}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xdhPQC5aKmyP"
      },
      "source": [
        "#### Evaluate models\n",
        "\n",
        "As an additional step, you can evaluate the tuned model. To evaluate the model you compare models qualitatively and quantitatively.\n",
        "\n",
        "In one case, you compare responses generated by the base Gemma model with the ones generated by the tuned Gemma model. In the other case, you calculate ROUGE metrics and its improvements which gives you an idea of how well the tuned models is able to reproduce the reference summaries correctly with respect to the base model.\n",
        "\n",
        "Evaluate models by comparing generated summaries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nmbCCJpWora7"
      },
      "outputs": [],
      "source": [
        "gemma_pipeline = pipeline(\n",
        "    \"text-generation\", model=base_model, tokenizer=tokenizer, max_new_tokens=50\n",
        ")\n",
        "\n",
        "generated_gemma_summary = gemma_pipeline(\n",
        "    prompt, do_sample=True, temperature=0.1, add_special_tokens=True\n",
        ")[0][\"generated_text\"][len(prompt) :]\n",
        "\n",
        "print(f\"Reference summary: {reference_summary}\")\n",
        "print(\"-\" * 100)\n",
        "print(f\"Base generated summary: {generated_gemma_summary}\")\n",
        "print(\"-\" * 100)\n",
        "print(f\"Tuned generated summary: {generated_tuned_gemma_summary}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RLt-3ovZq7y1"
      },
      "source": [
        "Evaluate models by computing ROUGE metrics and its improvements."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FRSPxJVpfJ3E"
      },
      "outputs": [],
      "source": [
        "rouge = evaluate.load(\"rouge\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iEO5qqtsMC75"
      },
      "outputs": [],
      "source": [
        "gemma_results = rouge.compute(\n",
        "    predictions=[generated_gemma_summary],\n",
        "    references=[reference_summary],\n",
        "    rouge_types=[\"rouge1\", \"rouge2\", \"rougeL\", \"rougeLsum\"],\n",
        "    use_aggregator=True,\n",
        "    use_stemmer=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sYC5t4I863pn"
      },
      "outputs": [],
      "source": [
        "tuned_gemma_results = rouge.compute(\n",
        "    predictions=[generated_tuned_gemma_summary],\n",
        "    references=[reference_summary],\n",
        "    rouge_types=[\"rouge1\", \"rouge2\", \"rougeL\", \"rougeLsum\"],\n",
        "    use_aggregator=True,\n",
        "    use_stemmer=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QnWVuRe394gh"
      },
      "outputs": [],
      "source": [
        "improvements = {}\n",
        "for rouge_metric, gemma_rouge in gemma_results.items():\n",
        "    tuned_gemma_rouge = tuned_gemma_results[rouge_metric]\n",
        "    if gemma_rouge != 0:\n",
        "        improvement = ((tuned_gemma_rouge - gemma_rouge) / gemma_rouge) * 100\n",
        "    else:\n",
        "        improvement = None\n",
        "    improvements[rouge_metric] = improvement\n",
        "\n",
        "print(\"Base Gemma vs Tuned Gemma - ROUGE improvements\")\n",
        "for rouge_metric, improvement in improvements.items():\n",
        "    print(f\"{rouge_metric}: {improvement:.3f}%\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G51Xhb4H0fa1"
      },
      "source": [
        "### Batch prediction with Ray Data\n",
        "\n",
        "To generate batch prediction with the tuned model using Ray Data on Ray on Vertex AI, you need a dataset to generate predictions and the tuned model stored in the Cloud bucket.\n",
        "\n",
        "Then, you can leverage Ray Data which provides an easy-to-use API for offline batch inference."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JjV1oHCTCe1r"
      },
      "source": [
        "#### Upload the tuned model\n",
        "\n",
        "Upload the tuned model on the Cloud storage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OFwD2p8jCjgH"
      },
      "outputs": [],
      "source": [
        "! gsutil -q cp -r {models_path} {MODELS_PATH}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eIwOKSnrpxkW"
      },
      "source": [
        "#### Prepare the batch prediction training script\n",
        "\n",
        "Prepare `src/batch_predict.py` file which is the Python script for executing the Ray batch prediction job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ekL5tV31pxkX"
      },
      "outputs": [],
      "source": [
        "batch_predictor_script = \"\"\"\n",
        "# General\n",
        "import argparse\n",
        "import os\n",
        "from huggingface_hub import login\n",
        "\n",
        "# Serving\n",
        "import datasets\n",
        "import transformers\n",
        "import numpy as np\n",
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "from transformers.pipelines import pipeline\n",
        "\n",
        "# Ray\n",
        "import ray\n",
        "\n",
        "# Settings\n",
        "datasets.disable_progress_bar()\n",
        "\n",
        "# Variables\n",
        "base_model_path = \"google/gemma-2b-it\"\n",
        "\n",
        "\n",
        "# helpers\n",
        "def get_args():\n",
        "    parser = argparse.ArgumentParser(description='Batch prediction with Gemma on Ray on Vertex AI')\n",
        "    parser.add_argument('--tuned_model_path', type=str, help='path of adapter model')\n",
        "    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus')\n",
        "    parser.add_argument('--batch_size', type=int, default=8, help='batch size')\n",
        "    parser.add_argument('--sample_size', type=int, default=20, help='number of articles to summarize')\n",
        "    parser.add_argument('--temperature', type=float, default=0.1, help='temperature for generating summaries')\n",
        "    parser.add_argument('--max_new_tokens', type=int, default=50, help='max new token for generating summaries')\n",
        "    parser.add_argument('--output_dir', type=str, help='output directory for predictions')\n",
        "    args = parser.parse_args()\n",
        "    return args\n",
        "\n",
        "def main():\n",
        "\n",
        "    # Set configuration\n",
        "    args = get_args()\n",
        "    config = vars(args)\n",
        "\n",
        "    # Setting training\n",
        "    login(token=os.environ['HF_TOKEN'], add_to_git_credential=True)\n",
        "    transformers.set_seed(8)\n",
        "\n",
        "    # Load dataset\n",
        "    dataset_id = \"xsum\"\n",
        "    sample_size = config[\"sample_size\"]\n",
        "    input_data = datasets.load_dataset(dataset_id, split=\"validation\", trust_remote_code=True)\n",
        "    input_data = input_data.select(range(sample_size))\n",
        "    ray_input_data = ray.data.from_huggingface(input_data)\n",
        "\n",
        "    # Generate predictions\n",
        "\n",
        "    class Summarizer:\n",
        "\n",
        "      def __init__(self):\n",
        "          self.tokenizer = AutoTokenizer.from_pretrained(base_model_path)\n",
        "          self.tokenizer.padding_side = \"right\"\n",
        "\n",
        "          self.tuned_model = AutoModelForCausalLM.from_pretrained(config[\"tuned_model_path\"],\n",
        "                                                                  device_map='auto',\n",
        "                                                                  torch_dtype=torch.float16)\n",
        "\n",
        "          self.pipeline = pipeline(\"text-generation\",\n",
        "                                    model=self.tuned_model,\n",
        "                                    tokenizer=self.tokenizer,\n",
        "                                    max_new_tokens=config[\"max_new_tokens\"])\n",
        "\n",
        "      def __call__(self, batch: np.ndarray):\n",
        "\n",
        "          # prepare dataset\n",
        "          messages = [{\"role\": \"user\",\n",
        "                      \"content\": f\"Summarize the following ARTICLE in one sentence.\\\\n###ARTICLE: {document}\"}\n",
        "                      for document in batch[\"document\"]]\n",
        "\n",
        "          batch['prompt'] = [self.tokenizer.apply_chat_template([message], tokenize=False, add_generation_prompt=True)\n",
        "                             for message in messages]\n",
        "\n",
        "          # generate\n",
        "          batch['generated_summary'] = [self.pipeline(prompt,\n",
        "                                                    do_sample=True,\n",
        "                                                    temperature=config[\"temperature\"],\n",
        "                                                    add_special_tokens=True)[0][\"generated_text\"][len(prompt):]\n",
        "                                                    for prompt in batch['prompt']]\n",
        "\n",
        "          return batch\n",
        "\n",
        "\n",
        "    predictions_data = ray_input_data.map_batches(\n",
        "        Summarizer,\n",
        "        concurrency=config[\"num_gpus\"],\n",
        "        num_gpus=1,\n",
        "        batch_size=config['batch_size'])\n",
        "\n",
        "    # Store resulting predictions\n",
        "    predictions_data.write_json(config[\"output_dir\"], try_create_dir=True)\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()\n",
        "\"\"\"\n",
        "\n",
        "with open(src_path / \"batch_predictor.py\", \"w\") as f:\n",
        "    f.write(batch_predictor_script)\n",
        "f.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_xynt8impxkX"
      },
      "source": [
        "####  Submit a Ray job using the Ray Jobs API\n",
        "\n",
        "Submit the script to the Ray on Vertex AI cluster using the Ray Jobs API via  the public Ray dashboard address.\n",
        "\n",
        "Initiate the client to submit the job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hY96qSQFpxkX"
      },
      "outputs": [],
      "source": [
        "client = JobSubmissionClient(\n",
        "    address=\"vertex_ray://{}\".format(ray_cluster.dashboard_address)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jQ6EqvIF5q9F"
      },
      "source": [
        "Set some job configuration including model path, job id, prediction entrypoint and more."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SDycAWy45siE"
      },
      "outputs": [],
      "source": [
        "batch_predict_id = \"\".join(random.choices(string.ascii_lowercase   string.digits, k=4))\n",
        "batch_predict_submission_id = f\"ray-job-{batch_predict_id}\"\n",
        "tuned_model_uri_path = str(MODELS_PATH / \"xsum-tuned-gemma-it\").replace(\n",
        "    \"gs://\", \"/gcs/\"\n",
        ")\n",
        "batch_predict_entrypoint = f\"python3 batch_predictor.py --tuned_model_path={tuned_model_uri_path} --num_gpus=2 --output_dir={PREDICTIONS_FOLDER_URI}\"\n",
        "batch_predict_runtime_env = {\n",
        "    \"working_dir\": str(src_path),\n",
        "    \"env_vars\": {\"HF_TOKEN\": HF_TOKEN},\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NcaMEKz6pxkX"
      },
      "source": [
        "Submit the job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qADWE5LOpxkX"
      },
      "outputs": [],
      "source": [
        "batch_predict_job_id = client.submit_job(\n",
        "    submission_id=batch_predict_submission_id,\n",
        "    entrypoint=batch_predict_entrypoint,\n",
        "    runtime_env=batch_predict_runtime_env,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LxrOwuVf6f3R"
      },
      "source": [
        "Check the status of the job using the `monitor_job` helper function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SXTF8pp9KJS8"
      },
      "outputs": [],
      "source": [
        "batch_predict_job_status = monitor_job(client, batch_predict_job_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_2Iafitj7wtH"
      },
      "source": [
        "#### Get generated summaries\n",
        "\n",
        "Have a quick view of generated summaries using a Pandas DataFrame."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "em4WQokY7050"
      },
      "outputs": [],
      "source": [
        "predictions_df = read_json_files(prefix=\"predictions/\", bucket_name=BUCKET_NAME)\n",
        "predictions_df = predictions_df[\n",
        "    [\"id\", \"document\", \"prompt\", \"summary\", \"generated_summary\"]\n",
        "]\n",
        "predictions_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "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) that you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources that you created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "delete_tensorboards = False\n",
        "delete_experiments = False\n",
        "delete_ray_clusters = False\n",
        "delete_image_repo = False\n",
        "delete_bucket = False\n",
        "delete_tutorial = False\n",
        "\n",
        "# Delete tensorboard\n",
        "if delete_tensorboards:\n",
        "    tensorboard_list = vertex_ai.Tensorboard.list()\n",
        "    for tensorboard in tensorboard_list:\n",
        "        tensorboard.delete()\n",
        "\n",
        "# Delete experiments\n",
        "if delete_experiments:\n",
        "    experiment_list = vertex_ai.Experiment.list()\n",
        "    for experiment in experiment_list:\n",
        "        experiment.delete()\n",
        "\n",
        "# Delete ray on vertex cluster\n",
        "if delete_ray_clusters:\n",
        "    ray_cluster_list = vertex_ray.list_ray_clusters()\n",
        "    for ray_cluster in ray_cluster_list:\n",
        "        vertex_ray.delete_ray_cluster(ray_cluster.cluster_resource_name)\n",
        "\n",
        "if delete_image_repo:\n",
        "    ! gcloud artifacts repositories delete {REPO_NAME}\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "if delete_bucket:\n",
        "    ! gsutil -q -m rm -r {BUCKET_URI}\n",
        "\n",
        "# Delete tutorial folder\n",
        "if delete_tutorial:\n",
        "    shutil.rmtree(tutorial_path)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_gemma_fine_tuning_batch_deployment_on_rov.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}