scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(router): free a least-busy pick when a vector store create ends

Browse source
This commit is contained in:
mateo-berri 2026-10-30 04:26:41 -07:00
parent 5cdcb8954d
commit 252012a469
2 changed files with 37 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/router.py
View file

@ -7828,6 +7828,10 @@ class Router(DeprecatedRouterMembers):
if model is not None:
self.fail_calls[model] += 1
raise e
finally:
release_unstarted_picks(
self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))
)
def _override_vector_store_methods_for_router(self):
"""

36
tests/unit/router_strategy/test_least_busy.py
View file

@ -798,7 +798,7 @@ async def _drain_logging() -> None:
await GLOBAL_LOGGING_WORKER.flush()
def _batch_deployment(route: respx.Route) -> str:
def _called_deployment(route: respx.Route) -> str:
return route.calls.last.request.headers["authorization"].removeprefix("Bearer key-")
@ -829,7 +829,37 @@ async def test_a_batch_create_does_not_keep_its_deployment_busy(
)
await _drain_logging()
assert _router_pick(router) == _batch_deployment(route)
assert _router_pick(router) == _called_deployment(route)
_VECTOR_STORE: Final = {
"id": "vs_least_busy",
"object": "vector_store",
"created_at": 1,
"name": "least-busy",
"usage_bytes": 0,
"file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
"status": "completed",
}
_REJECTED_VECTOR_STORE: Final = {"error": {"message": "Invalid name", "type": "invalid_request_error"}}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"provider_response", [httpx.Response(200, json=_VECTOR_STORE), httpx.Response(400, json=_REJECTED_VECTOR_STORE)]
)
async def test_a_vector_store_create_does_not_keep_its_deployment_busy(
provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post("https://api.openai.com/v1/vector_stores").mock(return_value=provider_response)
router: Final = _least_busy_router()
with contextlib.suppress(openai.BadRequestError):
await router.avector_store_create(model=ROUTER_GROUP, name="least-busy", litellm_call_id="vector-store-call")
await _drain_logging()
assert _router_pick(router) == _called_deployment(route)
@pytest.mark.asyncio
@ -844,7 +874,7 @@ async def test_a_batch_cancel_does_not_keep_its_deployment_busy(
await router.acancel_batch(model=ROUTER_GROUP, batch_id="batch_least_busy")
assert _router_pick(router) == _batch_deployment(route)
assert _router_pick(router) == _called_deployment(route)
@pytest.mark.asyncio