Merge branch 'main' of https://github.com/GoogleCloudPlatform/ai-on-gke…

… into makefile-and-readme-change
GoogleCloudPlatform · Feb 24, 2024 · 0eba732 · 0eba732
2 parents e629cdd + 9016cde
commit 0eba732
Show file tree

Hide file tree

Showing 43 changed files with 955 additions and 783 deletions.
diff --git a/applications/jupyter/README.md b/applications/jupyter/README.md
@@ -49,16 +49,16 @@ To use GCS, create a bucket with your username. For example, when authenticating
 If using this with the Ray module (`applications/ray/`), it is recommended to use the same k8s namespace
 for both i.e. set this to the same namespace as `applications/ray/workloads.tfvars`.
 
-| Variable               | Description       | Required |
-| ---------------------- |-------------------|:--------:|
-| project_id             | GCP Project Id    | Yes      |
-| cluster_name           | GKE Cluster Name  | Yes      |
-| cluster_location       | GCP Region        | Yes      |
-| cluster_membership_id  | Fleet membership name for GKE cluster. <br /> Required when using private clusters with Anthos Connect Gateway   | |
-| namespace              | The namespace that Jupyterhub and rest of the other resources will be installed in.        | Yes      |
-| gcs_bucket             | GCS bucket to be used for Jupyter storage        |       |
-| create_service_account | Create service accounts used for Workload Identity mapping       | Yes      |
-| gcp_service_account    | GCP service account used for Workload Identity mapping      | Yes      |
+| Variable                    | Description                                                                                                    | Required |
+|-----------------------------|----------------------------------------------------------------------------------------------------------------|:--------:|
+| project_id                  | GCP Project Id                                                                                                 | Yes      |
+| cluster_name                | GKE Cluster Name                                                                                               | Yes      |
+| cluster_location            | GCP Region                                                                                                     | Yes      |
+| cluster_membership_id       | Fleet membership name for GKE cluster. <br /> Required when using private clusters with Anthos Connect Gateway | |
+| namespace                   | The namespace that Jupyterhub and rest of the other resources will be installed in.                            | Yes      |
+| gcs_bucket                  | GCS bucket to be used for Jupyter storage                                                                      |       |
+| create_service_account      | Create service accounts used for Workload Identity mapping                                                     | Yes      |
+| gcp_and_k8s_service_account | GCP service account used for Workload Identity mapping and k8s sa attached with workload                       | Yes      |
 
 For variables under `Jupyterhub with IAP`, please see the section below 
 

diff --git a/applications/jupyter/main.tf b/applications/jupyter/main.tf
@@ -77,7 +77,7 @@ module "jupyterhub" {
 
   namespace                     = var.namespace
   create_service_account        = var.create_service_account
-  gcp_service_account           = var.gcp_service_account
+  gcp_and_k8s_service_account   = var.gcp_and_k8s_service_account
   gcs_bucket                    = var.gcs_bucket
 
   # IAP Auth parameters

diff --git a/applications/jupyter/variables.tf b/applications/jupyter/variables.tf
@@ -42,7 +42,7 @@ variable "create_service_account" {
   default     = true
 }
 
-variable "gcp_service_account" {
+variable "gcp_and_k8s_service_account" {
   type        = string
   description = "gcp service account"
 }
@@ -59,9 +59,10 @@ variable "default_backend_service" {
 }
 
 variable "members_allowlist" {
-  type    = string
-  default = ""
+  type    = list(string)
+  default = []
 }
+
 variable "add_auth" {
   type        = bool
   description = "Enable iap authentication on jupyterhub"

diff --git a/applications/jupyter/workloads.tfvars b/applications/jupyter/workloads.tfvars
@@ -35,7 +35,7 @@ create_service_account        = true
 gcp_service_account           = "jupyter-service-account"
 
 # Jupyterhub with IAP
-add_auth                = true
+add_auth                = false
 brand                   = "projects/<prj-number>/brands/<prj-number>"
 support_email           = "<email>"
 default_backend_service = "proxy-public"
@@ -45,4 +45,4 @@ url_domain_addr   = ""
 url_domain_name   = ""
 client_id         = ""
 client_secret     = ""
-members_allowlist = "user:<email>"
+members_allowlist = ["user:<email>"]
diff --git a/applications/rag/main.tf b/applications/rag/main.tf
@@ -66,39 +66,39 @@ provider "helm" {
 data "kubernetes_all_namespaces" "allns" {}
 
 module "kuberay-operator" {
-  source                  = "../../modules/kuberay-operator"
-  project_id              = var.project_id
-  create_namespace        = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.kubernetes_namespace)
-  namespace               = var.kubernetes_namespace
-  name                    = "kuberay-operator"
-  google_service_account  = var.ray_service_account
-  create_service_account  = var.create_ray_service_account
-  enable_autopilot        = data.google_container_cluster.default.enable_autopilot
+  source                 = "../../modules/kuberay-operator"
+  project_id             = var.project_id
+  create_namespace       = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.kubernetes_namespace)
+  namespace              = var.kubernetes_namespace
+  name                   = "kuberay-operator"
+  google_service_account = var.ray_service_account
+  create_service_account = var.create_ray_service_account
+  enable_autopilot       = data.google_container_cluster.default.enable_autopilot
 }
 
 module "gcs" {
-  source     = "../../modules/gcs"
-  project_id = var.project_id
+  source      = "../../modules/gcs"
+  project_id  = var.project_id
   bucket_name = var.gcs_bucket
 }
 
 module "cloudsql" {
   source     = "../../modules/cloudsql"
-  depends_on = [module.gcs,module.kuberay-operator]
+  depends_on = [module.gcs, module.kuberay-operator]
   project_id = var.project_id
   namespace  = var.kubernetes_namespace
 }
 
 module "jupyterhub" {
-  source      = "../../modules/jupyter"
+  source     = "../../modules/jupyter"
   depends_on = [module.kuberay-operator]
-  namespace   = var.kubernetes_namespace
-  project_id  = var.project_id
-  gcs_bucket  = var.gcs_bucket
-  add_auth    = false # TODO: Replace with IAP.
+  namespace  = var.kubernetes_namespace
+  project_id = var.project_id
+  gcs_bucket = var.gcs_bucket
+  add_auth   = false # TODO: Replace with IAP.
 
-  gcp_service_account    = var.jupyter_service_account
-  create_service_account = var.create_jupyter_service_account
+  gcp_and_k8s_service_account    = var.jupyter_service_account
+  create_service_account         = var.create_jupyter_service_account
 
   # IAP Auth parameters
   brand                   = var.brand
@@ -119,25 +119,26 @@ module "kuberay-logging" {
 }
 
 module "kuberay-cluster" {
-  source     = "../../modules/kuberay-cluster"
-  project_id  = var.project_id
-  depends_on = [module.kuberay-operator, module.gcs, module.kuberay-monitoring]
-  namespace  = var.kubernetes_namespace
-  gcs_bucket = var.gcs_bucket 
-  create_namespace    = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.kubernetes_namespace)
-  enable_tpu          = data.google_container_cluster.default.enable_tpu
-  enable_autopilot    = data.google_container_cluster.default.enable_autopilot
+  source                 = "../../modules/kuberay-cluster"
+  project_id             = var.project_id
+  depends_on             = [module.kuberay-operator, module.gcs, module.kuberay-monitoring]
+  namespace              = var.kubernetes_namespace
+  gcs_bucket             = var.gcs_bucket
+  create_namespace       = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.kubernetes_namespace)
+  enable_tpu             = data.google_container_cluster.default.enable_tpu
+  enable_autopilot       = data.google_container_cluster.default.enable_autopilot
   google_service_account = var.ray_service_account
-  grafana_host        = module.kuberay-monitoring.grafana_uri
+  grafana_host           = module.kuberay-monitoring.grafana_uri
 }
 
 module "kuberay-monitoring" {
-  source     = "../../modules/kuberay-monitoring"
-  depends_on = [module.kuberay-operator]
-  project_id = var.project_id
-  namespace  = var.kubernetes_namespace
-  create_namespace    = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.kubernetes_namespace)
-  k8s_service_account = var.ray_service_account
+  source                          = "../../modules/kuberay-monitoring"
+  depends_on                      = [module.kuberay-operator]
+  project_id                      = var.project_id
+  namespace                       = var.kubernetes_namespace
+  create_namespace                = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.kubernetes_namespace)
+  enable_grafana_on_ray_dashboard = var.enable_grafana_on_ray_dashboard
+  k8s_service_account             = var.ray_service_account
 }
 
 module "inference-server" {
@@ -147,15 +148,15 @@ module "inference-server" {
 }
 
 module "frontend" {
-  source            = "./frontend"
-  depends_on = [module.cloudsql, module.gcs, module.inference-server]
-  project_id = var.project_id
-  create_service_account = var.create_rag_service_account
-  google_service_account = var.rag_service_account
-  namespace  = var.kubernetes_namespace
-  inference_service_name = module.inference-server.inference_service_name
-  inference_service_namespace = module.inference-server.inference_service_namespace
-  db_secret_name = module.cloudsql.db_secret_name
-  db_secret_namespace = module.cloudsql.db_secret_namespace
+  source                        = "./frontend"
+  depends_on                    = [module.cloudsql, module.gcs, module.inference-server]
+  project_id                    = var.project_id
+  create_service_account        = var.create_rag_service_account
+  google_service_account        = var.rag_service_account
+  namespace                     = var.kubernetes_namespace
+  inference_service_name        = module.inference-server.inference_service_name
+  inference_service_namespace   = module.inference-server.inference_service_namespace
+  db_secret_name                = module.cloudsql.db_secret_name
+  db_secret_namespace           = module.cloudsql.db_secret_namespace
   dataset_embeddings_table_name = var.dataset_embeddings_table_name
 }
diff --git a/applications/rag/variables.tf b/applications/rag/variables.tf
@@ -49,6 +49,11 @@ variable "jupyter_service_account" {
   default     = "jupyter-system-account"
 }
 
+variable "enable_grafana_on_ray_dashboard" {
+  type        = bool
+  description = "Add option to enable or disable grafana for the ray dashboard. Enabling requires anonymous access."
+  default     = false
+}
 variable "create_ray_service_account" {
   type        = bool
   description = "Creates a google IAM service account & k8s service account & configures workload identity"
@@ -89,9 +94,10 @@ variable "default_backend_service" {
 }
 
 variable "members_allowlist" {
-  type    = string
-  default = ""
+  type    = list(string)
+  default = []
 }
+
 variable "add_auth" {
   type        = bool
   description = "Enable iap authentication on jupyterhub"
@@ -138,4 +144,4 @@ variable "client_secret" {
   description = "Client secret used for enabling IAP"
   default     = ""
   sensitive   = false
-}
+}
diff --git a/applications/rag/workloads.tfvars b/applications/rag/workloads.tfvars
@@ -12,31 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-project_id = "<project_id>"
+project_id = "<your project ID>"
 
 ## this is required for terraform to connect to GKE master and deploy workloads
 cluster_name     = "<cluster_name>"
 cluster_location = "us-central1"
 
 ## GKE environment variables
-kubernetes_namespace      = "rag"
-gcs_bucket                = "rag-data-xyzu" # Choose a globally unique bucket name.
+kubernetes_namespace = "rag"
+gcs_bucket           = "rag-data-xyzu" # Choose a globally unique bucket name.
 
 ## Service accounts
 # Creates a google service account & k8s service account & configures workload identity with appropriate permissions.
 # Set to false & update the variable `ray_service_account` to use an existing IAM service account.
-create_ray_service_account = true
-ray_service_account        = "ray-system-account"
-
+create_ray_service_account      = true
+ray_service_account             = "ray-system-account"
+enable_grafana_on_ray_dashboard = false
 # Creates a google service account & k8s service account & configures workload identity with appropriate permissions.
 # Set to false & update the variable `rag_service_account` to use an existing IAM service account.
 create_rag_service_account = true
 rag_service_account        = "rag-system-account"
 
 # Creates a google service account & k8s service account & configures workload identity with appropriate permissions.
 # Set to false & update the variable `jupyter_service_account` to use an existing IAM service account.
-create_jupyter_service_account  = true
-jupyter_service_account = "jupyter-system-account"
+create_jupyter_service_account = true
+jupyter_service_account        = "jupyter-system-account"
 
 ## Embeddings table name - change this to the TABLE_NAME used in the notebook.
 dataset_embeddings_table_name = "googlemaps_reviews_db"
@@ -52,4 +52,4 @@ url_domain_addr   = ""
 url_domain_name   = ""
 client_id         = ""
 client_secret     = ""
-members_allowlist = "allAuthenticatedUsers,user:<email>"
+members_allowlist = ["allAuthenticatedUsers", "user:<email>"]
diff --git a/applications/ray/main.tf b/applications/ray/main.tf
@@ -68,50 +68,51 @@ provider "helm" {
 data "kubernetes_all_namespaces" "allns" {}
 
 module "kuberay-operator" {
-  source           = "../../modules/kuberay-operator"
-  name             = "kuberay-operator"
-  create_namespace = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.ray_namespace)
-  namespace        = var.ray_namespace
-  project_id       = var.project_id
-  enable_autopilot = data.google_container_cluster.default.enable_autopilot
+  source                 = "../../modules/kuberay-operator"
+  name                   = "kuberay-operator"
+  create_namespace       = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.ray_namespace)
+  namespace              = var.ray_namespace
+  project_id             = var.project_id
+  enable_autopilot       = data.google_container_cluster.default.enable_autopilot
   google_service_account = var.gcp_service_account
   create_service_account = var.create_service_account
 }
 
 module "kuberay-logging" {
-  source     = "../../modules/kuberay-logging"
-  namespace  = var.ray_namespace
+  source    = "../../modules/kuberay-logging"
+  namespace = var.ray_namespace
 
   depends_on = [module.kuberay-operator]
 }
 
 module "kuberay-monitoring" {
-  count               = var.create_ray_cluster == true ? 1 : 0
-  source              = "../../modules/kuberay-monitoring"
-  project_id          = var.project_id
-  namespace           = var.ray_namespace
-  create_namespace    = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.ray_namespace)
-  k8s_service_account = var.gcp_service_account
-  depends_on = [module.kuberay-operator]
+  count                           = var.create_ray_cluster ? 1 : 0
+  source                          = "../../modules/kuberay-monitoring"
+  project_id                      = var.project_id
+  namespace                       = var.ray_namespace
+  create_namespace                = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.ray_namespace)
+  enable_grafana_on_ray_dashboard = var.enable_grafana_on_ray_dashboard
+  k8s_service_account             = var.gcp_service_account
+  depends_on                      = [module.kuberay-operator]
 }
 
 module "gcs" {
-  source     = "../../modules/gcs"
-  project_id = var.project_id
+  source      = "../../modules/gcs"
+  project_id  = var.project_id
   bucket_name = var.gcs_bucket
 }
 
 module "kuberay-cluster" {
-  count                   = var.create_ray_cluster == true ? 1 : 0
-  source                  = "../../modules/kuberay-cluster"
-  create_namespace        = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.ray_namespace)
-  namespace               = var.ray_namespace
-  project_id              = var.project_id
-  enable_tpu              = data.google_container_cluster.default.enable_tpu
-  gcs_bucket              = var.gcs_bucket
-  enable_autopilot        = data.google_container_cluster.default.enable_autopilot
-  google_service_account  = var.gcp_service_account
-  grafana_host            = module.kuberay-monitoring[0].grafana_uri
-  depends_on              = [module.kuberay-monitoring, module.gcs]
+  count                  = var.create_ray_cluster == true ? 1 : 0
+  source                 = "../../modules/kuberay-cluster"
+  create_namespace       = !contains(data.kubernetes_all_namespaces.allns.namespaces, var.ray_namespace)
+  namespace              = var.ray_namespace
+  project_id             = var.project_id
+  enable_tpu             = data.google_container_cluster.default.enable_tpu
+  gcs_bucket             = var.gcs_bucket
+  enable_autopilot       = data.google_container_cluster.default.enable_autopilot
+  google_service_account = var.gcp_service_account
+  grafana_host           = var.enable_grafana_on_ray_dashboard ? module.kuberay-monitoring[0].grafana_uri : ""
+  depends_on             = [module.kuberay-monitoring, module.gcs]
 }
 
diff --git a/...rain-examples/raytrain-with-gcsfusecsi/kuberaytf/user/modules/kuberay/kuberay-values.yaml b/...rain-examples/raytrain-with-gcsfusecsi/kuberaytf/user/modules/kuberay/kuberay-values.yaml
@@ -56,6 +56,7 @@ head:
     #     memory: "512Mi"
   labels:
     cloud.google.com/gke-ray-node-type: head
+    created-by: ray-on-gke
   serviceAccountName: my-ksa
   rayStartParams:
     dashboard-host: '0.0.0.0'
@@ -130,6 +131,7 @@ worker:
   type: worker
   labels:
     cloud.google.com/gke-ray-node-type: worker
+    created-by: ray-on-gke
   serviceAccountName: my-ksa
   rayStartParams:
     block: 'true'

diff --git a/applications/ray/variables.tf b/applications/ray/variables.tf
@@ -42,6 +42,12 @@ variable "ray_namespace" {
   default     = "myray"
 }
 
+variable "enable_grafana_on_ray_dashboard" {
+  type    = bool
+  description = "Add option to enable or disable grafana for the ray dashboard. Enabling requires anonymous access."
+  default = false  
+}
+
 variable "gcs_bucket" {
   type = string
 }