scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

feat(proxy): elect one owner per auxiliary DB job and add a worker role
Some checks failed
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Details

Browse source 
Every proxy process registers the same scheduler, and it registers it once per
uvicorn worker process rather than once per pod, so a job with a shared side
effect runs replicas times processes. The jobs that did elect an owner each
hand-rolled it, and the dialects disagree on the one case that matters:
acquire_lock reports contention and an unreachable Redis identically, so the
PTU rollup ran unguarded while key rotation, spend log cleanup and UI session
cleanup skipped, meaning a Redis outage silently stopped those three
everywhere at once.

Adds run_as_single_owner and claim_once_per_window, which name that decision
as WhenLockUnavailable and carry the two lease shapes the proxy actually uses:
a mutex held for one run, and a done-marker sized to a reporting window. The
lease is now renewed while the body runs, so the TTL is a failover deadline
rather than a run budget. On a live proxy that is the difference between a
spend log cleanup losing its 60s lock partway through a multi-minute sweep,
which a challenger could take at 75s, 120s and 180s, and holding it to
completion.

LITELLM_JOB_ROLE=serving registers none of these jobs, so an operator can run
them on a dedicated worker deployment instead of on every replica taking
traffic. A serving pod keeps the jobs that drain its own in-memory queues and
refresh its own model registry, so it still writes spend. The default is
unchanged, and an unrecognised value falls back to it with a warning.

Both Helm charts gain an opt-in worker Deployment, off by default and pinned
to one process. Live-Redis integration tests cover election, failover,
renewal, rolling restart and once-per-window claims.

Also fixes the ownership gauge, which only ever fired on a reentrant
re-acquisition and so was driven to 0 by every release and never back up.
This commit is contained in:
Yassin Kortam 2026-08-11 18:13:05 -07:00 • committed by yassin
parent 035a3227ac
commit 2ad9e53d90
29 changed files with 2766 additions and 461 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

105
.github/workflows/test-integration-redis.yml vendored Normal file
View file

@ -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

6
helm/litellm-helm/README.md
View file

@ -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` |

8
helm/litellm-helm/templates/NOTES.txt
View file

@ -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.*
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 }}

286
helm/litellm-helm/templates/_deployment.tpl Normal file
View file

@ -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 -}}

33
helm/litellm-helm/templates/_helpers.tpl
View file

@ -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

251
helm/litellm-helm/templates/deployment.yaml
View file

@ -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) }}

38
helm/litellm-helm/templates/worker-deployment.yaml Normal file
View file

@ -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 }}

316
helm/litellm-helm/tests/worker_deployment_tests.yaml Normal file
View file

@ -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

32
helm/litellm-helm/values.yaml
View file

@ -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

10
helm/litellm/templates/_helpers.tpl
View file

@ -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.

4
helm/litellm/templates/backend/deployment.yaml
View file

@ -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:

4
helm/litellm/templates/gateway/deployment.yaml
View file

@ -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:

104
helm/litellm/templates/worker/deployment.yaml Normal file
View file

@ -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 }}

269
helm/litellm/tests/worker_deployment_tests.yaml Normal file
View file

@ -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

62
helm/litellm/values.yaml
View file

@ -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

7
litellm/constants.py
View file

@ -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")))

44
litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py
View file

@ -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(

59
litellm/proxy/common_utils/key_rotation_manager.py
View file

@ -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]:
"""

220
litellm/proxy/common_utils/single_owner_job.py Normal file
View file

@ -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

139
litellm/proxy/db/db_transaction_queue/pod_lock_manager.py
View file

@ -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):

114
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -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")

74
litellm/proxy/proxy_server.py
View file

@ -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()

59
litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py
View file

@ -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(

348
tests/redis_integration_tests/test_single_owner_job_redis.py Normal file
View file

@ -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"

5
tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py
View file

@ -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()

480
tests/test_litellm/proxy/common_utils/test_single_owner_job.py Normal file
View file

@ -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

52
tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py
View file

@ -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

12
tests/test_litellm/proxy/proxy_server/test_lifecycle.py
View file

@ -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):

86
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -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)."""