diff --git a/gateway/spend_worker.py b/gateway/spend_worker.py new file mode 100644 index 00000000000..b437425d8c3 --- /dev/null +++ b/gateway/spend_worker.py @@ -0,0 +1,43 @@ +"""Spend sidecar entrypoint for the gateway image. + +Assembles ``DATABASE_URL`` the way ``gateway.launch`` does, but instead of starting a PgBouncer it +points the URL at the one the gateway container already runs on the pod's loopback (the sidecar +must see the same ``DATABASE_URL`` inputs and ``LITELLM_PGBOUNCER_*`` values as the gateway +container), then hands off to ``litellm.proxy.spend_worker``. + +Run with: + python -m gateway.spend_worker [--address unix:///var/run/litellm/spend-worker.sock] +""" + +import os +import sys +from collections.abc import Mapping, Sequence +from typing import Final + +from litellm.proxy.db.db_url_settings import DatabaseURLSettings +from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, pooled_database_url +from litellm.proxy.spend_worker import main as spend_worker_main + + +def pod_pgbouncer_database_url(pgbouncer: PgBouncerSettings, environ: Mapping[str, str]) -> str | PgBouncerError | None: + """The gateway container's PgBouncer URL for ``environ["DATABASE_URL"]``, or None when PgBouncer is off.""" + if not pgbouncer.enabled: + return None + upstream_url: Final = environ.get("DATABASE_URL") + if upstream_url is None: + return PgBouncerError("LITELLM_PGBOUNCER_ENABLED is set but no DATABASE_URL could be assembled") + return pooled_database_url(upstream_url, pgbouncer) + + +def main(argv: Sequence[str]) -> None: + DatabaseURLSettings.from_env().apply_to_env() + pooled: Final = pod_pgbouncer_database_url(PgBouncerSettings(), os.environ) + if isinstance(pooled, PgBouncerError): + sys.exit(f"LiteLLM spend worker: cannot use the pod's pgbouncer: {pooled.reason}") + if pooled is not None: + os.environ["DATABASE_URL"] = pooled + spend_worker_main(argv) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/helm/litellm-helm/templates/_helpers.tpl b/helm/litellm-helm/templates/_helpers.tpl index 8f2acb20fce..f0787f28929 100644 --- a/helm/litellm-helm/templates/_helpers.tpl +++ b/helm/litellm-helm/templates/_helpers.tpl @@ -161,3 +161,156 @@ taken before the change, which by that point no longer exists. {{- fail (printf "postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got %q). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore." $tag) -}} {{- end -}} {{- end -}} + +{{/* +Environment shared by the proxy container and the opt-in spend worker sidecar: +database, pgbouncer, master key, redis, user envVars. Both containers must see +the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer +and the same spend transaction buffer. +*/}} +{{- define "litellm.proxyEnv" -}} +- 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 }} +{{- if .Values.db.connectionPool.enabled }} +- name: LITELLM_PGBOUNCER_ENABLED + value: "true" +- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: {{ .Values.db.connectionPool.maxDbConnections | quote }} +- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: {{ .Values.db.connectionPool.maxClientConn | 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 . }} +{{- end }} +{{- if .Values.billingMetrics.enabled }} +{{ include "litellm.billingMetricsEnv" . }} +{{- end }} +{{- if .Values.metricsServer.enabled }} +{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }} +{{- fail "metricsServer.port must differ from service.port" }} +{{- end }} +- name: PROMETHEUS_METRICS_PORT + value: {{ .Values.metricsServer.port | quote }} +{{- 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 }} +{{- end -}} + +{{/* +Directory of the spend worker's unix socket, shared between the two containers +through an emptyDir. Empty when the sidecar is off or uses 127.0.0.1 TCP. +*/}} +{{- define "litellm.spendWorker.socketDir" -}} +{{- if and .Values.spendWorker.enabled (hasPrefix "unix://" .Values.spendWorker.address) -}} +{{- dir (trimPrefix "unix://" .Values.spendWorker.address) -}} +{{- end -}} +{{- end -}} + +{{- define "litellm.spendWorkerEnv" -}} +- name: LITELLM_SPEND_WORKER_ENABLED + value: "true" +- name: LITELLM_SPEND_WORKER_ADDRESS + value: {{ .Values.spendWorker.address | quote }} +- name: LITELLM_SPEND_WORKER_BUFFER_SIZE + value: {{ .Values.spendWorker.bufferSize | quote }} +- name: LITELLM_SPEND_WORKER_ON_UNAVAILABLE + value: {{ .Values.spendWorker.onUnavailable | quote }} +- name: LITELLM_SPEND_WORKER_DRAIN_TIMEOUT_SECONDS + value: {{ .Values.spendWorker.drainTimeoutSeconds | quote }} +{{- end -}} diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index 834071eb9b2..8083699f981 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -56,126 +56,9 @@ spec: 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 }} - {{- if .Values.db.connectionPool.enabled }} - - name: LITELLM_PGBOUNCER_ENABLED - value: "true" - - name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS - value: {{ .Values.db.connectionPool.maxDbConnections | quote }} - - name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN - value: {{ .Values.db.connectionPool.maxClientConn | 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.metricsServer.enabled }} - {{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }} - {{- fail "metricsServer.port must differ from service.port" }} - {{- end }} - - name: PROMETHEUS_METRICS_PORT - value: {{ .Values.metricsServer.port | quote }} - {{- 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" + {{- include "litellm.proxyEnv" . | nindent 12 }} + {{- if .Values.spendWorker.enabled }} + {{- include "litellm.spendWorkerEnv" . | nindent 12 }} {{- end }} envFrom: {{- range .Values.environmentSecrets }} @@ -253,6 +136,10 @@ spec: {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} {{- end }} + {{- if include "litellm.spendWorker.socketDir" . }} + - name: spend-worker-socket + mountPath: {{ include "litellm.spendWorker.socketDir" . }} + {{- end }} {{- with .Values.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -260,6 +147,51 @@ spec: lifecycle: {{- toYaml . | nindent 12 }} {{- end }} + {{- if .Values.spendWorker.enabled }} + - name: {{ include "litellm.name" . }}-spend-worker + securityContext: + {{- toYaml .Values.securityContext | nindent 12 }} + image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + command: {{ toYaml .Values.spendWorker.command | nindent 12 }} + env: + {{- include "litellm.proxyEnv" . | nindent 12 }} + {{- include "litellm.spendWorkerEnv" . | nindent 12 }} + - name: LITELLM_JOB_ROLE + value: spend_worker + {{- if not (hasKey (default dict .Values.envVars) "CONFIG_FILE_PATH") }} + - name: CONFIG_FILE_PATH + value: /etc/litellm/config.yaml + {{- end }} + envFrom: + {{- range .Values.environmentSecrets }} + - secretRef: + name: {{ . }} + {{- end }} + {{- range .Values.environmentConfigMaps }} + - configMapRef: + name: {{ . }} + {{- end }} + resources: + {{- toYaml .Values.spendWorker.resources | nindent 12 }} + volumeMounts: + - name: litellm-config + mountPath: /etc/litellm/config.yaml + subPath: config.yaml + {{- if include "litellm.spendWorker.socketDir" . }} + - name: spend-worker-socket + mountPath: {{ include "litellm.spendWorker.socketDir" . }} + {{- end }} + {{ if .Values.securityContext.readOnlyRootFilesystem }} + - name: tmp + mountPath: /tmp + - name: cache + mountPath: /.cache + {{- end }} + {{- with .Values.volumeMounts }} + {{- toYaml . | nindent 12 }} + {{- end }} + {{- end }} {{- with .Values.extraContainers }} {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} @@ -288,6 +220,11 @@ spec: {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} {{- end }} + {{- if include "litellm.spendWorker.socketDir" . }} + - name: spend-worker-socket + emptyDir: + sizeLimit: 1Mi + {{- end }} {{- with .Values.volumes }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/helm/litellm-helm/templates/hpa.yaml b/helm/litellm-helm/templates/hpa.yaml index ddfef56e75c..dd1441f50b9 100644 --- a/helm/litellm-helm/templates/hpa.yaml +++ b/helm/litellm-helm/templates/hpa.yaml @@ -18,6 +18,15 @@ spec: {{- end }} metrics: {{- if .Values.autoscaling.targetCPUUtilizationPercentage }} + {{- if and .Values.spendWorker.enabled .Values.spendWorker.scaleOnProxyContainerCpu }} + - type: ContainerResource + containerResource: + name: cpu + container: {{ include "litellm.name" . }} + target: + type: Utilization + averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }} + {{- else }} - type: Resource resource: name: cpu @@ -25,6 +34,7 @@ spec: type: Utilization averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }} {{- end }} + {{- end }} {{- if .Values.autoscaling.targetMemoryUtilizationPercentage }} - type: Resource resource: diff --git a/helm/litellm-helm/tests/spend_worker_tests.yaml b/helm/litellm-helm/tests/spend_worker_tests.yaml new file mode 100644 index 00000000000..83a3534fce3 --- /dev/null +++ b/helm/litellm-helm/tests/spend_worker_tests.yaml @@ -0,0 +1,214 @@ +suite: test spend worker sidecar +templates: + - deployment.yaml + - hpa.yaml + - configmap-litellm.yaml +tests: + - it: should run the proxy alone with no spend worker env by default + template: deployment.yaml + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 1 + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_SPEND_WORKER_ENABLED + value: "true" + - notContains: + path: spec.template.spec.volumes + content: + name: spend-worker-socket + any: true + + - it: should add the sidecar on the same image and point both containers at the unix socket + template: deployment.yaml + set: + image.tag: test + db.connectionPool.enabled: true + spendWorker.enabled: true + spendWorker.resources: + requests: + cpu: 500m + memory: 1Gi + limits: + cpu: "1" + memory: 2Gi + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 2 + - equal: + path: spec.template.spec.containers[1].name + value: litellm-spend-worker + - equal: + path: spec.template.spec.containers[1].image + value: ghcr.io/berriai/litellm:test + - equal: + path: spec.template.spec.containers[1].command + value: [python, -m, gateway.spend_worker] + - equal: + path: spec.template.spec.containers[1].resources.requests.cpu + value: 500m + - equal: + path: spec.template.spec.containers[1].resources.limits.memory + value: 2Gi + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_SPEND_WORKER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_SPEND_WORKER_ADDRESS + value: unix:///var/run/litellm/spend-worker.sock + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_SPEND_WORKER_BUFFER_SIZE + value: "1000" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_SPEND_WORKER_ON_UNAVAILABLE + value: fallback + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: spend_worker + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_JOB_ROLE + value: spend_worker + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_SPEND_WORKER_ADDRESS + value: unix:///var/run/litellm/spend-worker.sock + - contains: + path: spec.template.spec.containers[1].env + content: + name: CONFIG_FILE_PATH + value: /etc/litellm/config.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: DATABASE_HOST + value: RELEASE-NAME-postgresql + - contains: + path: spec.template.spec.containers[1].env + content: + name: DATABASE_PASSWORD + valueFrom: + secretKeyRef: + name: RELEASE-NAME-litellm-dbcredentials + key: password + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "20" + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: spend-worker-socket + mountPath: /var/run/litellm + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: spend-worker-socket + mountPath: /var/run/litellm + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: litellm-config + mountPath: /etc/litellm/config.yaml + subPath: config.yaml + - contains: + path: spec.template.spec.volumes + content: + name: spend-worker-socket + emptyDir: + sizeLimit: 1Mi + + - it: should skip the socket volume and pass the policy through on tcp transport + template: deployment.yaml + set: + spendWorker.enabled: true + spendWorker.address: tcp://127.0.0.1:4100 + spendWorker.onUnavailable: drop + spendWorker.bufferSize: 50 + envVars: + CONFIG_FILE_PATH: /custom/config.yaml + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 2 + - notContains: + path: spec.template.spec.containers[1].env + content: + name: CONFIG_FILE_PATH + value: /etc/litellm/config.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: CONFIG_FILE_PATH + value: /custom/config.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_SPEND_WORKER_ADDRESS + value: tcp://127.0.0.1:4100 + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_SPEND_WORKER_ON_UNAVAILABLE + value: drop + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_SPEND_WORKER_BUFFER_SIZE + value: "50" + - notContains: + path: spec.template.spec.volumes + content: + name: spend-worker-socket + any: true + + - it: should keep the pod-wide cpu metric unless asked to scale on the proxy container + template: hpa.yaml + set: + autoscaling.enabled: true + spendWorker.enabled: true + asserts: + - equal: { path: "spec.metrics[0].type", value: Resource } + - equal: { path: "spec.metrics[0].resource.name", value: cpu } + + - it: should scale on the proxy container's cpu only when opted in + template: hpa.yaml + set: + autoscaling.enabled: true + spendWorker.enabled: true + spendWorker.scaleOnProxyContainerCpu: true + asserts: + - equal: { path: "spec.metrics[0].type", value: ContainerResource } + - equal: { path: "spec.metrics[0].containerResource.name", value: cpu } + - equal: { path: "spec.metrics[0].containerResource.container", value: litellm } + - equal: { path: "spec.metrics[0].containerResource.target.averageUtilization", value: 60 } + - isNull: { path: "spec.metrics[0].resource" } + + - it: should not switch to the container metric while the sidecar is off + template: hpa.yaml + set: + autoscaling.enabled: true + spendWorker.scaleOnProxyContainerCpu: true + asserts: + - equal: { path: "spec.metrics[0].type", value: Resource } diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 1c22a5165b2..9f4c52736e9 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -190,6 +190,48 @@ metricsServer: enabled: false port: 4001 +# Opt-in sidecar that runs the post-response spend pipeline (cost calculation, +# spend logs, spend counters, budget reservation reconciliation) so the proxy's +# uvicorn workers only serialise a compact typed event and go back to serving +# inference. Same image and tag as the proxy, second container in the same pod, +# fed over loopback (a unix socket on a shared emptyDir, or 127.0.0.1 TCP). It +# reuses the pod's in-container pgbouncer (db.connectionPool) and the same Redis +# spend transaction buffer, so the per-pod DB connection budget is unchanged. +# Delivery is at-most-once inside the pod: events already handed to the sidecar +# are lost if it crashes before writing them; events the workers could not hand +# over follow onUnavailable. Both containers drain on SIGTERM within +# terminationGracePeriodSeconds +spendWorker: + enabled: false + # unix:////.sock (the becomes a shared emptyDir) or tcp://127.0.0.1: + address: unix:///var/run/litellm/spend-worker.sock + # Events each uvicorn worker holds in memory while the sidecar is slow or restarting + bufferSize: 1000 + # fallback: run the pipeline in the worker when the sidecar is unreachable or the + # buffer is full (spend stays exact, that request costs proxy CPU again) + # drop: count and discard the event instead (spend under-reports) + onUnavailable: fallback + # How long the workers keep pushing buffered events on shutdown, and how long the + # sidecar keeps serving its open connections after SIGTERM + drainTimeoutSeconds: 10 + command: + - python + - -m + - gateway.spend_worker + # Sized independently of the proxy container; the pipeline is CPU bound + resources: {} + # requests: + # cpu: 500m + # memory: 1Gi + # limits: + # cpu: "1" + # memory: 2Gi + # When autoscaling.enabled, swap the pod-wide cpu Resource metric for an + # autoscaling/v2 ContainerResource metric on the proxy container only, so the + # sidecar's CPU never scales inference replicas. Needs Kubernetes 1.30+ (or the + # HPAContainerMetrics feature gate on 1.27 to 1.29) + scaleOnProxyContainerCpu: false + resources: {} # Unset by default so the chart installs on small clusters such as Minikube, and so an diff --git a/litellm/proxy/db/pgbouncer.py b/litellm/proxy/db/pgbouncer.py index 85542c66ef7..5ad21f8de0e 100644 --- a/litellm/proxy/db/pgbouncer.py +++ b/litellm/proxy/db/pgbouncer.py @@ -236,6 +236,12 @@ def plan_pgbouncer( return PgBouncerPlan(ini=ini, userlist=userlist, pooled_url=pooled_url) +def pooled_database_url(upstream_url: str, settings: PgBouncerSettings) -> str | PgBouncerError: + """The loopback URL of a PgBouncer another container in the pod already runs for ``upstream_url``.""" + plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir=Path("/nonexistent"), run_as_user=None) + return plan if isinstance(plan, PgBouncerError) else plan.pooled_url + + def write_pgbouncer_files(plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None) -> Path: """Write the ini and userlist (both hold the password, so mode 0600) and return the ini path. diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c4fba8ecf9e..fc04882244a 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -27,12 +27,23 @@ from litellm.proxy.db.db_spend_update_writer import ( get_llm_router, ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.spend_event import ( + ObjectMapping, + SpendEventBuildError, + SpendEventDecodeError, + build_spend_event, + decode_spend_event, + is_offloadable_success, + spend_event_callback_args, +) +from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer from litellm.proxy.spend_tracking.spend_log_error_logger import ( should_suppress_spend_log_tracebacks, spend_log_error, ) from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, + _should_store_prompts_and_responses_in_spend_logs, get_request_model_access_groups, ) from litellm.proxy.utils import ProxyUpdateSpend @@ -71,8 +82,30 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( class _ProxyDBLogger(CustomLogger): - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) + def __init__(self, spend_event_producer: SpendEventProducer | None = None) -> None: + super().__init__() + self.spend_event_producer = spend_event_producer + + async def async_log_success_event( + self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if self.spend_event_producer is None or not is_offloadable_success(response_obj): + await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) + return + event: Final = build_spend_event( + kwargs, + response_obj, + start_time, + end_time, + store_bodies=_should_store_prompts_and_responses_in_spend_logs(), + ) + if isinstance(event, SpendEventBuildError): + verbose_proxy_logger.warning( + "spend worker: tracking cost in-process, event not buildable: %s", event.reason + ) + await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) + return + await self.spend_event_producer.publish(event) async def async_post_call_failure_hook( self, @@ -503,6 +536,16 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value +async def run_spend_event(line: bytes) -> None: + """Run the unchanged cost pipeline on a serialized spend event (sidecar consumer and in-process fallback).""" + event: Final = decode_spend_event(line) + if isinstance(event, SpendEventDecodeError): + verbose_proxy_logger.error("spend worker: discarding undecodable spend event: %s", event.reason) + return + args: Final = spend_event_callback_args(event) + await _ProxyDBLogger()._PROXY_track_cost_callback(args.kwargs, args.response_obj, args.start_time, args.end_time) + + def _is_unbilled_interaction_response(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import missing_usage_is_expected from litellm.types.interactions import InteractionsAPIResponse diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 09e43eb74e1..30c289bb1e8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -453,7 +453,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( from litellm.proxy.hooks.prompt_injection_detection import ( _OPTIONAL_PromptInjectionDetection, ) -from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger +from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger, run_spend_event from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.list_api.common import ( PROBLEM_TYPE_BASE, @@ -576,6 +576,11 @@ from litellm.proxy.plugin_routes import ( from litellm.proxy.plugin_routes import ( router as plugin_router, ) +from litellm.proxy.spend_tracking.spend_event_producer import ( + SpendEventProducer, + SpendWorkerSettings, + build_spend_event_producer, +) from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail try: @@ -1362,6 +1367,8 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: except Exception as e: verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e) + await _drain_spend_event_producer_on_shutdown() + await _flush_spend_logs_queue_on_shutdown() await proxy_config.stop_config_sync_subscriber() @@ -2441,16 +2448,27 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): ) +spend_event_producer: SpendEventProducer | None = None + + def cost_tracking(): - global prisma_client + global prisma_client, spend_event_producer if prisma_client is not None: from litellm.integrations.shadow_eval_logger import ShadowEvalLogger - litellm.logging_callback_manager.add_litellm_callback(_ProxyDBLogger()) - litellm.logging_callback_manager.add_litellm_async_success_callback(_ProxyDBLogger()) + spend_event_producer = build_spend_event_producer(SpendWorkerSettings(), fallback=run_spend_event) + litellm.logging_callback_manager.add_litellm_callback(_ProxyDBLogger(spend_event_producer)) + litellm.logging_callback_manager.add_litellm_async_success_callback(_ProxyDBLogger(spend_event_producer)) litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger()) +async def _drain_spend_event_producer_on_shutdown() -> None: + if spend_event_producer is None: + return + await spend_event_producer.close(drain_timeout=SpendWorkerSettings().drain_timeout_seconds) + verbose_proxy_logger.info("spend worker: producer drained on shutdown. stats=%s", spend_event_producer.stats()) + + # Bounds authoritative DB re-reads when enforcing a budget against a # stale-low spend counter: at most one DB read per counter per window. SPEND_DB_FLOOR_CACHE_TTL_SECONDS: Final = 5 diff --git a/litellm/proxy/spend_tracking/spend_event.py b/litellm/proxy/spend_tracking/spend_event.py new file mode 100644 index 00000000000..07c659eea6d --- /dev/null +++ b/litellm/proxy/spend_tracking/spend_event.py @@ -0,0 +1,418 @@ +"""Compact, typed success event handed from an inference worker to the spend sidecar. + +``build_spend_event`` runs on the inference worker right after ``Logging.async_success_handler`` +has built the ``standard_logging_object`` (so the cost is already known). It validates the success +callback's ``kwargs`` into the projection ``_PROXY_track_cost_callback`` and +``DBSpendUpdateWriter.update_database`` actually read: identities and metadata, timings, usage, the +standard logging payload without its prompt/response bodies, and the tool names. The request +messages, the raw ``proxy_server_request`` body and the full response travel only when spend logs +are configured to store prompts and responses. The cache key is the preset key the caching layer +already computed, never a fresh hash over the request body. + +``spend_event_callback_args`` rebuilds the ``(kwargs, response_obj, start_time, end_time)`` tuple +the existing cost pipeline consumes, so the sidecar runs the unchanged pipeline against the event. +""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.spend_log_tool_index import response_tool_call_names +from litellm.types.interactions import InteractionsAPIResponse +from litellm.types.utils import LiteLLMBatch, Usage + +SPEND_EVENT_VERSION: Final = 1 +CACHE_OFF_KEY: Final = "Cache OFF" + +ObjectMapping: TypeAlias = Mapping[str, object] + +_UNSERIALIZABLE_METADATA_KEYS: Final = frozenset({"user_api_key_auth", "litellm_parent_otel_span"}) +_STANDARD_LOGGING_BODY_KEYS: Final = frozenset({"messages", "response"}) +_STANDARD_LOGGING_DROPPED_KEYS: Final = frozenset({"model_parameters"}) +_NOT_OFFLOADED_RESPONSE_TYPES: Final = (LiteLLMBatch, InteractionsAPIResponse) + + +class _LitellmParams(TypedDict, total=False): + api_base: ReadOnly[str | None] + custom_llm_provider: ReadOnly[str | None] + litellm_call_id: ReadOnly[str | None] + user_api_key_end_user_id: ReadOnly[str | None] + metadata: ReadOnly[ObjectMapping | None] + litellm_metadata: ReadOnly[ObjectMapping | None] + proxy_server_request: ReadOnly[ObjectMapping | None] + preset_cache_key: ReadOnly[str | None] + + +class _DynamicParams(TypedDict, total=False): + turn_off_message_logging: ReadOnly[bool | None] + + +class _RequestBody(TypedDict, total=False): + tools: ReadOnly[Sequence[ObjectMapping] | None] + + +class _PassthroughPayload(TypedDict, total=False): + request_body: ReadOnly[_RequestBody | None] + + +class _ToolCallFunction(TypedDict): + name: ReadOnly[str] + arguments: ReadOnly[str] + + +class _ToolCall(TypedDict): + id: ReadOnly[str | None] + type: ReadOnly[Literal["function"]] + function: ReadOnly[_ToolCallFunction] + + +class _ToolCallMessage(TypedDict): + role: ReadOnly[Literal["assistant"]] + content: ReadOnly[None] + tool_calls: ReadOnly[Sequence[_ToolCall]] + + +class _ToolCallChoice(TypedDict): + index: ReadOnly[int] + finish_reason: ReadOnly[Literal["tool_calls"]] + message: ReadOnly[_ToolCallMessage] + + +class CompactResponse(TypedDict, total=False): + """What the spend pipeline reads off a response: its id, usage and which tools it called.""" + + id: ReadOnly[object] + model: ReadOnly[object] + usage: ReadOnly[object] + usage_info: ReadOnly[object] + status: ReadOnly[object] + background: ReadOnly[object] + choices: ReadOnly[Sequence[_ToolCallChoice]] + + +class _SuccessKwargs(TypedDict, total=False): + """The success callback's ``kwargs`` (``Logging.model_call_details``), validated and projected.""" + + litellm_call_id: ReadOnly[str | None] + call_type: ReadOnly[str | None] + model: ReadOnly[str | None] + custom_llm_provider: ReadOnly[str | None] + stream: ReadOnly[bool | None] + complete_streaming_response: ReadOnly[object] + cache_hit: ReadOnly[bool | None] + response_cost: ReadOnly[float | None] + completion_start_time: ReadOnly[datetime | None] + agent_id: ReadOnly[str | None] + litellm_trace_id: ReadOnly[str | None] + litellm_params: ReadOnly[_LitellmParams] + standard_logging_object: ReadOnly[ObjectMapping | None] + standard_callback_dynamic_params: ReadOnly[_DynamicParams | None] + combined_usage_object: ReadOnly[Usage | None] + realtime_tools: ReadOnly[Sequence[object] | None] + realtime_tool_calls: ReadOnly[Sequence[object] | None] + tools: ReadOnly[Sequence[ObjectMapping] | None] + passthrough_logging_payload: ReadOnly[_PassthroughPayload | None] + + +class _FunctionToolFunction(TypedDict): + name: ReadOnly[str] + + +class _FunctionTool(TypedDict): + type: ReadOnly[Literal["function"]] + function: ReadOnly[_FunctionToolFunction] + + +class SpendCallbackKwargs(TypedDict): + """The ``kwargs`` handed to ``_PROXY_track_cost_callback`` on the sidecar.""" + + litellm_call_id: ReadOnly[str | None] + call_type: ReadOnly[str | None] + model: ReadOnly[str | None] + custom_llm_provider: ReadOnly[str | None] + stream: ReadOnly[bool | None] + cache_hit: ReadOnly[bool | None] + response_cost: ReadOnly[float | None] + completion_start_time: ReadOnly[datetime | None] + agent_id: ReadOnly[str | None] + litellm_trace_id: ReadOnly[str | None] + litellm_params: ReadOnly[_LitellmParams] + standard_logging_object: ReadOnly[ObjectMapping | None] + standard_callback_dynamic_params: ReadOnly[_DynamicParams | None] + combined_usage_object: ReadOnly[Usage | None] + realtime_tools: ReadOnly[Sequence[object] | None] + realtime_tool_calls: ReadOnly[Sequence[object] | None] + tools: ReadOnly[Sequence[_FunctionTool] | None] + complete_streaming_response: NotRequired[ReadOnly[CompactResponse | None]] + + +_NO_LITELLM_PARAMS: Final[_LitellmParams] = {} +_SUCCESS_KWARGS: Final = TypeAdapter(_SuccessKwargs) +_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) +_COMPACT_RESPONSE: Final = TypeAdapter(CompactResponse) + + +class SpendEvent(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + version: Literal[1] + litellm_call_id: str | None + call_type: str | None + model: str | None + custom_llm_provider: str | None + stream: bool | None + complete_streaming_response: bool + cache_hit: bool | None + response_cost: float | None + start_time: datetime + end_time: datetime + completion_start_time: datetime | None + agent_id: str | None + litellm_trace_id: str | None + litellm_params: _LitellmParams + standard_logging_object: ObjectMapping | None + standard_callback_dynamic_params: _DynamicParams | None + response: CompactResponse | None + combined_usage: ObjectMapping | None + realtime_tools: Sequence[object] | None + realtime_tool_calls: Sequence[object] | None + request_tool_names: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class SpendEventCallbackArgs: + kwargs: SpendCallbackKwargs + response_obj: CompactResponse | None + start_time: datetime + end_time: datetime + + +@dataclass(frozen=True, slots=True) +class SpendEventBuildError: + reason: str + + +@dataclass(frozen=True, slots=True) +class SpendEventDecodeError: + reason: str + + +def is_offloadable_success(response_obj: object) -> bool: + """Batch retrieves and interaction polls branch on the concrete response class, so they stay in-process.""" + return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) + + +def _json_fallback(value: object) -> str: + return str(value) + + +def _mapping_or_none(value: object) -> ObjectMapping | None: + try: + return _OBJECT_MAPPING.validate_python(value) + except ValidationError: + return None + + +def _drop_keys(mapping: ObjectMapping, keys: frozenset[str]) -> ObjectMapping: + return MappingProxyType({key: value for key, value in mapping.items() if key not in keys}) + + +def _budget_reservation(metadata: ObjectMapping) -> ObjectMapping | None: + """The admission-time reservation, wherever the request setup left it, so the sidecar can reconcile it.""" + direct: Final = _mapping_or_none(metadata.get("user_api_key_budget_reservation")) + if direct is not None: + return direct + auth: Final = metadata.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth): + return auth.budget_reservation + auth_mapping: Final = _mapping_or_none(auth) + return _mapping_or_none(auth_mapping.get("budget_reservation")) if auth_mapping is not None else None + + +def _metadata_for_event( + metadata: ObjectMapping | None, budget_reservation: ObjectMapping | None +) -> ObjectMapping | None: + if metadata is None: + return None + kept: Final = _drop_keys(metadata, _UNSERIALIZABLE_METADATA_KEYS) + if budget_reservation is None: + return kept + return MappingProxyType({**kept, "user_api_key_budget_reservation": budget_reservation}) + + +def _litellm_params_for_event( + litellm_params: _LitellmParams, cache_key: str | None, store_bodies: bool +) -> _LitellmParams: + metadata: Final = litellm_params.get("metadata") + litellm_metadata: Final = litellm_params.get("litellm_metadata") + budget_reservation: Final = next( + ( + reservation + for source in (litellm_metadata, metadata) + if source is not None and (reservation := _budget_reservation(source)) is not None + ), + None, + ) + projected: Final[_LitellmParams] = { + "api_base": litellm_params.get("api_base"), + "custom_llm_provider": litellm_params.get("custom_llm_provider"), + "litellm_call_id": litellm_params.get("litellm_call_id"), + "user_api_key_end_user_id": litellm_params.get("user_api_key_end_user_id"), + "metadata": _metadata_for_event(metadata, budget_reservation), + "litellm_metadata": _metadata_for_event(litellm_metadata, budget_reservation), + "proxy_server_request": litellm_params.get("proxy_server_request") if store_bodies else None, + "preset_cache_key": cache_key, + } + return projected + + +def _standard_logging_for_event(sl_object: ObjectMapping | None, store_bodies: bool) -> ObjectMapping | None: + if sl_object is None: + return None + dropped: Final = ( + _STANDARD_LOGGING_DROPPED_KEYS if store_bodies else _STANDARD_LOGGING_DROPPED_KEYS | _STANDARD_LOGGING_BODY_KEYS + ) + return _drop_keys(sl_object, dropped) + + +def _tool_call(name: str) -> _ToolCall: + tool_call: Final[_ToolCall] = {"id": None, "type": "function", "function": {"name": name, "arguments": "{}"}} + return tool_call + + +def _tool_call_choice(names: Sequence[str]) -> _ToolCallChoice: + choice: Final[_ToolCallChoice] = { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": tuple(_tool_call(name) for name in names)}, + } + return choice + + +def _compact_response(response_obj: object) -> CompactResponse | None: + """Usage, identity and tool calls of the response, in chat-completions shape, without the content.""" + dumped: Final = response_obj.model_dump() if isinstance(response_obj, BaseModel) else _mapping_or_none(response_obj) + if dumped is None: + return None + scalars: Final = _COMPACT_RESPONSE.validate_python(_drop_keys(dumped, frozenset({"choices"}))) + tool_call_names: Final = response_tool_call_names(response_obj) + if not tool_call_names: + return scalars + with_tool_calls: Final[CompactResponse] = {**scalars, "choices": (_tool_call_choice(tool_call_names),)} + return with_tool_calls + + +def _tool_name(tool: ObjectMapping) -> str | None: + """Chat tools nest the name under ``function``; Anthropic and Responses API tools keep it at the top.""" + function: Final = _mapping_or_none(tool.get("function")) + name: Final = function.get("name") if function is not None else tool.get("name") + return name.strip() if isinstance(name, str) and name.strip() else None + + +def _request_tool_names(kwargs: _SuccessKwargs) -> tuple[str, ...]: + passthrough: Final = kwargs.get("passthrough_logging_payload") + request_body: Final = passthrough.get("request_body") if passthrough is not None else None + passthrough_tools: Final = request_body.get("tools") if request_body is not None else None + return tuple( + name + for source in (kwargs.get("tools"), passthrough_tools) + if source is not None + for tool in source + if (name := _tool_name(tool)) is not None + ) + + +def preset_spend_log_cache_key(litellm_params: _LitellmParams) -> str | None: + """The key the caching layer already stored in ``litellm_params``, or ``Cache OFF``; never hashes the body.""" + if litellm.cache is None: + return CACHE_OFF_KEY + return litellm_params.get("preset_cache_key") + + +def _function_tool(name: str) -> _FunctionTool: + tool: Final[_FunctionTool] = {"type": "function", "function": {"name": name}} + return tool + + +def build_spend_event( + raw_kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime, store_bodies: bool +) -> bytes | SpendEventBuildError: + """Validate the success callback's kwargs and serialize the event once, as a single JSON line.""" + try: + kwargs: Final = _SUCCESS_KWARGS.validate_python(raw_kwargs) + except ValidationError as error: + return SpendEventBuildError(reason=str(error)) + litellm_params: Final = kwargs.get("litellm_params", _NO_LITELLM_PARAMS) + sl_object: Final = kwargs.get("standard_logging_object") + cache_key: Final = preset_spend_log_cache_key(litellm_params) + response_cost: Final = sl_object.get("response_cost") if sl_object is not None else kwargs.get("response_cost") + combined_usage: Final = kwargs.get("combined_usage_object") + event: Final = SpendEvent( + version=SPEND_EVENT_VERSION, + litellm_call_id=kwargs.get("litellm_call_id"), + call_type=kwargs.get("call_type"), + model=kwargs.get("model"), + custom_llm_provider=kwargs.get("custom_llm_provider"), + stream=kwargs.get("stream"), + complete_streaming_response="complete_streaming_response" in kwargs, + cache_hit=kwargs.get("cache_hit"), + response_cost=response_cost if isinstance(response_cost, (int, float)) else None, + start_time=start_time, + end_time=end_time, + completion_start_time=kwargs.get("completion_start_time"), + agent_id=kwargs.get("agent_id"), + litellm_trace_id=kwargs.get("litellm_trace_id"), + litellm_params=_litellm_params_for_event(litellm_params, cache_key, store_bodies), + standard_logging_object=_standard_logging_for_event(sl_object, store_bodies), + standard_callback_dynamic_params=kwargs.get("standard_callback_dynamic_params"), + response=_compact_response(response_obj), + combined_usage=combined_usage.model_dump() if combined_usage is not None else None, + realtime_tools=kwargs.get("realtime_tools"), + realtime_tool_calls=kwargs.get("realtime_tool_calls"), + request_tool_names=_request_tool_names(kwargs), + ) + return event.model_dump_json(fallback=_json_fallback).encode() + b"\n" + + +def decode_spend_event(line: bytes) -> SpendEvent | SpendEventDecodeError: + try: + return SpendEvent.model_validate_json(line) + except ValidationError as error: + return SpendEventDecodeError(reason=str(error)) + + +def spend_event_callback_args(event: SpendEvent) -> SpendEventCallbackArgs: + """The ``(kwargs, response_obj, start_time, end_time)`` the in-process cost callback receives.""" + tools: Final = tuple(_function_tool(name) for name in event.request_tool_names) + kwargs: Final[SpendCallbackKwargs] = { + "litellm_call_id": event.litellm_call_id, + "call_type": event.call_type, + "model": event.model, + "custom_llm_provider": event.custom_llm_provider, + "stream": event.stream, + "cache_hit": event.cache_hit, + "response_cost": event.response_cost, + "completion_start_time": event.completion_start_time, + "agent_id": event.agent_id, + "litellm_trace_id": event.litellm_trace_id, + "litellm_params": event.litellm_params, + "standard_logging_object": event.standard_logging_object, + "standard_callback_dynamic_params": event.standard_callback_dynamic_params, + "combined_usage_object": Usage.model_validate(event.combined_usage) + if event.combined_usage is not None + else None, + "realtime_tools": event.realtime_tools, + "realtime_tool_calls": event.realtime_tool_calls, + "tools": tools or None, + } + if not event.complete_streaming_response: + return SpendEventCallbackArgs(kwargs, event.response, event.start_time, event.end_time) + streaming_kwargs: Final[SpendCallbackKwargs] = {**kwargs, "complete_streaming_response": event.response} + return SpendEventCallbackArgs(streaming_kwargs, event.response, event.start_time, event.end_time) diff --git a/litellm/proxy/spend_tracking/spend_event_producer.py b/litellm/proxy/spend_tracking/spend_event_producer.py new file mode 100644 index 00000000000..ca417132048 --- /dev/null +++ b/litellm/proxy/spend_tracking/spend_event_producer.py @@ -0,0 +1,276 @@ +"""Fire-and-forget push of serialized spend events from an inference worker to the pod-local sidecar. + +``LITELLM_SPEND_WORKER_ENABLED=true`` turns the push on in the gateway; the sidecar process sets +``LITELLM_JOB_ROLE=spend_worker`` and always runs the pipeline in-process. Events queue in a bounded +in-memory buffer that a single writer task flushes over a unix socket or loopback TCP connection. +When the sidecar is unreachable, the buffer is full, or the connection breaks mid-write, each affected +event follows ``LITELLM_SPEND_WORKER_ON_UNAVAILABLE``: ``fallback`` runs the existing cost pipeline in +the worker, ``drop`` counts it and moves on. Transitions are logged with the counters, so a sidecar +outage is visible without scraping anything. + +Delivery is at-most-once: a sidecar crash loses the events already handed to its socket. Events from +one uvicorn worker are handled in the order it produced them; events from different workers interleave, +exactly like the in-process callbacks do today. +""" + +import asyncio +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias +from urllib.parse import urlsplit + +from pydantic import AliasChoices, Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from litellm._logging import verbose_proxy_logger + +SPEND_WORKER_ENV_PREFIX: Final = "LITELLM_SPEND_WORKER_" +SPEND_WORKER_JOB_ROLE: Final = "spend_worker" +DEFAULT_SPEND_WORKER_ADDRESS: Final = "unix:///var/run/litellm/spend-worker.sock" +RECONNECT_BACKOFF_SECONDS: Final = 1.0 +DROP_LOG_EVERY: Final = 1000 + +UnavailablePolicy: TypeAlias = Literal["fallback", "drop"] +PublishOutcome: TypeAlias = Literal["queued", "fallback", "dropped"] + + +class SpendWorkerSettings(BaseSettings): + """``LITELLM_SPEND_WORKER_*`` env vars, shared by the gateway producer and the sidecar consumer.""" + + model_config = SettingsConfigDict( + env_prefix=SPEND_WORKER_ENV_PREFIX, case_sensitive=False, extra="ignore", frozen=True, populate_by_name=True + ) + + enabled: bool = False + address: str = DEFAULT_SPEND_WORKER_ADDRESS + buffer_size: int = Field(default=1000, ge=1) + on_unavailable: UnavailablePolicy = "fallback" + drain_timeout_seconds: float = Field(default=10.0, gt=0) + connect_timeout_seconds: float = Field(default=1.0, gt=0) + job_role: str | None = Field(default=None, validation_alias=AliasChoices("LITELLM_JOB_ROLE")) + + @property + def produces(self) -> bool: + return self.enabled and self.job_role != SPEND_WORKER_JOB_ROLE + + +@dataclass(frozen=True, slots=True) +class UnixAddress: + path: str + + +@dataclass(frozen=True, slots=True) +class TcpAddress: + host: str + port: int + + +@dataclass(frozen=True, slots=True) +class AddressError: + reason: str + + +SpendWorkerAddress: TypeAlias = UnixAddress | TcpAddress + + +def parse_spend_worker_address(address: str) -> SpendWorkerAddress | AddressError: + """``unix:///path/to.sock`` or ``tcp://127.0.0.1:port``.""" + parsed: Final = urlsplit(address) + if parsed.scheme == "unix" and parsed.path: + return UnixAddress(path=parsed.path) + if parsed.scheme == "tcp" and parsed.hostname and parsed.port is not None: + return TcpAddress(host=parsed.hostname, port=parsed.port) + return AddressError(reason=f"expected unix:///path or tcp://host:port, got {address!r}") + + +async def open_spend_worker_connection( + address: SpendWorkerAddress, timeout: float +) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + match address: + case UnixAddress(path=path): + return await asyncio.wait_for(asyncio.open_unix_connection(path), timeout) + case TcpAddress(host=host, port=port): + return await asyncio.wait_for(asyncio.open_connection(host, port), timeout) + + +def build_spend_event_producer( + settings: SpendWorkerSettings, fallback: Callable[[bytes], Awaitable[None]] +) -> "SpendEventProducer | None": + """The gateway producer for these settings, or ``None`` when the pipeline stays in-process.""" + if not settings.produces: + return None + address: Final = parse_spend_worker_address(settings.address) + if isinstance(address, AddressError): + verbose_proxy_logger.error("spend worker: %s; running the spend pipeline in-process", address.reason) + return None + verbose_proxy_logger.info( + "spend worker: offloading spend tracking to %s (buffer=%d, on_unavailable=%s)", + settings.address, + settings.buffer_size, + settings.on_unavailable, + ) + return SpendEventProducer( + address=address, + on_unavailable=settings.on_unavailable, + buffer_size=settings.buffer_size, + connect_timeout=settings.connect_timeout_seconds, + fallback=fallback, + ) + + +@dataclass(frozen=True, slots=True) +class SpendEventProducerStats: + queued: int + sent: int + fallback: int + dropped: int + connected: bool + + +class SpendEventProducer: + """Bounded buffer plus one writer task per process; see the module docstring for the contract.""" + + def __init__( + self, + address: SpendWorkerAddress, + on_unavailable: UnavailablePolicy, + buffer_size: int, + connect_timeout: float, + fallback: Callable[[bytes], Awaitable[None]], + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._address = address + self._on_unavailable = on_unavailable + self._buffer_size = buffer_size + self._connect_timeout = connect_timeout + self._fallback = fallback + self._clock = clock + self._queue: asyncio.Queue[bytes] | None = None + self._writer_task: asyncio.Task[None] | None = None + self._writer: asyncio.StreamWriter | None = None + self._closing = False + self._next_connect_at = 0.0 + self._queued = 0 + self._sent = 0 + self._fallback_count = 0 + self._dropped = 0 + + def stats(self) -> SpendEventProducerStats: + return SpendEventProducerStats( + queued=self._queued, + sent=self._sent, + fallback=self._fallback_count, + dropped=self._dropped, + connected=self._writer is not None, + ) + + async def publish(self, line: bytes) -> PublishOutcome: + """Hand one serialized event to the writer task, or apply the unavailable policy right away.""" + if self._closing or self._clock() < self._next_connect_at: + return await self._unavailable(line, "sidecar unreachable") + queue: Final = self._ensure_writer() + try: + queue.put_nowait(line) + except asyncio.QueueFull: + return await self._unavailable(line, "buffer full") + self._queued += 1 + return "queued" + + async def close(self, drain_timeout: float) -> None: + """Flush the buffer for up to ``drain_timeout`` seconds, then apply the unavailable policy to the rest.""" + self._closing = True + queue: Final = self._queue + task: Final = self._writer_task + if queue is None or task is None: + return + try: + await asyncio.wait_for(queue.join(), drain_timeout) + except asyncio.TimeoutError: + verbose_proxy_logger.warning( + "spend worker: %s events still buffered after %.1fs drain timeout", queue.qsize(), drain_timeout + ) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + while not queue.empty(): + await self._unavailable(queue.get_nowait(), "shutdown") + await self._disconnect() + + def _ensure_writer(self) -> asyncio.Queue[bytes]: + if self._queue is None: + self._queue = asyncio.Queue(maxsize=self._buffer_size) + if self._writer_task is None or self._writer_task.done(): + self._writer_task = asyncio.get_running_loop().create_task(self._run_writer(self._queue)) + return self._queue + + async def _run_writer(self, queue: asyncio.Queue[bytes]) -> None: + while True: + line = await queue.get() + try: + await self._send(line) + finally: + queue.task_done() + + async def _send(self, line: bytes) -> None: + writer: Final = await self._connect() + if writer is None: + await self._unavailable(line, "sidecar unreachable") + return + try: + writer.write(line) + await writer.drain() + except (ConnectionError, OSError) as error: + await self._disconnect() + self._next_connect_at = self._clock() + RECONNECT_BACKOFF_SECONDS + await self._unavailable(line, f"write failed: {error}") + return + self._sent += 1 + + async def _connect(self) -> asyncio.StreamWriter | None: + if self._writer is not None: + return self._writer + if self._clock() < self._next_connect_at: + return None + try: + _, writer = await open_spend_worker_connection(self._address, self._connect_timeout) + except (ConnectionError, OSError, asyncio.TimeoutError) as error: + self._next_connect_at = self._clock() + RECONNECT_BACKOFF_SECONDS + verbose_proxy_logger.warning( + "spend worker: cannot reach %s (%s); applying %s policy for %.0fs. stats=%s", + self._address, + error, + self._on_unavailable, + RECONNECT_BACKOFF_SECONDS, + self.stats(), + ) + return None + self._writer = writer + verbose_proxy_logger.info("spend worker: connected to %s. stats=%s", self._address, self.stats()) + return writer + + async def _disconnect(self) -> None: + writer: Final = self._writer + self._writer = None + if writer is None: + return + writer.close() + try: + await writer.wait_closed() + except (ConnectionError, OSError): + pass + + async def _unavailable(self, line: bytes, reason: str) -> PublishOutcome: + if self._on_unavailable == "fallback": + self._fallback_count += 1 + try: + await self._fallback(line) + except Exception: + verbose_proxy_logger.exception("spend worker: in-process fallback failed (%s)", reason) + return "fallback" + self._dropped += 1 + if self._dropped % DROP_LOG_EVERY == 1: + verbose_proxy_logger.warning("spend worker: dropping spend event (%s). stats=%s", reason, self.stats()) + return "dropped" diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a21d761996f..0d589e6f329 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -533,10 +533,12 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs additional_usage_values["cache_creation_input_tokens"] = cache_write_tokens clean_metadata["additional_usage_values"] = additional_usage_values - if litellm.cache is not None: - cache_key = litellm.cache.get_cache_key(**kwargs) - else: + if litellm.cache is None: cache_key = "Cache OFF" + elif litellm_params.get("preset_cache_key") is not None: + cache_key = litellm_params["preset_cache_key"] + else: + cache_key = litellm.cache.get_cache_key(**kwargs) if cache_hit is True: import time diff --git a/litellm/proxy/spend_worker.py b/litellm/proxy/spend_worker.py new file mode 100644 index 00000000000..e90d9638009 --- /dev/null +++ b/litellm/proxy/spend_worker.py @@ -0,0 +1,160 @@ +"""Spend sidecar: consume spend events from the pod's inference workers and run the cost pipeline. + +Runs the proxy startup lifespan (config, Prisma, Redis transaction buffer, scheduled spend flushes) +without serving HTTP, then listens on ``LITELLM_SPEND_WORKER_ADDRESS`` for newline-delimited spend +events. Each event goes through the unchanged ``_ProxyDBLogger._PROXY_track_cost_callback``, so +spend logs, spend counters, budget reservation reconciliation and cache updates happen exactly as +they would in-process, just in this container. Events are handled in order per producer connection +(one per uvicorn worker); a slow pipeline fills the socket buffer and the producer's bounded queue, +which is the backpressure that triggers its fallback or drop policy. ``SIGTERM`` stops accepting +connections, finishes the events already sent, then runs the proxy shutdown (which flushes the +buffered spend transactions). + + LITELLM_JOB_ROLE=spend_worker python -m litellm.proxy.spend_worker [--address unix:///path.sock] +""" + +import asyncio +import os +import signal +import sys +from collections.abc import Awaitable, Callable, Sequence +from pathlib import Path +from typing import Final + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.spend_tracking.spend_event_producer import ( + SPEND_WORKER_JOB_ROLE, + AddressError, + SpendWorkerAddress, + SpendWorkerSettings, + TcpAddress, + UnixAddress, + parse_spend_worker_address, +) + +MAX_EVENT_BYTES: Final = 64 * 1024 * 1024 + + +class SpendEventConsumer: + """Accepts producer connections and runs ``handler`` on every line each one sends, in order.""" + + def __init__(self, handler: Callable[[bytes], Awaitable[None]]) -> None: + self._handler = handler + self._open_connections = 0 + self._idle = asyncio.Event() + self._idle.set() + self._received = 0 + self._handled = 0 + self._failed = 0 + + @property + def received(self) -> int: + return self._received + + @property + def handled(self) -> int: + return self._handled + + @property + def failed(self) -> int: + return self._failed + + async def serve(self, address: SpendWorkerAddress) -> asyncio.Server: + match address: + case UnixAddress(path=path): + socket_path: Final = Path(path) + socket_path.parent.mkdir(parents=True, exist_ok=True) + socket_path.unlink(missing_ok=True) + return await asyncio.start_unix_server(self._on_connection, path=path, limit=MAX_EVENT_BYTES) + case TcpAddress(host=host, port=port): + return await asyncio.start_server(self._on_connection, host=host, port=port, limit=MAX_EVENT_BYTES) + + async def _on_connection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + self._open_connections += 1 + self._idle.clear() + try: + while line := await reader.readline(): + if not line.endswith(b"\n"): + verbose_proxy_logger.error("spend worker: discarding truncated spend event (%d bytes)", len(line)) + break + self._received += 1 + await self._handle(line) + except (ConnectionError, asyncio.IncompleteReadError, asyncio.LimitOverrunError) as error: + verbose_proxy_logger.warning("spend worker: producer connection ended abnormally: %s", error) + finally: + writer.close() + self._open_connections -= 1 + if self._open_connections == 0: + self._idle.set() + + async def _handle(self, line: bytes) -> None: + try: + await self._handler(line) + self._handled += 1 + except Exception: + self._failed += 1 + verbose_proxy_logger.exception("spend worker: spend event failed") + + async def drain(self, timeout: float) -> int: + """Keep serving the open connections until every producer hangs up, or ``timeout`` seconds pass. + + Returns how many producer connections were still open when the timeout hit. + """ + try: + await asyncio.wait_for(self._idle.wait(), timeout) + except TimeoutError: + pass + return self._open_connections + + +def _install_stop_signals(loop: asyncio.AbstractEventLoop, stop: asyncio.Event) -> None: + for signum in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(signum, stop.set) + + +async def run_spend_worker(address: SpendWorkerAddress, drain_timeout: float) -> None: + from fastapi import FastAPI + + from litellm.proxy.hooks.proxy_track_cost_callback import run_spend_event + from litellm.proxy.proxy_server import proxy_startup_event + + stop: Final = asyncio.Event() + _install_stop_signals(asyncio.get_running_loop(), stop) + consumer: Final = SpendEventConsumer(handler=run_spend_event) + async with proxy_startup_event(FastAPI()): + server: Final = await consumer.serve(address) + verbose_proxy_logger.info("spend worker: listening on %s", address) + await stop.wait() + server.close() + still_open: Final = await consumer.drain(drain_timeout) + verbose_proxy_logger.info( + "spend worker: stopping. received=%d handled=%d failed=%d connections_cut=%d", + consumer.received, + consumer.handled, + consumer.failed, + still_open, + ) + + +def _address_argument(argv: Sequence[str], default: str) -> str | AddressError: + match tuple(argv): + case (): + return default + case ("--address", value): + return value + case _: + return AddressError(f"usage: python -m litellm.proxy.spend_worker [--address ADDRESS], got {tuple(argv)}") + + +def main(argv: Sequence[str]) -> None: + os.environ.setdefault("LITELLM_JOB_ROLE", SPEND_WORKER_JOB_ROLE) + settings: Final = SpendWorkerSettings() + raw_address: Final = _address_argument(argv, default=settings.address) + address: Final = raw_address if isinstance(raw_address, AddressError) else parse_spend_worker_address(raw_address) + if isinstance(address, AddressError): + sys.exit(f"LiteLLM spend worker: {address.reason}") + asyncio.run(run_spend_worker(address, drain_timeout=settings.drain_timeout_seconds)) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/litellm/utils.py b/litellm/utils.py index 06ee3dd88af..edc727f79cc 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1843,6 +1843,9 @@ def client(original_function): elif _caching_handler_response.embedding_all_elements_cache_hit is True: return _caching_handler_response.final_embedding_cached_response + if _llm_caching_handler.preset_cache_key is not None: + logging_obj.litellm_params["preset_cache_key"] = _llm_caching_handler.preset_cache_key + # CHECK MAX TOKENS if ( kwargs.get("max_tokens", None) is not None diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index eff892f2d80..cddb69b03f7 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -5,14 +5,20 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth +from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( _get_budget_reservation_from_metadata, _ProxyDBLogger, _should_track_cost_callback, _update_database_and_spend_counters, + run_spend_event, ) -from litellm.types.utils import CallTypes, Usage +from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event +from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.spend_worker import SpendEventConsumer +from litellm.types.utils import CallTypes, ModelResponse, Usage @pytest.mark.asyncio @@ -1875,9 +1881,15 @@ async def test_track_cost_callback_keeps_guardrail_cost_on_cache_hit(): } with ( - patch("litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock) as mock_increment, # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam - patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), # test-quality-ok: same function-body import, no injection seam - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, # test-quality-ok: same function-body import, no injection seam + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() @@ -2096,3 +2108,234 @@ async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_un ) assert charged == ("premium", "tier0") + + +def _offload_kwargs() -> dict: + big_prompt = "x" * 10_000 + reservation = {"reserved_cost": 0.5, "entries": [{"counter_key": "key:hash-1", "reserved_cost": 0.5}]} + return { + "litellm_call_id": "call-1", + "call_type": "acompletion", + "model": "gpt-4o", + "custom_llm_provider": "openai", + "stream": False, + "cache_hit": None, + "response_cost": 0.0125, + "completion_start_time": datetime(2026, 1, 1, 0, 0, 1), + "messages": [{"role": "user", "content": big_prompt}], + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + "litellm_params": { + "api_base": "https://api.openai.com", + "preset_cache_key": None, + "proxy_server_request": {"body": {"messages": [{"role": "user", "content": big_prompt}]}}, + "metadata": { + "user_api_key": "hash-1", + "user_api_key_hash": "hash-1", + "user_api_key_alias": "alias-1", + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-1", + "user_api_key_org_id": "org-1", + "user_api_key_end_user_id": "end-user-1", + "user_api_key_auth": UserAPIKeyAuth(api_key="hash-1", budget_reservation=reservation), + "model_group": "gpt-4o", + "model_info": {"id": "deployment-1"}, + "tags": ["tag-a"], + }, + }, + "standard_logging_object": { + "id": "chatcmpl-1", + "trace_id": "trace-1", + "response_cost": 0.0125, + "model": "gpt-4o-2024-08-06", + "model_id": "deployment-1", + "model_group": "gpt-4o", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "prompt_tokens": 5000, + "completion_tokens": 4000, + "total_tokens": 9000, + "request_tags": ["tag-a"], + "request_model_access_groups": ["premium"], + "messages": [{"role": "user", "content": big_prompt}], + "response": {"choices": [{"message": {"content": "y" * 10_000}}]}, + "model_parameters": {"temperature": 0.1}, + "metadata": { + "user_api_key_hash": "hash-1", + "user_api_key_end_user_id": "end-user-1", + "usage_object": {"prompt_tokens": 5000, "completion_tokens": 4000, "total_tokens": 9000}, + }, + "hidden_params": {"litellm_overhead_time_ms": 3}, + "model_map_information": {}, + "cost_breakdown": {"input_cost": 0.0125, "output_cost": 0.0}, + }, + } + + +def _offload_response() -> ModelResponse: + return ModelResponse( + id="chatcmpl-1", + model="gpt-4o-2024-08-06", + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call-1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + "finish_reason": "tool_calls", + } + ], + usage=Usage(prompt_tokens=5000, completion_tokens=4000, total_tokens=9000), + ) + + +class _RecordingHandler: + def __init__(self) -> None: + self.lines: list[bytes] = [] # mutable-ok: test double records the events the sidecar received + + async def __call__(self, line: bytes) -> None: + self.lines.append(line) + + +async def _no_fallback(line: bytes) -> None: + raise AssertionError("the sidecar was reachable, nothing should fall back") + + +@pytest.mark.asyncio +async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_skips_the_pipeline(tmp_path): + handler = _RecordingHandler() + consumer = SpendEventConsumer(handler) + address = UnixAddress(path=str(tmp_path / "spend.sock")) + server = await consumer.serve(address) + producer = SpendEventProducer( + address=address, on_unavailable="fallback", buffer_size=10, connect_timeout=1.0, fallback=_no_fallback + ) + logger = _ProxyDBLogger(producer) + + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as counters, + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + await logger.async_log_success_event(_offload_kwargs(), _offload_response(), datetime.now(), datetime.now()) + await producer.close(drain_timeout=5.0) + server.close() + assert await consumer.drain(timeout=5.0) == 0 + + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_awaited() + counters.assert_not_awaited() + assert producer.stats().sent == 1 + assert len(handler.lines) == 1 + assert len(handler.lines[0]) < 4_000 + event = decode_spend_event(handler.lines[0]) + assert not isinstance(event, SpendEventDecodeError) + assert event.litellm_params["metadata"]["user_api_key_team_id"] == "team-1" + assert event.response_cost == 0.0125 + + +@pytest.mark.asyncio +async def test_async_log_success_event_keeps_batch_retrieves_in_process(): + producer = SpendEventProducer( + address=UnixAddress(path="/nonexistent/spend.sock"), + on_unavailable="drop", + buffer_size=10, + connect_timeout=1.0, + fallback=_no_fallback, + ) + logger = _ProxyDBLogger(producer) + kwargs = _batch_retrieve_kwargs(CallTypes.aretrieve_batch.value) + + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + await logger.async_log_success_event( + kwargs, _retrieved_batch("in_progress", output_file_id=None), datetime.now(), datetime.now() + ) + + assert producer.stats().queued == 0 + + +async def _spend_row_written_by(run) -> tuple[SpendLogsPayload, dict, tuple[str, ...]]: + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as counters, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True) + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + await run() + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + written = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs + counters.assert_awaited_once() + counted = dict(counters.await_args.kwargs) + + row = get_logging_payload( + kwargs=written["kwargs"], + response_obj=written["completion_response"], + start_time=written["start_time"], + end_time=written["end_time"], + ) + return row, counted, response_tool_call_names(written["completion_response"]) + + +@pytest.mark.asyncio +async def test_sidecar_writes_the_same_spend_row_and_counters_as_the_in_process_path(): + start_time = datetime(2026, 1, 1, 0, 0, 0) + end_time = datetime(2026, 1, 1, 0, 0, 2) + + async def in_process() -> None: + await _ProxyDBLogger()._PROXY_track_cost_callback( + _offload_kwargs(), _offload_response(), start_time=start_time, end_time=end_time + ) + + async def via_sidecar() -> None: + line = build_spend_event(_offload_kwargs(), _offload_response(), start_time, end_time, store_bodies=False) + assert isinstance(line, bytes) + await run_spend_event(line) + + in_process_row, in_process_counters, in_process_tools = await _spend_row_written_by(in_process) + sidecar_row, sidecar_counters, sidecar_tools = await _spend_row_written_by(via_sidecar) + + assert sidecar_row == in_process_row + assert in_process_row["spend"] == 0.0125 + assert in_process_row["team_id"] == "team-1" + assert in_process_row["end_user"] == "end-user-1" + assert in_process_row["total_tokens"] == 9000 + assert in_process_row["model_id"] == "deployment-1" + assert in_process_row["request_tags"] == '["tag-a"]' + assert in_process_row["messages"] == "{}" + assert in_process_row["response"] == "{}" + assert sidecar_counters == in_process_counters + assert in_process_counters["token"] == "hash-1" + assert in_process_counters["response_cost"] == 0.0125 + assert in_process_counters["budget_reservation"]["reserved_cost"] == 0.5 + assert in_process_counters["model_access_groups"] == ("premium",) + assert sidecar_tools == in_process_tools == ("get_weather",) + + +@pytest.mark.asyncio +async def test_sidecar_ignores_an_undecodable_event(): # test-quality-ok: a discarded event has no observable output other than the DB writer never being reached + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + await run_spend_event(b"garbage\n") + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_awaited() diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_event.py b/tests/test_litellm/proxy/spend_tracking/test_spend_event.py new file mode 100644 index 00000000000..ff449235582 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_event.py @@ -0,0 +1,213 @@ +import json +from datetime import datetime +from typing import Final + +import pytest + +import litellm +from litellm.caching.caching import Cache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.spend_tracking.spend_event import ( + CACHE_OFF_KEY, + SpendEventBuildError, + SpendEventDecodeError, + build_spend_event, + decode_spend_event, + is_offloadable_success, + spend_event_callback_args, +) +from litellm.types.utils import LiteLLMBatch, ModelResponse, Usage + +_BIG_PROMPT: Final = "x" * 20_000 +_RESERVATION: Final = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "key:hash", "reserved_cost": 0.5}], + "finalized": False, + "input_cost": 0.1, + "input_tokens": 5000, +} + + +def _response(tool_name: str | None = None) -> ModelResponse: + tool_calls: Final = ( + [{"id": "call-1", "type": "function", "function": {"name": tool_name, "arguments": "{}"}}] + if tool_name is not None + else None + ) + return ModelResponse( + id="chatcmpl-1", + model="gpt-4o-2024-08-06", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "y" * 20_000, "tool_calls": tool_calls}, + "finish_reason": "tool_calls" if tool_name else "stop", + } + ], + usage=Usage(prompt_tokens=5000, completion_tokens=4000, total_tokens=9000), + ) + + +def _success_kwargs(preset_cache_key: str | None = "preset-key") -> dict: + return { + "litellm_call_id": "call-1", + "call_type": "acompletion", + "model": "gpt-4o", + "custom_llm_provider": "openai", + "stream": False, + "cache_hit": None, + "response_cost": 0.0125, + "completion_start_time": datetime(2026, 1, 1, 0, 0, 1), + "messages": [{"role": "user", "content": _BIG_PROMPT}], + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + "litellm_params": { + "api_base": "https://api.openai.com", + "preset_cache_key": preset_cache_key, + "proxy_server_request": {"body": {"messages": [{"role": "user", "content": _BIG_PROMPT}]}}, + "metadata": { + "user_api_key": "hash-1", + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-1", + "user_api_key_org_id": "org-1", + "user_api_key_end_user_id": "end-user-1", + "user_api_key_auth": UserAPIKeyAuth(api_key="hash-1", budget_reservation=dict(_RESERVATION)), + "model_group": "gpt-4o", + "model_info": {"id": "deployment-1"}, + "tags": ["tag-a"], + "litellm_parent_otel_span": object(), + }, + }, + "standard_logging_object": { + "response_cost": 0.0125, + "model": "gpt-4o-2024-08-06", + "model_id": "deployment-1", + "request_tags": ["tag-a"], + "request_model_access_groups": ["premium"], + "messages": [{"role": "user", "content": _BIG_PROMPT}], + "response": {"choices": [{"message": {"content": "y" * 20_000}}]}, + "model_parameters": {"temperature": 0.1}, + "metadata": {"user_api_key_hash": "hash-1", "usage_object": {"prompt_tokens": 5000}}, + "hidden_params": {"litellm_overhead_time_ms": 3}, + "model_map_information": {}, + }, + } + + +def _build(kwargs: dict, response: object, store_bodies: bool = False) -> bytes: + line: Final = build_spend_event( + kwargs, response, datetime(2026, 1, 1), datetime(2026, 1, 1, 0, 0, 2), store_bodies=store_bodies + ) + assert isinstance(line, bytes) + return line + + +def test_event_is_compact_and_omits_bodies_by_default(): + line: Final = _build(_success_kwargs(), _response(tool_name="get_weather")) + assert line.endswith(b"\n") + assert len(line) < 4_000 + assert _BIG_PROMPT.encode() not in line + assert b"yyyy" not in line + decoded: Final = json.loads(line) + assert "messages" not in decoded["standard_logging_object"] + assert "response" not in decoded["standard_logging_object"] + assert decoded["litellm_params"]["proxy_server_request"] is None + + +def test_event_carries_bodies_when_spend_logs_store_them(): + line: Final = _build(_success_kwargs(), _response(), store_bodies=True) + decoded: Final = json.loads(line) + assert decoded["standard_logging_object"]["messages"][0]["content"] == _BIG_PROMPT + assert decoded["standard_logging_object"]["response"]["choices"][0]["message"]["content"] == "y" * 20_000 + assert decoded["litellm_params"]["proxy_server_request"]["body"]["messages"][0]["content"] == _BIG_PROMPT + + +def test_round_trip_preserves_identity_usage_reservation_and_tools(): + line: Final = _build(_success_kwargs(), _response(tool_name="get_weather")) + event: Final = decode_spend_event(line) + assert not isinstance(event, SpendEventDecodeError) + args: Final = spend_event_callback_args(event) + + metadata: Final = args.kwargs["litellm_params"]["metadata"] + assert metadata is not None + assert (metadata["user_api_key"], metadata["user_api_key_team_id"], metadata["user_api_key_org_id"]) == ( + "hash-1", + "team-1", + "org-1", + ) + assert metadata["user_api_key_budget_reservation"] == _RESERVATION + assert "user_api_key_auth" not in metadata + assert "litellm_parent_otel_span" not in metadata + assert args.kwargs["standard_logging_object"]["request_model_access_groups"] == ["premium"] + assert args.kwargs["standard_logging_object"]["response_cost"] == 0.0125 + assert args.kwargs["tools"] == ({"type": "function", "function": {"name": "get_weather"}},) + assert args.kwargs["completion_start_time"] == datetime(2026, 1, 1, 0, 0, 1) + assert (args.start_time, args.end_time) == (datetime(2026, 1, 1), datetime(2026, 1, 1, 0, 0, 2)) + assert args.response_obj is not None + assert args.response_obj["id"] == "chatcmpl-1" + assert args.response_obj["usage"]["prompt_tokens"] == 5000 + assert args.response_obj["usage"]["completion_tokens"] == 4000 + tool_calls: Final = args.response_obj["choices"][0]["message"]["tool_calls"] + assert [call["function"]["name"] for call in tool_calls] == ["get_weather"] + assert "complete_streaming_response" not in args.kwargs + + +def test_streaming_event_reconstructs_complete_streaming_response(): + kwargs: Final = {**_success_kwargs(), "stream": True, "complete_streaming_response": _response()} + event: Final = decode_spend_event(_build(kwargs, _response())) + assert not isinstance(event, SpendEventDecodeError) + args: Final = spend_event_callback_args(event) + assert args.kwargs["stream"] is True + assert args.kwargs["complete_streaming_response"] == args.response_obj + + +class _HashingCache(Cache): + def __init__(self) -> None: + pass + + def get_cache_key(self, **kwargs) -> str: + raise AssertionError("the fast path must not hash the request body") + + +@pytest.mark.parametrize( + ("cache", "preset", "expected"), + [ + (None, "preset-key", CACHE_OFF_KEY), + (_HashingCache(), "preset-key", "preset-key"), + (_HashingCache(), None, None), + ], +) +def test_event_reuses_preset_cache_key_and_never_hashes(monkeypatch, cache, preset, expected): + monkeypatch.setattr(litellm, "cache", cache) + decoded: Final = json.loads(_build(_success_kwargs(preset_cache_key=preset), _response())) + assert decoded["litellm_params"]["preset_cache_key"] == expected + + +def test_unbuildable_kwargs_fall_back_to_in_process_tracking(): + kwargs: Final = {**_success_kwargs(), "response_cost": "not-a-number"} + assert isinstance( + build_spend_event(kwargs, _response(), datetime.now(), datetime.now(), False), SpendEventBuildError + ) + + +def test_undecodable_line_is_an_error_value(): + assert isinstance(decode_spend_event(b'{"version": 2}\n'), SpendEventDecodeError) + assert isinstance(decode_spend_event(b"not json\n"), SpendEventDecodeError) + + +def test_batch_retrieves_stay_in_process(): + assert is_offloadable_success(_response()) is True + assert is_offloadable_success(None) is True + assert ( + is_offloadable_success( + LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + ) + ) + is False + ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_event_producer.py b/tests/test_litellm/proxy/spend_tracking/test_spend_event_producer.py new file mode 100644 index 00000000000..644b896747c --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_event_producer.py @@ -0,0 +1,165 @@ +import asyncio +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.spend_tracking.spend_event_producer import ( + AddressError, + SpendEventProducer, + SpendWorkerSettings, + TcpAddress, + UnixAddress, + build_spend_event_producer, + parse_spend_worker_address, +) + + +class _Sidecar: + """A unix-socket server that records every line it receives, standing in for the spend worker.""" + + def __init__(self, path: Path) -> None: + self.path = path + self.lines: list[bytes] = [] # mutable-ok: test double records what the producer sent + self._server: asyncio.Server | None = None + + async def __aenter__(self) -> "_Sidecar": + self._server = await asyncio.start_unix_server(self._on_connection, path=str(self.path)) + return self + + async def __aexit__(self, *exc: object) -> None: + assert self._server is not None + self._server.close() + await self._server.wait_closed() + + async def _on_connection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + while line := await reader.readline(): + self.lines.append(line) + writer.close() + + +class _Fallback: + def __init__(self) -> None: + self.lines: list[bytes] = [] # mutable-ok: test double records what fell back to in-process + + async def __call__(self, line: bytes) -> None: + self.lines.append(line) + + +def _producer(path: Path, fallback: _Fallback, on_unavailable="fallback", buffer_size: int = 100) -> SpendEventProducer: + return SpendEventProducer( + address=UnixAddress(path=str(path)), + on_unavailable=on_unavailable, + buffer_size=buffer_size, + connect_timeout=1.0, + fallback=fallback, + ) + + +def test_parse_spend_worker_address(): + assert parse_spend_worker_address("unix:///var/run/litellm/spend-worker.sock") == UnixAddress( + path="/var/run/litellm/spend-worker.sock" + ) + assert parse_spend_worker_address("tcp://127.0.0.1:4100") == TcpAddress(host="127.0.0.1", port=4100) + assert isinstance(parse_spend_worker_address("redis://localhost:6379"), AddressError) + assert isinstance(parse_spend_worker_address("tcp://127.0.0.1"), AddressError) + + +def test_gateway_produces_only_when_enabled_and_not_the_sidecar_itself(): + fallback: Final = _Fallback() + assert build_spend_event_producer(SpendWorkerSettings(enabled=False), fallback) is None + assert build_spend_event_producer(SpendWorkerSettings(enabled=True, job_role="spend_worker"), fallback) is None + assert build_spend_event_producer(SpendWorkerSettings(enabled=True, address="redis://x"), fallback) is None + assert isinstance(build_spend_event_producer(SpendWorkerSettings(enabled=True), fallback), SpendEventProducer) + + +def test_settings_read_the_documented_env(monkeypatch): + monkeypatch.setenv("LITELLM_SPEND_WORKER_ENABLED", "true") + monkeypatch.setenv("LITELLM_SPEND_WORKER_ADDRESS", "tcp://127.0.0.1:4100") + monkeypatch.setenv("LITELLM_SPEND_WORKER_BUFFER_SIZE", "50") + monkeypatch.setenv("LITELLM_SPEND_WORKER_ON_UNAVAILABLE", "drop") + monkeypatch.setenv("LITELLM_JOB_ROLE", "spend_worker") + settings: Final = SpendWorkerSettings() + assert (settings.enabled, settings.address, settings.buffer_size, settings.on_unavailable) == ( + True, + "tcp://127.0.0.1:4100", + 50, + "drop", + ) + assert settings.produces is False + + +@pytest.mark.asyncio +async def test_events_reach_the_sidecar_once_and_in_order(tmp_path: Path): + fallback: Final = _Fallback() + async with _Sidecar(tmp_path / "spend.sock") as sidecar: + producer: Final = _producer(sidecar.path, fallback) + outcomes: Final = [await producer.publish(f"event-{i}\n".encode()) for i in range(20)] + await producer.close(drain_timeout=5.0) + await asyncio.sleep(0.05) + + assert outcomes == ["queued"] * 20 + assert sidecar.lines == [f"event-{i}\n".encode() for i in range(20)] + assert fallback.lines == [] + stats: Final = producer.stats() + assert (stats.queued, stats.sent, stats.fallback, stats.dropped) == (20, 20, 0, 0) + + +@pytest.mark.asyncio +async def test_unreachable_sidecar_falls_back_in_process_and_backs_off(tmp_path: Path): + fallback: Final = _Fallback() + producer: Final = _producer(tmp_path / "missing.sock", fallback) + first: Final = await producer.publish(b"event-1\n") + await asyncio.sleep(0.05) + second: Final = await producer.publish(b"event-2\n") + await producer.close(drain_timeout=5.0) + + assert first == "queued" + assert second == "fallback" + assert fallback.lines == [b"event-1\n", b"event-2\n"] + stats: Final = producer.stats() + assert (stats.sent, stats.fallback, stats.dropped, stats.connected) == (0, 2, 0, False) + + +@pytest.mark.asyncio +async def test_drop_policy_counts_instead_of_running_in_process(tmp_path: Path): + fallback: Final = _Fallback() + producer: Final = _producer(tmp_path / "missing.sock", fallback, on_unavailable="drop") + await producer.publish(b"event-1\n") + await producer.close(drain_timeout=5.0) + assert await producer.publish(b"event-2\n") == "dropped" + + assert fallback.lines == [] + assert producer.stats().dropped == 2 + + +@pytest.mark.asyncio +async def test_full_buffer_applies_the_unavailable_policy_immediately(tmp_path: Path): + fallback: Final = _Fallback() + async with _Sidecar(tmp_path / "spend.sock") as sidecar: + producer: Final = _producer(sidecar.path, fallback, buffer_size=2) + outcomes: Final = [await producer.publish(f"event-{i}\n".encode()) for i in range(3)] + await producer.close(drain_timeout=5.0) + await asyncio.sleep(0.05) + + assert outcomes == ["queued", "queued", "fallback"] + assert fallback.lines == [b"event-2\n"] + assert sidecar.lines == [b"event-0\n", b"event-1\n"] + + +@pytest.mark.asyncio +async def test_close_flushes_buffered_events_then_refuses_new_ones(tmp_path: Path): + fallback: Final = _Fallback() + async with _Sidecar(tmp_path / "spend.sock") as sidecar: + producer: Final = _producer(sidecar.path, fallback) + for i in range(50): + await producer.publish(f"event-{i}\n".encode()) + assert sidecar.lines == [] + await producer.close(drain_timeout=5.0) + await asyncio.sleep(0.05) + after_close: Final = await producer.publish(b"late\n") + + assert len(sidecar.lines) == 50 + assert after_close == "fallback" + assert fallback.lines == [b"late\n"] + assert producer.stats().connected is False diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 5a79560b972..22ab3a2d64c 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1,8 +1,8 @@ import asyncio import datetime import json -from datetime import timezone from collections.abc import Mapping +from datetime import timezone from typing import Any, Final, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -19,7 +19,7 @@ from litellm.constants import ( SESSION_ID_OMITTED_METADATA_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_messages_for_spend_logs_payload, @@ -39,7 +39,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_logging_payload, get_spend_logs_id, ) -from litellm.proxy._types import SpendLogsPayload from litellm.proxy.utils import hash_token from litellm.types.utils import ( StandardLoggingHiddenParams, @@ -81,6 +80,48 @@ def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_token assert additional_usage_values["prompt_tokens_details"]["cached_tokens"] == 123 +class _HashingCache(litellm.Cache): + def __init__(self) -> None: + pass + + def get_cache_key(self, **kwargs) -> str: + raise AssertionError("a preset cache key must be reused instead of hashing the request") + + +def _cache_key_in_spend_log(monkeypatch: pytest.MonkeyPatch, cache: litellm.Cache | None, preset: str | None) -> str: + monkeypatch.setattr(litellm, "cache", cache) + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "x" * 10_000}], + "litellm_params": {"metadata": {"user_api_key": "test-key"}, "preset_cache_key": preset}, + }, + response_obj=litellm.ModelResponse(id="chatcmpl-test", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + return payload["cache_key"] + + +def test_get_logging_payload_reuses_the_preset_cache_key_instead_of_hashing_the_body(monkeypatch): + assert _cache_key_in_spend_log(monkeypatch, _HashingCache(), "preset-key") == "preset-key" + + +def test_get_logging_payload_records_cache_off_without_hashing(monkeypatch): + assert _cache_key_in_spend_log(monkeypatch, None, None) == "Cache OFF" + + +def test_get_logging_payload_still_hashes_when_caching_is_on_and_no_preset_key_exists(monkeypatch): + class _RecordingCache(litellm.Cache): + def __init__(self) -> None: + pass + + def get_cache_key(self, **kwargs) -> str: + return "hashed-from-" + kwargs["model"] + + assert _cache_key_in_spend_log(monkeypatch, _RecordingCache(), None) == "hashed-from-gpt-4o-mini" + + _TRACE_ONLY_STANDARD_LOGGING: Final = cast( StandardLoggingPayload, { @@ -139,9 +180,7 @@ def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id(): from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1") - cost_session: Final = _get_batch_trace_session_id( - call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost" - ) + cost_session: Final = _get_batch_trace_session_id(call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost") assert create_session == cost_session == "batch-uid-1" @@ -4252,7 +4291,7 @@ ANTHROPIC_MESSAGES_SSE_CHUNKS: Final = ( 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' '"usage":{"output_tokens":4}}\n\n', - "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', ) @@ -4290,9 +4329,7 @@ def test_spend_log_request_id_is_the_message_id_a_non_streaming_messages_caller_ """ logging_obj = _anthropic_messages_logging_obj(stream=False) - logged_response = logging_obj._handle_anthropic_messages_response_logging( - result=ANTHROPIC_MESSAGES_RESPONSE - ) + logged_response = logging_obj._handle_anthropic_messages_response_logging(result=ANTHROPIC_MESSAGES_RESPONSE) assert logged_response.id == "msg_01Lit6806NonStreaming" assert ( @@ -4368,9 +4405,7 @@ def test_spend_log_request_id_still_falls_back_to_litellm_call_id_without_a_prov end_time=datetime.datetime.now(timezone.utc), logging_obj=logging_obj, ) - assert logging_obj.model_call_details["complete_streaming_response"].id == ( - "6806cafe-0000-4000-8000-000000000001" - ) + assert logging_obj.model_call_details["complete_streaming_response"].id == ("6806cafe-0000-4000-8000-000000000001") def test_spend_log_request_id_for_chat_completions_is_untouched(): diff --git a/tests/test_litellm/proxy/test_spend_worker.py b/tests/test_litellm/proxy/test_spend_worker.py new file mode 100644 index 00000000000..79d7f0815f1 --- /dev/null +++ b/tests/test_litellm/proxy/test_spend_worker.py @@ -0,0 +1,97 @@ +import asyncio +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.spend_tracking.spend_event_producer import ( + AddressError, + SpendEventProducer, + TcpAddress, + UnixAddress, + open_spend_worker_connection, +) +from litellm.proxy.spend_worker import SpendEventConsumer, _address_argument + + +class _Handler: + def __init__(self, fail_on: bytes | None = None) -> None: + self.lines: list[bytes] = [] # mutable-ok: test double records the events the consumer handed over + self._fail_on = fail_on + + async def __call__(self, line: bytes) -> None: + if line == self._fail_on: + raise RuntimeError("pipeline failed") + self.lines.append(line) + + +async def _no_fallback(line: bytes) -> None: + raise AssertionError(f"unexpected fallback for {line!r}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["unix", "tcp"]) +async def test_consumer_handles_each_producer_line_once_in_order(tmp_path: Path, transport: str): + handler: Final = _Handler(fail_on=b"event-3\n") + consumer: Final = SpendEventConsumer(handler) + server: Final = await consumer.serve( + UnixAddress(path=str(tmp_path / "spend.sock")) if transport == "unix" else TcpAddress("127.0.0.1", 0) + ) + address: Final = ( + UnixAddress(path=str(tmp_path / "spend.sock")) + if transport == "unix" + else TcpAddress("127.0.0.1", server.sockets[0].getsockname()[1]) + ) + producer: Final = SpendEventProducer( + address=address, on_unavailable="fallback", buffer_size=100, connect_timeout=1.0, fallback=_no_fallback + ) + for i in range(6): + await producer.publish(f"event-{i}\n".encode()) + await producer.close(drain_timeout=5.0) + + server.close() + assert await consumer.drain(timeout=5.0) == 0 + assert handler.lines == [f"event-{i}\n".encode() for i in range(6) if i != 3] + assert (consumer.received, consumer.handled, consumer.failed) == (6, 5, 1) + + +@pytest.mark.asyncio +async def test_consumer_discards_a_truncated_trailing_event(tmp_path: Path): + handler: Final = _Handler() + consumer: Final = SpendEventConsumer(handler) + address: Final = UnixAddress(path=str(tmp_path / "spend.sock")) + server: Final = await consumer.serve(address) + _, writer = await open_spend_worker_connection(address, timeout=1.0) + writer.write(b"whole\npartial-without-newline") + await writer.drain() + writer.close() + await writer.wait_closed() + await asyncio.sleep(0.05) + + server.close() + assert await consumer.drain(timeout=5.0) == 0 + assert handler.lines == [b"whole\n"] + assert consumer.received == 1 + + +@pytest.mark.asyncio +async def test_drain_reports_producers_still_connected_after_the_timeout(tmp_path: Path): + consumer: Final = SpendEventConsumer(_Handler()) + address: Final = UnixAddress(path=str(tmp_path / "spend.sock")) + server: Final = await consumer.serve(address) + _, writer = await open_spend_worker_connection(address, timeout=1.0) + await asyncio.sleep(0.05) + + server.close() + assert await consumer.drain(timeout=0.1) == 1 + writer.close() + await writer.wait_closed() + assert await consumer.drain(timeout=5.0) == 0 + + +def test_address_argument(): + assert _address_argument((), default="unix:///tmp/x.sock") == "unix:///tmp/x.sock" + assert _address_argument(("--address", "tcp://127.0.0.1:4100"), default="unix:///tmp/x.sock") == ( + "tcp://127.0.0.1:4100" + ) + assert isinstance(_address_argument(("--listen", "x"), default="unix:///tmp/x.sock"), AddressError)