mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
test(proxy): isolate _runtime_registered_model_cost in distributed reload test
A price data reload replays runtime registrations on top of the freshly adopted catalog. When a prior test in the worker had persisted a 'gpt-3.5-turbo' entry, this test's exact-equality assertion on litellm.model_cost['gpt-3.5-turbo'] would fail because register_model merges the built-in schema over the sparse mocked entry, producing over 100 fields instead of the single input_cost_per_token. Snapshot _runtime_registered_model_cost, clear it during the test, and restore it in finally. This matches the pattern already used in test_scheduled_reload_replays_runtime_registrations and test_swap_in_model_cost_map_counts_the_fetched_catalog_only in the same class. Co-authored-by: Krrish Dholakia <krrish-berri-2@users.noreply.github.com>
This commit is contained in:
parent
b66d4e6965
commit
bac9a36ca1
1 changed files with 9 additions and 0 deletions
|
|
@ -4006,6 +4006,7 @@ class TestPriceDataReloadIntegration:
|
|||
A revision this pod has not applied takes effect here even one minute into a 6h
|
||||
interval; a missing row is a no-op
|
||||
"""
|
||||
from litellm import utils as litellm_utils
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
|
|
@ -4034,6 +4035,12 @@ class TestPriceDataReloadIntegration:
|
|||
from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
# A reload replays runtime registrations, so any gpt-3.5-turbo entry a
|
||||
# prior test in this worker persisted would be re-registered on top of
|
||||
# the mocked catalog and defeat the exact-equality check below. Isolate
|
||||
# the registry for the duration of this test.
|
||||
original_registry = dict(litellm_utils._runtime_registered_model_cost)
|
||||
litellm_utils._runtime_registered_model_cost.clear()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map,
|
||||
|
|
@ -4053,6 +4060,8 @@ class TestPriceDataReloadIntegration:
|
|||
assert proxy_config.model_cost_map_applied_revision == 4
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
litellm_utils._runtime_registered_model_cost.clear()
|
||||
litellm_utils._runtime_registered_model_cost.update(original_registry)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_distributed_reload_ignores_already_applied_request(self):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue