From 091511f3bd2b5a25bf370bccec15edadd6b2aef9 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 13 Feb 2026 17:00:22 -0800 Subject: [PATCH] perf: cache Deployment objects in Router.get_deployment() (~25x faster) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cache Deployment Pydantic objects to avoid reconstructing them on every call. Deployment.__init__ drops from 12,011 calls to 1 per model_id. Per-call cost: ~28µs → ~1.1µs. Also removes dead get_deployment call in common_request_processing.py. --- litellm/proxy/common_request_processing.py | 2 - litellm/router.py | 32 +++- .../test_router_deployment_cache.py | 158 ++++++++++++++++++ 3 files changed, 188 insertions(+), 4 deletions(-) create mode 100644 tests/router_unit_tests/test_router_deployment_cache.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a02bc7f9e51..cbe7ee1a514 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -799,8 +799,6 @@ class ProxyBaseLLMRequestProcessing: ) # Post Call Processing - if llm_router is not None: - self.data["deployment"] = llm_router.get_deployment(model_id=model_id) asyncio.create_task( proxy_logging_obj.update_request_status( litellm_call_id=self.data.get("litellm_call_id", ""), status="success" diff --git a/litellm/router.py b/litellm/router.py index 888c97ca0b1..6ccecbce4fd 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -456,6 +456,7 @@ class Router: # Initialize model name to deployment indices mapping for O(1) lookups # Maps model_name -> list of indices in model_list self.model_name_to_deployment_indices: Dict[str, List[int]] = {} + self._deployment_cache: Dict[str, Deployment] = {} if model_list is not None: # set_model_list will build indices automatically @@ -6211,6 +6212,7 @@ class Router: self.model_list = [] self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index + self._invalidate_deployment_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works for model in original_model_list: @@ -6514,6 +6516,7 @@ class Router: """ idx = len(self.model_list) self.model_list.append(model) + self._invalidate_deployment_cache() # Update model_id index for O(1) lookup if model_id is not None: @@ -6561,6 +6564,7 @@ class Router: if removal_idx is not None: self.model_list.pop(removal_idx) + self._invalidate_deployment_cache() self._update_deployment_indices_after_removal( model_id=deployment_id, removal_idx=removal_idx ) @@ -6594,6 +6598,7 @@ class Router: if deployment_idx is not None: # Pop the item from the list first item = self.model_list.pop(deployment_idx) + self._invalidate_deployment_cache() self._update_deployment_indices_after_removal( model_id=id, removal_idx=deployment_idx ) @@ -6609,12 +6614,19 @@ class Router: Raise Exception -> if model found in invalid format """ + # Check cache first + cached = self._deployment_cache.get(model_id) + if cached is not None: + return cached + # Use O(1) lookup via model_id_to_deployment_index_map only if model_id in self.model_id_to_deployment_index_map: idx = self.model_id_to_deployment_index_map[model_id] model = self.model_list[idx] if isinstance(model, dict): - return Deployment(**model) + deployment = Deployment(**model) + self._deployment_cache[model_id] = deployment + return deployment elif isinstance(model, Deployment): return model else: @@ -6650,7 +6662,15 @@ class Router: # Return first deployment for this model_name model = self.model_list[indices[0]] if isinstance(model, dict): - return Deployment(**model) + model_id = model.get("model_info", {}).get("id") + if model_id: + cached = self._deployment_cache.get(model_id) + if cached is not None: + return cached + deployment = Deployment(**model) + if model_id: + self._deployment_cache[model_id] = deployment + return deployment elif isinstance(model, Deployment): return model else: @@ -7328,6 +7348,7 @@ class Router: """ # First populate the model_list self.model_list = [] + self._invalidate_deployment_cache() for _, model in enumerate(model_list): # Extract model_info from the model dict model_info = model.get("model_info", {}) @@ -7664,6 +7685,13 @@ class Router: return returned_models + def _invalidate_deployment_cache(self) -> None: + """Invalidate the cached Deployment objects. + + Call this whenever self.model_list is modified to ensure deployments are rebuilt. + """ + self._deployment_cache.clear() + def get_model_access_groups( self, model_name: Optional[str] = None, diff --git a/tests/router_unit_tests/test_router_deployment_cache.py b/tests/router_unit_tests/test_router_deployment_cache.py new file mode 100644 index 00000000000..f1196bb94d2 --- /dev/null +++ b/tests/router_unit_tests/test_router_deployment_cache.py @@ -0,0 +1,158 @@ +import sys +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import pytest +from litellm import Router +from litellm.types.router import Deployment, LiteLLM_Params + + +@pytest.fixture +def model_list(): + return [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "sk-test1", + }, + "model_info": { + "id": "model-id-1", + }, + }, + { + "model_name": "gpt-4o", + "litellm_params": { + "model": "gpt-4o", + "api_key": "sk-test2", + }, + "model_info": { + "id": "model-id-2", + }, + }, + ] + + +@pytest.fixture +def router(model_list): + return Router(model_list=model_list) + + +def test_get_deployment_returns_correct_object(router): + """get_deployment returns correct Deployment object.""" + deployment = router.get_deployment("model-id-1") + assert deployment is not None + assert isinstance(deployment, Deployment) + assert deployment.model_name == "gpt-3.5-turbo" + assert deployment.litellm_params.model == "gpt-3.5-turbo" + + +def test_get_deployment_returns_cached_object(router): + """get_deployment returns same cached object on second call.""" + result1 = router.get_deployment("model-id-1") + result2 = router.get_deployment("model-id-1") + assert result1 is result2 # same object reference = cache hit + + +def test_get_deployment_by_model_group_name_returns_correct_object(router): + """get_deployment_by_model_group_name returns correct Deployment.""" + deployment = router.get_deployment_by_model_group_name("gpt-3.5-turbo") + assert deployment is not None + assert isinstance(deployment, Deployment) + assert deployment.model_name == "gpt-3.5-turbo" + + +def test_add_deployment_invalidates_cache(router): + """add_deployment invalidates cache so new deployments are found.""" + # Populate cache + router.get_deployment("model-id-1") + + # Add a new deployment + new_deployment = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params(model="gpt-4", api_key="sk-test3"), + model_info={"id": "model-id-3"}, + ) + router.add_deployment(deployment=new_deployment) + + # New deployment should be found + result = router.get_deployment("model-id-3") + assert result is not None + assert result.model_name == "gpt-4" + + +def test_delete_deployment_invalidates_cache(router): + """delete_deployment invalidates cache so deleted deployments return None.""" + # Populate cache + result = router.get_deployment("model-id-1") + assert result is not None + + # Delete the deployment + router.delete_deployment(id="model-id-1") + + # Should return None now + result = router.get_deployment("model-id-1") + assert result is None + + +def test_upsert_deployment_invalidates_cache(router): + """upsert_deployment invalidates cache so updated deployments are returned.""" + # Populate cache + original = router.get_deployment("model-id-1") + assert original is not None + assert original.litellm_params.model == "gpt-3.5-turbo" + + # Upsert with changed api_key + updated_deployment = Deployment( + model_name="gpt-3.5-turbo", + litellm_params=LiteLLM_Params( + model="gpt-3.5-turbo", api_key="sk-test-updated" + ), + model_info={"id": "model-id-1"}, + ) + router.upsert_deployment(deployment=updated_deployment) + + # Should return updated deployment + result = router.get_deployment("model-id-1") + assert result is not None + assert result.litellm_params.api_key == "sk-test-updated" + + +def test_set_model_list_invalidates_cache(router): + """set_model_list invalidates cache completely.""" + # Populate cache + router.get_deployment("model-id-1") + router.get_deployment("model-id-2") + + # Set completely new model list + router.set_model_list( + [ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4", + "api_key": "sk-new", + }, + "model_info": { + "id": "new-model-id", + }, + } + ] + ) + + # Old model IDs should return None + assert router.get_deployment("model-id-1") is None + assert router.get_deployment("model-id-2") is None + + # New model ID should work + result = router.get_deployment("new-model-id") + assert result is not None + assert result.model_name == "gpt-4" + + +def test_get_deployment_invalid_model_id_returns_none(router): + """get_deployment with invalid model_id returns None.""" + result = router.get_deployment("nonexistent-model-id") + assert result is None