diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py index 867bd553f..5b5852a18 100644 --- a/python/sglang/srt/model_executor/runner_utils/pool.py +++ b/python/sglang/srt/model_executor/runner_utils/pool.py @@ -58,9 +58,11 @@ def set_graph_pool_borrow_runs(runs: list[tuple[int, int]]) -> None: borrowing. """ global _borrow_static_runs - _borrow_static_runs = sorted(runs, key=lambda run: run[1], reverse=True)[ + static_runs = sorted(runs, key=lambda run: run[1], reverse=True)[ : BumpArenaStub.MAX_EXTENTS ] + _teardown_borrow_pool() + _borrow_static_runs = static_runs logger.info( "Graph pool borrow runs pinned: runs=%d free=%d", len(_borrow_static_runs), diff --git a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py index ab86a4b13..a88b0e529 100644 --- a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py +++ b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py @@ -122,6 +122,20 @@ class TestGraphPoolBorrow(CustomTestCase): pool.disable_graph_pool_borrow("graph storage is externally managed") self.assertFalse(pool.graph_pool_borrow_enabled()) + def test_setting_static_runs_retires_previous_borrow_pool(self): + runs = [(0x1000, 4096), (0x2000, 8192)] + + def reset_static_runs(): + pool._borrow_static_runs = None + + with patch.object( + pool, "_teardown_borrow_pool", side_effect=reset_static_runs + ) as teardown: + pool.set_graph_pool_borrow_runs(runs) + + teardown.assert_called_once_with() + self.assertEqual(pool._borrow_static_runs, [(0x2000, 8192), (0x1000, 4096)]) + def test_eagle_non_greedy_probabilities_do_not_borrow_graph_pool(self): def fake_sampling(**kwargs): kwargs["predicts"].fill_(3)