fix: Reorganize the way values are passed (#266)

wandb · Nov 25, 2024 · 20a3879 · 20a3879
1 parent 634003d
commit 20a3879
Show file tree

Hide file tree

Showing 14 changed files with 162 additions and 17 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.1
+version: 0.19.2
 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
@@ -209,7 +209,7 @@ spec:
             {{- end }}
             {{- end }}
             - name: BUCKET
-              value: "{{ include "app.bucket" . }}"
+              value: {{ include "app.bucket" . | quote}}
             - name: AWS_REGION
               value: {{ .Values.global.bucket.region | default .Values.global.defaultBucket.region }}
             - name: AWS_S3_KMS_ID
@@ -239,7 +239,7 @@ spec:
               value: "otlp+grpc://{{ .Release.Name }}-otel-daemonset:4317?trace_ratio={{ .Values.traceRatio }}"
             {{- end }}
             - name: OVERFLOW_BUCKET_ADDR
-              value: "{{ include "app.bucket" .}}"
+              value: {{ include "app.bucket" . | quote }}
             {{- if not .Values.global.pubSub.enabled}}
             - name: KAFKA_BROKER_HOST
               value: "{{ include "wandb.kafka.brokerHost" . }}"
@@ -256,18 +256,23 @@ spec:
               value: {{ include "wandb.kafka.runUpdatesShadowTopic" .}}
             - name: KAFKA_RUN_UPDATE_SHADOW_QUEUE_NUM_PARTITIONS
               value: "{{ include "wandb.kafka.runUpdatesShadowNumPartitions" .}}"
+            {{- end }}
             - name: GORILLA_RUN_UPDATE_SHADOW_QUEUE
               value: >
                 {
                   "overflow-bucket": {
-                    "store": "{{ include "app.bucket" .}}",
+                    "store": {{ include "app.bucket" . | quote}},
                     "name": "wandb",
                     "prefix": "wandb-overflow"
                   },
-                  "addr": "kafka://$(KAFKA_CLIENT_USER):$(KAFKA_CLIENT_PASSWORD)@$(KAFKA_BROKER_HOST):$(KAFKA_BROKER_PORT)/$(KAFKA_TOPIC_RUN_UPDATE_SHADOW_QUEUE)?producer_batch_bytes=1048576&num_partitions=$(KAFKA_RUN_UPDATE_SHADOW_QUEUE_NUM_PARTITIONS)&replication_factor=3"
+                  "addr": {{ include "app.runUpdateShadowTopic" . | quote }}
                 }
-            {{- end }}
-
+            - name: GORILLA_HISTORY_STORE
+              value: {{ include "app.historyStore" . | quote }}
+            - name: GORILLA_PARQUET_LIVE_HISTORY_STORE
+              value: {{ include "app.liveHistoryStore" . | quote }}
+            - name: GORILLA_FILE_STREAM_STORE_ADDRESS
+              value: {{ include "app.fileStreamStore" . | quote }}
             - name: GORILLA_ARTIFACTS_GC_BATCH_SIZE
               value: {{ .Values.artifactsGc.BatchSize | quote }}
             - name: GORILLA_ARTIFACTS_GC_NUM_WORKERS

diff --git a/charts/operator-wandb/charts/app/templates/_helpers.tpl b/charts/operator-wandb/charts/app/templates/_helpers.tpl
@@ -139,3 +139,37 @@ app deployments.
 {{- end -}}
 }'
 {{- end -}}
+
+{{- define "app.runUpdateShadowTopic" -}}
+{{- if .Values.global.pubSub.enabled -}}
+pubsub:/{{ .Values.global.pubSub.project }}/{{ .Values.global.pubSub.runUpdateShadowTopic }}
+{{- else }}
+kafka://$(KAFKA_CLIENT_USER):$(KAFKA_CLIENT_PASSWORD)@$(KAFKA_BROKER_HOST):$(KAFKA_BROKER_PORT)/$(KAFKA_TOPIC_RUN_UPDATE_SHADOW_QUEUE)?producer_batch_bytes=1048576&num_partitions=$(KAFKA_RUN_UPDATE_SHADOW_QUEUE_NUM_PARTITIONS)&replication_factor=3
+{{- end -}}
+{{- 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 -}}
+{{- 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 -}}
+{{- end -}}
+mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred
+{{- end -}}
+
+{{/* TODO(dpanzella) - Probably need to make this support kafka as well*/}}
+{{- define "app.fileStreamStore" -}}
+{{- if .Values.global.pubSub.enabled -}}
+pubsub:/{{ .Values.global.pubSub.project }}/{{ .Values.global.pubSub.filestreamTopic }}
+{{- else -}}
+mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred
+{{- end -}}
+{{- end -}}
diff --git a/charts/operator-wandb/charts/executor/templates/_helpers.tpl b/charts/operator-wandb/charts/executor/templates/_helpers.tpl
@@ -127,3 +127,29 @@ Create the name of the service account to use
 {{- end -}}
 {{- trimSuffix "/" $bucket -}}
 {{- 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 -}}
+{{- 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 -}}
+{{- end -}}
+mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred
+{{- end -}}
+
+{{/* TODO(dpanzella) - Probably need to make this support kafka as well*/}}
+{{- define "executor.fileStreamStore" -}}
+{{- if .Values.global.pubSub.enabled -}}
+pubsub:/{{ .Values.global.pubSub.project }}/{{ .Values.global.pubSub.filestreamTopic }}
+{{- else -}}
+mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred
+{{- end -}}
+{{- end -}}
diff --git a/charts/operator-wandb/charts/executor/templates/deployment.yaml b/charts/operator-wandb/charts/executor/templates/deployment.yaml
@@ -107,6 +107,13 @@ spec:
             - name: GORILLA_TASK_QUEUE_WORKER_ENABLED
               value: "true"
 
+            - name: GORILLA_HISTORY_STORE
+              value: {{ include "executor.historyStore" . | quote }}
+            - name: GORILLA_PARQUET_LIVE_HISTORY_STORE
+              value: {{ include "executor.liveHistoryStore" .| quote }}
+            - name: GORILLA_FILE_STREAM_STORE_ADDRESS
+              value: {{ include "executor.fileStreamStore" . | quote }}
+
             - name: BUCKET
               value: "{{ include "executor.bucket" . }}"
             - name: AWS_REGION

diff --git a/charts/operator-wandb/charts/filestream/templates/_helpers.tpl b/charts/operator-wandb/charts/filestream/templates/_helpers.tpl
@@ -126,4 +126,18 @@ Create the name of the service account to use
 {{- end -}}
 {{- end -}}
 {{- trimSuffix "/" $bucket -}}
-{{- end -}}
+{{- end -}}
+
+{{- define "filestream.fileStreamWorkerSource" -}}
+{{- if .Values.global.pubSub.enabled -}}
+pubsub:/{{ .Values.global.pubSub.project }}/{{ .Values.global.pubSub.filestreamTopic }}/{{ .Values.pubSub.subscription }}
+{{- 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 }}
diff --git a/charts/operator-wandb/charts/filestream/templates/deployment.yaml b/charts/operator-wandb/charts/filestream/templates/deployment.yaml
@@ -135,9 +135,9 @@ spec:
             {{- end }}
             {{- end }}
             - name: GORILLA_FILE_STREAM_WORKER_SOURCE_ADDRESS
-              value: "{{ required "fileStreamWorkerSource required" .Values.env.fileStreamWorkerSource }}"
+              value: {{ include "filestream.fileStreamWorkerSource" . | quote }}
             - name: GORILLA_FILE_STREAM_WORKER_STORE_ADDRESS
-              value: "{{ required "fileStreamWorkerStore required" .Values.env.fileStreamWorkerStore }}"
+              value: {{ include "filestream.fileStreamWorkerStore" . | quote }}
             - name: GORILLA_FILE_STREAM_WORKER_STORE_TIMEOUT
               value: "{{ .Values.env.fileStreamWorkerStoreTimeout }}"
             - name: GORILLA_METADATA_STORE

diff --git a/charts/operator-wandb/charts/filestream/values.yaml b/charts/operator-wandb/charts/filestream/values.yaml
@@ -25,10 +25,11 @@ serviceAccount:
   name: ""
 
 env:
-  fileStreamWorkerSource: ""
-  fileStreamWorkerStore: ""
   fileStreamWorkerStoreTimeout: "30s"
 
+pubSub:
+  subscription: ""
+
 mysql:
   max-idle-conns: 10
   max-open-conns: 10

diff --git a/charts/operator-wandb/charts/flat-run-fields-updater/templates/_helpers.tpl b/charts/operator-wandb/charts/flat-run-fields-updater/templates/_helpers.tpl
@@ -126,4 +126,14 @@ Create the name of the service account to use
 {{- end -}}
 {{- end -}}
 {{- trimSuffix "/" $bucket -}}
-{{- end -}}
+{{- end -}}
+
+{{- define "flat-run-fields-updater.runUpdateShadowQueue" -}}
+{{- if .Values.global.pubSub.enabled -}}
+pubsub:/{{ .Values.global.pubSub.project }}/{{ .Values.global.pubSub.runUpdateShadowTopic }}/{{ .Values.pubSub.subscription }}
+{{- else }}
+kafka://$(KAFKA_CLIENT_USER):$(KAFKA_CLIENT_PASSWORD)@wandb-kafka:9092/$(KAFKA_TOPIC_RUN_UPDATE_SHADOW_QUEUE)?consumer_group_id=default-group&num_partitions=$(KAFKA_RUN_UPDATE_SHADOW_QUEUE_NUM_PARTITIONS)&replication_factor=3
+{{- end }}
+{{- 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
@@ -122,6 +122,9 @@ spec:
               value: {{ include "wandb.kafka.runUpdatesShadowTopic" .}}
             - name: KAFKA_RUN_UPDATE_SHADOW_QUEUE_NUM_PARTITIONS
               value: "{{ include "wandb.kafka.runUpdatesShadowNumPartitions" .}}"
+            {{- end }}
+            - name: BUCKET
+              value: "{{ include "flat-run-fields-updater.bucket" .}}"
             - name: GORILLA_RUN_UPDATE_SHADOW_QUEUE
               value: >
                 {
@@ -131,12 +134,9 @@ spec:
                     "prefix": "wandb-overflow"
                   },
                   "subscriptions": {
-                    "flatRunFieldsUpdater": "kafka://$(KAFKA_CLIENT_USER):$(KAFKA_CLIENT_PASSWORD)@wandb-kafka:9092/$(KAFKA_TOPIC_RUN_UPDATE_SHADOW_QUEUE)?consumer_group_id=default-group&num_partitions=$(KAFKA_RUN_UPDATE_SHADOW_QUEUE_NUM_PARTITIONS)&replication_factor=3"
+                    "flatRunFieldsUpdater": "{{ include "flat-run-fields-updater.runUpdateShadowQueue" .}}"
                   }
                 }
-            {{- end }}
-            - name: BUCKET
-              value: "{{ include "flat-run-fields-updater.bucket" .}}"
             - name: AWS_REGION
               value: {{ .Values.global.bucket.region | default .Values.global.defaultBucket.region }}
             - name: AWS_S3_KMS_ID

diff --git a/charts/operator-wandb/charts/flat-run-fields-updater/values.yaml b/charts/operator-wandb/charts/flat-run-fields-updater/values.yaml
@@ -28,6 +28,9 @@ serviceAccount:
   # If not set and create is true, a name is generated using the fullname template
   name: ""
 
+pubSub:
+  subscription: ""
+
 pod:
   labels: {}
   annotations: {}

diff --git a/charts/operator-wandb/charts/parquet/templates/_helpers.tpl b/charts/operator-wandb/charts/parquet/templates/_helpers.tpl
@@ -131,3 +131,29 @@ app deployments.
 {{- end -}}
 {{- trimSuffix "/" $bucket -}}
 {{- 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 -}}
+{{- 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 -}}
+{{- end -}}
+mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred
+{{- end -}}
+
+{{/* TODO(dpanzella) - Probably need to make this support kafka as well*/}}
+{{- define "parquet.fileStreamStore" -}}
+{{- if .Values.global.pubSub.enabled -}}
+pubsub:/{{ .Values.global.pubSub.project }}/{{ .Values.global.pubSub.filestreamTopic }}
+{{- else -}}
+mysql://$(MYSQL_USER):$(MYSQL_PASSWORD)@$(MYSQL_HOST):$(MYSQL_PORT)/$(MYSQL_DATABASE)?tls=preferred
+{{- end -}}
+{{- end -}}
diff --git a/charts/operator-wandb/charts/parquet/templates/deployment.yaml b/charts/operator-wandb/charts/parquet/templates/deployment.yaml
@@ -120,6 +120,13 @@ spec:
               value: "false"
             {{- end }}
 
+            - name: GORILLA_HISTORY_STORE
+              value: {{ include "parquet.historyStore" . | quote }}
+            - name: GORILLA_PARQUET_LIVE_HISTORY_STORE
+              value: {{ include "parquet.liveHistoryStore" . | quote }}
+            - name: GORILLA_FILE_STREAM_STORE_ADDRESS
+              value: {{ include "parquet.fileStreamStore" . | quote }}
+
             - name: BUCKET
               value: "{{ include "parquet.bucket" . }}"
             - name: AWS_REGION

diff --git a/charts/operator-wandb/values.yaml b/charts/operator-wandb/values.yaml
@@ -122,8 +122,16 @@ global:
     secretName: ""
     secretKey: "REDIS_PASSWORD"
 
+  bigTable:
+    enabled: false
+    project: ""
+    instance: ""
+
   pubSub:
     enabled: false
+    project: ""
+    filestreamTopic: ""
+    runUpdateShadowTopic: ""
 
   kafka:
     # The following values are anchored here, and referenced by alias later for
@@ -228,12 +236,16 @@ flat-run-fields-updater:
   image:
     repository: wandb/local
     tag: latest
+  pubSub:
+    subscription: ""
 
 filestream:
   install: false
   image:
     repository: wandb/local-dev
     tag: 0.63.0-danielpanzella-filestream-local.1
+  pubSub:
+    subscription: ""
 
 executor:
   install: false