notebooks/official/ray_on_vertex_ai/ray_cluster_management.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": [
        "# Ray on Vertex AI cluster management\n",
        "\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ray_on_vertex_ai/ray_cluster_management.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> <br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fray_on_vertex_ai%2Fray_cluster_management.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"> <br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/ray_on_vertex_ai/ray_cluster_management.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"> <br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "<td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ray_on_vertex_ai/ray_cluster_management.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"> <br>\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use Ray on Vertex AI SDK for cluster management.\n",
        "\n",
        "Learn more about [Ray on Vertex AI overview](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 learn how to create a cluster, list existing clusters, get a cluster, update (manually scaling) a cluster, and delete a cluster.\n",
        "\n",
        "This tutorial uses the following Vertex AI services and resources:\n",
        "\n",
        "- [Ray on Vertex AI](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/overview)\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create a cluster.\n",
        "- List existing clusters.\n",
        "- Get a cluster.\n",
        "- Manually scale up the cluster, then scale down the cluster.\n",
        "- Delete existing clusters."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aed92deeb4a0"
      },
      "source": [
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "\n",
        "Learn about [Ray on Vertex AI pricing](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/overview#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": "8925ff9e165e"
      },
      "source": [
        "## Getting Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform[ray]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "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": "f200f10a1da3"
      },
      "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": "7b49231643e4"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7176ea64999b"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7de6ef0fac42"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "# Retrieve the project number\n",
        "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
        "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Set network information\n",
        "\n",
        "[Set up a VPC peering network](https://cloud.google.com/vertex-ai/docs/general/vpc-peering) and private services connection to access Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3927074343e3"
      },
      "outputs": [],
      "source": [
        "VPC_NETWORK = \"[your-network-name]\"  # @param {type:\"string\"}\n",
        "VPC_NETWORK_FULL = \"projects/{}/global/networks/{}\".format(PROJECT_NUMBER, VPC_NETWORK)\n",
        "VPC_NETWORK_FULL"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import vertex_ray"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4481de089cc2"
      },
      "source": [
        "## Create a cluster\n",
        "\n",
        "Note that within the same VPC network, IP ranges restrict the numbers of clusters and nodes you can create."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6031842560d2"
      },
      "outputs": [],
      "source": [
        "head_node_type = vertex_ray.Resources(\n",
        "    machine_type=\"n1-standard-16\",\n",
        "    node_count=1,\n",
        ")\n",
        "\n",
        "worker_node_types = [\n",
        "    vertex_ray.Resources(\n",
        "        machine_type=\"n1-standard-8\",\n",
        "        node_count=2,  # Can be > 1\n",
        "        accelerator_type=\"NVIDIA_TESLA_T4\",\n",
        "        accelerator_count=1,\n",
        "    )\n",
        "]\n",
        "\n",
        "cluster_resource_name = vertex_ray.create_ray_cluster(\n",
        "    head_node_type=head_node_type,\n",
        "    network=VPC_NETWORK_FULL,\n",
        "    worker_node_types=worker_node_types,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "95da29c3e54b"
      },
      "source": [
        "## List existing clusters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f35e48bf1b35"
      },
      "outputs": [],
      "source": [
        "clusters = vertex_ray.list_ray_clusters()\n",
        "clusters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "be2e79a1af08"
      },
      "source": [
        "## Update an existing cluster (manually scaling)\n",
        "\n",
        "\n",
        "Note that the maximum number of worker nodes you can scale up depends on the initial node counts (more details are in these [formulas](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/set-up)) and is restricted by IP ranges within the same VPC network.\n",
        "\n",
        "Get the cluster you want to scale."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7b0d96c88191"
      },
      "outputs": [],
      "source": [
        "cluster = vertex_ray.get_ray_cluster(cluster_resource_name)\n",
        "cluster"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f1a20a92c80b"
      },
      "source": [
        "Scale down workers from 2 nodes to 1 node."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "03d3780f8b5f"
      },
      "outputs": [],
      "source": [
        "new_worker_node_types = []\n",
        "for worker_node_type in cluster.worker_node_types:\n",
        "    worker_node_type.node_count = 1\n",
        "    new_worker_node_types.append(worker_node_type)\n",
        "\n",
        "cluster_resource_name = vertex_ray.update_ray_cluster(\n",
        "    cluster_resource_name=cluster_resource_name,\n",
        "    worker_node_types=new_worker_node_types,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b025fdd97348"
      },
      "source": [
        "Verify if the cluster is successfully scaled down."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "604a124c425f"
      },
      "outputs": [],
      "source": [
        "cluster = vertex_ray.get_ray_cluster(cluster_resource_name)\n",
        "cluster"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "72dc8445c9cd"
      },
      "source": [
        "Scale up to 2 worker nodes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "750293506627"
      },
      "outputs": [],
      "source": [
        "new_worker_node_types = []\n",
        "for worker_node_type in cluster.worker_node_types:\n",
        "    worker_node_type.node_count = 2\n",
        "    new_worker_node_types.append(worker_node_type)\n",
        "\n",
        "cluster_resource_name = vertex_ray.update_ray_cluster(\n",
        "    cluster_resource_name=cluster_resource_name,\n",
        "    worker_node_types=new_worker_node_types,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e62b56ac6429"
      },
      "source": [
        "Verify that the cluster is successfully scaled up."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "be90cef02cf2"
      },
      "outputs": [],
      "source": [
        "cluster = vertex_ray.get_ray_cluster(cluster_resource_name)\n",
        "cluster"
      ]
    },
    {
      "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) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the cluster you created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "# Delete the cluster\n",
        "vertex_ray.delete_ray_cluster(cluster.cluster_resource_name)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "ray_cluster_management.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}