From 2d48c6ae82b24b49cbd5da1385b5870eebe6ba0c Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 15:09:45 -0700 Subject: [PATCH] fix(adaptive_router): add the persisted delta to the cold-start prior on load load_state_from_db assigned a DB row's (alpha, beta) straight into the bandit cell, discarding the cold-start prior _init_cold_start_cells had already put there. AdaptiveRouterUpdateQueue.flush_state_to_db only ever persists accumulated deltas (its upsert creates a row with the raw delta as the initial value, then increments it), never a full posterior, so a cell whose first flush sees only one kind of signal persists a one-sided row: e.g. alpha=1.0, beta=0.0. Loading that row as the whole cell hands thompson_sample() a Beta(alpha, 0), and random.betavariate raises 'gammavariate: alpha and beta must be > 0.0' on every draw from that cell from then on, surviving restarts since the bad row stays in place. Fix: add the row on top of a freshly computed prior instead of replacing the cell with it. Deltas are never negative, so both parameters stay positive. Fixes #35590. Fixes #29397. --- .../adaptive_router/adaptive_router.py | 19 +++++++- .../adaptive_router/test_adaptive_router.py | 43 ++++++++++++++++--- .../test_e2e_adaptive_router.py | 12 ++++-- 3 files changed, 63 insertions(+), 11 deletions(-) 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