diff --git a/helm/litellm-helm/templates/migrations-job.yaml b/helm/litellm-helm/templates/migrations-job.yaml index 5a873cbb965..b199ad2d8f3 100644 --- a/helm/litellm-helm/templates/migrations-job.yaml +++ b/helm/litellm-helm/templates/migrations-job.yaml @@ -6,6 +6,9 @@ metadata: name: {{ include "litellm.fullname" . }}-migrations labels: {{- include "litellm.labels" . | nindent 4 }} + {{- with .Values.migrationJob.jobLabels }} + {{- toYaml . | nindent 4 }} + {{- end }} annotations: {{- if .Values.migrationJob.hooks.argocd.enabled }} argocd.argoproj.io/hook: PreSync @@ -17,6 +20,9 @@ metadata: helm.sh/hook-weight: {{ .Values.migrationJob.hooks.helm.weight | default "1" | quote }} {{- end }} checksum/config: {{ toYaml .Values | sha256sum }} + {{- with .Values.migrationJob.jobAnnotations }} + {{- toYaml . | nindent 4 }} + {{- end }} spec: template: metadata: @@ -25,6 +31,9 @@ spec: {{- with .Values.podLabels }} {{- toYaml . | nindent 8 }} {{- end }} + {{- with .Values.migrationJob.podLabels }} + {{- toYaml . | nindent 8 }} + {{- end }} annotations: {{- with .Values.migrationJob.annotations }} {{- toYaml . | nindent 8 }} @@ -47,7 +56,16 @@ spec: imagePullPolicy: {{ .Values.image.pullPolicy }} securityContext: {{- toYaml .Values.securityContext | nindent 12 }} + {{- if .Values.migrationJob.command }} + command: {{ toYaml .Values.migrationJob.command | nindent 12 }} + {{- else }} command: ["python", "litellm/proxy/prisma_migration.py"] + {{- end }} + {{- if .Values.migrationJob.args }} + args: {{ toYaml .Values.migrationJob.args | nindent 12 }} + {{- else if .Values.migrationJob.command }} + args: [] + {{- end }} workingDir: "/app" env: {{- if .Values.db.useExisting }} diff --git a/helm/litellm-helm/tests/migrations-job_tests.yaml b/helm/litellm-helm/tests/migrations-job_tests.yaml index dd4276ac60f..bd5dd457f95 100644 --- a/helm/litellm-helm/tests/migrations-job_tests.yaml +++ b/helm/litellm-helm/tests/migrations-job_tests.yaml @@ -360,3 +360,73 @@ tests: asserts: - notExists: path: spec.activeDeadlineSeconds + + - it: should set custom jobLabels and jobAnnotations on Job metadata + template: migrations-job.yaml + set: + migrationJob: + enabled: true + jobLabels: + environment: production + team: platform + jobAnnotations: + example.com/cost-center: "1234" + asserts: + - equal: + path: metadata.labels.environment + value: production + - equal: + path: metadata.labels.team + value: platform + - equal: + path: metadata.annotations['example.com/cost-center'] + value: "1234" + + - it: should set custom podLabels on Pod template + template: migrations-job.yaml + set: + migrationJob: + enabled: true + podLabels: + custom.io/pod-role: migration + asserts: + - equal: + path: spec.template.metadata.labels['custom.io/pod-role'] + value: migration + + - it: should override container command and args + template: migrations-job.yaml + set: + migrationJob: + enabled: true + command: + - sh + args: + - -c + - echo migrating + asserts: + - equal: + path: spec.template.spec.containers[0].command + value: + - sh + - equal: + path: spec.template.spec.containers[0].args + value: + - -c + - echo migrating + + - it: should clear container args when only command is specified + template: migrations-job.yaml + set: + migrationJob: + enabled: true + command: + - sh + asserts: + - equal: + path: spec.template.spec.containers[0].command + value: + - sh + - equal: + path: spec.template.spec.containers[0].args + value: [] diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 83dbb3c5aa0..42ed777e6bc 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -572,6 +572,11 @@ migrationJob: # In that case, pre-install/pre-upgrade hooks run before normal resources, so this defaults to "default". serviceAccountName: "" annotations: {} + jobLabels: {} # Custom labels for the Job metadata + jobAnnotations: {} # Custom annotations for the Job metadata + podLabels: {} # Custom labels for the Job pod template + command: [] # Override container command (defaults to ["python", "litellm/proxy/prisma_migration.py"]) + args: [] # Override container args ttlSecondsAfterFinished: 120 resources: {} # Unset by default. This job runs the database migration and exits, so it does not