Chart refactoring

stackhpc · Nov 11, 2024 · fecd477 · fecd477
1 parent 5ea88b3
commit fecd477
Show file tree

Hide file tree

Showing 11 changed files with 86 additions and 38 deletions.
diff --git a/.gitignore b/.gitignore
@@ -8,7 +8,7 @@ __pycache__/
 **/.ruff_cache
 
 # Ignore local dev helpers
-test-values.y[a]ml
+**/dev-values.yml
 **venv*/
 
 # Helm chart stuff

diff --git a/charts/flux-image-gen/Chart.yaml b/charts/flux-image-gen/Chart.yaml
@@ -21,4 +21,4 @@ version: 0.1.0
 # incremented each time you make changes to the application. Versions are not expected to
 # follow Semantic Versioning. They should reflect the version the application is using.
 # It is recommended to use it with quotes.
-appVersion: "4339cc5"
+appVersion: "02aeef3"
diff --git a/charts/flux-image-gen/ci/test-values.yaml b/charts/flux-image-gen/ci/test-values.yaml
@@ -0,0 +1,13 @@
+models:
+  - flux-dev
+  - flux-schnell
+api:
+  # Run in dev mode so that we skip
+  # the image gen step and can therefore
+  # test in a kind cluster
+  commandOverride:
+  - fastapi
+  - dev
+  - stackhpc-app/api_server.py
+  - --host
+  - "0.0.0.0"
diff --git a/charts/flux-image-gen/templates/_helpers.tpl b/charts/flux-image-gen/templates/_helpers.tpl
@@ -51,12 +51,20 @@ app.kubernetes.io/instance: {{ .Release.Name }}
 {{- end }}
 
 {{/*
-Selector labels
+Model selector labels
 */}}
 {{- define "flux-image-gen.modelLabels" -}}
 app.kubernetes.io/component: {{ . }}-api
 {{- end }}
 
+{{/*
+UI selector labels
+*/}}
+{{- define "flux-image-gen.uiLabels" -}}
+app.kubernetes.io/component: {{ .Release.Name }}-ui
+{{- end }}
+
+
 {{/*
 Create the name of the service account to use
 */}}

diff --git a/charts/flux-image-gen/templates/api/deployment.yaml b/charts/flux-image-gen/templates/api/deployment.yaml
@@ -42,6 +42,12 @@ spec:
             {{- toYaml $.Values.api.securityContext | nindent 12 }}
           image: "{{ $.Values.image.repository }}:{{ $.Values.image.tag | default $.Chart.AppVersion }}"
           imagePullPolicy: {{ $.Values.image.pullPolicy }}
+          {{- with $.Values.api.commandOverride }}
+          {{- if . }}
+          command:
+          {{- . | toYaml | nindent 12 }}
+          {{- end }}
+          {{- end }}
           ports:
             - name: http
               containerPort: {{ $.Values.api.service.port }}

diff --git a/charts/flux-image-gen/templates/tests/gradio-api.yaml b/charts/flux-image-gen/templates/tests/gradio-api.yaml
@@ -0,0 +1,22 @@
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: gradio-client-test
+  annotations:
+    "helm.sh/hook": test
+spec:
+  template:
+    spec:
+      containers:
+      - name: gradio-client
+        image: "{{ $.Values.image.repository }}:{{ $.Values.image.tag | default $.Chart.AppVersion }}"
+        command:
+        - python
+        - stackhpc-app/test_client.py
+        env:
+        - name: GRADIO_HOST
+          value: {{ printf "http://%s-ui.%s.svc:%v" (include "flux-image-gen.fullname" .) .Release.Namespace .Values.ui.service.port }}
+        - name: FLUX_MODEL
+          value: {{ .Values.models | first }}
+      restartPolicy: Never
+  backoffLimit: 3
diff --git a/charts/flux-image-gen/templates/tests/test-connection.yaml b/charts/flux-image-gen/templates/tests/test-connection.yaml
diff --git a/charts/flux-image-gen/templates/ui/configmap.yml b/charts/flux-image-gen/templates/ui/configmap.yml
@@ -0,0 +1,14 @@
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: {{ .Release.Name }}-ui-config
+  labels:
+    {{- include "flux-image-gen.labels" . | nindent 4 }}
+data:
+  gradio_config.yaml: |
+    models:
+    {{- range $model := .Values.models }}
+      - name: {{ . }}
+        address: {{ printf "http://%s.%s.svc:%v" ( printf "%s-%s-api" (include "flux-image-gen.fullname" $) . ) $.Release.Namespace $.Values.api.service.port }}
+    {{- end -}}
diff --git a/charts/flux-image-gen/templates/ui/deployment.yaml b/charts/flux-image-gen/templates/ui/deployment.yaml
@@ -9,6 +9,7 @@ spec:
   selector:
     matchLabels:
       {{- include "flux-image-gen.selectorLabels" . | nindent 6 }}
+      {{- include "flux-image-gen.uiLabels" . | nindent 6 }}
   template:
     metadata:
       {{- with .Values.ui.podAnnotations }}
@@ -17,6 +18,7 @@ spec:
       {{- end }}
       labels:
         {{- include "flux-image-gen.labels" . | nindent 8 }}
+        {{- include "flux-image-gen.uiLabels" . | nindent 8 }}
         {{- with .Values.ui.podLabels }}
         {{- toYaml . | nindent 8 }}
         {{- end }}
@@ -35,13 +37,7 @@ spec:
           imagePullPolicy: {{ .Values.image.pullPolicy }}
           command:
           - python
-          - gradio_ui.py
-          args:
-          - --models
-          - {{ .Values.models | join "," }}
-          {{- if .Values.ui.extraArgs }}
-          {{- .Values.ui.extraArgs | toYaml | nindent 10 }}
-          {{- end }}
+          - stackhpc-app/gradio_ui.py
           ports:
             - name: http
               containerPort: {{ .Values.ui.service.port }}
@@ -58,26 +54,29 @@ spec:
           readinessProbe:
             {{- toYaml .Values.ui.readinessProbe | nindent 12 }}
           {{- end }}
-          {{- with .Values.ui.volumeMounts }}
           volumeMounts:
+            - name: app-config
+              mountPath: /etc/gradio-app/
+          {{- with .Values.ui.volumeMounts }}
+          {{- if . -}}
             {{- toYaml . | nindent 12 }}
+          {{- end -}}
           {{- end }}
           # Make stdout from python visible in k8s logs
           tty: true
           env:
           - name: PYTHONUNBUFFERED
             value: "1"
-          - name: HELM_RELEASE_NAME
-            value: {{ include "flux-image-gen.fullname" . }}
-          - name: HELM_RELEASE_NAMESPACE
-            value: {{ .Release.Namespace }}
-          {{- with .Values.ui.envFrom }}
-          envFrom:
-            {{- toYaml . | nindent 12 }}
-          {{- end }}
-      {{- with .Values.ui.volumes }}
+          - name: GRADIO_SERVER_NAME
+            value: 0.0.0.0
       volumes:
+      - name: app-config
+        configMap:
+          name: {{ .Release.Name }}-ui-config
+      {{- with .Values.ui.volumes }}
+      {{- if . -}}
         {{- toYaml . | nindent 8 }}
+      {{- end -}}
       {{- end }}
       {{- with .Values.ui.nodeSelector }}
       nodeSelector:

diff --git a/charts/flux-image-gen/templates/ui/service.yaml b/charts/flux-image-gen/templates/ui/service.yaml
@@ -5,6 +5,7 @@ metadata:
   name: {{ include "flux-image-gen.fullname" . }}-ui
   labels:
     {{- include "flux-image-gen.labels" $ | nindent 4 }}
+    {{- include "flux-image-gen.uiLabels" . | nindent 4 }}
 spec:
   type: {{ .Values.ui.service.type }}
   ports:
@@ -14,3 +15,4 @@ spec:
       name: http
   selector:
     {{- include "flux-image-gen.selectorLabels" . | nindent 4 }}
+    {{- include "flux-image-gen.uiLabels" . | nindent 4 }}
diff --git a/charts/flux-image-gen/values.yaml b/charts/flux-image-gen/values.yaml
@@ -21,9 +21,6 @@ nameOverride: ""
 fullnameOverride: ""
 
 ui:
-
-  extraArgs:
-
   # This is for setting Kubernetes Annotations to a Pod.
   # For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/annotations/
   podAnnotations: {}
@@ -132,6 +129,8 @@ api:
 
   affinity: {}
 
+  commandOverride:
+
   # This is for setting Kubernetes Annotations to a Pod.
   # For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/annotations/
   podAnnotations: {}