phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix - locking in memory leads to failing tests

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-28 20:10:58 -07:00
parent fb98bd3a4a
commit 93c20b8d2d
4 changed files with 25 additions and 38 deletions
Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/db/db_spend_update_writer.py
View file

@ -172,10 +172,9 @@ class DBSpendUpdateWriter:
f"track_cost_callback: {entity_type.value}_id is None. Not tracking spend for {entity_type.value}" f"track_cost_callback: {entity_type.value}_id is None. Not tracking spend for {entity_type.value}"
) )
return False return False
async with prisma_client.in_memory_transaction_lock: transaction_list[entity_id] = response_cost + transaction_list.get(
transaction_list[entity_id] = response_cost + transaction_list.get( entity_id, 0
entity_id, 0 )
)
return True return True
except Exception as e: except Exception as e:

47
litellm/proxy/db/redis_update_buffer.py
View file

@ -80,35 +80,32 @@ class RedisUpdateBuffer:
"redis_cache is None, skipping store_in_memory_spend_updates_in_redis" "redis_cache is None, skipping store_in_memory_spend_updates_in_redis"
) )
return return
async with prisma_client.in_memory_transaction_lock: db_spend_update_transactions: DBSpendUpdateTransactions = (
db_spend_update_transactions: DBSpendUpdateTransactions = ( DBSpendUpdateTransactions(
DBSpendUpdateTransactions( user_list_transactions=prisma_client.user_list_transactions,
user_list_transactions=prisma_client.user_list_transactions, end_user_list_transactions=prisma_client.end_user_list_transactions,
end_user_list_transactions=prisma_client.end_user_list_transactions, key_list_transactions=prisma_client.key_list_transactions,
key_list_transactions=prisma_client.key_list_transactions, team_list_transactions=prisma_client.team_list_transactions,
team_list_transactions=prisma_client.team_list_transactions, team_member_list_transactions=prisma_client.team_member_list_transactions,
team_member_list_transactions=prisma_client.team_member_list_transactions, org_list_transactions=prisma_client.org_list_transactions,
org_list_transactions=prisma_client.org_list_transactions,
)
) )
)
# only store in redis if there are any updates to commit # only store in redis if there are any updates to commit
if ( if (
self._number_of_transactions_to_store_in_redis( self._number_of_transactions_to_store_in_redis(db_spend_update_transactions)
db_spend_update_transactions == 0
) ):
== 0 return
):
return
list_of_transactions = [safe_dumps(db_spend_update_transactions)] list_of_transactions = [safe_dumps(db_spend_update_transactions)]
await self.redis_cache.async_rpush( await self.redis_cache.async_rpush(
key=REDIS_UPDATE_BUFFER_KEY, key=REDIS_UPDATE_BUFFER_KEY,
values=list_of_transactions, values=list_of_transactions,
) )
# clear the in-memory spend updates # clear the in-memory spend updates
RedisUpdateBuffer._clear_all_in_memory_spend_updates(prisma_client) RedisUpdateBuffer._clear_all_in_memory_spend_updates(prisma_client)
@staticmethod @staticmethod
def _number_of_transactions_to_store_in_redis( def _number_of_transactions_to_store_in_redis(

8
litellm/proxy/proxy_config.yaml
View file

@ -5,11 +5,3 @@ model_list:
api_key: fake-key api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/ api_base: https://exampleopenaiendpoint-production.up.railway.app/
general_settings:
use_redis_transaction_buffer: True
litellm_settings:
cache: true
cache_params:
type: redis
supported_call_types: []

1
litellm/proxy/utils.py
View file

@ -1121,7 +1121,6 @@ class PrismaClient:
self.iam_token_db_auth: Optional[bool] = str_to_bool( self.iam_token_db_auth: Optional[bool] = str_to_bool(
os.getenv("IAM_TOKEN_DB_AUTH") os.getenv("IAM_TOKEN_DB_AUTH")
) )
self.in_memory_transaction_lock = asyncio.Lock()
verbose_proxy_logger.debug("Creating Prisma Client..") verbose_proxy_logger.debug("Creating Prisma Client..")
try: try:
from prisma import Prisma # type: ignore from prisma import Prisma # type: ignore