fix: Additional bigtable and pubsub fixes (#268)

Co-authored-by: Annirudh Prasad <[email protected]>
wandb · Dec 2, 2024 · 4103c08 · 4103c08
1 parent ddb0dab
commit 4103c08
Show file tree

Hide file tree

Showing 12 changed files with 129 additions and 42 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -122,8 +122,11 @@ global:
     secretName: ""
     secretKey: "REDIS_PASSWORD"
 
-  bigTable:
-    enabled: false
+  bigtable:
+    v3:
+      enabled: false
+    v2:
+      enabled: false
     project: ""
     instance: ""