diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 02188f48496..12ccacbbc1d 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -123,16 +123,11 @@ class AdaptiveRouter: self._cells[(rt, model)] = initial_cell(prefs, rt) async def load_state_from_db(self, prisma_client: Any) -> None: - """Add persisted deltas on top of the cold-start prior for every cell with a row. + """Add each row's persisted delta to a freshly computed cold-start prior. - A DB row holds accumulated deltas only (AdaptiveRouterUpdateQueue.flush_state_to_db - creates the row with the raw delta as its initial value, then increments it), never - the prior. Assigning `row.alpha`/`row.beta` straight into the cell would silently drop - the cold-start prior _init_cold_start_cells already put there, and the first flush after - a cell sees only one kind of signal persists a one-sided row (e.g. alpha=1, beta=0) - as - a bare Beta(alpha, beta) that zeroes out one shape parameter, which is invalid and 500s - on every later thompson_sample() draw for that cell. Adding the row on top of a freshly - computed prior keeps both parameters positive, since deltas are never negative. + A row holds an accumulated delta, not a full posterior, and can be one-sided + (e.g. beta=0) - assigning it straight into the cell would zero out a Beta shape + parameter and crash thompson_sample() on every later draw for that cell. """ if prisma_client is None: return diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py index 8637dbcc06e..f36443db1e5 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py @@ -270,10 +270,8 @@ async def test_record_turn_bounds_feedback_contexts_and_evicts_least_recent_sess @pytest.mark.asyncio async def test_load_state_from_db_adds_the_persisted_delta_to_the_cold_start_prior(): - """A DB row holds an accumulated delta, not a full posterior (AdaptiveRouterUpdateQueue - creates the row with the raw delta and increments it from there) - loading it must add - that delta on top of the same cold-start prior _init_cold_start_cells already computed, - not replace the cell outright.""" + """A row holds an accumulated delta, not a full posterior; loading must add it to the + cold-start prior, not replace the cell outright.""" r = _make_router() cold = r._cells[(RequestType.GENERAL, "fast")] @@ -293,11 +291,8 @@ async def test_load_state_from_db_adds_the_persisted_delta_to_the_cold_start_pri @pytest.mark.asyncio async def test_load_state_from_db_keeps_a_one_sided_delta_row_sampleable(): - """Regression: a cell whose only DB activity is one signal type persists a one-sided row - (e.g. delta_beta=0.0, per AdaptiveRouterUpdateQueue.flush_state_to_db's create branch). - Loading that row must not zero out a Beta shape parameter - thompson_sample() raises - `ValueError: gammavariate: alpha and beta must be > 0.0` on a zeroed side, bricking every - request for that cell until the process restarts.""" + """A cell whose only DB activity is one signal type persists a one-sided row (e.g. + beta=0.0); loading it must not zero out a Beta shape parameter and crash thompson_sample().""" from litellm.router_strategy.adaptive_router.bandit import thompson_sample r = _make_router() @@ -321,7 +316,8 @@ async def test_load_state_from_db_keeps_a_one_sided_delta_row_sampleable(): @pytest.mark.asyncio async def test_load_state_from_db_handles_unknown_request_type(): r = _make_router() - cold = r._cells[(RequestType.GENERAL, "fast")] + cold_general = r._cells[(RequestType.GENERAL, "fast")] + cold_writing = r._cells[(RequestType.WRITING, "fast")] bad_row = MagicMock() bad_row.request_type = "nonexistent_type_v999" @@ -341,9 +337,9 @@ async def test_load_state_from_db_handles_unknown_request_type(): # Unknown skipped; good added to the cold-start prior. new_general = r._cells[(RequestType.GENERAL, "fast")] - assert new_general.alpha == cold.alpha + 7.0 - # Other request types kept their cold-start values. - assert r._cells[(RequestType.WRITING, "fast")] == cold + assert new_general.alpha == cold_general.alpha + 7.0 + # Other request types kept their own cold-start values. + assert r._cells[(RequestType.WRITING, "fast")] == cold_writing # ---- Session state eviction --------------------------------------------- diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py index 322866936a1..23fc859d4a6 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py @@ -187,10 +187,8 @@ async def test_failure_signal_increments_beta_after_flush(): @pytest.mark.asyncio async def test_load_state_from_db_adds_persisted_delta_to_cold_start(): - """A DB row is an accumulated delta, not a full posterior, so loading it must add onto the - same cold-start prior _init_cold_start_cells already computed, not replace the cell outright - (see test_adaptive_router.py's version of this test, and the one-sided create row - test_failure_signal_increments_beta_after_flush above asserts, for why).""" + """A row holds an accumulated delta, not a full posterior; loading must add it to the + cold-start prior, not replace the cell outright.""" router = _make_router() cold = router._cells[(RequestType.GENERAL, "gpt-4o")]