From 93c20b8d2d6549735f51cb3a2bc5cc13d35ab0a5 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 28 Mar 2025 20:10:58 -0700
Subject: [PATCH] fix - locking in memory leads to failing tests

---
 litellm/proxy/db/db_spend_update_writer.py |  7 ++--
 litellm/proxy/db/redis_update_buffer.py    | 47 ++++++++++------------
 litellm/proxy/proxy_config.yaml            |  8 ----
 litellm/proxy/utils.py                     |  1 -
 4 files changed, 25 insertions(+), 38 deletions(-)

diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 72692339c8..f46b03b57a 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -172,10 +172,9 @@ class DBSpendUpdateWriter:
                     f"track_cost_callback: {entity_type.value}_id is None. Not tracking spend for {entity_type.value}"
                 )
                 return False
-            async with prisma_client.in_memory_transaction_lock:
-                transaction_list[entity_id] = response_cost + transaction_list.get(
-                    entity_id, 0
-                )
+            transaction_list[entity_id] = response_cost + transaction_list.get(
+                entity_id, 0
+            )
             return True
 
         except Exception as e:
diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py
index 3d27e41d15..f842fc85d2 100644
--- a/litellm/proxy/db/redis_update_buffer.py
+++ b/litellm/proxy/db/redis_update_buffer.py
@@ -80,35 +80,32 @@ class RedisUpdateBuffer:
                 "redis_cache is None, skipping store_in_memory_spend_updates_in_redis"
             )
             return
-        async with prisma_client.in_memory_transaction_lock:
-            db_spend_update_transactions: DBSpendUpdateTransactions = (
-                DBSpendUpdateTransactions(
-                    user_list_transactions=prisma_client.user_list_transactions,
-                    end_user_list_transactions=prisma_client.end_user_list_transactions,
-                    key_list_transactions=prisma_client.key_list_transactions,
-                    team_list_transactions=prisma_client.team_list_transactions,
-                    team_member_list_transactions=prisma_client.team_member_list_transactions,
-                    org_list_transactions=prisma_client.org_list_transactions,
-                )
+        db_spend_update_transactions: DBSpendUpdateTransactions = (
+            DBSpendUpdateTransactions(
+                user_list_transactions=prisma_client.user_list_transactions,
+                end_user_list_transactions=prisma_client.end_user_list_transactions,
+                key_list_transactions=prisma_client.key_list_transactions,
+                team_list_transactions=prisma_client.team_list_transactions,
+                team_member_list_transactions=prisma_client.team_member_list_transactions,
+                org_list_transactions=prisma_client.org_list_transactions,
             )
+        )
 
-            # only store in redis if there are any updates to commit
-            if (
-                self._number_of_transactions_to_store_in_redis(
-                    db_spend_update_transactions
-                )
-                == 0
-            ):
-                return
+        # only store in redis if there are any updates to commit
+        if (
+            self._number_of_transactions_to_store_in_redis(db_spend_update_transactions)
+            == 0
+        ):
+            return
 
-            list_of_transactions = [safe_dumps(db_spend_update_transactions)]
-            await self.redis_cache.async_rpush(
-                key=REDIS_UPDATE_BUFFER_KEY,
-                values=list_of_transactions,
-            )
+        list_of_transactions = [safe_dumps(db_spend_update_transactions)]
+        await self.redis_cache.async_rpush(
+            key=REDIS_UPDATE_BUFFER_KEY,
+            values=list_of_transactions,
+        )
 
-            # clear the in-memory spend updates
-            RedisUpdateBuffer._clear_all_in_memory_spend_updates(prisma_client)
+        # clear the in-memory spend updates
+        RedisUpdateBuffer._clear_all_in_memory_spend_updates(prisma_client)
 
     @staticmethod
     def _number_of_transactions_to_store_in_redis(
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index 9ca63f18ea..106003f996 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -5,11 +5,3 @@ model_list:
       api_key: fake-key
       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: []
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 66dd224dff..d0d4ea9b4e 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -1121,7 +1121,6 @@ class PrismaClient:
         self.iam_token_db_auth: Optional[bool] = str_to_bool(
             os.getenv("IAM_TOKEN_DB_AUTH")
         )
-        self.in_memory_transaction_lock = asyncio.Lock()
         verbose_proxy_logger.debug("Creating Prisma Client..")
         try:
             from prisma import Prisma  # type: ignore