feat: Add support for env vars on Launch Agent (#69)

wandb · Jan 30, 2024 · 630ef56 · 630ef56
1 parent 4803685
commit 630ef56
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 2 deletions.
diff --git a/charts/launch-agent/Chart.yaml b/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.11.7
+version: 0.11.8
 maintainers:
   - name: wandb
     email: support@wandb.com

diff --git a/charts/launch-agent/README.md b/charts/launch-agent/README.md
@@ -57,3 +57,5 @@ The table below describes all the available variables in the chart:
 | `gitCreds`                     | mutiline string | No                   | `null`                          | Contents of a git credentials file.                                                                                                              |
 | `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.                                                                                |
+| `additionalSecretEnvVars`      | map(string)     | No                   | {}                              | Map with environment variables to be stored in the `launch-agent-secret-env-vars` secret and set in the Launch Agent Pod                         |
diff --git a/charts/launch-agent/templates/deployment.yaml b/charts/launch-agent/templates/deployment.yaml
@@ -6,6 +6,7 @@ metadata:
   namespace: {{ .Values.namespace }}
   annotations:
     checksum/config: {{ include (print $.Template.BasePath "/configmap.yaml") . | sha256sum }}
+    checksum/config: {{ include (print $.Template.BasePath "/secret.yaml") . | sha256sum }}
 spec:
   replicas: 1
   strategy:
@@ -57,6 +58,21 @@ spec:
               value: "{{ .Values.agent.startTimeout }}"
             - name: WANDB_RELEASE_NAME
               value: "{{ .Release.Name }}"
+            {{- if .Values.additionalEnvVars }}
+            {{- range .Values.additionalEnvVars }}
+            - name: {{ .name }}
+              value: {{ .value }}
+            {{- end }}
+            {{- end }}
+            {{- if .Values.additionalSecretEnvVars }}
+            {{- range .Values.additionalSecretEnvVars }}
+            - name: {{ .name }}
+              valueFrom:
+                secretKeyRef:
+                  name: launch-agent-secret-env-vars
+                  key: {{ .name }}
+            {{- end }}
+            {{- end }}
           volumeMounts:
             - name: wandb-launch-config
               mountPath: /home/launch_agent/.config/wandb

diff --git a/charts/launch-agent/templates/secret.yaml b/charts/launch-agent/templates/secret.yaml
@@ -43,4 +43,17 @@ metadata:
 type: kubernetes.io/basic-auth
 stringData:
   password: {{ .Values.azureStorageAccessKey }}
-{{ end }}
+{{ end }}
+{{- if .Values.additionalSecretEnvVars }}
+---
+apiVersion: v1
+kind: Secret
+metadata:
+  name: launch-agent-secret-env-vars
+  namespace: wandb
+type: Opaque
+data:
+{{- range .Values.additionalSecretEnvVars }}
+  {{ .name }}: {{ .value | b64enc | quote }}
+{{- end }}
+{{- end }}
diff --git a/charts/launch-agent/values.yaml b/charts/launch-agent/values.yaml
@@ -38,6 +38,16 @@ launchConfig: |
   builder:  # Builder config. Type `noop` and `kaniko` can be used in a kubernetes cluster.
     type: noop
 
+# # Additional environment variables to set on the agent container.
+additionalEnvVars: {}
+  # - name: MY_ENV_VAR
+  #   value: "my-value"
+
+# # Additional secret environment variables to set on the agent container.
+additionalSecretEnvVars: {}
+  # - name: MY_SECRET_ENV_VAR
+  #   value: "my-secret-value"
+
 # The contents of a git credentials file. This will be stored in a k8s secret
 # and mounted into the agent container. Set this if you want to clone private
 # repos.