Include override configuration for Prefect Agent (#813)

* prefect override config * prefect override config * format files * update overrides option for Prefect helm chart * Add agent auth credentials override option * Add agent auth credentials override option * Add prefect overrides specification doc * allow arbitrary environment variables to be set on the prefect agent * Add prefect overrides extraVar option * Remove duplicate loop * Update extraEnvVars and update docs * don't require base64encoding fix * Fix missing envVars and add secretEnvVars field * Add overrides ability to create secrets * rm commented code * Update docs * Update docs * Update docs Co-authored-by: Adam-D-Lewis <> Co-authored-by: leej3 <johnleenimh@gmail.com>
nebari-dev · Sep 29, 2021 · c10358a · c10358a
1 parent bc80e7c
commit c10358a
Show file tree

Hide file tree

Showing 10 changed files with 115 additions and 19 deletions.
diff --git a/docs/source/admin_guide/prefect.md b/docs/source/admin_guide/prefect.md
@@ -33,7 +33,6 @@ The `TF_VAR_prefect_token` API key is set as `PREFECT__CLOUD__AGENT__AUTH_TOKEN`
 environment variable in the agent. It is used while deploying Prefect Agent so that
 it can connect to Prefect Cloud and query flows.
 
-
 ## Prefect Cloud
 
 Prefect Cloud is a fully hosted, production-ready backend for Prefect Core.
@@ -53,6 +52,57 @@ runs, and allocates resources for them on their respective deployment platforms.
 When you enable prefect via `qhub-config.yml` prefect agent is deployed on the
 QHub's kubernetes cluster, which querys the Prefect Cloud for flow runs.
 
+## Agent configuration overrides
+You can override your agent configuration without having to modify the helm files directly.  The extra variable `overrides` makes this
+possible by changing the default values for the Agent chart according to the settings presented on your qhub-config.yaml file.
+
+The current variables, originaly available in the [Agent helm chart](https://github.com/PrefectHQ/server/blob/master/helm/prefect-server/templates/agent/deployment.yaml) that can be overridden include:
+
+```
+- IMAGE_PULL_SECRETS
+- PREFECT__CLOUD__AGENT__LABELS
+- JOB_MEM_REQUEST
+- JOB_MEM_LIMIT
+- JOB_CPU_REQUEST
+- JOB_CPU_LIMIT
+- IMAGE_PULL_POLICY
+```
+
+For example, if you just want to override the amount of CPU limits for each job, you would need to craft a declarative configuration,
+in you qhub-config.yaml file, as follows:
+
+```yaml
+prefect:
+ enabled: true
+ overrides:
+     prefect_agent:
+       job:
+         resources:
+           limit:
+             cpu: 4
+```
+Also, if you would like to include an extra variable to the agent environment configuration, that was not previcosly in the helm chart, you can do it by including it under
+the `envVars` field in the overrides block. For example, if you would like to add `MY_VAR: "<value>"` to you agent environment, you can do so by adding the following to your qhub-config
+
+```yaml
+prefect:
+ enabled: true
+ overrides:
+     envVars:
+       MY_VAR: "<value>"
+```
+### Adding secrets to you Agent configuration
+Overrides also allow you to define extra secrets to pass through your agent configuration, for example, when using [default secrets](https://docs.prefect.io/core/concepts/secrets.html#default-secrets) to automatically authenticate your flow with the listed service. In the Google cloud case, for `GCP_CREDENTIALS` context secret, you can do it by adding that specific key value pair into your configuration:
+
+```yalm
+prefect:
+ enabled: true
+ overrides:
+   secretEnvVars:
+       PREFECT__CONTEXT__SECRETS__GCP_CREDENTIALS: '<Your value>'
+```
+This secret will then be stored as a [kubernetes secret](https://kubernetes.io/docs/concepts/configuration/secret/) variable into you QHub secrets volume.
+
 ## Flows
 
 Prefect agent can only orchestrate your flows, you need an actual flow to run via

diff --git a/qhub/render/jinja.py b/qhub/render/jinja.py
@@ -12,7 +12,7 @@ def __init__(self, environment):
         super().__init__(environment)
 
         def yamlify(obj):
-            return yaml.dump(obj)
+            return yaml.dump(obj, default_flow_style=False)
 
         environment.filters["yamlify"] = yamlify
 

diff --git a/qhub/schema.py b/qhub/schema.py
@@ -77,6 +77,7 @@ class ClearML(Base):
 class Prefect(Base):
     enabled: bool
     image: typing.Optional[str]
+    overrides: typing.Optional[typing.Dict]
 
 
 # ============= Terraform ===============

diff --git a/qhub/template/cookiecutter.json b/qhub/template/cookiecutter.json
@@ -22,7 +22,8 @@
     },
     "prefect": {
         "enabled": null,
-        "image": null
+        "image": null,
+        "overrides": null
     },
     "security": null,
     "default_images": null,

diff --git a/qhub/template/{{ cookiecutter.repo_directory }}/infrastructure/kubernetes.tf b/qhub/template/{{ cookiecutter.repo_directory }}/infrastructure/kubernetes.tf
@@ -235,6 +235,12 @@ module "prefect" {
   {% if cookiecutter.prefect.image is defined -%}
   image                = "{{ cookiecutter.prefect.image }}"
   {% endif -%}
+  {% if cookiecutter.prefect.overrides is defined %}
+  overrides            = [<<EOT
+{{ cookiecutter.prefect.overrides|yamlify -}}
+    EOT
+    ]
+  {% endif %}
 }
 {% endif -%}
 

diff --git a/...ectory }}/infrastructure/modules/kubernetes/services/prefect/chart/templates/prefect.yaml b/...ectory }}/infrastructure/modules/kubernetes/services/prefect/chart/templates/prefect.yaml
@@ -18,11 +18,22 @@ spec:
       serviceAccountName: {{ .Values.serviceAccount }}
       containers:
         - args:
-            - prefect agent start kubernetes -e JUPYTERHUB_API_TOKEN=$JUPYTERHUB_API_TOKEN
+            - prefect agent start kubernetes -e JUPYTERHUB_API_TOKEN=$JUPYTERHUB_API_TOKEN {{- range $k, $v := .Values.secretEnvVars }} {{ printf " -e %s=$%s" $k  $k -}} {{ end }}
           command:
             - /bin/bash
             - -c
           env:
+            {{- range $k, $v := .Values.envVars }}
+            - name: {{ $k }}
+              value: {{ squote $v -}}
+            {{ end }}
+            {{- range $i, $k := keys .Values.secretEnvVars }}
+            - name: {{ $k }}
+              valueFrom:
+                secretKeyRef:
+                  name: prefect-envsecret-{{ $i }}
+                  key: prefectEnvSecret{{ $i -}}
+            {{ end }}
             - name: PREFECT__CLOUD__AGENT__AUTH_TOKEN
               valueFrom:
                 secretKeyRef:
@@ -37,20 +48,6 @@ spec:
               value: {{ .Values.cloudApi }}
             - name: NAMESPACE
               value: {{ .Values.namespace }}
-            - name: IMAGE_PULL_SECRETS
-              value: ''
-            - name: PREFECT__CLOUD__AGENT__LABELS
-              value: '[]'
-            - name: JOB_MEM_REQUEST
-              value: ''
-            - name: JOB_MEM_LIMIT
-              value: ''
-            - name: JOB_CPU_REQUEST
-              value: ''
-            - name: JOB_CPU_LIMIT
-              value: ''
-            - name: IMAGE_PULL_POLICY
-              value: ''
             - name: SERVICE_ACCOUNT_NAME
               value: {{ .Values.serviceAccount }}
             - name: PREFECT__BACKEND
@@ -111,4 +108,4 @@ roleRef:
   name: prefect-agent-rbac
 subjects:
   - kind: ServiceAccount
-    name: {{ .Values.serviceAccount }}
+    name: {{ .Values.serviceAccount }}
diff --git a/...rectory }}/infrastructure/modules/kubernetes/services/prefect/chart/templates/secret.yaml b/...rectory }}/infrastructure/modules/kubernetes/services/prefect/chart/templates/secret.yaml
@@ -12,3 +12,18 @@ metadata:
   name: jupyterhub-token
 data:
   jupyterHubToken: "{{ .Values.jupyterHubToken | b64enc }}"
+
+---
+{{ $index := dict "index" "0" }}
+{{- range $k, $v := .Values.secretEnvVars }}
+{{ $i := get $index "index" }}
+{{ $ni := add1 $i }}
+apiVersion: v1
+kind: Secret
+metadata:
+  name: prefect-envsecret-{{ $i }}
+data:
+  prefectEnvSecret{{ $i }}: "{{ $v | b64enc }}"
+{{ $_ := set $index "index" $ni }}
+---
+{{ end }}
diff --git a/...cookiecutter.repo_directory }}/infrastructure/modules/kubernetes/services/prefect/main.tf b/...cookiecutter.repo_directory }}/infrastructure/modules/kubernetes/services/prefect/main.tf
@@ -3,6 +3,9 @@ resource "helm_release" "prefect" {
   namespace = var.namespace
   chart     = "${path.module}/chart"
 
+  values = concat([
+  file("${path.module}/values.yaml")], var.overrides)
+
   set_sensitive {
     name  = "prefectToken"
     value = var.prefect_token

diff --git a/...iecutter.repo_directory }}/infrastructure/modules/kubernetes/services/prefect/values.yaml b/...iecutter.repo_directory }}/infrastructure/modules/kubernetes/services/prefect/values.yaml
@@ -0,0 +1,17 @@
+
+envVars: {}
+
+secretEnvVars: {}
+
+prefect_agent:
+  job:
+    imagePullSecrets: ''
+    imagePullPolicy: ''
+    resources:
+      requests:
+        memory: ''
+        cpu: ''
+      limits:
+        memory: ''
+        cpu: ''
+  prefectLabels: '[]'
diff --git a/...ecutter.repo_directory }}/infrastructure/modules/kubernetes/services/prefect/variables.tf b/...ecutter.repo_directory }}/infrastructure/modules/kubernetes/services/prefect/variables.tf
@@ -21,3 +21,9 @@ variable "cloud_api" {
   type    = string
   default = "https://api.prefect.io"
 }
+
+variable "overrides" {
+  description = "Prefect helm chart list of overrides"
+  type        = list(string)
+  default     = []
+}