Deploy GPU-accelerated Ray for AI workloads on GKE

Create a cluster in GKE Autopilot mode

In Autopilot mode, Google manages your cluster configuration, including scaling, security, and other preconfigured settings. Autopilot clusters are optimized to run most production workloads and provision compute resources based on your Kubernetes manifests.

1. In the terminal, run the following command:

   gcloud container clusters create-auto my-ray-enabled-cluster \
       --enable-ray-operator \
       --enable-ray-cluster-monitoring \
       --enable-ray-cluster-logging \
       --location=us-central1
   

   This Google Cloud CLI command creates an Autopilot cluster named my-ray-enabled-cluster, with these settings:

   • --enable-ray-operator: installs the Ray Operator on the cluster, which simplifies managing Ray clusters within Kubernetes.
   • --enable-ray-cluster-monitoring: sets up integration with Cloud Monitoring for Ray cluster metrics.
   • --enable-ray-cluster-logging: configures integration with Cloud Logging to collect logs from your Ray cluster.
   • --location=us-central1: specifies the geographical region where the GKE cluster will be created. Be aware that accelerator availability and supported types vary by location. For more information, see GPU regions and zones.

   Autopilot clusters automatically provisions and scales nodes (including GPU-enabled ones that use compute classes) based on workload demands, optimizing both utilization and cost. This automatic provisioning and scaling of resources in Autopilot clusters simplifies resource management, and lets you focus on your models instead of infrastructure.

   If you use a GKE Standard cluster, the flags and settings for enabling Ray might be different. For more information, see Enable the Ray Operator on GKE.

2. Configure kubectl to communicate with your cluster:

   gcloud container clusters get-credentials my-ray-enabled-cluster \
       --region=us-central1
   

Install the KubeRay kubectl ray plugin

The kubectl ray plugin simplifies common Ray on Kubernetes workflows. For the latest installation instructions, see the KubeRay documentation.

1. Go to the releases page and download the binary for your platform. For example, to install kubectl ray plugin version 1.3.2 on Linux amd64, run the following commands:

   curl -LO https://github.com/ray-project/kuberay/releases/download/v1.3.2/kubectl-ray_v1.3.2_linux_amd64.tar.gz
   tar -xvf kubectl-ray_v1.3.2_linux_amd64.tar.gz
   cp kubectl-ray ~/.local/bin
   

   Replace ~/.local/bin with the directory in your PATH.

2. Verify your installation:

   kubectl ray version
   

   If you see a warning like KubeRay operator installation cannot be found: no KubeRay operator deployments found in any namespace., you can safely ignore it. If the installation was successful, the output is the version of the kubectl ray plugin.

Create the custom compute class in your cluster

A custom compute class in GKE lets you define specific hardware requirements, like the GPU type and count, by priority. Autopilot uses this priority when provisioning nodes for your workloads.

To create the example custom compute class for this guide, follow these steps:

1. Save the following compute class specification as nvidia-l4-compute-class.yaml:

   apiVersion: cloud.google.com/v1
   kind: ComputeClass
   metadata:
     name: gpu-compute-class
   spec:
     priorities:
     - gpu:
         type: nvidia-l4
         count: 1
       spot: true
     - gpu:
         type: nvidia-l4
         count: 4
       spot: false
     nodePoolAutoCreation:
       enabled: true
     whenUnsatisfiable: DoNotScaleUp
   

2. Create the compute class by by running this command:

   kubectl apply -f nvidia-l4-compute-class.yaml
   

   What happens after you apply the custom compute class? After applying the compute class YAML file, Autopilot will automatically provision nodes with one NVIDIA L4 GPU when your Ray cluster requests them. It will prioritize Spot VMs as defined in your manifest. The whenUnsatisfiable: DoNotScaleUp setting instructs GKE to not create nodes that don't meet the specified requirements.

Create and verify the Ray cluster

You can now create your Ray cluster on your Autopilot cluster.

1. Create the Ray cluster. Use the kubectl ray create cluster command to define and create your Ray cluster. This command simplifies the process by handling the underlying Kubernetes resource creation.

     kubectl ray create cluster my-ray-cluster \
         --worker-replicas=1 \
         --worker-cpu=2 \
         --worker-memory=4Gi \
         --worker-gpu=1 \
         --worker-node-selectors="cloud.google.com/compute-class=gpu-compute-class"
   

   The kubectl ray create cluster command uses the KubeRay plugin to translate your specifications into a RayCluster custom resource definition and sends it to the Kubernetes API to create the RayCluster object. This object acts as a blueprint, telling the Ray Operator how to provision and manage the actual Ray cluster components (head and worker nodes) within GKE. By default, GKE schedules the head Pod on a node from the E2 machine series.

   The example command uses these settings:

   • kubectl ray create cluster my-ray-cluster: specifies the name of the Ray cluster you want to create.
   • --worker-node-selectors: tells Kubernetes to schedule worker Pods on nodes that are labeled with cloud.google.com/compute-class=gpu-compute-class.

2. Verify the Ray cluster status. Use these commands to verify that your Ray cluster is up and running.

   • Check the RayCluster resource by running this command:

     kubectl ray get cluster
     

     This command lists all the Ray clusters in your Kubernetes namespace. You should see my-ray-cluster listed. It might take a few minutes for the cluster to initialize.

   • Check the Kubernetes Pods by running this command:

     kubectl get pods
     

     This command lists all the Pods that are running in your Kubernetes namespace. You should see Pods related to your Ray cluster:

     • A Pod for the Ray head node, with a name like my-ray-cluster-head-0.
     • One or more Pods for the Ray worker nodes, with names like my-ray-cluster-worker-group-0-xxxxx. The number of worker Pods depends on the initial configuration of your Ray cluster and any autoscaling that might have occurred.

Submit a Ray job

After your Ray cluster is ready, you can submit a Ray job to your running Ray cluster on GKE.

To do this, use the kubectl ray session command to start an interactive session, and the ray job submit command to start job execution.

1. Start a Ray interactive session. To establish a local connection to your Ray cluster, run this command:

     kubectl ray session my-ray-cluster
   

   This command initiates port forwarding between your local machine and the Ray head node in your GKE cluster. Note that your terminal will be occupied while this session is active; to proceed, open a separate terminal instance.

2. Create the sample code to run. Save the following example code to a file named sample_code.py.

     import ray
     import os
     import requests
   
     ray.init()
   
     @ray.remote
     class Counter:
         def __init__(self):
             # Used to verify runtimeEnv
             self.name = os.getenv("counter_name")
             assert self.name == "test_counter"
             self.counter = 0
   
         def inc(self):
             self.counter += 1
   
         def get_counter(self):
             return "{} got {}".format(self.name, self.counter)
   
     counter = Counter.remote()
   
     for _ in range(5):
             ray.get(counter.inc.remote())
             print(ray.get(counter.get_counter.remote()))
   
     # Verify that the correct runtime env was used for the job.
     assert requests.__version__ == "2.26.0"
   

3. Open a new terminal and submit the Ray job.

   1. Open a new terminal on your local machine. If you created a virtual environment for your Ray installation, ensure you activate it in this new terminal as well.
   2. Navigate to the directory where you saved the sample_code.py file.

   3. Run the following command:

        ray job submit \
            --working-dir=. \
            --runtime-env-json='{"pip": ["requests==2.26.0"], "env_vars": {"counter_name": "test_counter"}}' \
            --address http://localhost:8265 python sample_code.py
      

      The example command uses these settings:

      • --working-dir=.: specifies the current directory as the working directory to be made available to the Ray job on the cluster. Ideally, this directory contains just the application source code you want to run on Ray. In our example, this is the sample_code.py file.
      • --runtime-env-json='{"pip": ["requests==2.26.0"], "env_vars": {"counter_name": "test_counter"}}': defines the runtime environment for the job. It specifies the required Python library version (requests==2.26.0) and sets an environment variable (counter_name) for the job's execution.

4. View the Ray job. You can monitor the job in one of these ways:

   • Ray dashboard: open your web browser and go to http://localhost:8265. This URL opens the Ray dashboard, where you can see information about your running job, including its status, logs, and resource usage.
   • Metrics Explorer: use the prebuilt metrics dashboards in the Google Cloud console. For more information, see Collect and view logs and metrics for Ray clusters on GKE.

5. Terminate the Ray session. After you finish interacting with the Ray cluster or monitoring your job, you can terminate the interactive session by pressing Ctrl+C in the terminal where you ran the kubectl ray session command.