From 2d48c6ae82b24b49cbd5da1385b5870eebe6ba0c Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 15:09:45 -0700 Subject: [PATCH 1/2] 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 From 1d86efde9c1cb67ec00cff257580f63db620d387 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 15:32:48 -0700 Subject: [PATCH 2/2] address review: trim verbose comments, fix wrong-request-type assertion test_load_state_from_db_handles_unknown_request_type compared the WRITING cell after load against a cold-start value captured for GENERAL. They happened to be equal for this fixture (the fast model's empty strengths list makes every request type's prior identical), which hid that the assertion was comparing the wrong baseline. Capture each request type's own cold-start value instead. --- .../adaptive_router/adaptive_router.py | 13 ++++------- .../adaptive_router/test_adaptive_router.py | 22 ++++++++----------- .../test_e2e_adaptive_router.py | 6 ++--- 3 files changed, 15 insertions(+), 26 deletions(-) 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")]