diff --git a/.github/workflows/test-integration-redis.yml b/.github/workflows/test-integration-redis.yml new file mode 100644 index 00000000000..a74bcd99ce6 --- /dev/null +++ b/.github/workflows/test-integration-redis.yml @@ -0,0 +1,105 @@ +name: "Integration Tests: Redis" + +on: + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" + push: + branches: + - main + - litellm_internal_staging + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +# The unit shards run against fakeredis and mocks, which cannot fail a +# compare-and-set, a TTL expiry or a Lua compare-and-delete. Everything under +# tests/redis_integration_tests/ needs the real server, so it gets its own job +# with a Redis service rather than a service on the shared unit-test base +# workflow, which would start one for every unrelated shard. +jobs: + redis-integration: + name: redis-integration + runs-on: ubuntu-latest + timeout-minutes: 25 + + services: + redis: + image: redis:7-alpine@sha256:7aec734b2bb298a1d769fd8729f13b8514a41bf90fcdd1f38ec52267fbaa8ee6 + ports: + - 6379:6379 + options: >- + --health-cmd "redis-cli ping" + --health-interval 5s + --health-timeout 3s + --health-retries 10 + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + timeout-minutes: 3 + with: + persist-credentials: false + + - name: Detect relevant changes + id: changes + timeout-minutes: 2 + uses: ./.github/actions/detect-changes + + - name: Set up Python + timeout-minutes: 3 + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + timeout-minutes: 3 + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache uv dependencies + timeout-minutes: 5 + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cache/uv + .venv + key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }} + restore-keys: | + ${{ runner.os }}-uv- + + - name: Install dependencies + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 8 + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + + - name: Cache Prisma binaries + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + uses: ./.github/actions/cache-prisma-binaries + + - name: Generate Prisma client + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Run tests + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 10 + env: + REDIS_HOST: 127.0.0.1 + REDIS_PORT: "6379" + LITELLM_REQUIRE_LIVE_REDIS: "1" + run: | + uv run --no-sync pytest tests/redis_integration_tests \ + --tb=short -vv \ + --durations=20 diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index b242373de5d..f034e3835c8 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -53,6 +53,12 @@ If `db.useStackgresOperator` is used (not yet implemented): | `pdb.maxUnavailable` | Maximum number/percentage of pods that can be unavailable during **voluntary** disruptions (choose **one** of minAvailable/maxUnavailable) | `null` | | `pdb.annotations` | Extra metadata annotations to add to the PDB | `{}` | | `pdb.labels` | Extra metadata labels to add to the PDB | `{}` | +| `worker.enabled` | Run the proxy's single-owner auxiliary jobs (budget resets, cost polls, cleanups) on a dedicated Deployment with `LITELLM_JOB_ROLE=worker`, and set `LITELLM_JOB_ROLE=serving` on the proxy Deployment. The worker gets no Service, HPA, or KEDA scaler. | `false` | +| `worker.replicaCount` | Worker pods to deploy. The jobs are single-owner, so raising this buys availability across a rollout, not throughput. | `1` | +| `worker.resources` | CPU/memory requests and limits for the worker container. Falls back to `resources` when empty. | `{}` | +| `worker.nodeSelector` | Node selector for the worker pods, e.g. to keep them off the serving node pool. Falls back to `nodeSelector` when empty. | `{}` | +| `worker.tolerations` | Tolerations for the worker pods. Falls back to `tolerations` when empty. | `[]` | +| `worker.affinity` | Affinity rules for the worker pods. Falls back to `affinity` when empty. | `{}` | | `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` | | `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` | diff --git a/helm/litellm-helm/templates/NOTES.txt b/helm/litellm-helm/templates/NOTES.txt index 017bbfa78bd..b958a10c421 100644 --- a/helm/litellm-helm/templates/NOTES.txt +++ b/helm/litellm-helm/templates/NOTES.txt @@ -20,4 +20,10 @@ echo "Visit http://127.0.0.1:8080 to use your application" kubectl --namespace {{ .Release.Namespace }} port-forward $POD_NAME 8080:$CONTAINER_PORT {{- end }} -PDB: {{ if .Values.pdb.enabled }}enabled{{ else }}disabled{{ end }}. Configure via .Values.pdb.* \ No newline at end of file +PDB: {{ if .Values.pdb.enabled }}enabled{{ else }}disabled{{ end }}. Configure via .Values.pdb.* +{{- if .Values.worker.enabled }} + +Worker: enabled. {{ include "litellm.worker.fullname" . }} owns the single-owner scheduled jobs +(budget resets, cost polls, cleanups) and the proxy pods run with LITELLM_JOB_ROLE=serving. +The proxy pods still flush their own spend and request queues, so spend keeps being written. +{{- end }} \ No newline at end of file diff --git a/helm/litellm-helm/templates/_deployment.tpl b/helm/litellm-helm/templates/_deployment.tpl new file mode 100644 index 00000000000..621c9515b2b --- /dev/null +++ b/helm/litellm-helm/templates/_deployment.tpl @@ -0,0 +1,286 @@ +{{/* +Pod spec shared by the proxy Deployment and the optional worker Deployment. + +Invoke with a dict: `(dict "root" $ "worker" false)`. Both run the same image +against the same config, database, and Redis; they differ only in +LITELLM_JOB_ROLE and in the sizing / scheduling overrides under `.Values.worker`, +which fall back to the top-level values when unset. + +The body carries the caller's indentation, so it is included directly under a +Deployment's pod `spec:` with no nindent. +*/}} +{{- define "litellm.podSpec" -}} +{{- $root := .root -}} +{{- $isWorker := .worker -}} +{{- $worker := $root.Values.worker -}} +{{- $resources := ternary (default $root.Values.resources $worker.resources) $root.Values.resources $isWorker -}} +{{- $nodeSelector := ternary (default $root.Values.nodeSelector $worker.nodeSelector) $root.Values.nodeSelector $isWorker -}} +{{- $tolerations := ternary (default $root.Values.tolerations $worker.tolerations) $root.Values.tolerations $isWorker -}} +{{- $affinity := ternary (default $root.Values.affinity $worker.affinity) $root.Values.affinity $isWorker -}} +{{- with $root -}} + {{- with .Values.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + serviceAccountName: {{ include "litellm.serviceAccountName" . }} + securityContext: + {{- toYaml .Values.podSecurityContext | nindent 8 }} + {{- with .Values.extraInitContainers }} + initContainers: + {{- tpl (toYaml .) $root | nindent 8 }} + {{- end }} + containers: + - name: {{ include "litellm.name" . }} + securityContext: + {{- toYaml .Values.securityContext | nindent 12 }} + image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + env: + - name: HOST + value: "{{ .Values.listen | default "0.0.0.0" }}" + - name: PORT + value: {{ .Values.service.port | quote}} + {{- if .Values.db.deployStandalone }} + - name: DATABASE_USERNAME + valueFrom: + secretKeyRef: + name: {{ include "litellm.fullname" . }}-dbcredentials + key: username + - name: DATABASE_PASSWORD + valueFrom: + secretKeyRef: + name: {{ include "litellm.fullname" . }}-dbcredentials + key: password + - name: DATABASE_HOST + value: {{ .Release.Name }}-postgresql + - name: DATABASE_NAME + value: litellm + {{- else if .Values.db.useExisting }} + - name: DATABASE_USERNAME + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.usernameKey }} + - name: DATABASE_PASSWORD + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.passwordKey }} + - name: DATABASE_HOST + {{- if .Values.db.secret.endpointKey }} + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.endpointKey }} + {{- else }} + value: {{ .Values.db.endpoint }} + {{- end }} + - name: DATABASE_NAME + value: {{ .Values.db.database }} + - name: DATABASE_URL + value: {{ .Values.db.url | quote }} + {{- end }} + {{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }} + - name: DATABASE_READER_HOST + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.readReplicaEndpointKey }} + {{- end }} + {{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }} + - name: DATABASE_URL_READ_REPLICA + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.readReplicaUrlKey }} + {{- else if .Values.db.readReplicaUrl }} + - name: DATABASE_URL_READ_REPLICA + value: {{ .Values.db.readReplicaUrl | quote }} + {{- end }} + - name: PROXY_MASTER_KEY + valueFrom: + secretKeyRef: + name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }} + key: {{ .Values.masterkeySecretKey | default "masterkey" }} + {{- if .Values.redis.enabled }} + - name: REDIS_HOST + value: {{ include "litellm.redis.serviceName" . }} + - name: REDIS_PORT + value: {{ include "litellm.redis.port" . | quote }} + - name: REDIS_PASSWORD + valueFrom: + secretKeyRef: + name: {{ include "redis.secretName" .Subcharts.redis }} + key: {{include "redis.secretPasswordKey" .Subcharts.redis }} + {{- end }} + {{- /* + Inject LITELLM_LOG only when envVars does not already define it. + */}} + {{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }} + - name: LITELLM_LOG + value: {{ .Values.logLevel | quote }} + {{- end }} + {{- if .Values.envVars }} + {{- range $key, $val := .Values.envVars }} + - name: {{ $key }} + value: {{ $val | quote }} + {{- end }} + {{- end }} + {{- with .Values.extraEnvVars }} + {{- toYaml . | nindent 12 }} + {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsEnv" . | nindent 12 }} + {{- end }} + {{- if .Values.migrationJob.enabled }} + # Schema updates are owned by the dedicated migrations Job; skip + # the proxy's startup `prisma db push` so N replicas don't race + # one DB on every rollout. Placed last (after envVars and + # extraEnvVars) so this override can't be silently shadowed by a + # user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env + # semantics — same pattern the migrations Job uses. + - name: DISABLE_SCHEMA_UPDATE + value: "true" + {{- end }} + {{- if $isWorker }} + - name: LITELLM_JOB_ROLE + value: "worker" + {{- else if .Values.worker.enabled }} + - name: LITELLM_JOB_ROLE + value: "serving" + {{- end }} + envFrom: + {{- range .Values.environmentSecrets }} + - secretRef: + name: {{ . }} + {{- end }} + {{- range .Values.environmentConfigMaps }} + - configMapRef: + name: {{ . }} + {{- end }} + {{- if .Values.command }} + command: {{ toYaml .Values.command | nindent 12 }} + {{- end }} + {{- if .Values.args }} + args: {{ toYaml .Values.args | nindent 12 }} + {{- else }} + args: + - --config + - /etc/litellm/config.yaml + {{- if $isWorker }} + # The scheduler is registered once per uvicorn worker process, not + # once per pod, so numWorkers copies of every job would run here. + - --num_workers + - "1" + {{- else }} + {{ if .Values.numWorkers }} + - --num_workers + - {{ .Values.numWorkers | quote }} + {{- end }} + {{- end }} + {{- end }} + ports: + - name: http + containerPort: {{ .Values.service.port }} + protocol: TCP + livenessProbe: + httpGet: + path: {{ .Values.livenessProbe.path | quote }} + port: "http" + initialDelaySeconds: {{ .Values.livenessProbe.initialDelaySeconds }} + periodSeconds: {{ .Values.livenessProbe.periodSeconds }} + timeoutSeconds: {{ .Values.livenessProbe.timeoutSeconds }} + successThreshold: {{ .Values.livenessProbe.successThreshold }} + failureThreshold: {{ .Values.livenessProbe.failureThreshold }} + readinessProbe: + httpGet: + path: {{ .Values.readinessProbe.path | quote }} + port: "http" + initialDelaySeconds: {{ .Values.readinessProbe.initialDelaySeconds }} + periodSeconds: {{ .Values.readinessProbe.periodSeconds }} + timeoutSeconds: {{ .Values.readinessProbe.timeoutSeconds }} + successThreshold: {{ .Values.readinessProbe.successThreshold }} + failureThreshold: {{ .Values.readinessProbe.failureThreshold }} + startupProbe: + httpGet: + path: {{ .Values.startupProbe.path | quote }} + port: "http" + initialDelaySeconds: {{ .Values.startupProbe.initialDelaySeconds }} + periodSeconds: {{ .Values.startupProbe.periodSeconds }} + timeoutSeconds: {{ .Values.startupProbe.timeoutSeconds }} + successThreshold: {{ .Values.startupProbe.successThreshold }} + failureThreshold: {{ .Values.startupProbe.failureThreshold }} + resources: + {{- toYaml $resources | nindent 12 }} + volumeMounts: + - name: litellm-config + mountPath: /etc/litellm/config.yaml + subPath: config.yaml + {{ if .Values.securityContext.readOnlyRootFilesystem }} + - name: tmp + mountPath: /tmp + - name: cache + mountPath: /.cache + - name: npm + mountPath: /.npm + {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} + {{- end }} + {{- with .Values.volumeMounts }} + {{- toYaml . | nindent 12 }} + {{- end }} + {{- with .Values.lifecycle }} + lifecycle: + {{- toYaml . | nindent 12 }} + {{- end }} + {{- with .Values.extraContainers }} + {{- tpl (toYaml .) $root | nindent 8 }} + {{- end }} + volumes: + {{ if .Values.securityContext.readOnlyRootFilesystem }} + - name: tmp + emptyDir: + sizeLimit: 500Mi + - name: cache + emptyDir: + sizeLimit: 500Mi + - name: npm + emptyDir: + sizeLimit: 500Mi + {{- end }} + - name: litellm-config + configMap: + {{- if .Values.proxyConfigMap.create }} + name: {{ include "litellm.fullname" . }}-config + {{- else }} + name: {{ .Values.proxyConfigMap.name }} + {{- end }} + items: + - key: {{ .Values.proxyConfigMap.key | default "config.yaml" }} + path: "config.yaml" + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} + {{- end }} + {{- with .Values.volumes }} + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with $nodeSelector }} + nodeSelector: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with $affinity }} + affinity: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with $tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} + terminationGracePeriodSeconds: {{ .Values.terminationGracePeriodSeconds | default 90 }} + {{- if .Values.topologySpreadConstraints }} + topologySpreadConstraints: + {{- toYaml .Values.topologySpreadConstraints | nindent 8 }} + {{- end }} +{{- end -}} +{{- end -}} diff --git a/helm/litellm-helm/templates/_helpers.tpl b/helm/litellm-helm/templates/_helpers.tpl index 8f2acb20fce..08ffeda59ba 100644 --- a/helm/litellm-helm/templates/_helpers.tpl +++ b/helm/litellm-helm/templates/_helpers.tpl @@ -50,6 +50,39 @@ app.kubernetes.io/name: {{ include "litellm.name" . }} app.kubernetes.io/instance: {{ .Release.Name }} {{- end }} +{{/* +Worker naming and labels. + +The worker takes a distinct `app.kubernetes.io/name` rather than the shared one +plus a component label. The Service, PodDisruptionBudget, and ServiceMonitor all +select on name + instance only, so a worker sharing the name would be picked up +as a proxy endpoint. Narrowing those selectors instead would mean editing the +Service selector of a running release, which drops every existing pod out of the +endpoints until the new ones roll. +*/}} +{{- define "litellm.worker.name" -}} +{{- printf "%s-worker" (include "litellm.name" .) | trunc 63 | trimSuffix "-" }} +{{- end }} + +{{- define "litellm.worker.fullname" -}} +{{- printf "%s-worker" (include "litellm.fullname" .) | trunc 63 | trimSuffix "-" }} +{{- end }} + +{{- define "litellm.worker.selectorLabels" -}} +app.kubernetes.io/name: {{ include "litellm.worker.name" . }} +app.kubernetes.io/instance: {{ .Release.Name }} +{{- end }} + +{{- define "litellm.worker.labels" -}} +helm.sh/chart: {{ include "litellm.chart" . }} +{{ include "litellm.worker.selectorLabels" . }} +app.kubernetes.io/component: worker +{{- if .Chart.AppVersion }} +app.kubernetes.io/version: {{ .Chart.AppVersion | quote }} +{{- end }} +app.kubernetes.io/managed-by: {{ .Release.Service }} +{{- end }} + {{/* Enterprise billable-request metering. The client certificate identifies the deployment to LiteLLM's collector, so it is mounted read-only from an existing diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index 52ffd117535..d8ce006254f 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -38,253 +38,4 @@ spec: {{- toYaml . | nindent 8 }} {{- end }} spec: - {{- with .Values.imagePullSecrets }} - imagePullSecrets: - {{- toYaml . | nindent 8 }} - {{- end }} - serviceAccountName: {{ include "litellm.serviceAccountName" . }} - securityContext: - {{- toYaml .Values.podSecurityContext | nindent 8 }} - {{- with .Values.extraInitContainers }} - initContainers: - {{- tpl (toYaml .) $ | nindent 8 }} - {{- end }} - containers: - - name: {{ include "litellm.name" . }} - securityContext: - {{- toYaml .Values.securityContext | nindent 12 }} - image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.image.pullPolicy }} - env: - - name: HOST - value: "{{ .Values.listen | default "0.0.0.0" }}" - - name: PORT - value: {{ .Values.service.port | quote}} - {{- if .Values.db.deployStandalone }} - - name: DATABASE_USERNAME - valueFrom: - secretKeyRef: - name: {{ include "litellm.fullname" . }}-dbcredentials - key: username - - name: DATABASE_PASSWORD - valueFrom: - secretKeyRef: - name: {{ include "litellm.fullname" . }}-dbcredentials - key: password - - name: DATABASE_HOST - value: {{ .Release.Name }}-postgresql - - name: DATABASE_NAME - value: litellm - {{- else if .Values.db.useExisting }} - - name: DATABASE_USERNAME - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.usernameKey }} - - name: DATABASE_PASSWORD - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.passwordKey }} - - name: DATABASE_HOST - {{- if .Values.db.secret.endpointKey }} - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.endpointKey }} - {{- else }} - value: {{ .Values.db.endpoint }} - {{- end }} - - name: DATABASE_NAME - value: {{ .Values.db.database }} - - name: DATABASE_URL - value: {{ .Values.db.url | quote }} - {{- end }} - {{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }} - - name: DATABASE_READER_HOST - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.readReplicaEndpointKey }} - {{- end }} - {{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }} - - name: DATABASE_URL_READ_REPLICA - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.readReplicaUrlKey }} - {{- else if .Values.db.readReplicaUrl }} - - name: DATABASE_URL_READ_REPLICA - value: {{ .Values.db.readReplicaUrl | quote }} - {{- end }} - - name: PROXY_MASTER_KEY - valueFrom: - secretKeyRef: - name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }} - key: {{ .Values.masterkeySecretKey | default "masterkey" }} - {{- if .Values.redis.enabled }} - - name: REDIS_HOST - value: {{ include "litellm.redis.serviceName" . }} - - name: REDIS_PORT - value: {{ include "litellm.redis.port" . | quote }} - - name: REDIS_PASSWORD - valueFrom: - secretKeyRef: - name: {{ include "redis.secretName" .Subcharts.redis }} - key: {{include "redis.secretPasswordKey" .Subcharts.redis }} - {{- end }} - {{- /* - Inject LITELLM_LOG only when envVars does not already define it. - */}} - {{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }} - - name: LITELLM_LOG - value: {{ .Values.logLevel | quote }} - {{- end }} - {{- if .Values.envVars }} - {{- range $key, $val := .Values.envVars }} - - name: {{ $key }} - value: {{ $val | quote }} - {{- end }} - {{- end }} - {{- with .Values.extraEnvVars }} - {{- toYaml . | nindent 12 }} - {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsEnv" . | nindent 12 }} - {{- end }} - {{- if .Values.migrationJob.enabled }} - # Schema updates are owned by the dedicated migrations Job; skip - # the proxy's startup `prisma db push` so N replicas don't race - # one DB on every rollout. Placed last (after envVars and - # extraEnvVars) so this override can't be silently shadowed by a - # user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env - # semantics — same pattern the migrations Job uses. - - name: DISABLE_SCHEMA_UPDATE - value: "true" - {{- end }} - envFrom: - {{- range .Values.environmentSecrets }} - - secretRef: - name: {{ . }} - {{- end }} - {{- range .Values.environmentConfigMaps }} - - configMapRef: - name: {{ . }} - {{- end }} - {{- if .Values.command }} - command: {{ toYaml .Values.command | nindent 12 }} - {{- end }} - {{- if .Values.args }} - args: {{ toYaml .Values.args | nindent 12 }} - {{- else }} - args: - - --config - - /etc/litellm/config.yaml - {{ if .Values.numWorkers }} - - --num_workers - - {{ .Values.numWorkers | quote }} - {{- end }} - {{- end }} - ports: - - name: http - containerPort: {{ .Values.service.port }} - protocol: TCP - livenessProbe: - httpGet: - path: {{ .Values.livenessProbe.path | quote }} - port: "http" - initialDelaySeconds: {{ .Values.livenessProbe.initialDelaySeconds }} - periodSeconds: {{ .Values.livenessProbe.periodSeconds }} - timeoutSeconds: {{ .Values.livenessProbe.timeoutSeconds }} - successThreshold: {{ .Values.livenessProbe.successThreshold }} - failureThreshold: {{ .Values.livenessProbe.failureThreshold }} - readinessProbe: - httpGet: - path: {{ .Values.readinessProbe.path | quote }} - port: "http" - initialDelaySeconds: {{ .Values.readinessProbe.initialDelaySeconds }} - periodSeconds: {{ .Values.readinessProbe.periodSeconds }} - timeoutSeconds: {{ .Values.readinessProbe.timeoutSeconds }} - successThreshold: {{ .Values.readinessProbe.successThreshold }} - failureThreshold: {{ .Values.readinessProbe.failureThreshold }} - startupProbe: - httpGet: - path: {{ .Values.startupProbe.path | quote }} - port: "http" - initialDelaySeconds: {{ .Values.startupProbe.initialDelaySeconds }} - periodSeconds: {{ .Values.startupProbe.periodSeconds }} - timeoutSeconds: {{ .Values.startupProbe.timeoutSeconds }} - successThreshold: {{ .Values.startupProbe.successThreshold }} - failureThreshold: {{ .Values.startupProbe.failureThreshold }} - resources: - {{- toYaml .Values.resources | nindent 12 }} - volumeMounts: - - name: litellm-config - mountPath: /etc/litellm/config.yaml - subPath: config.yaml - {{ if .Values.securityContext.readOnlyRootFilesystem }} - - name: tmp - mountPath: /tmp - - name: cache - mountPath: /.cache - - name: npm - mountPath: /.npm - {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} - {{- with .Values.volumeMounts }} - {{- toYaml . | nindent 12 }} - {{- end }} - {{- with .Values.lifecycle }} - lifecycle: - {{- toYaml . | nindent 12 }} - {{- end }} - {{- with .Values.extraContainers }} - {{- tpl (toYaml .) $ | nindent 8 }} - {{- end }} - volumes: - {{ if .Values.securityContext.readOnlyRootFilesystem }} - - name: tmp - emptyDir: - sizeLimit: 500Mi - - name: cache - emptyDir: - sizeLimit: 500Mi - - name: npm - emptyDir: - sizeLimit: 500Mi - {{- end }} - - name: litellm-config - configMap: - {{- if .Values.proxyConfigMap.create }} - name: {{ include "litellm.fullname" . }}-config - {{- else }} - name: {{ .Values.proxyConfigMap.name }} - {{- end }} - items: - - key: {{ .Values.proxyConfigMap.key | default "config.yaml" }} - path: "config.yaml" - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} - {{- end }} - {{- with .Values.volumes }} - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.nodeSelector }} - nodeSelector: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.affinity }} - affinity: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.tolerations }} - tolerations: - {{- toYaml . | nindent 8 }} - {{- end }} - terminationGracePeriodSeconds: {{ .Values.terminationGracePeriodSeconds | default 90 }} - {{- if .Values.topologySpreadConstraints }} - topologySpreadConstraints: - {{- toYaml .Values.topologySpreadConstraints | nindent 8 }} - {{- end }} + {{- include "litellm.podSpec" (dict "root" $ "worker" false) }} diff --git a/helm/litellm-helm/templates/worker-deployment.yaml b/helm/litellm-helm/templates/worker-deployment.yaml new file mode 100644 index 00000000000..8342d099c90 --- /dev/null +++ b/helm/litellm-helm/templates/worker-deployment.yaml @@ -0,0 +1,38 @@ +{{- if .Values.worker.enabled }} +apiVersion: apps/v1 +kind: Deployment +metadata: + annotations: + {{- toYaml .Values.deploymentAnnotations | nindent 4 }} + name: {{ include "litellm.worker.fullname" . }} + labels: + {{- include "litellm.worker.labels" . | nindent 4 }} + {{- if .Values.deploymentLabels }} + {{- toYaml .Values.deploymentLabels | nindent 4 }} + {{- end }} +spec: + replicas: {{ .Values.worker.replicaCount }} + {{- with .Values.strategy }} + strategy: + {{- toYaml . | nindent 4 }} + {{- end }} + selector: + matchLabels: + {{- include "litellm.worker.selectorLabels" . | nindent 6 }} + template: + metadata: + annotations: + {{- if .Values.proxyConfigMap.create }} + checksum/config: {{ include (print $.Template.BasePath "/configmap-litellm.yaml") . | sha256sum }} + {{- end }} + {{- with .Values.podAnnotations }} + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} + labels: + {{- include "litellm.worker.labels" . | nindent 8 }} + {{- with .Values.podLabels }} + {{- toYaml . | nindent 8 }} + {{- end }} + spec: + {{- include "litellm.podSpec" (dict "root" $ "worker" true) }} +{{- end }} diff --git a/helm/litellm-helm/tests/worker_deployment_tests.yaml b/helm/litellm-helm/tests/worker_deployment_tests.yaml new file mode 100644 index 00000000000..c8c411b9d1a --- /dev/null +++ b/helm/litellm-helm/tests/worker_deployment_tests.yaml @@ -0,0 +1,316 @@ +suite: test worker deployment disabled +templates: + - deployment.yaml + - worker-deployment.yaml + - configmap-litellm.yaml +tests: + - it: renders no worker Deployment by default + template: worker-deployment.yaml + asserts: + - hasDocuments: + count: 0 + + - it: adds nothing to the proxy Deployment env by default + template: deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "serving" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "worker" + - lengthEqual: + path: spec.template.spec.containers[0].env + count: 9 + +--- +suite: test worker deployment enabled +templates: + - deployment.yaml + - worker-deployment.yaml + - configmap-litellm.yaml +set: + worker.enabled: true +tests: + - it: renders exactly one worker Deployment with one replica + template: worker-deployment.yaml + asserts: + - hasDocuments: + count: 1 + - isKind: + of: Deployment + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-worker + - equal: + path: spec.replicas + value: 1 + + - it: runs the worker job role + template: worker-deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "worker" + + - it: flips the proxy Deployment to the serving job role + template: deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "serving" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "worker" + + - it: keeps the worker out of the proxy selector labels + template: worker-deployment.yaml + asserts: + - equal: + path: spec.selector.matchLabels + value: + app.kubernetes.io/name: litellm-worker + app.kubernetes.io/instance: RELEASE-NAME + - equal: + path: spec.template.metadata.labels["app.kubernetes.io/name"] + value: litellm-worker + - equal: + path: spec.template.metadata.labels["app.kubernetes.io/component"] + value: worker + + - it: reaches the same database, master key, and config as the proxy + template: worker-deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROXY_MASTER_KEY + valueFrom: + secretKeyRef: + name: RELEASE-NAME-litellm-masterkey + key: masterkey + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_HOST + value: RELEASE-NAME-postgresql + - contains: + path: spec.template.spec.volumes + content: + name: litellm-config + configMap: + name: RELEASE-NAME-litellm-config + items: + - key: config.yaml + path: "config.yaml" + + - it: shares the proxy image + template: worker-deployment.yaml + set: + image.tag: test + asserts: + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:test + + - it: runs one uvicorn worker process whatever the proxy runs + set: + numWorkers: 4 + asserts: + - equal: + path: spec.template.spec.containers[0].args + value: + - --config + - /etc/litellm/config.yaml + - --num_workers + - "1" + template: worker-deployment.yaml + - equal: + path: spec.template.spec.containers[0].args + value: + - --config + - /etc/litellm/config.yaml + - --num_workers + - "4" + template: deployment.yaml + + - it: runs one uvicorn worker process when the proxy sets none + template: worker-deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].args + value: + - --config + - /etc/litellm/config.yaml + - --num_workers + - "1" + +--- +suite: test worker deployment is not a serving target +templates: + - service.yaml + - hpa.yaml + - keda.yaml + - poddisruptionbudget.yaml + - worker-deployment.yaml + - configmap-litellm.yaml +set: + worker.enabled: true +tests: + - it: does not route Service traffic to worker pods + template: service.yaml + asserts: + - equal: + path: spec.selector + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + + - it: scales only the proxy Deployment with the HPA + template: hpa.yaml + set: + autoscaling.enabled: true + asserts: + - equal: + path: spec.scaleTargetRef.name + value: RELEASE-NAME-litellm + + - it: scales only the proxy Deployment with KEDA + template: keda.yaml + set: + keda.enabled: true + asserts: + - equal: + path: spec.scaleTargetRef.name + value: RELEASE-NAME-litellm + + - it: does not budget worker pod disruptions + template: poddisruptionbudget.yaml + set: + pdb.enabled: true + pdb.minAvailable: 1 + asserts: + - equal: + path: spec.selector.matchLabels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + +--- +suite: test worker deployment sizing and scheduling overrides +templates: + - deployment.yaml + - worker-deployment.yaml + - configmap-litellm.yaml +set: + worker.enabled: true +tests: + - it: inherits the proxy sizing and scheduling when unset + set: + resources: + requests: + cpu: "2" + nodeSelector: + pool: serving + tolerations: + - key: serving + operator: Exists + asserts: + - equal: + path: spec.template.spec.containers[0].resources.requests.cpu + value: "2" + template: worker-deployment.yaml + - equal: + path: spec.template.spec.nodeSelector.pool + value: serving + template: worker-deployment.yaml + - equal: + path: spec.template.spec.tolerations[0].key + value: serving + template: worker-deployment.yaml + + - it: overrides sizing and scheduling without touching the proxy + set: + resources: + requests: + cpu: "2" + nodeSelector: + pool: serving + tolerations: + - key: serving + operator: Exists + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: serving + operator: Exists + worker: + resources: + requests: + cpu: 100m + nodeSelector: + pool: jobs + tolerations: + - key: jobs + operator: Exists + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: jobs + operator: Exists + asserts: + - equal: + path: spec.template.spec.containers[0].resources.requests.cpu + value: 100m + template: worker-deployment.yaml + - equal: + path: spec.template.spec.nodeSelector.pool + value: jobs + template: worker-deployment.yaml + - equal: + path: spec.template.spec.tolerations[0].key + value: jobs + template: worker-deployment.yaml + - equal: + path: spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].key + value: jobs + template: worker-deployment.yaml + - equal: + path: spec.template.spec.containers[0].resources.requests.cpu + value: "2" + template: deployment.yaml + - equal: + path: spec.template.spec.nodeSelector.pool + value: serving + template: deployment.yaml + - equal: + path: spec.template.spec.tolerations[0].key + value: serving + template: deployment.yaml + - equal: + path: spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].key + value: serving + template: deployment.yaml + + - it: replicas follow worker.replicaCount + template: worker-deployment.yaml + set: + worker.replicaCount: 2 + asserts: + - equal: + path: spec.replicas + value: 2 diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index f8df98de102..4dfdf154bb2 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -257,6 +257,38 @@ tolerations: [] affinity: {} +# Optional dedicated worker Deployment for the proxy's single-owner auxiliary +# jobs (budget resets, cost polls, cleanups). Off by default, which leaves +# LITELLM_JOB_ROLE unset everywhere and every proxy pod eligible to own them, +# exactly as before. +# +# Enabling this sets LITELLM_JOB_ROLE=worker here and LITELLM_JOB_ROLE=serving on +# the proxy Deployment, so the jobs stop competing with inference traffic for the +# event loop and the database pool. The worker gets no Service, HPA, or KEDA +# scaler; nothing routes to it. Its pods also take their own +# `app.kubernetes.io/name`, so they stay out of the proxy Service, PDB, and +# ServiceMonitor selectors. +# +# Only the cluster-wide scheduled jobs move. A serving pod keeps flushing its own +# in-memory spend and request queues and keeps refreshing its own model registry +# and credentials, so it still writes spend exactly as it does today. +# +# The jobs are single-owner, so raising replicaCount buys availability across a +# rollout, not throughput. The worker also pins `--num_workers 1` rather than +# inheriting `numWorkers`: the scheduler is registered once per uvicorn worker +# process, so N processes would run N copies of every job in one pod. +# Everything else (image, config, probes, env, +# securityContext, topologySpreadConstraints) is inherited from the proxy; the +# keys below override sizing and scheduling only, and fall back to the top-level +# values when left empty. +worker: + enabled: false + replicaCount: 1 + resources: {} + nodeSelector: {} + tolerations: [] + affinity: {} + db: # Use an existing postgres server/cluster useExisting: false diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 72f7f74bcf6..791ac7f9974 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -27,6 +27,10 @@ Common naming + label helpers shared by gateway, backend, and ui templates. {{- printf "%s-ui" (include "litellm.fullname" .) | trunc 63 | trimSuffix "-" -}} {{- end -}} +{{- define "litellm.worker.fullname" -}} +{{- printf "%s-worker" (include "litellm.fullname" .) | trunc 63 | trimSuffix "-" -}} +{{- end -}} + {{- define "litellm.commonLabels" -}} app.kubernetes.io/name: {{ include "litellm.name" . }} app.kubernetes.io/instance: {{ .Release.Name }} @@ -106,6 +110,12 @@ app.kubernetes.io/instance: {{ .Release.Name }} app.kubernetes.io/component: ui {{- end -}} +{{- define "litellm.worker.selectorLabels" -}} +app.kubernetes.io/name: {{ include "litellm.name" . }} +app.kubernetes.io/instance: {{ .Release.Name }} +app.kubernetes.io/component: worker +{{- end -}} + {{/* Per-component ServiceAccount name helpers. diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 5c0431fc0bd..37cccd63892 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -58,6 +58,10 @@ spec: {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsEnv" . | nindent 12 }} {{- end }} + {{- if .Values.worker.enabled }} + - name: LITELLM_JOB_ROLE + value: "serving" + {{- end }} {{- include "litellm.envFrom" .Values.backend | nindent 10 }} {{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }} volumeMounts: diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index d5363d0096e..bb1e3912462 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -60,6 +60,10 @@ spec: {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsEnv" . | nindent 12 }} {{- end }} + {{- if .Values.worker.enabled }} + - name: LITELLM_JOB_ROLE + value: "serving" + {{- end }} {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }} volumeMounts: diff --git a/helm/litellm/templates/worker/deployment.yaml b/helm/litellm/templates/worker/deployment.yaml new file mode 100644 index 00000000000..a2f077ae7b7 --- /dev/null +++ b/helm/litellm/templates/worker/deployment.yaml @@ -0,0 +1,104 @@ +{{- if .Values.worker.enabled }} +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ include "litellm.worker.fullname" . }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: worker +spec: + replicas: {{ .Values.worker.replicaCount }} + selector: + matchLabels: + {{- include "litellm.worker.selectorLabels" . | nindent 6 }} + template: + metadata: + {{- if or .Values.gateway.config.create .Values.worker.podAnnotations }} + annotations: + {{- if .Values.gateway.config.create }} + checksum/config: {{ include (print $.Template.BasePath "/gateway/configmap.yaml") . | sha256sum }} + {{- end }} + {{- with .Values.worker.podAnnotations }} + {{- toYaml . | nindent 8 }} + {{- end }} + {{- end }} + labels: + {{- include "litellm.worker.selectorLabels" . | nindent 8 }} + {{- with .Values.worker.podLabels }} + {{- include "litellm.podLabels" (dict "podLabels" . "componentName" "worker") | nindent 8 }} + {{- end }} + spec: + serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }} + automountServiceAccountToken: {{ .Values.serviceAccounts.backend.automount }} + {{- with .Values.worker.podSecurityContext }} + securityContext: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + containers: + - name: worker + image: "{{ .Values.backend.image.repository }}:{{ .Values.backend.image.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.backend.image.pullPolicy }} + {{- with .Values.worker.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} + ports: + - name: http + containerPort: 4001 + protocol: TCP + env: + {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.worker) | nindent 12 }} + {{- if .Values.gateway.config.create }} + - name: CONFIG_FILE_PATH + value: /app/config/config.yaml + {{- end }} + # The scheduler is registered once per uvicorn worker process, not + # once per pod, so anything above 1 runs that many copies of every job. + - name: NUM_WORKERS + value: "1" + - name: LITELLM_JOB_ROLE + value: "worker" + {{- include "litellm.envFrom" .Values.worker | nindent 10 }} + {{- if .Values.gateway.config.create }} + volumeMounts: + - name: gateway-config + mountPath: /app/config/config.yaml + subPath: config.yaml + {{- end }} + {{- with .Values.worker.livenessProbe }} + livenessProbe: + {{- toYaml . | nindent 12 }} + {{- end }} + {{- with .Values.worker.readinessProbe }} + readinessProbe: + {{- toYaml . | nindent 12 }} + {{- end }} + resources: + {{- toYaml .Values.worker.resources | nindent 12 }} + {{- if .Values.gateway.config.create }} + volumes: + - name: gateway-config + configMap: + name: {{ include "litellm.gateway.fullname" . }}-config + {{- end }} + {{- with .Values.worker.nodeSelector }} + nodeSelector: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.worker.affinity }} + affinity: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.worker.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- $gracePeriod := .Values.worker.terminationGracePeriodSeconds }} + {{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }} + terminationGracePeriodSeconds: {{ $gracePeriod }} + {{- end }} +{{- end }} diff --git a/helm/litellm/tests/worker_deployment_tests.yaml b/helm/litellm/tests/worker_deployment_tests.yaml new file mode 100644 index 00000000000..5abed050f0f --- /dev/null +++ b/helm/litellm/tests/worker_deployment_tests.yaml @@ -0,0 +1,269 @@ +suite: test worker component disabled +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - worker/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: renders no worker Deployment by default + template: worker/deployment.yaml + asserts: + - hasDocuments: + count: 0 + + - it: adds nothing to the gateway env by default + template: gateway/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "serving" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "worker" + - lengthEqual: + path: spec.template.spec.containers[0].env + count: 10 + + - it: adds nothing to the backend env by default + template: backend/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "serving" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "worker" + - lengthEqual: + path: spec.template.spec.containers[0].env + count: 9 + +--- +suite: test worker component enabled +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - worker/deployment.yaml +values: + - ./values/required.yaml +set: + worker.enabled: true +tests: + - it: renders exactly one worker Deployment with one replica + template: worker/deployment.yaml + asserts: + - hasDocuments: + count: 1 + - isKind: + of: Deployment + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-worker + - equal: + path: spec.replicas + value: 1 + + - it: runs the worker job role + template: worker/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "worker" + + - it: flips the gateway to the serving job role + template: gateway/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "serving" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "worker" + + - it: flips the backend to the serving job role + template: backend/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "serving" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: "worker" + + - it: carries its own component selector + template: worker/deployment.yaml + asserts: + - equal: + path: spec.selector.matchLabels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: worker + - equal: + path: spec.template.metadata.labels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: worker + + - it: reaches the same database, master key, and config as the backend + template: worker/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_MASTER_KEY + valueFrom: + secretKeyRef: + name: litellm-master-key-secret + key: master-key + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_HOST + value: postgres.example.com + - contains: + path: spec.template.spec.volumes + content: + name: gateway-config + configMap: + name: RELEASE-NAME-litellm-gateway-config + - equal: + path: spec.template.spec.serviceAccountName + value: default + + - it: shares the backend image + template: worker/deployment.yaml + set: + backend.image.tag: test + asserts: + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm-backend:test + + - it: runs one uvicorn worker process whatever the gateway runs + set: + gateway.numWorkers: 4 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: NUM_WORKERS + value: "1" + template: worker/deployment.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: NUM_WORKERS + value: "4" + template: gateway/deployment.yaml + + - it: rejects podLabels that collide with its selector + template: worker/deployment.yaml + set: + worker.podLabels: + app.kubernetes.io/component: something-else + asserts: + - failedTemplate: + errorMessage: "worker.podLabels cannot set app.kubernetes.io/component: it is part of the Deployment's immutable selector" + +--- +suite: test worker component is not a serving target +templates: + - gateway/service.yaml + - gateway/hpa.yaml + - backend/service.yaml + - backend/hpa.yaml + - ui/service.yaml + - worker/deployment.yaml + - gateway/configmap.yaml +values: + - ./values/required.yaml +set: + worker.enabled: true +tests: + - it: does not route gateway Service traffic to worker pods + template: gateway/service.yaml + asserts: + - equal: + path: spec.selector["app.kubernetes.io/component"] + value: gateway + + - it: does not route backend Service traffic to worker pods + template: backend/service.yaml + asserts: + - equal: + path: spec.selector["app.kubernetes.io/component"] + value: backend + + - it: scales only the gateway and backend Deployments + asserts: + - equal: + path: spec.scaleTargetRef.name + value: RELEASE-NAME-litellm-gateway + template: gateway/hpa.yaml + - equal: + path: spec.scaleTargetRef.name + value: RELEASE-NAME-litellm-backend + template: backend/hpa.yaml + +--- +suite: test worker component sizing and scheduling +templates: + - worker/deployment.yaml + - gateway/configmap.yaml +values: + - ./values/required.yaml +set: + worker.enabled: true +tests: + - it: sizes and schedules independently of the gateway + template: worker/deployment.yaml + set: + worker.resources: + requests: + cpu: 100m + worker.nodeSelector: + pool: jobs + worker.tolerations: + - key: jobs + operator: Exists + asserts: + - equal: + path: spec.template.spec.containers[0].resources.requests.cpu + value: 100m + - equal: + path: spec.template.spec.nodeSelector.pool + value: jobs + - equal: + path: spec.template.spec.tolerations[0].key + value: jobs + + - it: replicas follow worker.replicaCount + template: worker/deployment.yaml + set: + worker.replicaCount: 2 + asserts: + - equal: + path: spec.replicas + value: 2 diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 998d225a317..00d6b3459b7 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -376,6 +376,68 @@ backend: # Same shape as gateway.topologySpreadConstraints. topologySpreadConstraints: [] +# ---------- worker (single-owner auxiliary jobs) ---------- +# +# Optional fourth component for the proxy's single-owner auxiliary jobs (budget +# resets, cost polls, cleanups). Off by default, which leaves LITELLM_JOB_ROLE +# unset on gateway and backend and every pod eligible to own them, exactly as +# before. +# +# Enabling this sets LITELLM_JOB_ROLE=worker here and LITELLM_JOB_ROLE=serving on +# gateway and backend, so the jobs stop competing with request traffic for the +# event loop and the database pool. The worker gets no Service and no HPA; +# nothing routes to it. +# +# Only the cluster-wide scheduled jobs move. Gateway and backend keep flushing +# their own in-memory spend and request queues and keep refreshing their own +# model registry and credentials, so they still write spend as they do today. +# +# The worker pins NUM_WORKERS=1 rather than taking gateway.numWorkers: the +# scheduler is registered once per uvicorn worker process, so N processes would +# run N copies of every job in one pod. +# +# It runs the backend image and the backend ServiceAccount deliberately: it does +# the same management-plane work against the same database, so it needs the same +# schema expectations and the same IAM database-auth identity. It also mounts the +# gateway config, since the proxy needs a model list to start. +# +# The jobs are single-owner, so raising replicaCount buys availability across a +# rollout, not throughput. +worker: + enabled: false + replicaCount: 1 + logLevel: INFO + extraEnv: [] + envConfigMaps: [] + envSecrets: [] + resources: + requests: + cpu: 500m + memory: 1Gi + limits: + cpu: "1" + memory: 2Gi + livenessProbe: + httpGet: { path: /health/liveliness, port: http } + initialDelaySeconds: 10 + periodSeconds: 15 + timeoutSeconds: 10 + failureThreshold: 6 + readinessProbe: + httpGet: { path: /health/readiness, port: http } + initialDelaySeconds: 5 + periodSeconds: 10 + timeoutSeconds: 10 + podAnnotations: {} + # Same shape as the gateway blocks of the same name. + podLabels: {} + podSecurityContext: {} + securityContext: {} + terminationGracePeriodSeconds: "" + nodeSelector: {} + tolerations: [] + affinity: {} + # ---------- ui (Next.js static dashboard) ---------- ui: enabled: true diff --git a/litellm/constants.py b/litellm/constants.py index 0657d89529f..c1c124a3738 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1516,6 +1516,13 @@ SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_QUEUE_MAX_BYT SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute +# Which scheduled jobs this process registers: "all" (default), "serving", or "worker". +# A "serving" pod registers no single-owner job, so an operator can run those in a +# dedicated deployment instead of on every replica taking traffic +LITELLM_JOB_ROLE: Final = os.getenv("LITELLM_JOB_ROLE") +# Renew a held lease this many times per TTL, so a renewal can be lost without +# the lease expiring under a healthy owner +SINGLE_OWNER_JOB_RENEWAL_DIVISOR: Final = 3 PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500"))) RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", "100"))) diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py index e314aec497f..3fff0f00466 100644 --- a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -9,13 +9,20 @@ from typing import Any, Final from litellm._logging import verbose_proxy_logger from litellm.constants import ( + DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, UI_SESSION_TOKEN_TEAM_ID, ) from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken, UserAPIKeyAuth +from litellm.proxy.common_utils.single_owner_job import ( + JobLease, + WhenLockUnavailable, + run_as_single_owner, +) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, @@ -35,7 +42,7 @@ class ExpiredUISessionKeyCleanupManager: self, prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, - pod_lock_manager=None, + pod_lock_manager: PodLockManager | None = None, ): self.prisma_client = prisma_client self.user_api_key_cache = user_api_key_cache @@ -44,25 +51,21 @@ class ExpiredUISessionKeyCleanupManager: async def cleanup_expired_keys(self) -> int: """ Main entry point for deleting expired UI session keys. - Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments. + Runs on one elected pod, holding a renewed lease for the whole cycle. """ - lock_acquired = False - try: - if self.pod_lock_manager and self.pod_lock_manager.redis_cache: - lock_acquired = ( - await self.pod_lock_manager.acquire_lock( - cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, - ) - or False - ) - if not lock_acquired: - verbose_proxy_logger.debug( - "Expired UI session key cleanup: another pod is already " - "running cleanup or Redis lock acquisition failed - " - "skipping this cycle." - ) - return 0 + deleted: Final = await run_as_single_owner( + pod_lock_manager=self.pod_lock_manager, + job_name=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ttl_seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, + # Two pods deleting the same batch make the loser's delete 404 and log + # an error for work that already succeeded + when_unavailable=WhenLockUnavailable.SKIP, + run=self._delete_expired_keys, + ) + return deleted or 0 + async def _delete_expired_keys(self, _lease: JobLease) -> int: + try: verbose_proxy_logger.info("Starting expired UI session key cleanup...") expired_keys: Final = await self._find_expired_ui_session_keys() @@ -103,11 +106,6 @@ class ExpiredUISessionKeyCleanupManager: return 0 verbose_proxy_logger.error("Expired UI session key cleanup failed: %s", e) return 0 - finally: - if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: - await self.pod_lock_manager.release_lock( - cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, - ) @staticmethod def _get_deleted_token_count( diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 839ff28c354..5e7fe12a883 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -9,6 +9,7 @@ from typing import Final from litellm._logging import verbose_proxy_logger from litellm.constants import ( + KEY_ROTATION_JOB_NAME, LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, LITELLM_KEY_ROTATION_GRACE_PERIOD, LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, @@ -18,6 +19,12 @@ from litellm.proxy._types import ( LiteLLM_VerificationToken, RegenerateKeyRequest, ) +from litellm.proxy.common_utils.single_owner_job import ( + JobLease, + WhenLockUnavailable, + run_as_single_owner, +) +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.key_management_endpoints import ( _calculate_key_rotation_time, @@ -37,42 +44,32 @@ class KeyRotationManager: Manages automated key rotation based on individual key rotation schedules. """ - def __init__(self, prisma_client: PrismaClient, pod_lock_manager=None): + def __init__(self, prisma_client: PrismaClient, pod_lock_manager: PodLockManager | None = None): self.prisma_client = prisma_client self.pod_lock_manager = pod_lock_manager - async def process_rotations(self): + async def process_rotations(self) -> None: """ Main entry point - find and rotate keys that are due for rotation. - Uses PodLockManager to ensure only one pod runs rotation in multi-pod deployments. + Runs on one elected pod; the lease is renewed for as long as the cycle + takes, so a deployment with many due keys cannot be joined mid-rotation. """ - from litellm.constants import KEY_ROTATION_JOB_NAME + # A dedicated lock TTL (default 600s) rather than the check interval, which + # defaults to 24h: the interval as a TTL would strand rotation for a day if + # the owner crashed before releasing + lock_ttl: Final = max(LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, 300) + await run_as_single_owner( + pod_lock_manager=self.pod_lock_manager, + job_name=KEY_ROTATION_JOB_NAME, + ttl_seconds=lock_ttl, + # Rotating one key twice hands out two replacements and invalidates the + # first, so a cycle skipped during a Redis outage is the cheaper failure + when_unavailable=WhenLockUnavailable.SKIP, + run=self._rotate_due_keys, + ) - lock_acquired = False + async def _rotate_due_keys(self, _lease: JobLease) -> None: try: - # If we have a pod lock manager with Redis, try to acquire the lock - if self.pod_lock_manager and self.pod_lock_manager.redis_cache: - # Use a dedicated lock TTL (default 600s) instead of the check interval - # (which defaults to 86400s / 24h). Using the check interval would create - # a 24-hour deadlock window if a pod crashes before releasing the lock. - lock_ttl: Final = max( - LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, 300 - ) # At least 5 minutes, configurable via LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS - lock_acquired = ( - await self.pod_lock_manager.acquire_lock( - cronjob_id=KEY_ROTATION_JOB_NAME, - ttl=lock_ttl, - ) - or False - ) - if not lock_acquired: - verbose_proxy_logger.warning( - "Key rotation: another pod is already running rotation " - "or Redis lock acquisition failed — skipping this cycle. " - "Keys will be rotated on the next cycle." - ) - return - verbose_proxy_logger.info("Starting scheduled key rotation check...") # Clean up expired deprecated keys first @@ -99,12 +96,6 @@ class KeyRotationManager: except Exception as e: verbose_proxy_logger.error("Key rotation process failed: %s", e) - finally: - # Only release the lock if it was actually acquired - if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: - await self.pod_lock_manager.release_lock( - cronjob_id=KEY_ROTATION_JOB_NAME, - ) async def _find_keys_needing_rotation(self) -> list[LiteLLM_VerificationToken]: """ diff --git a/litellm/proxy/common_utils/single_owner_job.py b/litellm/proxy/common_utils/single_owner_job.py new file mode 100644 index 00000000000..cf60bc32dbf --- /dev/null +++ b/litellm/proxy/common_utils/single_owner_job.py @@ -0,0 +1,220 @@ +"""Ownership rules for scheduled jobs that must run once per deployment. + +Every proxy process registers the same scheduler, so a job with a shared +side effect runs once per pod unless something elects an owner. +""" + +import asyncio +from collections.abc import Awaitable, Callable +from contextlib import suppress +from enum import Enum +from typing import Final, TypeVar + +from litellm._logging import verbose_proxy_logger +from litellm.constants import SINGLE_OWNER_JOB_RENEWAL_DIVISOR +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + +T = TypeVar("T") + + +class JobRole(Enum): + """Which scheduled jobs a process registers.""" + + ALL = "all" + SERVING = "serving" + WORKER = "worker" + + @classmethod + def from_env_value(cls, raw: str | None) -> "JobRole": + if raw is None or not raw.strip(): + return cls.ALL + try: + return cls(raw.strip().lower()) + except ValueError: + verbose_proxy_logger.warning( + "LITELLM_JOB_ROLE=%r is not one of %s; running every job as if it were %s", + raw, + tuple(role.value for role in cls), + cls.ALL.value, + ) + return cls.ALL + + @property + def runs_single_owner_jobs(self) -> bool: + return self is not JobRole.SERVING + + +class JobLease(Enum): + """This pod's standing for one run of a single-owner job.""" + + LEADER = "leader" + FOLLOWER = "follower" + UNGUARDED = "unguarded" + + +class WhenLockUnavailable(Enum): + """What a pod does when the lock can be neither taken nor read. + + ``acquire_lock`` reports contention and an unreachable Redis identically, so + each job has to say which way to resolve the ambiguity. Work whose duplicate + is merely wasteful picks ``RUN``; work whose duplicate reaches a person or an + external system picks ``SKIP``. + """ + + RUN = "run" + SKIP = "skip" + + +async def run_as_single_owner( + *, + pod_lock_manager: PodLockManager | None, + job_name: str, + ttl_seconds: int, + when_unavailable: WhenLockUnavailable, + run: Callable[[JobLease], Awaitable[T]], +) -> T | None: + """Run ``run`` on one pod, holding a lease that is renewed for its duration. + + Returns ``None`` without running when another pod owns this tick. + + Renewal is what makes the TTL a failover deadline rather than a run budget: a + healthy owner keeps the lease however long its work takes, and a crashed one + strands the job for at most ``ttl_seconds``. + + The lease is released when ``run`` returns, so it dedupes for the body's runtime + and not for the TTL. Two pods whose ticks are further apart than that both run, + which is what the scheduler's stagger makes ordinary rather than incidental, so + ``run`` has to be idempotent. Work that must happen once per period regardless of + when each pod fires wants ``claim_once_per_window`` instead, whose marker outlives + the run. + """ + manager: Final = pod_lock_manager + if manager is None or manager.redis_cache is None: + return await run(JobLease.UNGUARDED) + + lease: Final = await _acquire( + manager=manager, + job_name=job_name, + ttl_seconds=ttl_seconds, + when_unavailable=when_unavailable, + ) + match lease: + case JobLease.FOLLOWER: + return None + case JobLease.UNGUARDED: + return await run(lease) + case JobLease.LEADER: + return await _run_holding_lease( + manager=manager, + job_name=job_name, + ttl_seconds=ttl_seconds, + run=run, + ) + + +async def claim_once_per_window( + *, + pod_lock_manager: PodLockManager | None, + job_name: str, + window_seconds: int, +) -> bool: + """True when this pod may run the job for the current window. + + The lock is a done-marker rather than a lease: nobody releases it, so it + expires with the window and the next window's first firer takes it. Each pod + anchors its interval to its own boot time, so a lock that outlived only the + run would let a later pod repeat the window. + + A deployment with no Redis at all runs, having no peer to duplicate against. + A configured but unreachable Redis does not, because these windows end at + someone's inbox. + """ + if pod_lock_manager is None: + return True + claimed: Final = await pod_lock_manager.acquire_lock( + cronjob_id=job_name, + ttl=window_seconds, + allow_reentrant=False, + ) + return claimed is not False + + +async def _acquire( + *, + manager: PodLockManager, + job_name: str, + ttl_seconds: int, + when_unavailable: WhenLockUnavailable, +) -> JobLease: + if await manager.acquire_lock(cronjob_id=job_name, ttl=ttl_seconds): + return JobLease.LEADER + + if await _lease_is_held(manager=manager, job_name=job_name): + verbose_proxy_logger.debug("%s: another pod holds the lease, skipping this run", job_name) + return JobLease.FOLLOWER + + match when_unavailable: + case WhenLockUnavailable.RUN: + verbose_proxy_logger.warning( + "%s: could not take the lease and no other pod holds it, running unguarded rather than skipping", + job_name, + ) + return JobLease.UNGUARDED + case WhenLockUnavailable.SKIP: + verbose_proxy_logger.warning( + "%s: could not take the lease and no other pod holds it, skipping rather than risking a duplicate", + job_name, + ) + return JobLease.FOLLOWER + + +async def _lease_is_held(*, manager: PodLockManager, job_name: str) -> bool: + """An unreadable lease reports as unheld so the caller reaches its own + ``when_unavailable`` policy instead of silently taking the follower branch. + """ + if manager.redis_cache is None: + return False + try: + lock_key: Final = manager.get_redis_lock_key(job_name) + return bool(await manager.redis_cache.async_get_cache(lock_key)) + except Exception as exc: # noqa: BLE001 # an unreadable lease must not decide the run + verbose_proxy_logger.warning("%s: could not read the lease: %s", job_name, exc) + return False + + +async def _run_holding_lease( + *, + manager: PodLockManager, + job_name: str, + ttl_seconds: int, + run: Callable[[JobLease], Awaitable[T]], +) -> T: + verbose_proxy_logger.info("%s: pod %s owns this run", job_name, manager.pod_id) + renewer: Final = asyncio.create_task( + _renew_until_cancelled(manager=manager, job_name=job_name, ttl_seconds=ttl_seconds) + ) + try: + return await run(JobLease.LEADER) + finally: + renewer.cancel() + with suppress(asyncio.CancelledError): + await renewer + await manager.release_lock(cronjob_id=job_name) + + +async def _renew_until_cancelled(*, manager: PodLockManager, job_name: str, ttl_seconds: int) -> None: + """Hold the lease until cancelled, or until it belongs to someone else. + + Renewal stops on the first failure so the new owner keeps what it took. The + release on the way out compares owners, so it cannot steal the lease back. + """ + interval: Final = max(1.0, ttl_seconds / SINGLE_OWNER_JOB_RENEWAL_DIVISOR) + while True: + await asyncio.sleep(interval) + if not await manager.renew_lock(cronjob_id=job_name, ttl=ttl_seconds): + verbose_proxy_logger.warning( + "%s: pod %s lost the lease mid-run, a second pod may be running this job", + job_name, + manager.pod_id, + ) + return diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 4be1331e955..24dc910bfd3 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -1,6 +1,8 @@ import asyncio import json -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Awaitable, Callable, Sequence +from contextlib import suppress +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -15,6 +17,12 @@ else: ProxyLogging = Any +class _RegisteredScript(Protocol): + """A Lua script already registered with Redis, as returned by async_register_script.""" + + def __call__(self, keys: Sequence[str], args: Sequence[object]) -> Awaitable[object]: ... + + class PodLockManager: """ Manager for acquiring and releasing locks for cron jobs using Redis. @@ -28,12 +36,21 @@ if redis.call("get", KEYS[1]) == ARGV[1] then else return 0 end +""" + + _COMPARE_AND_EXPIRE_LOCK_SCRIPT = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("expire", KEYS[1], ARGV[2]) +else + return 0 +end """ def __init__(self, redis_cache: RedisCache | None = None): self.pod_id = str(uuid.uuid4()) self.redis_cache = redis_cache - self._release_lock_script: Any | None = None + self._release_lock_script: _RegisteredScript | None = None + self._renew_lock_script: _RegisteredScript | None = None @staticmethod def get_redis_lock_key(cronjob_id: str) -> str: @@ -85,7 +102,7 @@ end self.pod_id, cronjob_id, ) - + self._emit_acquired_lock_event(cronjob_id, self.pod_id) return True else: # Check if the current pod already holds the lock @@ -112,6 +129,45 @@ end verbose_proxy_logger.error("Error acquiring Redis lock for %s: %s", cronjob_id, e) return False + async def renew_lock(self, cronjob_id: str, ttl: int | None = None) -> bool: + """Extend this pod's lock by another TTL, if it still owns it. + + Lets a lease outlive a run that takes longer than its TTL without making + the TTL a failover deadline for the whole job. Returns False when the + lock is gone or another pod took it, so the caller learns it is no + longer the owner rather than extending someone else's lease. + + Renewal has no non-atomic fallback, unlike release. A GET-then-SET would + write unconditionally, so a lease that lapsed and was taken over between + the two calls would be handed back to the pod that lost it, leaving that + pod and its successor both believing they own the job. Where the + compare-and-expire cannot run, this reports False and the lease is left + to expire into the failover it already describes. + """ + cache: Final = self.redis_cache + if cache is None: + verbose_proxy_logger.debug("redis_cache is None, skipping renew_lock") + return False + lock_ttl: Final = ttl or DEFAULT_CRON_JOB_LOCK_TTL_SECONDS + lock_key: Final = PodLockManager.get_redis_lock_key(cronjob_id) + + result, self._renew_lock_script = await self._act_if_owner( + cache=cache, + lock_key=lock_key, + script=self._COMPARE_AND_EXPIRE_LOCK_SCRIPT, + script_handle=self._renew_lock_script, + script_args=(lock_ttl,), + fallback=None, + ) + if result: + verbose_proxy_logger.debug( + "Pod %s renewed Redis lock for cronjob_id=%s (ttl=%ds)", + self.pod_id, + cronjob_id, + lock_ttl, + ) + return bool(result) + async def release_lock( self, cronjob_id: str, @@ -123,7 +179,8 @@ end Falls back to GET + DEL for cache implementations that don't support script registration. """ - if self.redis_cache is None: + cache: Final = self.redis_cache + if cache is None: verbose_proxy_logger.debug("redis_cache is None, skipping release_lock") return try: @@ -133,7 +190,7 @@ end cronjob_id, ) lock_key: Final = PodLockManager.get_redis_lock_key(cronjob_id) - result: Final = await self._compare_and_delete_lock(lock_key=lock_key) + result: Final = await self._compare_and_delete_lock(cache=cache, lock_key=lock_key) if result == 1: verbose_proxy_logger.info( "Pod %s successfully released Redis lock for cronjob_id=%s", @@ -153,40 +210,70 @@ end except Exception as e: verbose_proxy_logger.error("Error releasing Redis lock for %s: %s", cronjob_id, e) - async def _compare_and_delete_lock(self, lock_key: str) -> int: + async def _compare_and_delete_lock(self, cache: RedisCache, lock_key: str) -> int: """ Atomically delete lock key only if current pod owns it. Falls back to get/delete for non-RedisCache implementations that do not expose Lua script registration. """ - script_register: Final = getattr(self.redis_cache, "async_register_script", None) + + async def _delete() -> int: + return int(await cache.async_delete_cache(lock_key) or 0) + + result, self._release_lock_script = await self._act_if_owner( + cache=cache, + lock_key=lock_key, + script=self._COMPARE_AND_DELETE_LOCK_SCRIPT, + script_handle=self._release_lock_script, + script_args=(), + fallback=_delete, + ) + return result + + async def _act_if_owner( + self, + *, + cache: RedisCache, + lock_key: str, + script: str, + script_handle: _RegisteredScript | None, + script_args: Sequence[object], + fallback: Callable[[], Awaitable[int]] | None, + ) -> tuple[int, _RegisteredScript | None]: + """Act on the lock only if this pod still owns it, atomically where possible. + + Returns the script's result and the handle to cache for next time, which is + None whenever the GET-then-act fallback answered instead, so a Redis restart + that cleared the loaded scripts re-registers rather than failing forever. + + ``fallback`` is None for actions with no safe non-atomic form, which report + 0 rather than racing the lock's owner. + """ + script_register: Final = getattr(cache, "async_register_script", None) if callable(script_register): - try: - if self._release_lock_script is None: - self._release_lock_script = script_register(self._COMPARE_AND_DELETE_LOCK_SCRIPT) + with suppress(Exception): # acquire_lock stores the pod_id via async_set_cache, which # JSON-encodes the value; compare against the same encoding so - # the Lua equality check matches and the lock is released - result = await self._release_lock_script(keys=[lock_key], args=[json.dumps(self.pod_id)]) - return int(result or 0) - except Exception: - # Lua execution failed (e.g. Redis restart cleared loaded scripts, - # or scripting is disabled). Reset cached script handle and fall - # through to the GET + DEL fallback so the lock is still released. - self._release_lock_script = None - verbose_proxy_logger.warning( - "Lua compare-and-delete failed for lock_key=%s, falling back to GET+DEL", - lock_key, - ) + # the Lua equality check matches + handle: Final = script_handle or script_register(script) + result: Final = await handle(keys=(lock_key,), args=(json.dumps(self.pod_id), *script_args)) + return int(result or 0), handle + # scripting is disabled, or a Redis restart cleared the loaded scripts + verbose_proxy_logger.warning( + "Lua compare-and-act failed for lock_key=%s, falling back to GET then act", + lock_key, + ) - current_value = await self.redis_cache.async_get_cache(lock_key) + if fallback is None: + return 0, None + + current_value = await cache.async_get_cache(lock_key) if isinstance(current_value, bytes): current_value = current_value.decode("utf-8") if current_value != self.pod_id: - return 0 - result = await self.redis_cache.async_delete_cache(lock_key) - return int(result or 0) + return 0, None + return await fallback(), None @staticmethod def _emit_acquired_lock_event(cronjob_id: str, pod_id: str): diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index e97e9f6e683..8af1208b2a8 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -9,6 +9,7 @@ from pydantic import BaseModel, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import ( + DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS, SPEND_LOG_CLEANUP_BATCH_SIZE, SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS, @@ -19,6 +20,11 @@ from litellm.constants import ( SPEND_LOG_RUN_LOOPS, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy.common_utils.single_owner_job import ( + JobLease, + WhenLockUnavailable, + run_as_single_owner, +) from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import ( RunOutcome, SpendLogCleanupMetrics, @@ -586,10 +592,10 @@ class SpendLogCleanup: async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None: """ Main cleanup function. Deletes old spend logs in batches. - If pod_lock_manager is available, ensures only one pod runs cleanup. - If no pod_lock_manager, runs cleanup without distributed locking. + Runs on one elected pod, holding a lease that is renewed for the whole + sweep so a cleanup slower than the lease TTL cannot be joined mid-run. + A deployment with no Redis-backed lock manager runs unguarded. """ - lock_acquired = False try: verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now()) self._refresh_bounds() @@ -612,63 +618,55 @@ class SpendLogCleanup: SpendLogCleanupMetrics.record_run("skipped_disabled") return - # If we have a pod lock manager, try to acquire the lock - if self.pod_lock_manager and self.pod_lock_manager.redis_cache: - lock_acquired = ( - await self.pod_lock_manager.acquire_lock( - cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME, + async def _sweep(_lease: JobLease) -> RunOutcome: + deadline: Final = time.monotonic() + self.run_budget_seconds + configured_group_count: Final = ( + int(delete_spend_logs and self.retention_seconds is not None) + + int(autorouter_retention_seconds is not None) + + int(health_check_retention_seconds is not None) + ) + + spend_log_results: Final = ( + await self._clean_spend_log_tables( + prisma_client, + self._group_deadline(deadline, configured_group_count), ) - or False + if delete_spend_logs and self.retention_seconds is not None + else () ) - verbose_proxy_logger.info( - "Lock acquisition attempt: %s at %s", "successful" if lock_acquired else "failed", datetime.now() + remaining_groups_after_spend_logs: Final = int(autorouter_retention_seconds is not None) + int( + health_check_retention_seconds is not None ) + session_results: Final = ( + await self._clean_session_rollup( + prisma_client, + autorouter_retention_seconds, + self._group_deadline(deadline, remaining_groups_after_spend_logs), + ) + if autorouter_retention_seconds is not None + else () + ) + health_check_results: Final = ( + await self._clean_health_checks( + prisma_client, + health_check_retention_seconds, + deadline, + ) + if health_check_retention_seconds is not None + else () + ) + return self._run_outcome(spend_log_results + session_results + health_check_results) - if not lock_acquired: - verbose_proxy_logger.info("Another pod is already running cleanup") - SpendLogCleanupMetrics.record_run("skipped_locked") - return - - deadline: Final = time.monotonic() + self.run_budget_seconds - configured_group_count: Final = ( - int(delete_spend_logs and self.retention_seconds is not None) - + int(autorouter_retention_seconds is not None) - + int(health_check_retention_seconds is not None) - ) - - spend_log_results: Final = ( - await self._clean_spend_log_tables( - prisma_client, - self._group_deadline(deadline, configured_group_count), - ) - if delete_spend_logs and self.retention_seconds is not None - else () - ) - remaining_groups_after_spend_logs: Final = int(autorouter_retention_seconds is not None) + int( - health_check_retention_seconds is not None - ) - session_results: Final = ( - await self._clean_session_rollup( - prisma_client, - autorouter_retention_seconds, - self._group_deadline(deadline, remaining_groups_after_spend_logs), - ) - if autorouter_retention_seconds is not None - else () - ) - health_check_results: Final = ( - await self._clean_health_checks( - prisma_client, - health_check_retention_seconds, - deadline, - ) - if health_check_retention_seconds is not None - else () - ) - - SpendLogCleanupMetrics.record_run( - self._run_outcome(spend_log_results + session_results + health_check_results) + outcome: Final = await run_as_single_owner( + pod_lock_manager=self.pod_lock_manager, + job_name=SPEND_LOG_CLEANUP_JOB_NAME, + ttl_seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, + # A second concurrent sweep deletes rows the first is already deleting, + # doubling the DB load the retention job exists to bound + when_unavailable=WhenLockUnavailable.SKIP, + run=_sweep, ) + SpendLogCleanupMetrics.record_run("skipped_locked" if outcome is None else outcome) except Exception as e: # .exception() captures the traceback; str(e) alone on a Prisma/DB @@ -679,9 +677,3 @@ class SpendLogCleanup: e, ) SpendLogCleanupMetrics.record_run("aborted") - return # Return after error handling - finally: - # Only release the lock if it was actually acquired - if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: - await self.pod_lock_manager.release_lock(cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME) - verbose_proxy_logger.info("Released cleanup lock") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 711775ee0cf..e9b8d2ebff0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -236,6 +236,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, GLOBAL_PROXY_SPEND_CACHE_KEY, + LITELLM_JOB_ROLE, LITELLM_PROXY_ADMIN_NAME, LITELLM_PROXY_BUDGET_NAME, MONTHLY_SPEND_REPORT_JOB_ID, @@ -358,6 +359,7 @@ from litellm.proxy.common_utils.scheduled_job_stagger import ( parse_stagger_settings, stagger_trigger, ) +from litellm.proxy.common_utils.single_owner_job import JobRole, claim_once_per_window from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_settings, @@ -8843,6 +8845,14 @@ class ProxyStartupEvent: timezone=None, ) + job_role: Final = JobRole.from_env_value(LITELLM_JOB_ROLE) + if not job_role.runs_single_owner_jobs: + verbose_proxy_logger.info( + "LITELLM_JOB_ROLE=%s: registering no single-owner background job on this process. " + "Run a deployment with LITELLM_JOB_ROLE=worker, or the default 'all', or these jobs never run", + job_role.value, + ) + # Use fixed intervals with small random offset instead of jitter # This avoids the expensive jitter calculations in APScheduler budget_interval: Final = proxy_budget_rescheduler_min_time + random.randint( @@ -8869,7 +8879,7 @@ class ProxyStartupEvent: ) ### RESET BUDGET ### - if general_settings.get("disable_reset_budget", False) is False: + if general_settings.get("disable_reset_budget", False) is False and job_role.runs_single_owner_jobs: budget_reset_job: Final = ResetBudgetJob( proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client, @@ -9044,16 +9054,17 @@ class ProxyStartupEvent: general_settings=general_settings, proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client, + job_role=job_role, ) - await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) + await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler, job_role=job_role) ### PTU DAILY ROLLUP ### from litellm.proxy.spend_tracking.ptu_feature_flag import ( is_ptu_cost_attribution_enabled, ) - if is_ptu_cost_attribution_enabled(): + if is_ptu_cost_attribution_enabled() and job_role.runs_single_owner_jobs: from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( PTU_ROLLUP_JOB_ID, run_scheduled_ptu_rollup, @@ -9090,7 +9101,7 @@ class ProxyStartupEvent: ) ### SPEND LOG CLEANUP ### - if ( + if job_role.runs_single_owner_jobs and ( general_settings.get("maximum_spend_logs_retention_period") is not None or general_settings.get("maximum_autorouter_session_retention_period") is not None or general_settings.get("maximum_health_check_retention_period") is not None @@ -9131,7 +9142,7 @@ class ProxyStartupEvent: except ValueError: verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") ### CHECK BATCH COST ### - if llm_router is not None and PROXY_BATCH_POLLING_ENABLED: + if llm_router is not None and PROXY_BATCH_POLLING_ENABLED and job_role.runs_single_owner_jobs: try: from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, @@ -9162,7 +9173,7 @@ class ProxyStartupEvent: ) ### CHECK RESPONSES COST ### - if llm_router is not None and PROXY_BATCH_POLLING_ENABLED: + if llm_router is not None and PROXY_BATCH_POLLING_ENABLED and job_role.runs_single_owner_jobs: try: from litellm_enterprise.proxy.common_utils.check_responses_cost import ( CheckResponsesCost, @@ -9211,7 +9222,7 @@ class ProxyStartupEvent: return worker_heartbeat @classmethod - async def _initialize_spend_tracking_background_jobs(cls, scheduler: AsyncIOScheduler): + async def _initialize_spend_tracking_background_jobs(cls, scheduler: AsyncIOScheduler, job_role: JobRole): """ Initialize the spend tracking and other background jobs 1. CloudZero Background Job @@ -9219,13 +9230,20 @@ class ProxyStartupEvent: 3. Prometheus Background Job 4. Key Rotation Background Job + Every job here exports or reconciles shared state, so a pod dedicated to + serving traffic registers none of them. + Args: scheduler: The scheduler to add the background jobs to + job_role: Which jobs this process registers """ global prisma_client global proxy_logging_obj global user_api_key_cache + if not job_role.runs_single_owner_jobs: + return + ######################################################## # CloudZero Background Job ######################################################## @@ -9395,8 +9413,16 @@ class ProxyStartupEvent: general_settings: dict, proxy_logging_obj: ProxyLogging, prisma_client: PrismaClient, + job_role: JobRole, ): - """Initialize Slack alerting background jobs for spend reports.""" + """Initialize Slack alerting background jobs for spend reports. + + Each report reaches a channel once per window, so a pod dedicated to + serving traffic registers none of them. + """ + if not job_role.runs_single_owner_jobs: + return + if ( proxy_logging_obj is not None and proxy_logging_obj.slack_alerting_instance.alerting is not None @@ -9413,24 +9439,21 @@ class ProxyStartupEvent: weekly_lock_ttl: Final = duration_in_seconds(spend_report_frequency) - 3600 async def _scheduled_weekly_spend_report() -> None: - # TTL spans the whole reporting window: each pod's interval anchor is its own - # boot time + jitter, so a shorter lock would let a later pod re-send the report. - # Minus an hour so the next window's first firer finds a free key - if ( - await pod_lock_manager.acquire_lock( - cronjob_id=WEEKLY_SPEND_REPORT_JOB_ID, ttl=weekly_lock_ttl, allow_reentrant=False - ) - is False + # Window is the reporting period minus an hour, so the next window's + # first firer finds a free key + if not await claim_once_per_window( + pod_lock_manager=pod_lock_manager, + job_name=WEEKLY_SPEND_REPORT_JOB_ID, + window_seconds=weekly_lock_ttl, ): return await proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report(spend_report_frequency) async def _scheduled_monthly_spend_report() -> None: - if ( - await pod_lock_manager.acquire_lock( - cronjob_id=MONTHLY_SPEND_REPORT_JOB_ID, ttl=3600, allow_reentrant=False - ) - is False + if not await claim_once_per_window( + pod_lock_manager=pod_lock_manager, + job_name=MONTHLY_SPEND_REPORT_JOB_ID, + window_seconds=3600, ): return await proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report() @@ -9457,11 +9480,10 @@ class ProxyStartupEvent: from zoneinfo import ZoneInfo async def _scheduled_fallback_stats() -> None: - if ( - await pod_lock_manager.acquire_lock( - cronjob_id=PROMETHEUS_FALLBACK_STATS_JOB_ID, ttl=3600, allow_reentrant=False - ) - is False + if not await claim_once_per_window( + pod_lock_manager=pod_lock_manager, + job_name=PROMETHEUS_FALLBACK_STATS_JOB_ID, + window_seconds=3600, ): return await proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus() diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index f1f7248c064..9a686e70e20 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -31,6 +31,11 @@ from litellm.constants import ( PTU_SENTINEL_API_KEY, ) from litellm.litellm_core_utils.ptu_pricing import ptu_terms +from litellm.proxy.common_utils.single_owner_job import ( + JobLease, + WhenLockUnavailable, + run_as_single_owner, +) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled if TYPE_CHECKING: @@ -601,10 +606,8 @@ async def run_scheduled_ptu_rollup( unguarded, as ``SpendLogCleanup`` does, and so does a run that cannot reach Redis at all: the lock exists to avoid duplicate work, so no lock problem may cost a day. - The lease is a fixed TTL with no renewal, so a long scan can outlive it. That costs - duplicate work rather than correctness: the upserts are idempotent on the sentinel - key and the prune reads only the row's own timestamp, so a second pod arriving - mid-run cannot corrupt the day. + Only the elected owner may prune. An unguarded run reconciles without pruning, since + a pod that cannot prove it is alone could delete a row another pod just wrote. Returns None without touching the database when PTU cost attribution is off. Proxy startup already skips scheduling the cron, so this guards the function itself rather @@ -614,41 +617,23 @@ async def run_scheduled_ptu_rollup( if not is_ptu_cost_attribution_enabled(): return None - if pod_lock_manager is None or pod_lock_manager.redis_cache is None: - return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False) - - if not await pod_lock_manager.acquire_lock(cronjob_id=PTU_ROLLUP_JOB_ID, ttl=PTU_ROLLUP_LOCK_TTL_SECONDS): - if await _lock_is_held(pod_lock_manager): - verbose_proxy_logger.info("PTU rollup: another pod holds the rollup lock, skipping this run") - return None - # acquire_lock reports contention and a Redis outage the same way, so an - # unreachable Redis would otherwise skip the day on every pod at once. The - # reconcile is safe to run concurrently, so losing the lock costs duplicate - # work; losing the day costs a team's charges - verbose_proxy_logger.warning( - "PTU rollup: could not take the rollup lock and no other pod holds it, " - "running unguarded rather than skipping the day" + async def _reconcile(lease: JobLease) -> RollupResult | None: + return await _run_and_alert( + prisma_client, + target_date=target_date, + alert=alert, + may_prune=lease is JobLease.LEADER, ) - return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False) - try: - return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=True) - finally: - await pod_lock_manager.release_lock(cronjob_id=PTU_ROLLUP_JOB_ID) - - -async def _lock_is_held(pod_lock_manager: "PodLockManager") -> bool: - """True only when the rollup lock is readable and someone is holding it. - - A Redis that cannot be read is reported as "not held" so the caller runs the day - rather than skipping it; the cost of being wrong here is a duplicate reconcile. - """ - try: - lock_key: Final = pod_lock_manager.get_redis_lock_key(PTU_ROLLUP_JOB_ID) - return bool(await pod_lock_manager.redis_cache.async_get_cache(lock_key)) - except Exception as exc: # noqa: BLE001 # an unreadable lock must not skip the day - verbose_proxy_logger.warning("PTU rollup: could not read the rollup lock: %s", exc) - return False + return await run_as_single_owner( + pod_lock_manager=pod_lock_manager, + job_name=PTU_ROLLUP_JOB_ID, + ttl_seconds=PTU_ROLLUP_LOCK_TTL_SECONDS, + # The reconcile is safe to repeat, so losing the lock costs duplicate work + # while losing the day costs a team's charges + when_unavailable=WhenLockUnavailable.RUN, + run=_reconcile, + ) async def _run_and_alert( diff --git a/tests/redis_integration_tests/test_single_owner_job_redis.py b/tests/redis_integration_tests/test_single_owner_job_redis.py new file mode 100644 index 00000000000..8614dd99faf --- /dev/null +++ b/tests/redis_integration_tests/test_single_owner_job_redis.py @@ -0,0 +1,348 @@ +"""Leader election for auxiliary DB jobs, exercised against a live Redis. + +Every test here stands up N PodLockManager instances with distinct pod ids +against one real Redis, which is the only place the compare-and-set, the TTL +expiry and the Lua compare-and-delete actually execute. fakeredis and mocks +cannot fail these. +""" + +from __future__ import annotations + +import asyncio +import os +import socket +import time +from collections.abc import Awaitable, Callable +from typing import Final + +import pytest + +from litellm._uuid import uuid +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.single_owner_job import ( + JobLease, + WhenLockUnavailable, + claim_once_per_window, + run_as_single_owner, +) +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + +pytestmark = pytest.mark.asyncio(loop_scope="module") + +POD_COUNT: Final = 5 + + +@pytest.fixture(scope="module") +def redis_cache() -> RedisCache: + host: Final = os.environ.get("REDIS_HOST") or "127.0.0.1" + port: Final = int(os.environ.get("REDIS_PORT") or "6379") + try: + with socket.create_connection((host, port), timeout=3): + pass + except OSError as exc: + unreachable: Final = f"Redis at {host}:{port} is not reachable ({exc})." + if os.environ.get("LITELLM_REQUIRE_LIVE_REDIS") == "1": + pytest.fail(f"{unreachable} This suite is meaningless without it.") + pytest.skip(f"{unreachable} Start one and set REDIS_HOST/REDIS_PORT to run this suite.") + return RedisCache(host=host, port=port, password=os.environ.get("REDIS_PASSWORD")) + + +@pytest.fixture +def job_name(request: pytest.FixtureRequest) -> str: + return f"lit5434-{request.node.name}-{uuid.uuid4()}" + + +@pytest.fixture +def pods(redis_cache: RedisCache) -> list[PodLockManager]: + fleet: Final = [PodLockManager(redis_cache=redis_cache) for _ in range(POD_COUNT)] + assert len({pod.pod_id for pod in fleet}) == POD_COUNT, "each pod must carry a distinct pod id" + return fleet + + +def _recording_body(log: list[JobLease]) -> Callable[[JobLease], Awaitable[str]]: + async def body(lease: JobLease) -> str: + log.append(lease) + return "ran" + + return body + + +async def _concurrent_tick( + pods: list[PodLockManager], + job_name: str, + *, + ttl_seconds: int, + when_unavailable: WhenLockUnavailable, +) -> tuple[list[JobLease], list[str | None]]: + """Fire one tick from every pod at once, holding whoever wins until all of them have tried to acquire. + + Without the gate a body that returns immediately lets the winner release before the next pod even + attempts, so the pods would take the lease in sequence and the tick would look exclusive by luck. + """ + leases: Final[list[JobLease]] = [] + gate: Final = asyncio.Event() + + async def body(lease: JobLease) -> str: + leases.append(lease) + await gate.wait() + return "ran" + + tick: Final = asyncio.gather( + *( + run_as_single_owner( + pod_lock_manager=pod, + job_name=job_name, + ttl_seconds=ttl_seconds, + when_unavailable=when_unavailable, + run=body, + ) + for pod in pods + ) + ) + await asyncio.sleep(0.5) + gate.set() + return leases, list(await tick) + + +async def _lease_ttl_ms(redis_cache: RedisCache, job_name: str) -> int: + client: Final = redis_cache.init_async_client() + return int(await client.pttl(PodLockManager.get_redis_lock_key(job_name))) + + +async def _lease_holder(redis_cache: RedisCache, job_name: str) -> str | None: + raw: Final = await redis_cache.async_get_cache(PodLockManager.get_redis_lock_key(job_name)) + return raw.decode("utf-8") if isinstance(raw, bytes) else raw + + +async def test_only_one_pod_of_five_runs_the_tick(pods: list[PodLockManager], job_name: str) -> None: + """Five pods firing the same tick concurrently: one body runs, four never start.""" + leases, results = await _concurrent_tick( + pods, job_name, ttl_seconds=10, when_unavailable=WhenLockUnavailable.SKIP + ) + + assert leases == [JobLease.LEADER], f"exactly one body may run, and it must hold the lease: {leases}" + assert results.count("ran") == 1, f"one call returns the body's value: {results}" + assert results.count(None) == POD_COUNT - 1, f"the other four skip without running: {results}" + + +async def test_successor_runs_only_after_the_dead_owner_lease_expires( + pods: list[PodLockManager], job_name: str +) -> None: + """A pod that grabs the lease and dies blocks the job until the TTL runs out, then one successor takes over.""" + ttl_seconds: Final = 4 + dead_owner: Final = pods[0] + survivors: Final = pods[1:] + + assert await dead_owner.acquire_lock(cronjob_id=job_name, ttl=ttl_seconds) is True + took_lease_at: Final = time.monotonic() + + before_expiry, early = await _concurrent_tick( + survivors, job_name, ttl_seconds=ttl_seconds, when_unavailable=WhenLockUnavailable.SKIP + ) + probed_after: Final = time.monotonic() - took_lease_at + assert probed_after < ttl_seconds, ( + f"the before-expiry probe took {probed_after:.1f}s on a {ttl_seconds}s lease, so it proves nothing; " + "raise ttl_seconds if this machine is that slow" + ) + assert before_expiry == [], "no survivor may run while the dead owner's lease is still live" + assert early == [None] * len(survivors), f"every survivor skips before expiry: {early}" + + await asyncio.sleep(ttl_seconds - probed_after + 0.5) + + after_expiry, late = await _concurrent_tick( + survivors, job_name, ttl_seconds=ttl_seconds, when_unavailable=WhenLockUnavailable.SKIP + ) + assert after_expiry == [JobLease.LEADER], f"exactly one survivor takes over after expiry: {after_expiry}" + assert late.count("ran") == 1, f"the successor's body value is returned once: {late}" + + +async def test_renewal_holds_the_lease_past_the_ttl_and_frees_it_on_completion( + redis_cache: RedisCache, pods: list[PodLockManager], job_name: str +) -> None: + """A body outliving its TTL keeps the lease the whole time, and gives it up the moment it finishes.""" + ttl_seconds: Final = 3 + owner: Final = pods[0] + challenger: Final = pods[1] + + async def slow_body(lease: JobLease) -> JobLease: + await asyncio.sleep(ttl_seconds * 2) + return lease + + owner_task: Final = asyncio.create_task( + run_as_single_owner( + pod_lock_manager=owner, + job_name=job_name, + ttl_seconds=ttl_seconds, + when_unavailable=WhenLockUnavailable.SKIP, + run=slow_body, + ) + ) + started: Final = time.monotonic() + await asyncio.sleep(0.5) + assert await _lease_holder(redis_cache, job_name) == owner.pod_id + + await asyncio.sleep(ttl_seconds + 0.5) + elapsed: Final = time.monotonic() - started + assert elapsed > ttl_seconds, "the challenger must fire strictly after the original lease would have lapsed" + + challenger_leases: Final[list[JobLease]] = [] + challenger_result: Final = await run_as_single_owner( + pod_lock_manager=challenger, + job_name=job_name, + ttl_seconds=ttl_seconds, + when_unavailable=WhenLockUnavailable.RUN, + run=_recording_body(challenger_leases), + ) + assert challenger_leases == [], ( + f"the lease must still belong to the running owner {elapsed:.1f}s in, well past its {ttl_seconds}s TTL, " + f"so a challenger that is willing to run unguarded must still be turned away: {challenger_leases}" + ) + assert challenger_result is None + assert await _lease_holder(redis_cache, job_name) == owner.pod_id + + assert await owner_task == JobLease.LEADER + assert await _lease_holder(redis_cache, job_name) is None, "the lease must be released once the body returns" + + handover_leases: Final[list[JobLease]] = [] + await run_as_single_owner( + pod_lock_manager=challenger, + job_name=job_name, + ttl_seconds=ttl_seconds, + when_unavailable=WhenLockUnavailable.SKIP, + run=_recording_body(handover_leases), + ) + assert handover_leases == [JobLease.LEADER], "the released lease is available to the next pod" + + +async def test_rolling_restart_never_skips_or_doubles_a_tick(redis_cache: RedisCache, job_name: str) -> None: + """Across a rolling restart every tick is run by exactly one live pod, and the replacement picks the job up.""" + drained: Final = PodLockManager(redis_cache=redis_cache) + kept: Final = PodLockManager(redis_cache=redis_cache) + replacement: Final = PodLockManager(redis_cache=redis_cache) + schedule: Final = ( + ("tick-0", (drained, kept)), + ("tick-1", (drained, kept)), + ("tick-2", (kept, replacement)), + ("tick-3", (replacement,)), + ) + + async def run_tick(live_pods: tuple[PodLockManager, ...]) -> list[str]: + leaders: list[str] = [] + gate: Final = asyncio.Event() + + async def claim(pod: PodLockManager) -> None: + async def body(lease: JobLease) -> None: + assert lease is JobLease.LEADER + leaders.append(pod.pod_id) + await gate.wait() + + await run_as_single_owner( + pod_lock_manager=pod, + job_name=job_name, + ttl_seconds=10, + when_unavailable=WhenLockUnavailable.SKIP, + run=body, + ) + + tick: Final = asyncio.gather(*(claim(pod) for pod in live_pods)) + await asyncio.sleep(0.5) + gate.set() + await tick + return leaders + + outcome: Final = [(label, live_pods, await run_tick(live_pods)) for label, live_pods in schedule] + + for label, live_pods, leaders in outcome: + assert len(leaders) == 1, f"{label} must be run by exactly one pod, got {len(leaders)}" + assert leaders[0] in {pod.pod_id for pod in live_pods}, f"{label} was run by a pod that was not live" + + after_restart: Final = {leaders[0] for _, _, leaders in outcome[2:]} + assert drained.pod_id not in after_restart, "the drained pod must not run a tick after it is gone" + assert outcome[-1][2] == [replacement.pod_id], "once only the replacement is left it must take the lease" + assert await _lease_holder(redis_cache, job_name) is None, "no lease may outlive the last tick" + + +async def test_window_claim_is_taken_once_even_by_the_pod_that_won_it( + pods: list[PodLockManager], job_name: str +) -> None: + """One pod claims the window, and nobody, including the winner, claims it again until it rolls over.""" + window_seconds: Final = 2 + + claims: Final = await asyncio.gather( + *( + claim_once_per_window(pod_lock_manager=pod, job_name=job_name, window_seconds=window_seconds) + for pod in pods + ) + ) + assert claims.count(True) == 1, f"exactly one pod may claim the window: {claims}" + + winner: Final = pods[claims.index(True)] + repeat: Final = await claim_once_per_window( + pod_lock_manager=winner, job_name=job_name, window_seconds=window_seconds + ) + assert repeat is False, "the winner re-firing inside its own window must be refused" + + await asyncio.sleep(window_seconds + 0.5) + + next_window: Final = await asyncio.gather( + *( + claim_once_per_window(pod_lock_manager=pod, job_name=job_name, window_seconds=window_seconds) + for pod in pods + ) + ) + assert next_window.count(True) == 1, f"the next window is claimable exactly once: {next_window}" + + +async def test_renew_lock_extends_the_lease_only_for_its_owner( + redis_cache: RedisCache, pods: list[PodLockManager], job_name: str +) -> None: + """renew_lock is a compare-and-expire: a non-owner is refused and cannot push the owner's expiry out.""" + owner: Final = pods[0] + intruder: Final = pods[1] + initial_ttl: Final = 5 + extended_ttl: Final = 60 + + assert await owner.acquire_lock(cronjob_id=job_name, ttl=initial_ttl) is True + before: Final = await _lease_ttl_ms(redis_cache, job_name) + assert 0 < before <= initial_ttl * 1000 + + assert await intruder.renew_lock(cronjob_id=job_name, ttl=extended_ttl) is False + after_intruder: Final = await _lease_ttl_ms(redis_cache, job_name) + assert after_intruder <= before, f"a non-owner renewal must not move the expiry: {before}ms -> {after_intruder}ms" + assert await _lease_holder(redis_cache, job_name) == owner.pod_id + + assert await owner.renew_lock(cronjob_id=job_name, ttl=extended_ttl) is True + after_owner: Final = await _lease_ttl_ms(redis_cache, job_name) + assert after_owner > initial_ttl * 1000, f"the owner's renewal must push the expiry out: {after_owner}ms" + + await owner.release_lock(cronjob_id=job_name) + assert await owner.renew_lock(cronjob_id=job_name, ttl=extended_ttl) is False, "a released lease cannot be renewed" + + +async def test_release_lock_never_frees_a_lease_another_pod_has_taken_over( + redis_cache: RedisCache, pods: list[PodLockManager], job_name: str +) -> None: + """A stale owner finishing after its lease lapsed must not delete the lease its successor now holds.""" + stale_owner: Final = pods[0] + successor: Final = pods[1] + bystanders: Final = pods[2:] + short_ttl: Final = 2 + + assert await stale_owner.acquire_lock(cronjob_id=job_name, ttl=short_ttl) is True + await asyncio.sleep(short_ttl + 0.5) + assert await successor.acquire_lock(cronjob_id=job_name, ttl=30) is True + + await stale_owner.release_lock(cronjob_id=job_name) + + assert ( + await _lease_holder(redis_cache, job_name) == successor.pod_id + ), "the stale owner's release must compare owners before deleting" + + leases, results = await _concurrent_tick( + bystanders, job_name, ttl_seconds=30, when_unavailable=WhenLockUnavailable.SKIP + ) + assert leases == [], "no pod may start while the successor still holds the lease" + assert results == [None] * len(bystanders), f"every bystander skips: {results}" + + await successor.release_lock(cronjob_id=job_name) + assert await _lease_holder(redis_cache, job_name) is None, "the owner's own release must free the lease" diff --git a/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py index 3efeeee9a27..dfd7c6fe2d2 100644 --- a/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py +++ b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py @@ -13,6 +13,7 @@ from fastapi import HTTPException, status sys.path.insert(0, os.path.abspath("../../../..")) from litellm.constants import ( + DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, @@ -305,6 +306,9 @@ class TestExpiredUISessionKeyCleanupManager: mock_cache = MagicMock() mock_pod_lock_manager = MagicMock() mock_pod_lock_manager.redis_cache = MagicMock() + # the lease is readable and another pod owns it, so this is contention + # rather than an outage + mock_pod_lock_manager.redis_cache.async_get_cache = AsyncMock(return_value="another-pod") mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False) mock_pod_lock_manager.release_lock = AsyncMock() @@ -320,6 +324,7 @@ class TestExpiredUISessionKeyCleanupManager: assert deleted_count == 0 mock_pod_lock_manager.acquire_lock.assert_called_once_with( cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ttl=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, ) manager._find_expired_ui_session_keys.assert_not_called() mock_pod_lock_manager.release_lock.assert_not_called() diff --git a/tests/test_litellm/proxy/common_utils/test_single_owner_job.py b/tests/test_litellm/proxy/common_utils/test_single_owner_job.py new file mode 100644 index 00000000000..f2a90ac0272 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_single_owner_job.py @@ -0,0 +1,480 @@ +import asyncio +import json +import time +from typing import Any + +import pytest + +from litellm.proxy.common_utils.single_owner_job import ( + JobLease, + JobRole, + WhenLockUnavailable, + claim_once_per_window, + run_as_single_owner, +) +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + +JOB = "test_single_owner_job" + + +class FakeRedisCache: + """In-memory cache with real SET NX / EX semantics and a real clock. + + ``supports_scripts`` selects between PodLockManager's Lua path and its + GET-then-write fallback, so both are exercised by the same assertions. + """ + + def __init__(self, supports_scripts: bool = True, fail: bool = False): + self.store: dict[str, tuple[Any, float | None]] = {} + self.fail = fail + if supports_scripts: + self.async_register_script = self._register_script + + def _live(self, key: str) -> tuple[Any, float | None] | None: + entry = self.store.get(key) + if entry is None: + return None + _, expires_at = entry + if expires_at is not None and time.monotonic() >= expires_at: + del self.store[key] + return None + return entry + + async def async_set_cache(self, key: str, value: Any, nx: bool = False, ttl: int | None = None) -> bool: + if self.fail: + raise ConnectionError("redis down") + if nx and self._live(key) is not None: + return False + self.store[key] = (value, None if ttl is None else time.monotonic() + ttl) + return True + + async def async_get_cache(self, key: str) -> Any: + if self.fail: + raise ConnectionError("redis down") + entry = self._live(key) + return None if entry is None else entry[0] + + async def async_delete_cache(self, key: str) -> int: + if self.fail: + raise ConnectionError("redis down") + return 1 if self.store.pop(key, None) is not None else 0 + + def ttl_remaining(self, key: str) -> float | None: + entry = self._live(key) + if entry is None or entry[1] is None: + return None + return entry[1] - time.monotonic() + + def _register_script(self, script: str): + async def _run(keys: list[str], args: list[Any]) -> int: + if self.fail: + raise ConnectionError("redis down") + key = keys[0] + entry = self._live(key) + # PodLockManager compares against the JSON encoding redis stores + if entry is None or json.dumps(entry[0]) != args[0]: + return 0 + if "del" in script: + del self.store[key] + return 1 + self.store[key] = (entry[0], time.monotonic() + int(args[1])) + return 1 + + return _run + + +def manager(**kwargs: Any) -> PodLockManager: + return PodLockManager(redis_cache=FakeRedisCache(**kwargs)) + + +def sharing(count: int, **kwargs: Any) -> list[PodLockManager]: + """N pods with distinct ids against one cache, as separate replicas would be.""" + cache = FakeRedisCache(**kwargs) + return [PodLockManager(redis_cache=cache) for _ in range(count)] + + +@pytest.mark.parametrize("supports_scripts", [True, False]) +@pytest.mark.asyncio +async def test_exactly_one_pod_runs_a_tick(supports_scripts: bool): + pods = sharing(5, supports_scripts=supports_scripts) + ran: list[str] = [] + + async def body(lease: JobLease) -> str: + ran.append(lease.value) + return "done" + + results = await asyncio.gather( + *( + run_as_single_owner( + pod_lock_manager=pod, + job_name=JOB, + ttl_seconds=60, + when_unavailable=WhenLockUnavailable.RUN, + run=body, + ) + for pod in pods + ) + ) + + assert ran == ["leader"] + assert results.count("done") == 1 + assert results.count(None) == 4 + + +@pytest.mark.asyncio +async def test_follower_does_not_run_and_leader_releases(): + pods = sharing(2) + leader, follower = pods + calls: list[JobLease] = [] + + async def body(lease: JobLease) -> int: + calls.append(lease) + return 1 + + assert await run_as_single_owner( + pod_lock_manager=leader, + job_name=JOB, + ttl_seconds=60, + when_unavailable=WhenLockUnavailable.RUN, + run=body, + ) == 1 + # the lease is released on the way out, so the next pod may take the next tick + assert await run_as_single_owner( + pod_lock_manager=follower, + job_name=JOB, + ttl_seconds=60, + when_unavailable=WhenLockUnavailable.RUN, + run=body, + ) == 1 + assert calls == [JobLease.LEADER, JobLease.LEADER] + + +@pytest.mark.asyncio +async def test_a_held_lease_blocks_a_second_pod_for_the_whole_run(): + pods = sharing(2) + holder, other = pods + started = asyncio.Event() + release = asyncio.Event() + other_leases: list[JobLease] = [] + + async def slow(_lease: JobLease) -> None: + started.set() + await release.wait() + + async def quick(lease: JobLease) -> None: + other_leases.append(lease) + + holding = asyncio.create_task( + run_as_single_owner( + pod_lock_manager=holder, + job_name=JOB, + ttl_seconds=60, + when_unavailable=WhenLockUnavailable.RUN, + run=slow, + ) + ) + await started.wait() + assert ( + await run_as_single_owner( + pod_lock_manager=other, + job_name=JOB, + ttl_seconds=60, + when_unavailable=WhenLockUnavailable.RUN, + run=quick, + ) + is None + ) + assert other_leases == [] + release.set() + await holding + + +@pytest.mark.parametrize("supports_scripts", [True]) +@pytest.mark.asyncio +async def test_renewal_keeps_a_run_that_outlives_its_ttl(supports_scripts: bool): + """A body slower than the TTL keeps the lease, so no second pod joins mid-run.""" + pods = sharing(2, supports_scripts=supports_scripts) + holder, other = pods + ttl = 2 + started = asyncio.Event() + joined: list[JobLease] = [] + + async def slow(_lease: JobLease) -> None: + started.set() + await asyncio.sleep(ttl * 1.75) + + async def joiner(lease: JobLease) -> None: + joined.append(lease) + + holding = asyncio.create_task( + run_as_single_owner( + pod_lock_manager=holder, + job_name=JOB, + ttl_seconds=ttl, + when_unavailable=WhenLockUnavailable.RUN, + run=slow, + ) + ) + await started.wait() + # strictly past the original TTL, so only a renewal can still be holding it + await asyncio.sleep(ttl * 1.25) + assert ( + await run_as_single_owner( + pod_lock_manager=other, + job_name=JOB, + ttl_seconds=ttl, + when_unavailable=WhenLockUnavailable.RUN, + run=joiner, + ) + is None + ), "a renewed lease must still be held past its original TTL" + assert joined == [] + + await holding + # released once the body finished, so the next tick elects freely + await run_as_single_owner( + pod_lock_manager=other, + job_name=JOB, + ttl_seconds=ttl, + when_unavailable=WhenLockUnavailable.RUN, + run=joiner, + ) + assert joined == [JobLease.LEADER] + + +@pytest.mark.asyncio +async def test_an_expired_lease_fails_over_to_another_pod(): + pods = sharing(2) + crashed, survivor = pods + ttl = 1 + # a pod that took the lease and died: no renewal, no release + assert await crashed.acquire_lock(cronjob_id=JOB, ttl=ttl) is True + + took: list[JobLease] = [] + + async def body(lease: JobLease) -> None: + took.append(lease) + + assert ( + await run_as_single_owner( + pod_lock_manager=survivor, + job_name=JOB, + ttl_seconds=ttl, + when_unavailable=WhenLockUnavailable.RUN, + run=body, + ) + is None + ), "the lease must still be honoured before it expires" + assert took == [] + + await asyncio.sleep(ttl * 1.5) + await run_as_single_owner( + pod_lock_manager=survivor, + job_name=JOB, + ttl_seconds=ttl, + when_unavailable=WhenLockUnavailable.RUN, + run=body, + ) + assert took == [JobLease.LEADER] + + +@pytest.mark.asyncio +async def test_unreachable_redis_runs_unguarded_or_skips_by_policy(): + seen: list[JobLease] = [] + + async def body(lease: JobLease) -> str: + seen.append(lease) + return "ran" + + assert ( + await run_as_single_owner( + pod_lock_manager=manager(fail=True), + job_name=JOB, + ttl_seconds=60, + when_unavailable=WhenLockUnavailable.RUN, + run=body, + ) + == "ran" + ) + assert seen == [JobLease.UNGUARDED] + + assert ( + await run_as_single_owner( + pod_lock_manager=manager(fail=True), + job_name=JOB, + ttl_seconds=60, + when_unavailable=WhenLockUnavailable.SKIP, + run=body, + ) + is None + ) + assert seen == [JobLease.UNGUARDED], "SKIP must not run the body during an outage" + + +@pytest.mark.parametrize("lock_manager", [None, PodLockManager(redis_cache=None)]) +@pytest.mark.asyncio +async def test_a_deployment_without_redis_runs_unguarded(lock_manager: PodLockManager | None): + seen: list[JobLease] = [] + + async def body(lease: JobLease) -> str: + seen.append(lease) + return "ran" + + assert ( + await run_as_single_owner( + pod_lock_manager=lock_manager, + job_name=JOB, + ttl_seconds=60, + when_unavailable=WhenLockUnavailable.SKIP, + run=body, + ) + == "ran" + ) + assert seen == [JobLease.UNGUARDED] + + +class StaleReadCache(FakeRedisCache): + """Reports a lease holder that the store has already moved past. + + This is what a GET-then-write renewal sees when the lease lapses and another + pod takes it in the gap between the two calls: the read still names the old + owner while the key already belongs to the successor. + """ + + def __init__(self, stale_holder: str): + super().__init__(supports_scripts=False) + self.stale_holder = stale_holder + + async def async_get_cache(self, key: str) -> Any: + return self.stale_holder + + def committed(self, key: str) -> Any: + entry = self._live(key) + return None if entry is None else entry[0] + + +@pytest.mark.asyncio +async def test_renewal_never_takes_a_lease_back_from_a_successor(): + """Renewal has no safe non-atomic form, so where compare-and-expire cannot run it + must report failure rather than write. + + A GET-then-SET writes unconditionally, so a lease taken over between the two calls + is handed back to the pod that lost it and both then believe they own the job. That + is the exclusivity this whole module exists to hold, so it must not be traded for a + renewal a Redis without scripting could not have done atomically anyway. + """ + successor_id = "successor-pod" + owner = PodLockManager() + owner.redis_cache = StaleReadCache(stale_holder=owner.pod_id) + key = PodLockManager.get_redis_lock_key(JOB) + # the successor already holds it; only the read is still reporting the old owner + owner.redis_cache.store[key] = (successor_id, None) + + assert await owner.renew_lock(cronjob_id=JOB, ttl=30) is False + assert owner.redis_cache.committed(key) == successor_id, ( + "renewal must not write the lease back to the pod that lost it" + ) + + +@pytest.mark.parametrize("supports_scripts", [True]) +@pytest.mark.asyncio +async def test_renew_lock_only_extends_this_pods_lease(supports_scripts: bool): + pods = sharing(2, supports_scripts=supports_scripts) + owner, intruder = pods + cache: FakeRedisCache = owner.redis_cache # type: ignore[assignment] + assert await owner.acquire_lock(cronjob_id=JOB, ttl=30) is True + + await asyncio.sleep(0.05) + before = cache.ttl_remaining(PodLockManager.get_redis_lock_key(JOB)) + assert before is not None and before < 30 + + assert await intruder.renew_lock(cronjob_id=JOB, ttl=30) is False + assert cache.ttl_remaining(PodLockManager.get_redis_lock_key(JOB)) == pytest.approx(before, abs=0.05) + + assert await owner.renew_lock(cronjob_id=JOB, ttl=30) is True + after = cache.ttl_remaining(PodLockManager.get_redis_lock_key(JOB)) + assert after is not None and after > before + + +@pytest.mark.asyncio +async def test_renew_lock_reports_a_lost_lease(): + pods = sharing(2) + owner, thief = pods + assert await owner.acquire_lock(cronjob_id=JOB, ttl=1) is True + await asyncio.sleep(1.2) + assert await thief.acquire_lock(cronjob_id=JOB, ttl=30) is True + + assert await owner.renew_lock(cronjob_id=JOB, ttl=30) is False + assert await owner.redis_cache.async_get_cache(PodLockManager.get_redis_lock_key(JOB)) == thief.pod_id + + +@pytest.mark.asyncio +async def test_releasing_after_losing_the_lease_leaves_the_new_owner_alone(): + pods = sharing(2) + loser, winner = pods + assert await loser.acquire_lock(cronjob_id=JOB, ttl=1) is True + await asyncio.sleep(1.2) + assert await winner.acquire_lock(cronjob_id=JOB, ttl=30) is True + + await loser.release_lock(cronjob_id=JOB) + + assert await winner.redis_cache.async_get_cache(PodLockManager.get_redis_lock_key(JOB)) == winner.pod_id + + +@pytest.mark.asyncio +async def test_one_pod_claims_a_window_and_nobody_repeats_it(): + pods = sharing(4) + claims = [ + await claim_once_per_window(pod_lock_manager=pod, job_name=JOB, window_seconds=1) for pod in pods + ] + assert claims.count(True) == 1 + + winner = pods[claims.index(True)] + # not reentrant: even the holder may not redo the window it already sent + assert await claim_once_per_window(pod_lock_manager=winner, job_name=JOB, window_seconds=1) is False + + await asyncio.sleep(1.2) + reclaims = [ + await claim_once_per_window(pod_lock_manager=pod, job_name=JOB, window_seconds=1) for pod in pods + ] + assert reclaims.count(True) == 1 + + +@pytest.mark.asyncio +async def test_window_claim_runs_without_redis_and_skips_during_an_outage(): + assert await claim_once_per_window(pod_lock_manager=None, job_name=JOB, window_seconds=60) is True + assert ( + await claim_once_per_window( + pod_lock_manager=PodLockManager(redis_cache=None), job_name=JOB, window_seconds=60 + ) + is True + ) + assert ( + await claim_once_per_window(pod_lock_manager=manager(fail=True), job_name=JOB, window_seconds=60) is False + ), "a report that reaches a channel must not be sent when the claim cannot be proved" + + +@pytest.mark.parametrize( + "raw, expected", + [ + (None, JobRole.ALL), + ("", JobRole.ALL), + (" ", JobRole.ALL), + ("all", JobRole.ALL), + ("serving", JobRole.SERVING), + ("worker", JobRole.WORKER), + (" WORKER ", JobRole.WORKER), + ("Serving", JobRole.SERVING), + ("bogus", JobRole.ALL), + ], +) +def test_job_role_parsing(raw: str | None, expected: JobRole): + assert JobRole.from_env_value(raw) is expected + + +@pytest.mark.parametrize( + "role, registers", + [(JobRole.ALL, True), (JobRole.WORKER, True), (JobRole.SERVING, False)], +) +def test_only_a_serving_pod_skips_single_owner_jobs(role: JobRole, registers: bool): + assert role.runs_single_owner_jobs is registers diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py index 7a1ab60c547..b6b2636db9a 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -320,7 +321,7 @@ async def test_release_lock_uses_atomic_compare_delete_script_when_available(pod lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") mock_redis.async_register_script.assert_called_once_with(PodLockManager._COMPARE_AND_DELETE_LOCK_SCRIPT) - script_callable.assert_called_once_with(keys=[lock_key], args=[json.dumps(pod_lock_manager.pod_id)]) + script_callable.assert_called_once_with(keys=(lock_key,), args=(json.dumps(pod_lock_manager.pod_id),)) mock_redis.async_get_cache.assert_not_called() mock_redis.async_delete_cache.assert_not_called() @@ -456,3 +457,52 @@ async def test_acquire_lock_own_lock_not_reentrant(pod_lock_manager, mock_redis) assert await pod_lock_manager.acquire_lock(cronjob_id="test_job", allow_reentrant=False) is False assert await pod_lock_manager.acquire_lock(cronjob_id="test_job") is True + + +@pytest.mark.asyncio +async def test_fresh_acquisition_reports_ownership(pod_lock_manager, mock_redis): + """LIT-5434: the ownership gauge has to rise when a pod first takes a lock. + + It only ever fired on a reentrant re-acquisition, so in a normal deployment + the gauge was driven to 0 by every release and never back up, leaving job + ownership invisible in metrics. + """ + mock_redis.async_set_cache.return_value = True + + with patch( + "litellm.proxy.db.db_transaction_queue.pod_lock_manager.service_logger_obj.async_service_success_hook", + new=AsyncMock(), + ) as emit: + assert await pod_lock_manager.acquire_lock(cronjob_id="test_job") is True + # the emit is fire-and-forget, so let its task reach the hook + await asyncio.sleep(0.05) + + emit.assert_awaited_once() + metadata = emit.await_args.kwargs["event_metadata"] + assert metadata["gauge_labels"] == f"test_job:{pod_lock_manager.pod_id}" + assert metadata["gauge_value"] == 1 + + +@pytest.mark.asyncio +async def test_renew_lock_extends_only_the_owning_pod(pod_lock_manager, mock_redis): + """A renewal from a pod that no longer owns the lock must not extend it.""" + script_calls = [] + + async def fake_script(keys, args): + script_calls.append((keys, args)) + return 1 if args[0] == json.dumps(pod_lock_manager.pod_id) else 0 + + mock_redis.async_register_script = MagicMock(return_value=fake_script) + + assert await pod_lock_manager.renew_lock(cronjob_id="test_job", ttl=120) is True + lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") + assert script_calls == [((lock_key,), (json.dumps(pod_lock_manager.pod_id), 120))] + + other_pod = PodLockManager(redis_cache=mock_redis) + assert await other_pod.renew_lock(cronjob_id="test_job", ttl=120) is False + + +@pytest.mark.asyncio +async def test_renew_lock_without_redis_reports_failure(mock_redis): + """A deployment with no Redis owns no lease, so it can renew nothing.""" + assert await PodLockManager(redis_cache=None).renew_lock(cronjob_id="test_job") is False diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a06e7142122..ffc884400b1 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -32,6 +32,7 @@ from pydantic import BaseModel from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps +from litellm.proxy.common_utils.single_owner_job import JobRole from litellm.proxy.proxy_server import ( ProxyStartupEvent, _initialize_shared_aiohttp_session, @@ -840,6 +841,7 @@ def _make_slack_alerting_proxy_logging(acquire_lock_result: bool | None) -> Magi async def _init_slack_alerting_jobs( acquire_lock_result: bool | None, spend_report_frequency: str = "7d", + job_role: JobRole = JobRole.ALL, ) -> tuple[SlackAlertingJobs, MagicMock]: scheduler = MagicMock() proxy_logging_obj = _make_slack_alerting_proxy_logging(acquire_lock_result) @@ -849,12 +851,22 @@ async def _init_slack_alerting_jobs( general_settings={"spend_report_frequency": spend_report_frequency}, proxy_logging_obj=proxy_logging_obj, prisma_client=MagicMock(), + job_role=job_role, ) jobs = {call.kwargs["id"]: call.args[0] for call in scheduler.add_job.call_args_list} return jobs, proxy_logging_obj +@pytest.mark.asyncio +async def test_serving_role_registers_no_spend_report_jobs(): + """LIT-5434: each report reaches a channel once per window, so a pod dedicated to + serving traffic must not schedule them at all.""" + jobs, _ = await _init_slack_alerting_jobs(acquire_lock_result=True, job_role=JobRole.SERVING) + + assert jobs == {} + + @pytest.mark.parametrize("spend_report_frequency", ["0d", "-1d", "7h"]) @pytest.mark.asyncio async def test_initialize_slack_alerting_jobs_invalid_frequency_raises(spend_report_frequency: str): diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 75aa716bb85..cb18d702f5a 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -27,6 +27,11 @@ from litellm.caching.caching import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache +from litellm.constants import ( + MONTHLY_SPEND_REPORT_JOB_ID, + PTU_ROLLUP_JOB_ID, + WEEKLY_SPEND_REPORT_JOB_ID, +) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app, initialize @@ -11000,7 +11005,7 @@ async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkey assert mock_client.health_check.await_count == 0 -async def _run_scheduled_background_jobs(): +async def _run_scheduled_background_jobs(general_settings=None, job_role=None): from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging @@ -11016,9 +11021,12 @@ async def _run_scheduled_background_jobs(): patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch("litellm.proxy.proxy_server.store_model_in_db", True), patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True), + patch("litellm.proxy.proxy_server.LITELLM_JOB_ROLE", job_role), + # the key rotation and UI cleanup jobs read the module global, not the argument + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), ): await ProxyStartupEvent.initialize_scheduled_background_jobs( - general_settings={}, + general_settings=general_settings if general_settings is not None else {}, prisma_client=mock_prisma_client, proxy_budget_rescheduler_min_time=1, proxy_budget_rescheduler_max_time=2, @@ -11032,6 +11040,80 @@ async def _run_scheduled_background_jobs(): return ps.scheduler +# Jobs whose side effect is shared across the deployment, so a pod dedicated to +# serving traffic must register none of them +SINGLE_OWNER_JOB_IDS = ( + "reset_budget_job", + "spend_log_cleanup_job", + "key_rotation_job", + WEEKLY_SPEND_REPORT_JOB_ID, + MONTHLY_SPEND_REPORT_JOB_ID, +) +# Jobs that drain this pod's own in-memory queues or refresh its own registry, +# so every pod must keep registering them whatever its role +PER_POD_JOB_IDS = ( + "update_spend_job", + "update_daily_tag_spend_job", + "update_gateway_requests_job", + "periodic_reload_job", + "add_deployment_job", + "get_credentials_job", +) + + +def _enable_single_owner_jobs(monkeypatch): + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + monkeypatch.setattr("litellm.constants.LITELLM_KEY_ROTATION_ENABLED", "true") + monkeypatch.setattr("litellm.proxy.proxy_server.PROXY_BATCH_POLLING_ENABLED", True) + return { + "maximum_spend_logs_retention_period": "30d", + "spend_report_frequency": "7d", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("job_role", [None, "all", "worker"]) +async def test_worker_and_default_roles_register_every_job(monkeypatch, job_role): + """LIT-5434: only a 'serving' pod opts out, so an unset role keeps today's behavior.""" + settings = _enable_single_owner_jobs(monkeypatch) + + scheduler = await _run_scheduled_background_jobs(general_settings=settings, job_role=job_role) + registered = {job.id for job in scheduler.get_jobs()} + + assert set(SINGLE_OWNER_JOB_IDS) <= registered + assert set(PER_POD_JOB_IDS) <= registered + assert PTU_ROLLUP_JOB_ID in registered + + +@pytest.mark.asyncio +@pytest.mark.parametrize("job_role", ["serving", " SERVING "]) +async def test_serving_role_registers_no_single_owner_job(monkeypatch, job_role): + """LIT-5434: an operator moving auxiliary work to a dedicated worker must be able to + verify the serving pods register none of it, while their own per-pod flushes keep running.""" + settings = _enable_single_owner_jobs(monkeypatch) + + scheduler = await _run_scheduled_background_jobs(general_settings=settings, job_role=job_role) + registered = {job.id for job in scheduler.get_jobs()} + + assert registered.isdisjoint(SINGLE_OWNER_JOB_IDS) + assert PTU_ROLLUP_JOB_ID not in registered + assert set(PER_POD_JOB_IDS) <= registered + + +@pytest.mark.asyncio +async def test_unrecognised_job_role_registers_every_job(monkeypatch): + """A typo must not silently stop the deployment's budget resets and cleanups.""" + settings = _enable_single_owner_jobs(monkeypatch) + + scheduler = await _run_scheduled_background_jobs(general_settings=settings, job_role="serving-pod") + registered = {job.id for job in scheduler.get_jobs()} + + assert set(SINGLE_OWNER_JOB_IDS) <= registered + + @pytest.mark.asyncio async def test_ptu_rollup_job_registered_at_startup(monkeypatch): """The PTU rollup cron is registered once an operator opts in; only models with PTU config accrue flat cost (asserted in test_ptu_flat_cost_rollup.py)."""