[Feature]: Batch update db cost tracking calls

BerriAI / litellm

Call all LLM APIs using the OpenAI format. Use Bedrock, Azure, OpenAI, Cohere, Anthropic, Ollama, Sagemaker, HuggingFace, Replicate (100+ LLMs)

https://docs.litellm.ai/docs/

Other

10.26k stars 1.14k forks source link

[Feature]: Batch update db cost tracking calls #1867

Open krrishdholakia opened 4 months ago

krrishdholakia commented 4 months ago

The Feature

in update_db(),

have the get spend call for user + key be 1 request
have the update call for user, key, spend logs table be a batched request

If any call fails, add them to an in-memory job queue for updating db on a periodic basis

Motivation, pitch

reduce prisma clients used / call

Twitter / LinkedIn details

No response

Manouchehri commented 4 months ago

Just a note on this, for serverless deployments, the batch job should immediately be triggered if LiteLLM gets a SIGTERM signal. Without this, there's a decent chance the queue will be lost and stuff will end up going missing.

https://cloud.google.com/run/docs/container-contract#forced-termination-jobs

https://github.com/aws-samples/graceful-shutdown-with-aws-lambda?tab=readme-ov-file#graceful-shutdown-with-aws-lambda

krrishdholakia commented 4 months ago

@Manouchehri can i just hook into the fastapi shutdown event for this?

Manouchehri commented 4 months ago

Likely, just test it locally with kill -SIGTERM pid_of_litellm to make sure it works.

krrishdholakia commented 4 months ago

I think we can combine update data to 1 batched query

UPDATE LiteLLM_UserTable
SET spend = spend + 100
WHERE user_id = 'some_user_id';

Also this - https://prisma-client-py.readthedocs.io/en/stable/reference/operations/#updating-atomic-fields

krrishdholakia commented 4 months ago

Pushed - https://github.com/BerriAI/litellm/pull/1875