From 4103c08a960a63f39f6c098c039fe252a4378d8e Mon Sep 17 00:00:00 2001
From: Daniel Panzella <daniel.panzella@wandb.com>
Date: Mon, 2 Dec 2024 11:52:58 -0700
Subject: [PATCH] fix: Additional bigtable and pubsub fixes (#268)

Co-authored-by: Annirudh Prasad <annirudh@wandb.com>
---
 charts/operator-wandb/Chart.yaml              |  2 +-
 .../charts/app/templates/_deployment.tpl      |  6 +++
 .../charts/app/templates/_helpers.tpl         | 40 ++++++++++++++-----
 .../charts/executor/templates/_helpers.tpl    | 40 ++++++++++++++-----
 .../charts/executor/templates/deployment.yaml |  4 +-
 .../charts/executor/values.yaml               |  1 +
 .../charts/filestream/templates/_helpers.tpl  | 23 +++++++----
 .../templates/deployment.yaml                 |  2 +
 .../charts/parquet/templates/_helpers.tpl     | 40 ++++++++++++++-----
 .../charts/parquet/templates/cron.yaml        |  2 +
 .../charts/parquet/templates/deployment.yaml  |  4 ++
 charts/operator-wandb/values.yaml             |  7 +++-
 12 files changed, 129 insertions(+), 42 deletions(-)

diff --git a/charts/operator-wandb/Chart.yaml b/charts/operator-wandb/Chart.yaml
index efcd92bc..7a2ec08b 100644
--- a/charts/operator-wandb/Chart.yaml
+++ b/charts/operator-wandb/Chart.yaml
@@ -2,7 +2,7 @@ apiVersion: v2
 name: operator-wandb
 description: A Helm chart for deploying W&B to Kubernetes
 type: application
-version: 0.19.2
+version: 0.19.3
 appVersion: 1.0.0
 icon: https://wandb.ai/logo.svg
 
diff --git a/charts/operator-wandb/charts/app/templates/_deployment.tpl b/charts/operator-wandb/charts/app/templates/_deployment.tpl
index 958bb9af..b1b7ec33 100644
--- a/charts/operator-wandb/charts/app/templates/_deployment.tpl
+++ b/charts/operator-wandb/charts/app/templates/_deployment.tpl
@@ -267,6 +267,10 @@ spec:
                   },
                   "addr": {{ include "app.runUpdateShadowTopic" . | quote }}
                 }
+            - name: GORILLA_SETTINGS_CACHE
+              value: "{{ include "app.redis" . | trim }}"
+            - name: GORILLA_METADATA_CACHE
+              value: "{{ include "app.redis" . | trim }}"
             - name: GORILLA_HISTORY_STORE
               value: {{ include "app.historyStore" . | quote }}
             - name: GORILLA_PARQUET_LIVE_HISTORY_STORE
@@ -287,6 +291,8 @@ spec:
               value: "10000"
             - name: GORILLA_TASK_QUEUE_WORKER_ENABLED
               value: "false"
+            - name: GORILLA_CLEAR_TASK_DEDUPE_KEY_ENABLED
+              value: "false"
             {{- end }}
 
             {{- if index .Values.global "weave-trace" "enabled" }}
diff --git a/charts/operator-wandb/charts/app/templates/_helpers.tpl b/charts/operator-wandb/charts/app/templates/_helpers.tpl
index b7929229..95ab004f 100644
--- a/charts/operator-wandb/charts/app/templates/_helpers.tpl
+++ b/charts/operator-wandb/charts/app/templates/_helpers.tpl
@@ -149,20 +149,40 @@ kafka://$(KAFKA_CLIENT_USER):$(KAFKA_CLIENT_PASSWORD)@$(KAFKA_BROKER_HOST):$(KAF
 {{- end -}}
 
 {{- define "app.historyStore" -}}
-{{- $historyStore := printf "http://%s-parquet:8087/_goRPC_" .Release.Name -}}
-{{- if .Values.global.bigTable.enabled }}
-{{- $historyStore = printf "%s,bigtablev3://%s/%s,bigtablev2://%s/%s" $historyStore .Values.global.bigTable.project .Values.global.bigTable.instance .Values.global.bigTable.project .Values.global.bigTable.instance -}}
-{{- else -}}
-{{- $historyStore = printf "%s,mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred" $historyStore }}
-{{- end -}}
-{{- $historyStore -}}
+    {{- $stores := list -}}
+    {{- $stores = append $stores (printf "http://%s-parquet:8087/_goRPC_" .Release.Name) -}}
+
+    {{- if .Values.global.bigtable.v3.enabled -}}
+        {{- $stores = append $stores (printf "bigtablev3://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+    {{- end -}}
+
+    {{- if .Values.global.bigtable.v2.enabled -}}
+        {{- $stores = append $stores (printf "bigtablev2://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+    {{- end -}}
+
+    {{- if not (or .Values.global.bigtable.v2.enabled .Values.global.bigtable.v3.enabled) -}}
+        {{- $stores = append $stores "mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred" -}}
+    {{- end -}}
+
+    {{- join "," $stores -}}
 {{- end -}}
 
 {{- define "app.liveHistoryStore" -}}
-{{- if .Values.global.bigTable.enabled }}
-{{- printf "bigtablev3://%s/%s,bigtablev2://%s/%s" .Values.global.bigTable.project .Values.global.bigTable.instance .Values.global.bigTable.project .Values.global.bigTable.instance -}}
+{{- $historyStore := printf "mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred" -}}
+{{- if or .Values.global.bigtable.v2.enabled .Values.global.bigtable.v3.enabled -}}
+  {{- $stores := list -}}
+
+  {{- if .Values.global.bigtable.v3.enabled -}}
+    {{- $stores = append $stores (printf "bigtablev3://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+  {{- end -}}
+
+  {{- if .Values.global.bigtable.v2.enabled -}}
+    {{- $stores = append $stores (printf "bigtablev2://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+  {{- end -}}
+
+  {{- $historyStore = join "," $stores -}}
 {{- end -}}
-mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred
+{{- $historyStore -}}
 {{- end -}}
 
 {{/* TODO(dpanzella) - Probably need to make this support kafka as well*/}}
diff --git a/charts/operator-wandb/charts/executor/templates/_helpers.tpl b/charts/operator-wandb/charts/executor/templates/_helpers.tpl
index 9d939c4f..503534ad 100644
--- a/charts/operator-wandb/charts/executor/templates/_helpers.tpl
+++ b/charts/operator-wandb/charts/executor/templates/_helpers.tpl
@@ -129,20 +129,40 @@ Create the name of the service account to use
 {{- end -}}
 
 {{- define "executor.historyStore" -}}
-{{- $historyStore := printf "http://%s-parquet:8087/_goRPC_" .Release.Name -}}
-{{- if .Values.global.bigTable.enabled -}}
-{{- $historyStore = printf "%s,bigtablev3://%s/%s,bigtablev2://%s/%s" $historyStore .Values.global.bigTable.project .Values.global.bigTable.instance .Values.global.bigTable.project .Values.global.bigTable.instance -}}
-{{- else -}}
-{{- $historyStore = printf "%s,mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred" $historyStore -}}
-{{- end -}}
-{{- $historyStore -}}
+    {{- $stores := list -}}
+    {{- $stores = append $stores (printf "http://%s-parquet:8087/_goRPC_" .Release.Name) -}}
+
+    {{- if .Values.global.bigtable.v3.enabled -}}
+        {{- $stores = append $stores (printf "bigtablev3://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+    {{- end -}}
+
+    {{- if .Values.global.bigtable.v2.enabled -}}
+        {{- $stores = append $stores (printf "bigtablev2://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+    {{- end -}}
+
+    {{- if not (or .Values.global.bigtable.v2.enabled .Values.global.bigtable.v3.enabled) -}}
+        {{- $stores = append $stores "mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred" -}}
+    {{- end -}}
+
+    {{- join "," $stores -}}
 {{- end -}}
 
 {{- define "executor.liveHistoryStore" -}}
-{{- if .Values.global.bigTable.enabled -}}
-{{- printf "bigtablev3://%s/%s,bigtablev2://%s/%s" .Values.global.bigTable.project .Values.global.bigTable.instance .Values.global.bigTable.project .Values.global.bigTable.instance -}}
+{{- $historyStore := printf "mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred" -}}
+{{- if or .Values.global.bigtable.v2.enabled .Values.global.bigtable.v3.enabled -}}
+  {{- $stores := list -}}
+
+  {{- if .Values.global.bigtable.v3.enabled -}}
+    {{- $stores = append $stores (printf "bigtablev3://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+  {{- end -}}
+
+  {{- if .Values.global.bigtable.v2.enabled -}}
+    {{- $stores = append $stores (printf "bigtablev2://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+  {{- end -}}
+
+  {{- $historyStore = join "," $stores -}}
 {{- end -}}
-mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred
+{{- $historyStore -}}
 {{- end -}}
 
 {{/* TODO(dpanzella) - Probably need to make this support kafka as well*/}}
diff --git a/charts/operator-wandb/charts/executor/templates/deployment.yaml b/charts/operator-wandb/charts/executor/templates/deployment.yaml
index e6f9cf7e..87387ce8 100644
--- a/charts/operator-wandb/charts/executor/templates/deployment.yaml
+++ b/charts/operator-wandb/charts/executor/templates/deployment.yaml
@@ -101,9 +101,11 @@ spec:
 
             - name: GORILLA_SETTINGS_CACHE
               value: "{{ include "executor.redis" . | trim }}"
+            - name: GORILLA_METADATA_CACHE
+              value: "{{ include "executor.redis" . | trim }}"
 
             - name: GORILLA_TASK_QUEUE
-              value: "{{ include "executor.redis" . | trim }}"
+              value: "{{ include "executor.redis" . | trim }}{{- if .Values.workerConcurrency }}&concurrency={{ .Values.workerConcurrency }}{{- end }}"
             - name: GORILLA_TASK_QUEUE_WORKER_ENABLED
               value: "true"
 
diff --git a/charts/operator-wandb/charts/executor/values.yaml b/charts/operator-wandb/charts/executor/values.yaml
index 5e37191e..035dbe94 100644
--- a/charts/operator-wandb/charts/executor/values.yaml
+++ b/charts/operator-wandb/charts/executor/values.yaml
@@ -14,6 +14,7 @@ imagePullSecrets: []
 nameOverride: ""
 fullnameOverride: ""
 traceRatio: 0
+workerConcurrency: 10
 
 serviceAccount:
   # Specifies whether a service account should be created
diff --git a/charts/operator-wandb/charts/filestream/templates/_helpers.tpl b/charts/operator-wandb/charts/filestream/templates/_helpers.tpl
index 3205e4b9..98355b0e 100644
--- a/charts/operator-wandb/charts/filestream/templates/_helpers.tpl
+++ b/charts/operator-wandb/charts/filestream/templates/_helpers.tpl
@@ -131,13 +131,20 @@ Create the name of the service account to use
 {{- define "filestream.fileStreamWorkerSource" -}}
 {{- if .Values.global.pubSub.enabled -}}
 pubsub:/{{ .Values.global.pubSub.project }}/{{ .Values.global.pubSub.filestreamTopic }}/{{ .Values.pubSub.subscription }}
-{{- else }}
-{{- end }}
-{{- end }}
+{{- else -}}
+{{- end -}}
+{{- end -}}
 
 {{- define "filestream.fileStreamWorkerStore" -}}
-{{- if .Values.global.bigTable.enabled -}}
-bigtablev3://{{ .Values.global.bigTable.project }}/{{ .Values.global.bigTable.instance }},bigtablev2://{{ .Values.global.bigTable.project }}/{{ .Values.global.bigTable.instance }}
-{{- else }}
-{{- end }}
-{{- end }}
+    {{- $stores := list -}}
+
+    {{- if .Values.global.bigtable.v2.enabled -}}
+        {{- $stores = append $stores (printf "bigtablev2://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+    {{- end -}}
+
+    {{- if .Values.global.bigtable.v3.enabled -}}
+        {{- $stores = append $stores (printf "bigtablev3://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+    {{- end -}}
+
+    {{- join "," $stores -}}
+{{- end -}}
diff --git a/charts/operator-wandb/charts/flat-run-fields-updater/templates/deployment.yaml b/charts/operator-wandb/charts/flat-run-fields-updater/templates/deployment.yaml
index cafea60c..20bd74cf 100644
--- a/charts/operator-wandb/charts/flat-run-fields-updater/templates/deployment.yaml
+++ b/charts/operator-wandb/charts/flat-run-fields-updater/templates/deployment.yaml
@@ -155,6 +155,8 @@ spec:
               value: "{{ include "flat-run-fields-updater.redis" . | trim }}"
             - name: GORILLA_SETTINGS_CACHE
               value: "{{ include "flat-run-fields-updater.redis" . | trim }}"
+            - name: GORILLA_METADATA_CACHE
+              value: "{{ include "flat-run-fields-updater.redis" . | trim }}"
             {{- if and .Values.global .Values.global.observability }}
             {{- if eq (default "custom" .Values.global.observability.mode) "otel" }}
             - name: GORILLA_STATSD_PORT
diff --git a/charts/operator-wandb/charts/parquet/templates/_helpers.tpl b/charts/operator-wandb/charts/parquet/templates/_helpers.tpl
index a6b5ee97..6de10860 100644
--- a/charts/operator-wandb/charts/parquet/templates/_helpers.tpl
+++ b/charts/operator-wandb/charts/parquet/templates/_helpers.tpl
@@ -133,20 +133,40 @@ app deployments.
 {{- end -}}
 
 {{- define "parquet.historyStore" -}}
-{{- $historyStore := printf "http://%s-parquet:8087/_goRPC_" .Release.Name -}}
-{{- if .Values.global.bigTable.enabled -}}
-{{- $historyStore = printf "%s,bigtablev3://%s/%s,bigtablev2://%s/%s" $historyStore .Values.global.bigTable.project .Values.global.bigTable.instance .Values.global.bigTable.project .Values.global.bigTable.instance -}}
-{{- else -}}
-{{- $historyStore = printf "%s,mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred" $historyStore -}}
-{{- end -}}
-{{- $historyStore -}}
+    {{- $stores := list -}}
+    {{- $stores = append $stores (printf "http://%s-parquet:8087/_goRPC_" .Release.Name) -}}
+
+    {{- if .Values.global.bigtable.v3.enabled -}}
+        {{- $stores = append $stores (printf "bigtablev3://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+    {{- end -}}
+
+    {{- if .Values.global.bigtable.v2.enabled -}}
+        {{- $stores = append $stores (printf "bigtablev2://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+    {{- end -}}
+
+    {{- if not (or .Values.global.bigtable.v2.enabled .Values.global.bigtable.v3.enabled) -}}
+        {{- $stores = append $stores "mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred" -}}
+    {{- end -}}
+
+    {{- join "," $stores -}}
 {{- end -}}
 
 {{- define "parquet.liveHistoryStore" -}}
-{{- if .Values.global.bigTable.enabled -}}
-{{- printf "bigtablev3://%s/%s,bigtablev2://%s/%s" .Values.global.bigTable.project .Values.global.bigTable.instance .Values.global.bigTable.project .Values.global.bigTable.instance -}}
+{{- $historyStore := printf "mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred" -}}
+{{- if or .Values.global.bigtable.v2.enabled .Values.global.bigtable.v3.enabled -}}
+  {{- $stores := list -}}
+
+  {{- if .Values.global.bigtable.v3.enabled -}}
+    {{- $stores = append $stores (printf "bigtablev3://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+  {{- end -}}
+
+  {{- if .Values.global.bigtable.v2.enabled -}}
+    {{- $stores = append $stores (printf "bigtablev2://%s/%s" .Values.global.bigtable.project .Values.global.bigtable.instance) -}}
+  {{- end -}}
+
+  {{- $historyStore = join "," $stores -}}
 {{- end -}}
-mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred
+{{- $historyStore -}}
 {{- end -}}
 
 {{/* TODO(dpanzella) - Probably need to make this support kafka as well*/}}
diff --git a/charts/operator-wandb/charts/parquet/templates/cron.yaml b/charts/operator-wandb/charts/parquet/templates/cron.yaml
index 7eb5e617..85e0c55c 100644
--- a/charts/operator-wandb/charts/parquet/templates/cron.yaml
+++ b/charts/operator-wandb/charts/parquet/templates/cron.yaml
@@ -105,6 +105,8 @@ spec:
                   value: "{{ include "app.redis" . | trim }}"
                 - name: GORILLA_TASK_QUEUE_WORKER_ENABLED
                   value: "false"
+                - name: GORILLA_CLEAR_TASK_DEDUPE_KEY_ENABLED
+                  value: "false"
                 {{- end }}
 
                 - name: WEAVE_SERVICE
diff --git a/charts/operator-wandb/charts/parquet/templates/deployment.yaml b/charts/operator-wandb/charts/parquet/templates/deployment.yaml
index f3eabb21..b4002288 100644
--- a/charts/operator-wandb/charts/parquet/templates/deployment.yaml
+++ b/charts/operator-wandb/charts/parquet/templates/deployment.yaml
@@ -112,12 +112,16 @@ spec:
 
             - name: GORILLA_SETTINGS_CACHE
               value: "{{ include "parquet.redis" . | trim }}"
+            - name: GORILLA_METADATA_CACHE
+              value: "{{ include "parquet.redis" . | trim }}"
 
             {{- if .Values.global.executor.enabled }}
             - name: GORILLA_TASK_QUEUE
               value: "{{ include "app.redis" . | trim }}"
             - name: GORILLA_TASK_QUEUE_WORKER_ENABLED
               value: "false"
+            - name: GORILLA_CLEAR_TASK_DEDUPE_KEY_ENABLED
+              value: "false"
             {{- end }}
 
             - name: GORILLA_HISTORY_STORE
diff --git a/charts/operator-wandb/values.yaml b/charts/operator-wandb/values.yaml
index b7199996..749d6cf1 100644
--- a/charts/operator-wandb/values.yaml
+++ b/charts/operator-wandb/values.yaml
@@ -122,8 +122,11 @@ global:
     secretName: ""
     secretKey: "REDIS_PASSWORD"
 
-  bigTable:
-    enabled: false
+  bigtable:
+    v3:
+      enabled: false
+    v2:
+      enabled: false
     project: ""
     instance: ""