mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
feat(proxy): enhance stale row expiration logic with terminal status guard
- Updated `_expire_stale_rows` method to include a terminal status guard in the `update_many` query, preventing updates on rows with certain statuses. - Adjusted the proxy server to use a calculated polling interval for scheduling the response cost check job, improving scheduling accuracy. - Added a new unit test to verify the correct behavior of the updated `_expire_stale_rows` method, ensuring it respects the terminal status guard.
This commit is contained in:
parent
b4e598fe36
commit
74a780467b
3 changed files with 60 additions and 4 deletions
|
|
@ -65,7 +65,19 @@ class CheckResponsesCost:
|
|||
return 0
|
||||
|
||||
await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={"id": {"in": stale_ids}},
|
||||
where={
|
||||
"id": {"in": stale_ids},
|
||||
"status": {
|
||||
"not_in": [
|
||||
"completed",
|
||||
"complete",
|
||||
"failed",
|
||||
"expired",
|
||||
"cancelled",
|
||||
"stale_expired",
|
||||
]
|
||||
},
|
||||
},
|
||||
data={"status": "stale_expired"},
|
||||
)
|
||||
return len(stale_ids)
|
||||
|
|
|
|||
|
|
@ -7188,11 +7188,13 @@ class ProxyStartupEvent:
|
|||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
_effective_responses_interval = (
|
||||
proxy_responses_polling_interval + random.randint(0, 30)
|
||||
)
|
||||
scheduler.add_job(
|
||||
check_responses_cost_job.check_responses_cost,
|
||||
"interval",
|
||||
seconds=proxy_responses_polling_interval
|
||||
+ random.randint(0, 30), # Add small random offset
|
||||
seconds=_effective_responses_interval,
|
||||
# REMOVED jitter parameter - major cause of memory leak
|
||||
id="check_responses_cost_job",
|
||||
replace_existing=True,
|
||||
|
|
@ -7200,7 +7202,7 @@ class ProxyStartupEvent:
|
|||
)
|
||||
verbose_proxy_logger.info(
|
||||
"Responses cost check job scheduled successfully (interval=%ss)",
|
||||
proxy_responses_polling_interval,
|
||||
_effective_responses_interval,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -104,6 +104,48 @@ class TestCheckResponsesCost:
|
|||
call_args = check_responses_cost_instance._expire_stale_rows.call_args
|
||||
assert call_args[0][1] == STALE_OBJECT_CLEANUP_BATCH_SIZE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expire_stale_rows_rechecks_terminal_status_on_update(
|
||||
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
|
||||
):
|
||||
"""update_many should include terminal-status guard to avoid race clobber."""
|
||||
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
||||
CheckResponsesCost,
|
||||
)
|
||||
|
||||
check_responses_cost_instance = CheckResponsesCost(
|
||||
proxy_logging_obj=mock_proxy_logging_obj,
|
||||
prisma_client=mock_prisma_client,
|
||||
llm_router=mock_llm_router,
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[MagicMock(id="row-1"), MagicMock(id="row-2")]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
marked = await check_responses_cost_instance._expire_stale_rows(
|
||||
cutoff=datetime.now(), batch_size=10
|
||||
)
|
||||
|
||||
assert marked == 2
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_awaited_once()
|
||||
where = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1][
|
||||
"where"
|
||||
]
|
||||
)
|
||||
assert where["id"]["in"] == ["row-1", "row-2"]
|
||||
assert where["status"]["not_in"] == [
|
||||
"completed",
|
||||
"complete",
|
||||
"failed",
|
||||
"expired",
|
||||
"cancelled",
|
||||
"stale_expired",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_completed_response(
|
||||
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue