From 630ef567f64b3fa143241fb3ac9cc307e5e7abd5 Mon Sep 17 00:00:00 2001
From: Flamarion Jorge <29267749+flamarion@users.noreply.github.com>
Date: Tue, 30 Jan 2024 12:16:13 +0100
Subject: [PATCH] feat: Add support for env vars on Launch Agent (#69)

---
 charts/launch-agent/Chart.yaml                |  2 +-
 charts/launch-agent/README.md                 |  2 ++
 charts/launch-agent/templates/deployment.yaml | 16 ++++++++++++++++
 charts/launch-agent/templates/secret.yaml     | 15 ++++++++++++++-
 charts/launch-agent/values.yaml               | 10 ++++++++++
 5 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/charts/launch-agent/Chart.yaml b/charts/launch-agent/Chart.yaml
index 24668396..e4ce5463 100644
--- 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
index bdc5d089..bd14972c 100644
--- 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
index 79b984eb..7c80ea87 100644
--- 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
index f194e89f..904a9357 100644
--- 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 }}
\ No newline at end of file
+{{ 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 }}
\ No newline at end of file
diff --git a/charts/launch-agent/values.yaml b/charts/launch-agent/values.yaml
index 0f8ba6a1..09a7f070 100644
--- 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.