Merge branch 'main' into aastha/restric-k8s-permissions

.github/CODEOWNERS

@@ -1,2 +1,2 @@
 * @gls4 @jsbroks @nfoucha @vanpelt
-/charts/launch-agent/ @bcsherma @gtarpenning @KyleGoyette @nickpenaranda @TimH98 @wandb-zacharyblasczyk
+/charts/launch-agent/ @bcsherma @KyleGoyette @TimH98

charts/launch-agent/Chart.yaml

@@ -3,7 +3,7 @@ name: launch-agent
 icon: https://em-content.zobj.net/thumbs/240/apple/354/rocket_1f680.png
 description: A Helm chart for running the W&B Launch Agent in Kubernetes
 type: application
-version: 0.13.5
+version: 0.13.10
 maintainers:
   - name: wandb
     email: [email protected]

charts/launch-agent/README.md

@@ -57,6 +57,7 @@ The table below describes all the available variables in the chart:
 | `launchConfig`                 | mutiline string | **Yes**              | `null`                          | This should be set to the literal contents of your launch agent config. See the agent setup docs for details: https://docs.wandb.ai/guides/launch/setup-agent-advanced                                                                          |
 | `volcano`                      | bool            | No                   | `true`                          | Controls whether the volcano scheduler should be installed in your cluster along with the agent. Set to `false` to disable volcano installation. |
 | `gitCreds`                     | mutiline string | No                   | `null`                          | Contents of a git credentials file.                                                                                                              |
+| `sshAuthSecrets`               | list(object)    | No                   | `[]`                              | Name of secret containing an ssh-auth kubernetes secret and the associated host for the ssh key. |
 | `serviceAccount.annotations`   | object          | No                   | `null`                          | Annotations for the wandb service account.                                                                                                       |
 | `azureStorageAccessKey`        | string          | No                   | ""                              | Azure storage access key required for kaniko to acces build contexts in azure blob storage.                                                      |
 | `additionalEnvVars`            | map(string)     | No                   | {}                              | Map with environment variables to be set in the Launch Agent pod.                                                                                |

charts/launch-agent/templates/configmap.yaml

@@ -8,4 +8,21 @@ kind: ConfigMap
 metadata:
   name: wandb-launch-configmap-{{ .Release.Name }}
   namespace: {{ .Values.namespace }}
-...
+---
+{{- if .Values.sshAuthSecrets }}
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: ssh-config-{{ .Release.Name }}
+  namespace: {{ .Values.namespace }}
+data:
+  config: |
+    {{- range $index, $secret := .Values.sshAuthSecrets }}
+    Host {{ .host }}
+      IdentityFile /home/launch_agent/.ssh/id_repo{{ $index }}
+      IdentitiesOnly yes
+      StrictHostKeyChecking no
+      UserKnownHostsFile /dev/null
+    {{- end }}
+{{- end }}

charts/launch-agent/templates/deployment.yaml

@@ -18,13 +18,23 @@ spec:
     metadata:
       labels:
         app: launch-agent-{{ .Release.Name }}
-{{- if .Values.agent.labels }}
-{{- toYaml .Values.agent.labels | trim | nindent 8 }}
-{{- end }}
+        {{- if .Values.agent.labels }}
+        {{- toYaml .Values.agent.labels | trim | nindent 8 }}
+        {{- end }}
     spec:
       serviceAccountName: wandb-launch-serviceaccount-{{ .Release.Name }}
-      {{- if .Values.kanikoPvcName }}
+      {{- if or .Values.sshAuthSecrets .Values.kanikoPvcName }}
       initContainers:
+      {{- end}}
+      {{- if .Values.sshAuthSecrets }}
+      - name: init-create-ssh-dir
+        image: {{ .Values.agent.image }}
+        command: ["sh", "-c", "mkdir -p /home/launch_agent/.ssh"]
+        volumeMounts:
+        - name: ssh-dir
+          mountPath: /home/launch_agent/.ssh
+      {{- end }}
+      {{- if .Values.kanikoPvcName }}
       - name: kaniko-volume-chown
         image: {{ .Values.agent.image }}
         command: ["sh", "-c"]
@@ -101,12 +111,31 @@ spec:
               value: {{ .Values.kanikoDockerConfigSecret }}
             {{- end }}
           volumeMounts:
+            - name: ssh-dir
+              mountPath: /home/launch_agent/.ssh
             - name: wandb-launch-config
               mountPath: /home/launch_agent/.config/wandb
               readOnly: true
             {{ if .Values.gitCreds}}
             - name: git-creds
-              mountPath: /home/launch_agent/
+              mountPath: /home/launch_agent/.gitconfig
+              subPath: .gitconfig
+              readOnly: true
+            - name: git-creds
+              mountPath: /home/launch_agent/.git-credentials
+              subPath: .git-credentials
+              readOnly: true
+            {{ end }}
+            {{ if .Values.sshAuthSecrets }}
+            {{- range $index, $secret := .Values.sshAuthSecrets }}
+            - name: git-ssh-key-secret-{{ $index }}
+              mountPath: /home/launch_agent/.ssh/id_repo{{ $index }}
+              subPath: id_repo{{ $index }}
+              readOnly: true
+            {{- end }}
+            - name: ssh-config
+              mountPath: /home/launch_agent/.ssh/config
+              subPath: config
               readOnly: true
             {{ end }}
             {{- if and .Values.customCABundle.configMap.name .Values.customCABundle.configMap.key }}
@@ -120,30 +149,45 @@ spec:
               mountPath: /home/launch_agent/kaniko
             {{ end }}
       volumes:
+        - name: ssh-dir
+          emptyDir: {}
         - name: wandb-launch-config
           configMap:
             name: wandb-launch-configmap-{{ .Release.Name }}
         {{ if .Values.gitCreds}}
         - name: git-creds
           secret:
             secretName: git-creds
-        - name: git-config
-          secret:
-            secretName: git-config
         {{ end}}
+        {{ if .Values.sshAuthSecrets }}
+        {{- range $index, $secret := .Values.sshAuthSecrets }}
+        - name: git-ssh-key-secret-{{ $index }}
+          secret:
+            secretName: {{ $secret.name }}
+            items: 
+              - key: ssh-privatekey
+                path: id_repo{{ $index }}
+        {{- end }}
+        - name: ssh-config
+          configMap:
+            name: ssh-config-{{ .Release.Name }}
+            items:
+              - key: config
+                path: config
+        {{ end }}
         {{- if and .Values.customCABundle.configMap.name .Values.customCABundle.configMap.key }}
         - name: custom-cabundle
           configMap:
             name: {{ .Values.customCABundle.configMap.name }}
             items:
               - key: {{ .Values.customCABundle.configMap.key }}
                 path: custom-ca.crt
-        {{- end}}
+        {{- end }}
         {{ if .Values.kanikoPvcName }}
         - name: kaniko-pvc
           persistentVolumeClaim:
             claimName: {{ .Values.kanikoPvcName }}
-        {{- end}}
+        {{- end }}
       nodeSelector:
         {{- toYaml .Values.agent.nodeSelector | nindent 8 }}
       tolerations:

charts/launch-agent/values.yaml

@@ -5,7 +5,7 @@ agent:
   # Providing API key can be done external to this chart
   useExternalWandbSecret: false
   # Container image to use for the agent.
-  image: wandb/launch-agent:0.17.1
+  image: wandb/launch-agent:0.17.3
   # Image pull policy for agent image.
   imagePullPolicy: Always
   # Resources block for the agent spec.
@@ -63,6 +63,14 @@ additionalSecretEnvVars:
 # repos. Example: https://username:[email protected]
 gitCreds: |
 
+# list of secrets for the agent to use for ssh auth
+# format is a list of secret names and hosts secrets
+# should be created ass ssh-auth secrets,
+# see: https://kubernetes.io/docs/concepts/configuration/secret/#ssh-authentication-secrets
+sshAuthSecrets:
+  # - name: secret-name
+  #   host: example.com
+
 # Annotations for the wandb service account. Useful when setting up workload identity on gcp.
 serviceAccount:
   annotations:

charts/operator-wandb/Chart.lock

@@ -8,6 +8,9 @@ dependencies:
 - name: weave
   repository: file://charts/weave
   version: 0.1.0
+- name: weave-trace
+  repository: file://charts/weave-trace
+  version: 0.1.0
 - name: parquet
   repository: file://charts/parquet
   version: 0.1.0
@@ -35,5 +38,8 @@ dependencies:
 - name: stackdriver
   repository: file://charts/stackdriver
   version: 0.1.0
-digest: sha256:9a6c69506deb6969686d5b220a0692b53cfa29642e059bdf27c440c5d7086bdb
-generated: "2024-06-05T11:04:02.508473-07:00"
+- name: yace
+  repository: file://charts/yace
+  version: 0.1.0
+digest: sha256:bca2b6781737da6806e4485605cf9ce87b1428944b14cb88f082024cc3500bbd
+generated: "2024-07-18T01:17:04.532871-04:00"

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.13.14
+version: 0.15.3
 appVersion: 1.0.0
 icon: https://wandb.ai/logo.svg
 
@@ -24,6 +24,10 @@ dependencies:
     version: "*.*.*"
     repository: file://charts/weave
     condition: weave.install
+  - name: weave-trace
+    version: "*.*.*"
+    repository: file://charts/weave-trace
+    condition: weave-trace.install
   - name: parquet
     version: "*.*.*"
     repository: file://charts/parquet
@@ -60,3 +64,7 @@ dependencies:
     version: "*.*.*"
     repository: file://charts/stackdriver
     condition: stackdriver.install
+  - name: yace
+    version: "*.*.*"
+    repository: file://charts/yace
+    condition: yace.install

charts/operator-wandb/charts/app/templates/_helpers.tpl

@@ -116,7 +116,7 @@ app deployments.
 {{- $bucket = printf "az://%s/%s" .Values.global.bucket.name .Values.global.bucket.path -}}
 {{- end -}}
 {{- if eq .Values.global.bucket.provider "gcs" -}}
-{{- $bucket = printf "gs://%s" .Values.global.bucket.name -}}
+{{- $bucket = printf "gs://%s/%s" .Values.global.bucket.name .Values.global.bucket.path -}}
 {{- end -}}
 {{- if eq .Values.global.bucket.provider "s3" -}}
 {{- if and .Values.global.bucket.accessKey .Values.global.bucket.secretKey -}}

charts/operator-wandb/charts/app/templates/deployment.yaml

@@ -63,7 +63,7 @@ spec:
                 secretKeyRef:
                   name: {{ include "wandb.mysql.passwordSecret" . }}
                   key:  MYSQL_PASSWORD
-          command: ['bash', '-c', "until mysql -h$MYSQL_HOST -u$MYSQL_USER -p$MYSQL_PASSWORD -D$MYSQL_DATABASE --execute=\"SELECT 1\"; do echo waiting for db; sleep 2; done"]
+          command: ['bash', '-c', "until mysql -h$MYSQL_HOST -u$MYSQL_USER -p$MYSQL_PASSWORD -D$MYSQL_DATABASE -P$MYSQL_PORT --execute=\"SELECT 1\"; do echo waiting for db; sleep 2; done"]
       containers:
         - name: {{ .Chart.Name }}
           image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}"
@@ -114,6 +114,10 @@ spec:
               value: "http://{{ .Release.Name }}-parquet:8087"
             - name: PARQUET_ENABLED
               value: "true"
+            {{- if index .Values.global "weave-trace" "enabled" }}
+            - name: WEAVE_TRACES_ENABLED
+              value: "true"
+            {{- end }}
 
             {{- if ne (include "wandb.redis.password" .) "" }}
             - name: REDIS_PASSWORD
@@ -227,6 +231,8 @@ spec:
                   key: KAFKA_CLIENT_PASSWORD
             - name: KAFKA_TOPIC_RUN_UPDATE_SHADOW_QUEUE
               value: {{ include "wandb.kafka.runUpdatesShadowTopic" .}}
+            - name: KAFKA_RUN_UPDATE_SHADOW_QUEUE_NUM_PARTITIONS
+              value: "{{ include "wandb.kafka.runUpdatesShadowNumPartitions" .}}"
             - name: OVERFLOW_BUCKET_ADDR
               value: "{{ include "app.bucket" .}}"
             - name: GORILLA_RUN_UPDATE_SHADOW_QUEUE
@@ -237,7 +243,7 @@ spec:
                     "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"
+                  "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)"
                 }
             {{- include "app.extraEnv" (dict "global" $.Values.global "local" .Values) | nindent 12 }}
             {{- include "wandb.extraEnvFrom" (dict "root" $ "local" .) | nindent 12 }}

charts/operator-wandb/charts/app/templates/serviceaccount.yaml

@@ -1,3 +1,4 @@
+{{- if .Values.serviceAccount.create -}}
 apiVersion: v1
 kind: ServiceAccount
 metadata:
@@ -14,3 +15,4 @@ metadata:
     {{- if .Values.serviceAccount.annotations -}}
     {{-   toYaml .Values.serviceAccount.annotations | nindent 4 }}
     {{- end }}
+  {{- end }}

charts/operator-wandb/charts/app/values.yaml

@@ -53,6 +53,7 @@ resources:
 
 serviceAccount:
   create: true
+  annotations: {}
 
 role: {}
 roleBinding: {}

charts/operator-wandb/charts/console/templates/deployment.yaml

@@ -34,7 +34,7 @@ spec:
         {{-   toYaml .Values.pod.annotations | nindent 4 }}
         {{- end }}
     spec:
-      serviceAccountName: {{ include "console.fullname" . }}
+      serviceAccountName: {{ include "console.serviceAccountName" . }}
       {{- if .tolerations }}
       tolerations:
         {{- toYaml .tolerations | nindent 8 }}

charts/operator-wandb/charts/console/templates/service.yaml

@@ -11,7 +11,10 @@ metadata:
     {{-   toYaml .Values.service.labels | nindent 4 }}
     {{- end }}
   annotations:
-    {{- include "wandb.deploymentAnnotations" $ | nindent 4 }}
+    {{- if eq .Values.global.cloudProvider "aws" }}
+    alb.ingress.kubernetes.io/healthcheck-path: /console/api/ready
+    {{- end }}
+    {{- include "wandb.serviceAnnotations" $ | nindent 4 }}
     {{- if .Values.service.annotations -}}
     {{-   toYaml .Values.service.annotations | nindent 4 }}
     {{- end }}

charts/operator-wandb/charts/console/templates/serviceaccount.yaml

@@ -1,3 +1,4 @@
+{{- if .Values.serviceAccount.create -}}
 apiVersion: v1
 kind: ServiceAccount
 metadata:
@@ -14,3 +15,4 @@ metadata:
     {{- if .Values.serviceAccount.annotations -}}
     {{-   toYaml .Values.serviceAccount.annotations | nindent 4 }}
     {{- end }}
+{{- end }}

charts/operator-wandb/charts/console/values.yaml

@@ -21,7 +21,9 @@ extraCors: []
 common:
   labels: {}
 deployment: {}
-serviceAccount: {}
+serviceAccount:
+  create: true
+  annotations: {}
 clusterRole: {}
 
 pod:

charts/operator-wandb/charts/flat-run-fields-updater/templates/_helpers.tpl

@@ -112,7 +112,7 @@ Create the name of the service account to use
 {{- $bucket = printf "az://%s/%s" .Values.global.bucket.name .Values.global.bucket.path -}}
 {{- end -}}
 {{- if eq .Values.global.bucket.provider "gcs" -}}
-{{- $bucket = printf "gs://%s" .Values.global.bucket.name -}}
+{{- $bucket = printf "gs://%s/%s" .Values.global.bucket.name .Values.global.bucket.path -}}
 {{- end -}}
 {{- if eq .Values.global.bucket.provider "s3" -}}
 {{- if and .Values.global.bucket.accessKey .Values.global.bucket.secretKey -}}