Add support for autopilot

Author: richardsliu

platform/main.tf

@@ -17,7 +17,7 @@ data "google_client_config" "provider" {}
 data "google_container_cluster" "ml_cluster" {
   name       = var.cluster_name
   location   = var.region
-  depends_on = [module.gke_cluster]
+  depends_on = [module.gke_autopilot, module.gke_standard]
 }
 
 provider "google" {
@@ -52,26 +52,38 @@ provider "helm" {
   }
 }
 
-module "gke_cluster" {
-  source = "./modules/gke_cluster"
+module "gke_autopilot" {
+  source = "./modules/gke_autopilot"
 
   project_id   = var.project_id
   region       = var.region
   cluster_name = var.cluster_name
+  enable_autopilot = var.enable_autopilot
+}
+
+
+module "gke_standard" {
+  source = "./modules/gke_standard"
+
+  project_id   = var.project_id
+  region       = var.region
+  cluster_name = var.cluster_name
+  enable_autopilot = var.enable_autopilot
 }
 
 module "kubernetes" {
   source = "./modules/kubernetes"
 
-  depends_on   = [module.gke_cluster]
+  depends_on   = [module.gke_standard]
   region       = var.region
   cluster_name = var.cluster_name
+  enable_autopilot = var.enable_autopilot
 }
 
 module "kuberay" {
   source = "./modules/kuberay"
 
-  depends_on   = [module.kubernetes]
+  depends_on   = [module.gke_autopilot, module.gke_standard]
   region       = var.region
   cluster_name = var.cluster_name
 }

platform/modules/gke_autopilot/main.tf

@@ -0,0 +1,52 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+provider "google" {
+  project = var.project_id
+  region  = var.region
+}
+
+
+# GKE cluster
+resource "google_container_cluster" "ml_cluster" {
+  name     = var.cluster_name
+  location = var.region
+  count    = var.enable_autopilot == true ? 1 : 0
+
+  initial_node_count = 1
+
+  logging_config {
+    enable_components = ["SYSTEM_COMPONENTS", "WORKLOADS"]
+  }
+
+  monitoring_config {
+    enable_components = ["SYSTEM_COMPONENTS"]
+    managed_prometheus {
+      enabled = "true"
+    }
+  }
+  ip_allocation_policy {
+    cluster_ipv4_cidr_block  = ""
+    services_ipv4_cidr_block = ""
+  }
+
+  enable_autopilot = true
+
+  release_channel {
+    channel = "RAPID"
+  }
+
+  min_master_version = "1.27"
+}
+

platform/modules/gke_cluster/output.tf platform/modules/gke_autopilot/output.tf

@@ -14,26 +14,26 @@
 
 output "project_id" {
   description = "GCP project id"
-  value       = resource.google_container_cluster.ml_cluster.project
+  value       = var.enable_autopilot ? resource.google_container_cluster.ml_cluster[0].project : null
 }
 
 output "region" {
   description = "GCP region"
-  value       = resource.google_container_cluster.ml_cluster.location
+  value       = var.enable_autopilot ? resource.google_container_cluster.ml_cluster[0].location : null
 }
 
 output "cluster_name" {
   description = "The name of the GKE cluster"
-  value       = resource.google_container_cluster.ml_cluster.name
+  value       = var.enable_autopilot ? resource.google_container_cluster.ml_cluster[0].name : null
 }
 
 output "kubernetes_host" {
   description = "Kubernetes cluster host"
-  value       = resource.google_container_cluster.ml_cluster.endpoint
+  value       = var.enable_autopilot ? resource.google_container_cluster.ml_cluster[0].endpoint : null
 }
 
 output "cluster_certicicate" {
   description = "Kubernetes cluster ca certificate"
-  value       = base64decode(resource.google_container_cluster.ml_cluster.master_auth[0].cluster_ca_certificate)
+  value       = var.enable_autopilot ? base64decode(resource.google_container_cluster.ml_cluster[0].master_auth[0].cluster_ca_certificate) : null
   sensitive   = true
 }

platform/modules/gke_cluster/variables.tf platform/modules/gke_autopilot/variables.tf

@@ -21,7 +21,7 @@ variable "project_id" {
 variable "region" {
   type        = string
   description = "GCP project region or zone"
-  default     = "us-central1-c"
+  default     = "us-central1"
 }
 
 variable "cluster_name" {
@@ -40,3 +40,9 @@ variable "num_gpu_nodes" {
   description = "Number of GPU nodes in the cluster"
   default     = 1
 }
+
+variable "enable_autopilot" {
+  type        = bool
+  description = "Set to true to enable GKE Autopilot clusters"
+  default     = false
+}

platform/modules/gke_cluster/versions.tf platform/modules/gke_autopilot/versions.tf

@@ -22,6 +22,7 @@ provider "google" {
 resource "google_container_cluster" "ml_cluster" {
   name     = var.cluster_name
   location = var.region
+  count    = var.enable_autopilot == false ? 1 : 0
 
   initial_node_count = 3
 
@@ -44,8 +45,9 @@ resource "google_container_cluster" "ml_cluster" {
 resource "google_container_node_pool" "gpu_pool" {
   name       = "gpu-pool"
   location   = var.region
-  cluster    = google_container_cluster.ml_cluster.name
+  cluster    = var.enable_autopilot == false ? google_container_cluster.ml_cluster[0].name : null
   node_count = var.num_gpu_nodes
+  count      = var.enable_autopilot == false ? 1 : 0
 
   autoscaling {
     min_node_count = "1"

platform/modules/gke_cluster/main.tf platform/modules/gke_standard/main.tf

@@ -0,0 +1,39 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+output "project_id" {
+  description = "GCP project id"
+  value       = var.enable_autopilot ? null : resource.google_container_cluster.ml_cluster[0].project
+}
+
+output "region" {
+  description = "GCP region"
+  value       = var.enable_autopilot ? null : resource.google_container_cluster.ml_cluster[0].location
+}
+
+output "cluster_name" {
+  description = "The name of the GKE cluster"
+  value       = var.enable_autopilot ? null : resource.google_container_cluster.ml_cluster[0].name
+}
+
+output "kubernetes_host" {
+  description = "Kubernetes cluster host"
+  value       = var.enable_autopilot ? null : resource.google_container_cluster.ml_cluster[0].endpoint
+}
+
+output "cluster_certicicate" {
+  description = "Kubernetes cluster ca certificate"
+  value       = var.enable_autopilot ? null : base64decode(resource.google_container_cluster.ml_cluster[0].master_auth[0].cluster_ca_certificate)
+  sensitive   = true
+}

platform/modules/gke_standard/output.tf

@@ -0,0 +1,48 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+variable "project_id" {
+  type        = string
+  description = "GCP project id"
+  default     = "ricliu-gke-dev"
+}
+
+variable "region" {
+  type        = string
+  description = "GCP project region or zone"
+  default     = "us-central1"
+}
+
+variable "cluster_name" {
+  type        = string
+  description = "GKE cluster name"
+  default     = "ml-cluster"
+}
+
+variable "namespace" {
+  type        = string
+  description = "Kubernetes namespace where resources are deployed"
+  default     = "ray"
+}
+
+variable "num_gpu_nodes" {
+  description = "Number of GPU nodes in the cluster"
+  default     = 1
+}
+
+variable "enable_autopilot" {
+  type        = bool
+  description = "Set to true to enable GKE Autopilot clusters"
+  default     = false
+}

platform/modules/gke_standard/variables.tf

@@ -0,0 +1,21 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+terraform {
+  required_providers {
+  }
+  provider_meta "google" {
+    module_name = "blueprints/terraform/terraform-google-kubernetes-engine:kuberay/v0.1.0"
+  }
+}

platform/modules/gke_standard/versions.tf

@@ -18,4 +18,5 @@ data "http" "nvidia_driver_installer_manifest" {
 
 resource "kubectl_manifest" "nvidia_driver_installer" {
   yaml_body = data.http.nvidia_driver_installer_manifest.response_body
+  count = var.enable_autopilot == false ? 1 : 0
 }

platform/modules/kubernetes/kubernetes.tf

@@ -29,3 +29,9 @@ variable "namespace" {
   description = "Kubernetes namespace where resources are deployed"
   default     = "ray"
 }
+
+variable "enable_autopilot" {
+  type        = bool
+  description = "Set to true to enable GKE Autopilot clusters"
+  default     = false
+}

platform/modules/kubernetes/variables.tf

@@ -21,11 +21,17 @@ variable "project_id" {
 variable "region" {
   type        = string
   description = "GCP project region or zone"
-  default     = "us-central1-c"
+  default     = "us-central1"
 }
 
 variable "cluster_name" {
   type        = string
   description = "GKE cluster name"
   default     = "ml-cluster"
 }
+
+variable "enable_autopilot" {
+  type        = bool
+  description = "Set to true to enable GKE Autopilot clusters"
+  default     = false
+}

platform/variables.tf

@@ -26,6 +26,24 @@ hub:
       password: password
     JupyterHub:
       authenticator_class: dummy
+  networkPolicy:
+    enabled: false
+
+prePuller:
+  hook:
+    enabled: false
+
+proxy:
+  chp:
+    networkPolicy:
+      enabled: false
+  traefik:
+    networkPolicy:
+      enabled: false
+
+scheduling:
+  userScheduler:
+    enabled: false
 
 singleuser:
   nodeSelector:
@@ -40,3 +58,7 @@ singleuser:
   # `cmd: null` allows the custom CMD of the Jupyter docker-stacks to be used
   # which performs further customization on startup.
   cmd: null
+  cloudMetadata:
+    blockWithIptables: false
+  networkPolicy:
+    enabled: false

user/modules/jupyterhub/jupyterhub-values.yaml

@@ -95,6 +95,7 @@ head:
   annotations: {}
   nodeSelector:
     iam.gke.io/gke-metadata-server-enabled: "true"
+    cloud.google.com/gke-accelerator: "nvidia-tesla-a100"
   tolerations: []
   affinity: {}
   # Ray container security context.
@@ -179,6 +180,7 @@ worker:
     key: value
   nodeSelector:
     iam.gke.io/gke-metadata-server-enabled: "true"
+    cloud.google.com/gke-accelerator: "nvidia-tesla-a100"
   tolerations: []
   affinity: {}
   # Ray container security context.