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:
harish-berri 2026-05-12 00:43:10 +00:00
parent b4e598fe36
commit 74a780467b
3 changed files with 60 additions and 4 deletions

View file

@ -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)

View file

@ -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:

View file

@ -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