feat!: Service Account Mapping (#151)

* feat!: Service Account Mapping * terraform-docs: automated action * fix: flat run not runs --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Daniel Panzella <[email protected]>
wandb · Jul 18, 2024 · 8930eaf · 8930eaf
1 parent d6b6ec7
commit 8930eaf
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 27 deletions.
diff --git a/README.md b/README.md
@@ -119,7 +119,6 @@ resources that lack official modules.
 | <a name="input_gke_machine_type"></a> [gke\_machine\_type](#input\_gke\_machine\_type) | Specifies the machine type to be allocated for the database | `string` | `"n1-standard-4"` | no |
 | <a name="input_gke_node_count"></a> [gke\_node\_count](#input\_gke\_node\_count) | n/a | `number` | `2` | no |
 | <a name="input_ilb_proxynetwork_cidr"></a> [ilb\_proxynetwork\_cidr](#input\_ilb\_proxynetwork\_cidr) | Internal load balancer proxy subnetwork | `string` | `"10.127.0.0/24"` | no |
-| <a name="input_kms_gcs_sa_name"></a> [kms\_gcs\_sa\_name](#input\_kms\_gcs\_sa\_name) | n/a | `string` | `"wandb-app"` | no |
 | <a name="input_labels"></a> [labels](#input\_labels) | Labels to apply to resources | `map(string)` | `{}` | no |
 | <a name="input_license"></a> [license](#input\_license) | Your wandb/local license | `string` | n/a | yes |
 | <a name="input_local_restore"></a> [local\_restore](#input\_local\_restore) | Restores W&B to a stable state if needed | `bool` | `false` | no |

diff --git a/main.tf b/main.tf
@@ -26,13 +26,18 @@ locals {
   url            = "${local.url_prefix}://${local.fqdn}"
   create_bucket  = var.bucket_name == ""
   create_network = var.network == null
+  k8s_sa_map = {
+    app       = "wandb-app"
+    parquet   = "wandb-parquet"
+    flat_runs = "wandb-flat-run-fields-updater"
+  }
 }
 
 module "service_accounts" {
   source                   = "./modules/service_accounts"
   namespace                = var.namespace
   bucket_name              = var.bucket_name
-  kms_gcs_sa_name          = var.kms_gcs_sa_name
+  kms_gcs_sa_list          = values(local.k8s_sa_map)
   create_workload_identity = var.create_workload_identity
   stackdriver_sa_name      = var.stackdriver_sa_name
   enable_stackdriver       = var.enable_stackdriver
@@ -212,9 +217,9 @@ module "wandb" {
   spec = {
     values = {
       global = {
-        pod = { labels = { workload_hash: local.workload_hash } }
-        host    = local.url
-        license = var.license
+        pod           = { labels = { workload_hash : local.workload_hash } }
+        host          = local.url
+        license       = var.license
         cloudProvider = "gcp"
         extraEnv = merge({
           "GORILLA_DISABLE_CODE_SAVING"          = var.disable_code_saving,
@@ -262,7 +267,7 @@ module "wandb" {
       app = {
         extraEnvs = var.app_wandb_env
         serviceAccount = var.create_workload_identity ? {
-          name        = var.kms_gcs_sa_name
+          name        = local.k8s_sa_map.app
           annotations = { "iam.gke.io/gcp-service-account" = module.service_accounts.sa_account_role }
           } : {
           name        = ""
@@ -348,6 +353,23 @@ module "wandb" {
 
       parquet = {
         extraEnvs = var.parquet_wandb_env
+        serviceAccount = var.create_workload_identity ? {
+          name        = local.k8s_sa_map.parquet
+          annotations = { "iam.gke.io/gcp-service-account" = module.service_accounts.sa_account_role }
+          } : {
+          name        = null
+          annotations = {}
+        }
+      }
+
+      flat-runs-fields-updater = {
+        serviceAccount = var.create_workload_identity ? {
+          name        = local.k8s_sa_map.flat_runs
+          annotations = { "iam.gke.io/gcp-service-account" = module.service_accounts.sa_account_role }
+          } : {
+          name        = null
+          annotations = {}
+        }
       }
     }
   }
@@ -418,4 +440,4 @@ module "private_link" {
   psc_subnetwork        = var.psc_subnetwork_cidr
   proxynetwork_cidr     = var.ilb_proxynetwork_cidr
   depends_on            = [google_compute_subnetwork.proxy, data.google_compute_forwarding_rules.all]
-}
+}
diff --git a/modules/app_gke/main.tf b/modules/app_gke/main.tf
@@ -77,14 +77,9 @@ resource "google_container_node_pool" "default" {
       "https://www.googleapis.com/auth/trace.append",
       "https://www.googleapis.com/auth/sqlservice.admin",
     ]
-
-     dynamic "workload_metadata_config" {
-       for_each = var.create_workload_identity == true ? [1] : []
-        content {
-          mode = "GKE_METADATA"
-        }
-      }
-
+    workload_metadata_config {
+      mode = var.create_workload_identity ? "GKE_METADATA" : "GCE_METADATA"
+    }
     shielded_instance_config {
       enable_secure_boot = true
     }

diff --git a/modules/service_accounts/main.tf b/modules/service_accounts/main.tf
@@ -120,13 +120,12 @@ resource "google_service_account_iam_member" "token_creator_binding" {
 }
 
 resource "google_service_account_iam_member" "workload_binding" {
-  count              = var.create_workload_identity == true ? 1 : 0
+  for_each           = var.create_workload_identity ? { for sa in var.kms_gcs_sa_list : sa => sa } : {}
   service_account_id = google_service_account.kms_gcs_sa[0].id
   role               = "roles/iam.workloadIdentityUser"
-  member             = "serviceAccount:${local.project_id}.svc.id.goog[default/${var.kms_gcs_sa_name}]"
+  member             = "serviceAccount:${local.project_id}.svc.id.goog[default/${each.value}]"
 }
 
-
 ### service account for stackdriver
 resource "google_service_account" "stackdriver" {
   count        = var.enable_stackdriver == true ? 1 : 0

diff --git a/modules/service_accounts/variables.tf b/modules/service_accounts/variables.tf
@@ -14,8 +14,8 @@ variable "create_workload_identity" {
   type        = bool
 }
 
-variable "kms_gcs_sa_name" {
-  type    = string
+variable "kms_gcs_sa_list" {
+  type = list(string)
 }
 
 variable "stackdriver_sa_name" {
@@ -26,4 +26,4 @@ variable "stackdriver_sa_name" {
 variable "enable_stackdriver" {
   description = "Flag to indicate whether to enable workload identity for the service account."
   type        = bool
-}
+}
diff --git a/variables.tf b/variables.tf
@@ -269,7 +269,7 @@ variable "create_private_link" {
 variable "public_access" {
   type        = bool
   description = "Whether to create a public endpoint for wandb access."
-  default = true
+  default     = true
 }
 
 variable "allowed_project_names" {
@@ -298,13 +298,9 @@ variable "create_workload_identity" {
   default     = false
 }
 
-variable "kms_gcs_sa_name" {
-  type    = string
-  default = "wandb-app"
-}
 
 variable "enable_stackdriver" {
-  type = bool
+  type    = bool
   default = false
 }