feat: Adding bigtable and pubsub support for GCP (#178)

* feat: Adding bigtable and pubsub support for GCP

* feat: CMEK permissions for bigtable and pubsub

* fix: Add edition for cloudsql and remove lots of hard coded values from the TF

* fix: add executor and filestream support

* feat: Remove unused public address, don't assign nodes public addresses

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Venkata <venkata@wandb.com>
Co-authored-by: Andrew Levin <andrew.levin@wandb.com>
Co-authored-by: Daniel Barnes <dabarnes2b@gmail.com>

README.md

@@ -23,6 +23,7 @@ preparation, however it does have the following pre-requisites:
 - Google Kubernetes Engine
 - Google Storage Bucket
 - Google PubSub
+- Google BigTable
 - Google Managed Certificates
 - Google Cloud DNS
 
@@ -95,6 +96,7 @@ resources that lack official modules.
 |------|--------|---------|
 | <a name="module_app_gke"></a> [app\_gke](#module\_app\_gke) | ./modules/app_gke | n/a |
 | <a name="module_app_lb"></a> [app\_lb](#module\_app\_lb) | ./modules/app_lb | n/a |
+| <a name="module_bigtable"></a> [bigtable](#module\_bigtable) | ./modules/bigtable | n/a |
 | <a name="module_clickhouse"></a> [clickhouse](#module\_clickhouse) | ./modules/clickhouse | n/a |
 | <a name="module_cloud_nat"></a> [cloud\_nat](#module\_cloud\_nat) | ./modules/cloud_nat | n/a |
 | <a name="module_database"></a> [database](#module\_database) | ./modules/database | n/a |
@@ -105,6 +107,7 @@ resources that lack official modules.
 | <a name="module_networking"></a> [networking](#module\_networking) | ./modules/networking | n/a |
 | <a name="module_private_link"></a> [private\_link](#module\_private\_link) | ./modules/private_link | n/a |
 | <a name="module_project_factory_project_services"></a> [project\_factory\_project\_services](#module\_project\_factory\_project\_services) | terraform-google-modules/project-factory/google//modules/project_services | ~> 14.0 |
+| <a name="module_pubsub"></a> [pubsub](#module\_pubsub) | ./modules/pubsub | n/a |
 | <a name="module_redis"></a> [redis](#module\_redis) | ./modules/redis | n/a |
 | <a name="module_service_accounts"></a> [service\_accounts](#module\_service\_accounts) | ./modules/service_accounts | n/a |
 | <a name="module_storage"></a> [storage](#module\_storage) | ./modules/storage | n/a |
@@ -123,6 +126,10 @@ resources that lack official modules.
 | <a name="input_allowed_inbound_cidrs"></a> [allowed\_inbound\_cidrs](#input\_allowed\_inbound\_cidrs) | Which IPv4 addresses/ranges to allow access. This must be explicitly provided, and by default is set to ["*"] | `list(string)` | <pre>[<br/>  "*"<br/>]</pre> | no |
 | <a name="input_allowed_project_names"></a> [allowed\_project\_names](#input\_allowed\_project\_names) | A map of allowed projects where each key is a project number and the value is the connection limit. | `map(number)` | `{}` | no |
 | <a name="input_app_wandb_env"></a> [app\_wandb\_env](#input\_app\_wandb\_env) | Extra environment variables for W&B | `map(string)` | `{}` | no |
+| <a name="input_bigtable_cpu_target"></a> [bigtable\_cpu\_target](#input\_bigtable\_cpu\_target) | The target CPU utilization for the Bigtable cluster. | `number` | `70` | no |
+| <a name="input_bigtable_max_nodes"></a> [bigtable\_max\_nodes](#input\_bigtable\_max\_nodes) | The maximum number of nodes for the Bigtable cluster. | `number` | `3` | no |
+| <a name="input_bigtable_min_nodes"></a> [bigtable\_min\_nodes](#input\_bigtable\_min\_nodes) | The minimum number of nodes for the Bigtable cluster. | `number` | `1` | no |
+| <a name="input_bigtable_storage_type"></a> [bigtable\_storage\_type](#input\_bigtable\_storage\_type) | The storage type for the Bigtable cluster. | `string` | `"SSD"` | no |
 | <a name="input_bucket_default_encryption"></a> [bucket\_default\_encryption](#input\_bucket\_default\_encryption) | Boolean to determine if a default bucket encryption key should be used. If true, a default key will be created. Takes precedence over `bucket_kms_key_id`. | `bool` | `false` | no |
 | <a name="input_bucket_kms_key_id"></a> [bucket\_kms\_key\_id](#input\_bucket\_kms\_key\_id) | ID of the customer-provided bucket KMS key. | `string` | `null` | no |
 | <a name="input_bucket_location"></a> [bucket\_location](#input\_bucket\_location) | Location of the bucket (US, EU, ASIA) | `string` | `"US"` | no |
@@ -132,9 +139,12 @@ resources that lack official modules.
 | <a name="input_clickhouse_region"></a> [clickhouse\_region](#input\_clickhouse\_region) | ClickHouse region (us-east1, us-central1, etc). | `string` | `""` | no |
 | <a name="input_clickhouse_subnetwork_cidr"></a> [clickhouse\_subnetwork\_cidr](#input\_clickhouse\_subnetwork\_cidr) | ClickHouse private service connect subnetwork | `string` | `"10.50.0.0/24"` | no |
 | <a name="input_controller_image_tag"></a> [controller\_image\_tag](#input\_controller\_image\_tag) | Tag of the controller image to deploy | `string` | `"1.14.0"` | no |
+| <a name="input_create_bigtable"></a> [create\_bigtable](#input\_create\_bigtable) | Boolean indicating whether to provision a bigtable instance (true) or not (false). | `bool` | `false` | no |
 | <a name="input_create_private_link"></a> [create\_private\_link](#input\_create\_private\_link) | Whether to create a private link service. | `bool` | `false` | no |
+| <a name="input_create_pubsub"></a> [create\_pubsub](#input\_create\_pubsub) | Boolean indicating whether to provision a bigtable instance (true) or not (false). | `bool` | `false` | no |
 | <a name="input_create_redis"></a> [create\_redis](#input\_create\_redis) | Boolean indicating whether to provision an redis instance (true) or not (false). | `bool` | `false` | no |
 | <a name="input_create_workload_identity"></a> [create\_workload\_identity](#input\_create\_workload\_identity) | Flag to indicate whether to create a workload identity for the service account. | `bool` | `false` | no |
+| <a name="input_database_edition"></a> [database\_edition](#input\_database\_edition) | The edition of the Cloud SQL instance. Can be either `STANDARD` or `ENTERPRISE` or `ENTERPRISE_PLUS`. | `string` | `"ENTERPRISE"` | no |
 | <a name="input_database_flags"></a> [database\_flags](#input\_database\_flags) | Flags to set for the database | `map(string)` | `{}` | no |
 | <a name="input_database_machine_type"></a> [database\_machine\_type](#input\_database\_machine\_type) | Specifies the machine type to be allocated for the database. Defaults to null and value from deployment-size.tf is used | `string` | `null` | no |
 | <a name="input_database_sort_buffer_size"></a> [database\_sort\_buffer\_size](#input\_database\_sort\_buffer\_size) | Specifies the sort\_buffer\_size value to set for the database | `number` | `67108864` | no |
@@ -164,7 +174,7 @@ resources that lack official modules.
 | <a name="input_oidc_client_id"></a> [oidc\_client\_id](#input\_oidc\_client\_id) | The Client ID of application in your identity provider | `string` | `""` | no |
 | <a name="input_oidc_issuer"></a> [oidc\_issuer](#input\_oidc\_issuer) | A url to your Open ID Connect identity provider, i.e. https://cognito-idp.us-east-1.amazonaws.com/us-east-1_uiIFNdacd | `string` | `""` | no |
 | <a name="input_oidc_secret"></a> [oidc\_secret](#input\_oidc\_secret) | The Client secret of application in your identity provider | `string` | `""` | no |
-| <a name="input_operator_chart_version"></a> [operator\_chart\_version](#input\_operator\_chart\_version) | Version of the operator chart to deploy | `string` | `"1.3.4"` | no |
+| <a name="input_operator_chart_version"></a> [operator\_chart\_version](#input\_operator\_chart\_version) | Version of the operator chart to deploy | `string` | `"1.3.6"` | no |
 | <a name="input_other_wandb_env"></a> [other\_wandb\_env](#input\_other\_wandb\_env) | Extra environment variables for W&B | `map(string)` | `{}` | no |
 | <a name="input_parquet_wandb_env"></a> [parquet\_wandb\_env](#input\_parquet\_wandb\_env) | Extra environment variables for W&B | `map(string)` | `{}` | no |
 | <a name="input_psc_subnetwork_cidr"></a> [psc\_subnetwork\_cidr](#input\_psc\_subnetwork\_cidr) | Private link service reserved subnetwork | `string` | `"192.168.0.0/24"` | no |

main.tf

@@ -27,8 +27,10 @@ locals {
   create_network = var.network == null
   k8s_sa_map = {
     app                    = "wandb-app"
+    executor               = "wandb-executor"
     parquet                = "wandb-parquet"
     flat_runs              = "wandb-flat-run-fields-updater"
+    filestream             = "wandb-filestream"
     weave                  = "wandb-weave"
     weave_trace            = "wandb-weave-trace"
     settings_migration_job = "wandb-settings-migration-job"
@@ -54,6 +56,10 @@ module "service_accounts" {
   skip_bucket_admin_role   = var.skip_bucket_admin_role
 }
 
+locals {
+  app_service_account = (var.create_workload_identity) ? module.service_accounts.sa_account_role : module.service_accounts.service_account.email
+}
+
 module "kms" {
   # KMS is currently only used to encrypt pubsub queue. Disable it if we dont use it.
   count               = var.use_internal_queue ? 0 : 1
@@ -63,21 +69,23 @@ module "kms" {
 }
 
 module "kms_default_bucket" {
-  count                          = var.bucket_default_encryption ? 1 : 0
-  source                         = "./modules/kms"
-  namespace                      = var.namespace
-  deletion_protection            = var.deletion_protection
-  key_location                   = lower(var.bucket_location)
-  bind_pubsub_service_to_kms_key = false
+  count                            = var.bucket_default_encryption ? 1 : 0
+  source                           = "./modules/kms"
+  namespace                        = var.namespace
+  deletion_protection              = var.deletion_protection
+  key_location                     = lower(var.bucket_location)
+  bind_pubsub_service_to_kms_key   = false
+  bind_bigtable_service_to_kms_key = false
 }
 
 module "kms_default_sql" {
-  count                          = var.sql_default_encryption ? 1 : 0
-  source                         = "./modules/kms"
-  namespace                      = var.namespace
-  deletion_protection            = var.deletion_protection
-  key_location                   = data.google_client_config.current.region
-  bind_pubsub_service_to_kms_key = false
+  count                            = var.sql_default_encryption ? 1 : 0
+  source                           = "./modules/kms"
+  namespace                        = var.namespace
+  deletion_protection              = var.deletion_protection
+  key_location                     = data.google_client_config.current.region
+  bind_pubsub_service_to_kms_key   = var.create_pubsub
+  bind_bigtable_service_to_kms_key = var.create_bigtable
 }
 locals {
   default_bucket_key = length(module.kms_default_bucket) > 0 ? module.kms_default_bucket[0].crypto_key.id : var.bucket_kms_key_id
@@ -162,6 +170,7 @@ module "database" {
   database_version    = var.database_version
   force_ssl           = var.force_ssl
   tier                = local.database_machine_type
+  edition             = var.database_edition
   database_flags      = var.database_flags
   sort_buffer_size    = var.database_sort_buffer_size
   network_connection  = local.network_connection
@@ -171,6 +180,34 @@ module "database" {
   depends_on          = [module.project_factory_project_services, module.kms_default_sql]
 }
 
+module "bigtable" {
+  source = "./modules/bigtable"
+  count  = var.create_bigtable ? 1 : 0
+
+  namespace             = var.namespace
+  deletion_protection   = var.deletion_protection
+  service_account_email = local.app_service_account
+  crypto_key            = local.default_sql_key
+  storage_type          = var.bigtable_storage_type
+  cpu_target            = var.bigtable_cpu_target
+  min_nodes             = var.bigtable_min_nodes
+  max_nodes             = var.bigtable_max_nodes
+
+  labels = var.labels
+}
+
+module "pubsub" {
+  source = "./modules/pubsub"
+  count  = var.create_pubsub ? 1 : 0
+
+  namespace             = var.namespace
+  deletion_protection   = var.deletion_protection
+  service_account_email = local.app_service_account
+  crypto_key            = local.default_sql_key
+
+  labels = var.labels
+}
+
 module "redis" {
   count     = var.create_redis ? 1 : 0
   source    = "./modules/redis"
@@ -291,6 +328,18 @@ module "wandb" {
           "TAG_CUSTOMER_NS"                      = var.namespace
         }, var.other_wandb_env, local.oidc_envs)
 
+        bigtable = {
+          project  = local.project_id
+          instance = var.create_bigtable ? module.bigtable[0].bigtable_instance_id : ""
+        }
+
+        pubSub = {
+          enabled              = var.create_pubsub
+          project              = local.project_id
+          filestreamTopic      = var.create_pubsub ? module.pubsub[0].filestream_topic_name : ""
+          runUpdateShadowTopic = var.create_pubsub ? module.pubsub[0].run_updates_shadow_topic_name : ""
+        }
+
         bucket = var.bucket_name != "" ? {
           provider = "gcs"
           name     = var.bucket_name
@@ -375,7 +424,9 @@ module "wandb" {
           projectId          = data.google_client_config.current.project
           serviceAccountName = var.stackdriver_sa_name
         }
-        serviceAccount = { annotations = { "iam.gke.io/gcp-service-account" = module.service_accounts.stackdriver_role } }
+        serviceAccount = {
+          annotations = { "iam.gke.io/gcp-service-account" = module.service_accounts.stackdriver_role }
+        }
         } : {
         install        = false
         stackdriver    = {}
@@ -417,6 +468,16 @@ module "wandb" {
         }
       }
 
+      executor = {
+        serviceAccount = var.create_workload_identity ? {
+          name        = local.k8s_sa_map.executor
+          annotations = { "iam.gke.io/gcp-service-account" = module.service_accounts.sa_account_role }
+          } : {
+          name        = null
+          annotations = {}
+        }
+      }
+
       settingsMigrationJob = {
         serviceAccount = var.create_workload_identity ? {
           annotations = { "iam.gke.io/gcp-service-account" = module.service_accounts.sa_account_role }
@@ -448,6 +509,9 @@ module "wandb" {
       }
 
       flat-run-fields-updater = {
+        pubSub = {
+          subscription = var.create_pubsub ? module.pubsub[0].flat_run_fields_updater_subscription_name : ""
+        }
         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 }
@@ -456,6 +520,19 @@ module "wandb" {
           annotations = {}
         }
       }
+
+      filestream = {
+        pubSub = {
+          subscription = var.create_pubsub ? module.pubsub[0].filestream_gorilla_subscription_name : ""
+        }
+        serviceAccount = var.create_workload_identity ? {
+          name        = local.k8s_sa_map.filestream
+          annotations = { "iam.gke.io/gcp-service-account" = module.service_accounts.sa_account_role }
+          } : {
+          name        = null
+          annotations = {}
+        }
+      }
     }
   }
 

modules/app_gke/main.tf

@@ -96,6 +96,7 @@ resource "google_container_node_pool" "default" {
       "https://www.googleapis.com/auth/trace.append",
       "https://www.googleapis.com/auth/sqlservice.admin",
     ]
+
     workload_metadata_config {
       mode = var.create_workload_identity ? "GKE_METADATA" : "GCE_METADATA"
     }

modules/app_lb/https/variables.tf

@@ -3,10 +3,6 @@ variable "namespace" {
   description = "Friendly name prefix used for tagging and naming AWS resources."
 }
 
-variable "ip_address" {
-  type = string
-}
-
 variable "fqdn" {
   type = string
 }

modules/app_lb/main.tf

@@ -1,18 +1,13 @@
-resource "google_compute_global_address" "default" {
-  name = "${var.namespace}-address"
-}
-
 resource "google_compute_global_address" "operator" {
   name = "${var.namespace}-operator-address"
 }
 
 module "https" {
   count = var.ssl ? 1 : 0
 
-  source     = "./https"
-  fqdn       = var.fqdn
-  namespace  = var.namespace
-  ip_address = google_compute_global_address.default.address
+  source    = "./https"
+  fqdn      = var.fqdn
+  namespace = var.namespace
 
   labels = var.labels
 }

modules/app_lb/outputs.tf

@@ -1,7 +1,3 @@
-output "address" {
-  value = google_compute_global_address.default.address
-}
-
 output "address_operator" {
   value = google_compute_global_address.operator.address
 }