chore: add the bigtable tables and updated the envvars

wandb · Nov 16, 2024 · a73c3bc · a73c3bc
1 parent 454b376
commit a73c3bc
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 5 deletions.
diff --git a/main.tf b/main.tf
@@ -270,10 +270,13 @@ locals {
     "OIDC_SECRET"      = var.oidc_secret
   } : {}
   internal_lb_name = "${var.namespace}-internal"
-  filestream_envs = (var.create_bigtable && var.create_pubsub) ? {
-    "GORILLA_PARQUET_LIVE_HISTORY_STORE" = "bigtable://${module.bigtable[0].bigtable_project_id}/${module.bigtable[0].bigtable_instance_id}"
-    "GORILLA_HISTORY_STORE"              = "http://wandb-parquet:8087/_goRPC_,bigtable://${module.bigtable[0].bigtable_project_id}/${module.bigtable[0].bigtable_instance_id}"
+  bigtable_url     = var.create_bigtable ? "bigtablev2://${module.bigtable[0].bigtable_project_id}/${module.bigtable[0].bigtable_instance_id}" : ""
+  filestream_envs = var.create_bigtable && var.create_pubsub ? {
+    "GORILLA_PARQUET_LIVE_HISTORY_STORE" = "bigtablev2://${module.bigtable[0].bigtable_project_id}/${module.bigtable[0].bigtable_instance_id}"
     "GORILLA_FILE_STREAM_STORE_ADDRESS"  = "pubsub:/${module.pubsub[0].filestream_project_id}/${module.pubsub[0].filestream_topic_name}"
+    "GORILLA_HISTORY_STORE"              = <<-EOF
+      ${join("\\,", ["http://wandb-parquet:8087/_goRPC_", local.bigtable_url])}
+    EOF
   } : {}
 }
 
@@ -433,10 +436,10 @@ module "wandb" {
       }
 
       filestream = {
-        enabled = var.create_pubsub && var.create_bigtable
+        install = var.create_pubsub && var.create_bigtable
         env = {
           fileStreamWorkerSource = "pubsub:/${module.pubsub[0].filestream_project_id}/${module.pubsub[0].filestream_topic_name}/${module.pubsub[0].filestream_gorilla_subscription_name}"
-          fileStreamWorkerStore  = "bigtable://${module.bigtable[0].bigtable_project_id}/${module.bigtable[0].bigtable_instance_id}"
+          fileStreamWorkerStore  = "bigtablev2://${module.bigtable[0].bigtable_project_id}/${module.bigtable[0].bigtable_instance_id}"
         }
         serviceAccount = var.create_workload_identity ? {
           name        = local.k8s_sa_map.filestream

diff --git a/modules/bigtable/main.tf b/modules/bigtable/main.tf
@@ -25,4 +25,73 @@ resource "google_bigtable_instance_iam_member" "default" {
   instance = google_bigtable_instance.default.name
   role     = "roles/bigtable.admin"
   member   = local.sa_member
+}
+
+resource "google_bigtable_table" "omni_history" {
+  name          = "omni-history"
+  instance_name = google_bigtable_instance.default.name
+
+  deletion_protection = (var.deletion_protection) ? "PROTECTED" : "UNPROTECTED"
+
+  column_family {
+    family = "x"
+  }
+
+  change_stream_retention = "0"
+}
+
+resource "google_bigtable_table" "omni_history_lookup" {
+  name          = "omni-history-lookup"
+  instance_name = google_bigtable_instance.default.name
+
+  deletion_protection = (var.deletion_protection) ? "PROTECTED" : "UNPROTECTED"
+
+  column_family {
+    family = "x"
+  }
+
+  change_stream_retention = "0"
+}
+
+resource "google_bigtable_table" "runs" {
+  name          = "runs"
+  instance_name = google_bigtable_instance.default.name
+
+  deletion_protection = (var.deletion_protection) ? "PROTECTED" : "UNPROTECTED"
+
+  column_family {
+    family = "data"
+  }
+
+  column_family {
+    family = "metadata"
+  }
+
+  change_stream_retention = "0"
+}
+
+resource "google_bigtable_table" "logs" {
+  name          = "logs"
+  instance_name = google_bigtable_instance.default.name
+
+  deletion_protection = (var.deletion_protection) ? "PROTECTED" : "UNPROTECTED"
+
+  column_family {
+    family = "line"
+  }
+
+  change_stream_retention = "0"
+}
+
+resource "google_bigtable_table" "history" {
+  name          = "history"
+  instance_name = google_bigtable_instance.default.name
+
+  deletion_protection = (var.deletion_protection) ? "PROTECTED" : "UNPROTECTED"
+
+  column_family {
+    family = "event"
+  }
+
+  change_stream_retention = "0"
 }