mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Merge remote-tracking branch 'origin/litellm_perf_offload_logging' into litellm_perf_1k_rps_sidecar
This commit is contained in:
commit
0dfb7627b6
19 changed files with 2225 additions and 147 deletions
43
gateway/spend_worker.py
Normal file
43
gateway/spend_worker.py
Normal file
|
|
@ -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:])
|
||||
|
|
@ -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 -}}
|
||||
|
|
|
|||
|
|
@ -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 }}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
214
helm/litellm-helm/tests/spend_worker_tests.yaml
Normal file
214
helm/litellm-helm/tests/spend_worker_tests.yaml
Normal file
|
|
@ -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 }
|
||||
|
|
@ -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:///<dir>/<file>.sock (the <dir> becomes a shared emptyDir) or tcp://127.0.0.1:<port>
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
418
litellm/proxy/spend_tracking/spend_event.py
Normal file
418
litellm/proxy/spend_tracking/spend_event.py
Normal file
|
|
@ -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)
|
||||
276
litellm/proxy/spend_tracking/spend_event_producer.py
Normal file
276
litellm/proxy/spend_tracking/spend_event_producer.py
Normal file
|
|
@ -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"
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
160
litellm/proxy/spend_worker.py
Normal file
160
litellm/proxy/spend_worker.py
Normal file
|
|
@ -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:])
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
213
tests/test_litellm/proxy/spend_tracking/test_spend_event.py
Normal file
213
tests/test_litellm/proxy/spend_tracking/test_spend_event.py
Normal file
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
@ -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():
|
||||
|
|
|
|||
97
tests/test_litellm/proxy/test_spend_worker.py
Normal file
97
tests/test_litellm/proxy/test_spend_worker.py
Normal file
|
|
@ -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)
|
||||
Loading…
Add table
Reference in a new issue