diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 1a33ea23bd4..02188f48496 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -123,7 +123,17 @@ class AdaptiveRouter: self._cells[(rt, model)] = initial_cell(prefs, rt) async def load_state_from_db(self, prisma_client: Any) -> None: - """Override cold-start cells with persisted state. Called once at startup.""" + """Add persisted deltas on top of the cold-start prior for every cell with a row. + + 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. + """ if prisma_client is None: return try: @@ -139,7 +149,12 @@ class AdaptiveRouter: continue if row.model_name not in self.config.available_models: continue - self._cells[(rt, row.model_name)] = BanditCell(alpha=row.alpha, beta=row.beta) + prefs = self.model_to_prefs.get(row.model_name) or _default_prefs() + prior = initial_cell(prefs, rt) + self._cells[(rt, row.model_name)] = BanditCell( + alpha=prior.alpha + row.alpha, + beta=prior.beta + row.beta, + ) loaded += 1 verbose_router_logger.info( "AdaptiveRouter[%s]: loaded %d cells from DB", 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 cbf5635a5ae..8637dbcc06e 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 @@ -269,7 +269,11 @@ async def test_record_turn_bounds_feedback_contexts_and_evicts_least_recent_sess @pytest.mark.asyncio -async def test_load_state_from_db_overrides_cold_start(): +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.""" r = _make_router() cold = r._cells[(RequestType.GENERAL, "fast")] @@ -284,8 +288,34 @@ async def test_load_state_from_db_overrides_cold_start(): await r.load_state_from_db(prisma) new_cell = r._cells[(RequestType.GENERAL, "fast")] - assert (new_cell.alpha, new_cell.beta) == (42.0, 13.0) - assert (new_cell.alpha, new_cell.beta) != (cold.alpha, cold.beta) + assert (new_cell.alpha, new_cell.beta) == (cold.alpha + 42.0, cold.beta + 13.0) + + +@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.""" + from litellm.router_strategy.adaptive_router.bandit import thompson_sample + + r = _make_router() + + one_sided_row = MagicMock() + one_sided_row.request_type = "general" + one_sided_row.model_name = "fast" + one_sided_row.alpha = 1.0 + one_sided_row.beta = 0.0 + + prisma = MagicMock() + prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[one_sided_row]) + await r.load_state_from_db(prisma) + + loaded_cell = r._cells[(RequestType.GENERAL, "fast")] + assert loaded_cell.alpha > 0.0 + assert loaded_cell.beta > 0.0 + thompson_sample(loaded_cell) # must not raise @pytest.mark.asyncio @@ -309,10 +339,11 @@ async def test_load_state_from_db_handles_unknown_request_type(): prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[bad_row, good_row]) await r.load_state_from_db(prisma) - # Unknown skipped; good applied. - assert r._cells[(RequestType.GENERAL, "fast")].alpha == 7.0 + # 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 or True + assert r._cells[(RequestType.WRITING, "fast")] == cold # ---- 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 3071f916ef1..322866936a1 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 @@ -186,8 +186,14 @@ async def test_failure_signal_increments_beta_after_flush(): @pytest.mark.asyncio -async def test_load_state_from_db_overrides_cold_start(): +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).""" router = _make_router() + cold = router._cells[(RequestType.GENERAL, "gpt-4o")] + fake_row = MagicMock() fake_row.request_type = RequestType.GENERAL.value fake_row.model_name = "gpt-4o" @@ -200,8 +206,8 @@ async def test_load_state_from_db_overrides_cold_start(): await router.load_state_from_db(prisma) cell = router._cells[(RequestType.GENERAL, "gpt-4o")] - assert cell.alpha == 90.0 - assert cell.beta == 10.0 + assert cell.alpha == cold.alpha + 90.0 + assert cell.beta == cold.beta + 10.0 @pytest.mark.asyncio