[Memory] Retire graph borrow pool before updating static runs (#37966)

Co-authored-by: Shiyan Deng <dsy842974287@meta.com>
This commit is contained in:
Lianmin Zheng
2026-09-04 02:59:52 -07:00
committed by GitHub
co-authored by Shiyan Deng
parent 2216697f90
commit d7f235daca
2 changed files with 17 additions and 1 deletions
@@ -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)