[Memory] Retire graph borrow pool before updating static runs (#37966)
Co-authored-by: Shiyan Deng <dsy842974287@meta.com>
This commit is contained in:
co-authored by
Shiyan Deng
parent
2216697f90
commit
d7f235daca
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user