[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.
|
borrowing.
|
||||||
"""
|
"""
|
||||||
global _borrow_static_runs
|
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
|
: BumpArenaStub.MAX_EXTENTS
|
||||||
]
|
]
|
||||||
|
_teardown_borrow_pool()
|
||||||
|
_borrow_static_runs = static_runs
|
||||||
logger.info(
|
logger.info(
|
||||||
"Graph pool borrow runs pinned: runs=%d free=%d",
|
"Graph pool borrow runs pinned: runs=%d free=%d",
|
||||||
len(_borrow_static_runs),
|
len(_borrow_static_runs),
|
||||||
|
|||||||
@@ -122,6 +122,20 @@ class TestGraphPoolBorrow(CustomTestCase):
|
|||||||
pool.disable_graph_pool_borrow("graph storage is externally managed")
|
pool.disable_graph_pool_borrow("graph storage is externally managed")
|
||||||
self.assertFalse(pool.graph_pool_borrow_enabled())
|
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 test_eagle_non_greedy_probabilities_do_not_borrow_graph_pool(self):
|
||||||
def fake_sampling(**kwargs):
|
def fake_sampling(**kwargs):
|
||||||
kwargs["predicts"].fill_(3)
|
kwargs["predicts"].fill_(3)
|
||||||
|
|||||||
Reference in New Issue
Block a user