Merge pull request #63 from stackhpc/feat/flux-image-gen

Add new flux image generator chart
stackhpc · Nov 12, 2024 · f8d17e0 · f8d17e0
2 parents 701ad09 + 2f4cffe
commit f8d17e0
Show file tree

Hide file tree

Showing 26 changed files with 1,027 additions and 8 deletions.
diff --git a/.github/workflows/build-push-images.yml b/.github/workflows/build-push-images.yml
@@ -18,6 +18,7 @@ jobs:
         include:
           - component: chat
           - component: image-analysis
+          - component: flux-image-gen
     permissions:
       contents: read
       id-token: write         # needed for signing the images with GitHub OIDC Token

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/.helmignore b/charts/flux-image-gen/.helmignore
@@ -0,0 +1,23 @@
+# Patterns to ignore when building packages.
+# This supports shell glob matching, relative path matching, and
+# negation (prefixed with !). Only one pattern per line.
+.DS_Store
+# Common VCS dirs
+.git/
+.gitignore
+.bzr/
+.bzrignore
+.hg/
+.hgignore
+.svn/
+# Common backup files
+*.swp
+*.bak
+*.tmp
+*.orig
+*~
+# Various IDEs
+.project
+.idea/
+*.tmproj
+.vscode/
diff --git a/charts/flux-image-gen/Chart.yaml b/charts/flux-image-gen/Chart.yaml
@@ -0,0 +1,9 @@
+apiVersion: v2
+name: flux-image-gen
+description: A Helm chart for running Flux image generation models on Kubernetes
+
+type: application
+
+# The version and appVersion are updated by the chart build script
+version: 0.1.0
+appVersion: local
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
+  - api_server.py
+  - --host
+  - "0.0.0.0"
diff --git a/charts/flux-image-gen/templates/NOTES.txt b/charts/flux-image-gen/templates/NOTES.txt
diff --git a/charts/flux-image-gen/templates/_helpers.tpl b/charts/flux-image-gen/templates/_helpers.tpl
@@ -0,0 +1,77 @@
+{{/*
+Expand the name of the chart.
+*/}}
+{{- define "flux-image-gen.name" -}}
+{{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }}
+{{- end }}
+
+{{/*
+Create a default fully qualified app name.
+We truncate at 63 chars because some Kubernetes name fields are limited to this (by the DNS naming spec).
+If release name contains chart name it will be used as a full name.
+*/}}
+{{- define "flux-image-gen.fullname" -}}
+{{- if .Values.fullnameOverride }}
+{{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" }}
+{{- else }}
+{{- $name := default .Chart.Name .Values.nameOverride }}
+{{- if contains $name .Release.Name }}
+{{- .Release.Name | trunc 63 | trimSuffix "-" }}
+{{- else }}
+{{- printf "%s-%s" .Release.Name $name | trunc 63 | trimSuffix "-" }}
+{{- end }}
+{{- end }}
+{{- end }}
+
+{{/*
+Create chart name and version as used by the chart label.
+*/}}
+{{- define "flux-image-gen.chart" -}}
+{{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }}
+{{- end }}
+
+{{/*
+Common labels
+*/}}
+{{- define "flux-image-gen.labels" -}}
+helm.sh/chart: {{ include "flux-image-gen.chart" . }}
+{{ include "flux-image-gen.selectorLabels" . }}
+{{- if .Chart.AppVersion }}
+app.kubernetes.io/version: {{ .Chart.AppVersion | quote }}
+{{- end }}
+app.kubernetes.io/managed-by: {{ .Release.Service }}
+{{- end }}
+
+{{/*
+Selector labels
+*/}}
+{{- define "flux-image-gen.selectorLabels" -}}
+app.kubernetes.io/name: {{ include "flux-image-gen.name" . }}
+app.kubernetes.io/instance: {{ .Release.Name }}
+{{- end }}
+
+{{/*
+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
+*/}}
+{{- define "flux-image-gen.serviceAccountName" -}}
+{{- if .Values.serviceAccount.create }}
+{{- default (include "flux-image-gen.fullname" .) .Values.serviceAccount.name }}
+{{- else }}
+{{- default "default" .Values.serviceAccount.name }}
+{{- end }}
+{{- end }}
diff --git a/charts/flux-image-gen/templates/api/deployment.yaml b/charts/flux-image-gen/templates/api/deployment.yaml
@@ -0,0 +1,104 @@
+{{- range $model := .Values.models }}
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: {{ printf "%s-%s-api" (include "flux-image-gen.fullname" $) $model }}
+  labels:
+    {{- include "flux-image-gen.labels" $ | nindent 4 }}
+    {{- include "flux-image-gen.modelLabels" . | nindent 4 }}
+spec:
+  replicas: {{ $.Values.api.replicaCount }}
+  {{- with $.Values.api.deploymentStrategy }}
+  strategy:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+  selector:
+    matchLabels:
+      {{- include "flux-image-gen.selectorLabels" $ | nindent 6 }}
+      {{- include "flux-image-gen.modelLabels" . | nindent 6 }}
+  template:
+    metadata:
+      {{- with $.Values.api.podAnnotations }}
+      annotations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      labels:
+        {{- include "flux-image-gen.labels" $ | nindent 8 }}
+        {{- include "flux-image-gen.modelLabels" . | nindent 8 }}
+        {{- with $.Values.api.podLabels }}
+        {{- toYaml . | nindent 8 }}
+        {{- end }}
+    spec:
+      {{- with $.Values.api.imagePullSecrets }}
+      imagePullSecrets:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      securityContext:
+        {{- toYaml $.Values.api.podSecurityContext | nindent 8 }}
+      containers:
+        - name: {{ $.Chart.Name }}
+          securityContext:
+            {{- 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 }}
+              protocol: TCP
+          {{- if $.Values.api.startupProbe }}
+          startupProbe:
+            {{- toYaml $.Values.api.startupProbe | nindent 12 }}
+          {{- end }}
+          {{- if $.Values.api.livenessProbe }}
+          livenessProbe:
+            {{- toYaml $.Values.api.livenessProbe | nindent 12 }}
+          {{- end }}
+          {{- if $.Values.api.readinessProbe }}
+          readinessProbe:
+            {{- toYaml $.Values.api.readinessProbe | nindent 12 }}
+          {{- end }}
+          resources:
+            {{- toYaml $.Values.api.resources | nindent 12 }}
+          {{- with $.Values.api.volumeMounts }}
+          volumeMounts:
+            {{- toYaml . | nindent 12 }}
+          {{- end }}
+          # Make stdout from python visible in k8s logs
+          tty: true
+          env:
+          - name: FLUX_MODEL_NAME
+            value: {{ $model }}
+          - name: PYTHONUNBUFFERED
+            value: "1"
+          {{- if $.Values.api.huggingfaceToken }}
+          - name: HUGGING_FACE_HUB_TOKEN
+            value: {{ quote $.Values.api.huggingfaceToken }}
+          {{- end }}
+          {{- with $.Values.api.envFrom }}
+          envFrom:
+            {{- toYaml . | nindent 12 }}
+          {{- end }}
+      {{- with $.Values.api.volumes }}
+      volumes:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with $.Values.api.nodeSelector }}
+      nodeSelector:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with $.Values.api.affinity }}
+      affinity:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with $.Values.api.tolerations }}
+      tolerations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+{{- end -}}
diff --git a/charts/flux-image-gen/templates/api/service.yaml b/charts/flux-image-gen/templates/api/service.yaml
@@ -0,0 +1,20 @@
+{{- range $model := .Values.models }}
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: {{ printf "%s-%s-api" (include "flux-image-gen.fullname" $) $model }}
+  labels:
+    {{- include "flux-image-gen.labels" $ | nindent 4 }}
+    {{- include "flux-image-gen.modelLabels" . | nindent 4 }}
+spec:
+  type: {{ $.Values.api.service.type }}
+  ports:
+    - port: {{ $.Values.api.service.port }}
+      targetPort: http
+      protocol: TCP
+      name: http
+  selector:
+    {{- include "flux-image-gen.selectorLabels" $ | nindent 4 }}
+    {{- include "flux-image-gen.modelLabels" . | nindent 4 }}
+{{- end -}}
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
+        - 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/ui/configmap.yaml b/charts/flux-image-gen/templates/ui/configmap.yaml
@@ -0,0 +1,16 @@
+---
+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 }}
+    example_prompt: |
+      {{- .Values.examplePrompt | nindent 6 -}}