From d8f8e89ffb7b0f9818cb5063f789fe0165bdb1f4 Mon Sep 17 00:00:00 2001 From: EvanderFBoop Date: Sat, 13 Jun 2026 11:38:57 +0800 Subject: [PATCH] [Test] Add unit tests for srt/mem_cache/utils.py (#27913) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../unit/mem_cache/test_mem_cache_utils.py | 326 ++++++++++++++++++ 1 file changed, 326 insertions(+) create mode 100644 test/registered/unit/mem_cache/test_mem_cache_utils.py diff --git a/test/registered/unit/mem_cache/test_mem_cache_utils.py b/test/registered/unit/mem_cache/test_mem_cache_utils.py new file mode 100644 index 000000000..73b61b5c4 --- /dev/null +++ b/test/registered/unit/mem_cache/test_mem_cache_utils.py @@ -0,0 +1,326 @@ +"""Unit tests for mem_cache/utils.py — no server, no model loading.""" + +import hashlib +import unittest +from unittest.mock import MagicMock, patch + +from sglang.srt.mem_cache.evict_policy import ( + FIFOStrategy, + FILOStrategy, + LFUStrategy, + LRUStrategy, + MRUStrategy, + PriorityStrategy, + SLRUStrategy, +) +from sglang.srt.mem_cache.utils import ( + compute_node_hash_values, + get_eviction_strategy, + get_hash_str, + hash_str_to_int64, + maybe_init_custom_mem_pool, + split_node_hash_value, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=8, suite="base-a-test-cpu") + + +class TestGetEvictionStrategy(CustomTestCase): + def test_lru(self): + self.assertIsInstance(get_eviction_strategy("lru"), LRUStrategy) + + def test_lfu(self): + self.assertIsInstance(get_eviction_strategy("lfu"), LFUStrategy) + + def test_fifo(self): + self.assertIsInstance(get_eviction_strategy("fifo"), FIFOStrategy) + + def test_mru(self): + self.assertIsInstance(get_eviction_strategy("mru"), MRUStrategy) + + def test_filo(self): + self.assertIsInstance(get_eviction_strategy("filo"), FILOStrategy) + + def test_priority(self): + self.assertIsInstance(get_eviction_strategy("priority"), PriorityStrategy) + + def test_slru(self): + self.assertIsInstance(get_eviction_strategy("slru"), SLRUStrategy) + + def test_case_insensitive(self): + self.assertIsInstance(get_eviction_strategy("LRU"), LRUStrategy) + self.assertIsInstance(get_eviction_strategy("Lru"), LRUStrategy) + self.assertIsInstance(get_eviction_strategy("FIFO"), FIFOStrategy) + + def test_unknown_policy_raises_valueerror(self): + with self.assertRaises(ValueError) as ctx: + get_eviction_strategy("nonexistent") + msg = str(ctx.exception) + self.assertIn("Unknown eviction policy", msg) + self.assertIn("lru", msg) + self.assertIn("lfu", msg) + self.assertIn("fifo", msg) + self.assertIn("mru", msg) + self.assertIn("filo", msg) + self.assertIn("priority", msg) + self.assertIn("slru", msg) + + def test_each_call_creates_new_instance(self): + s1 = get_eviction_strategy("lru") + s2 = get_eviction_strategy("lru") + self.assertIsNot(s1, s2) + + +class TestMaybeInitCustomMemPool(CustomTestCase): + @patch("sglang.srt.mem_cache.utils.envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get") + def test_disabled_by_default(self, mock_env_get): + mock_env_get.return_value = None + enabled, pool, pool_type = maybe_init_custom_mem_pool("cuda:0") + self.assertFalse(enabled) + self.assertIsNone(pool) + self.assertIsNone(pool_type) + + @patch("sglang.srt.mem_cache.utils.envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get") + @patch("sglang.srt.disaggregation.mooncake.utils.init_mooncake_custom_mem_pool") + def test_enabled_via_env(self, mock_init, mock_env_get): + mock_env_get.return_value = "enabled" + mock_init.return_value = (True, "mock_pool_instance", "mooncake") + + enabled, pool, pool_type = maybe_init_custom_mem_pool("cuda:0") + self.assertTrue(enabled) + self.assertEqual(pool, "mock_pool_instance") + self.assertEqual(pool_type, "mooncake") + mock_init.assert_called_once_with("cuda:0") + + +class TestGetHashStr(CustomTestCase): + def test_empty_list(self): + result = get_hash_str([]) + self.assertIsInstance(result, str) + self.assertEqual(len(result), 64) + expected = hashlib.sha256().hexdigest() + self.assertEqual(result, expected) + + def test_different_sequence_different_hash(self): + h1 = get_hash_str([1, 2, 3]) + h2 = get_hash_str([3, 2, 1]) + self.assertNotEqual(h1, h2) + + def test_different_values_different_hash(self): + h1 = get_hash_str([100]) + h2 = get_hash_str([200]) + self.assertNotEqual(h1, h2) + + def test_bigram_vs_flat_token_equivalent(self): + h_flat = get_hash_str([1, 2]) + h_bigram = get_hash_str([(1, 2)]) + self.assertEqual(h_flat, h_bigram) + + def test_bigram_order_matters(self): + h1 = get_hash_str([(1, 2)]) + h2 = get_hash_str([(2, 1)]) + self.assertNotEqual(h1, h2) + + def test_prior_hash_chaining(self): + chained = get_hash_str([3, 4], prior_hash=get_hash_str([1, 2])) + # prior_hash must fold into the digest, so chaining differs from + # hashing [3, 4] alone... + self.assertNotEqual(chained, get_hash_str([3, 4])) + # ...and a different prior_hash must yield a different chained digest. + self.assertNotEqual( + chained, get_hash_str([3, 4], prior_hash=get_hash_str([9, 9])) + ) + + def test_prior_hash_single_step(self): + step1 = get_hash_str([1]) + step2 = get_hash_str([2], prior_hash=step1) + direct = get_hash_str([1, 2]) + self.assertNotEqual(step2, direct) + + def test_returns_64_char_hex(self): + for tokens in [[], [1], [1, 2, 3], [(1, 2)], [1, 2, 3, 4, 5]]: + result = get_hash_str(tokens) + self.assertRegex(result, r"^[0-9a-f]{64}$") + + +class TestHashStrToInt64(CustomTestCase): + def test_zero_hash(self): + result = hash_str_to_int64("0" * 64) + self.assertEqual(result, 0) + + def test_small_positive_value(self): + result = hash_str_to_int64("0000000000000001" + "0" * 48) + self.assertEqual(result, 1) + + def test_large_positive_value(self): + result = hash_str_to_int64("7fffffffffffffff" + "0" * 48) + self.assertEqual(result, 2**63 - 1) + + def test_negative_overflow(self): + result = hash_str_to_int64("8000000000000000" + "0" * 48) + self.assertEqual(result, -(2**63)) + + def test_max_unsigned_maps_to_minus_one(self): + result = hash_str_to_int64("f" * 16 + "0" * 48) + self.assertEqual(result, -1) + + def test_only_first_16_chars_matter(self): + h1 = hash_str_to_int64("a" * 16 + "0" * 48) + h2 = hash_str_to_int64("a" * 16 + "f" * 48) + self.assertEqual(h1, h2) + + def test_roundtrip_with_get_hash_str(self): + hash_hex = get_hash_str([42, 99]) + int64_val = hash_str_to_int64(hash_hex) + self.assertIsInstance(int64_val, int) + self.assertTrue(-(2**63) <= int64_val < 2**63) + + +class TestComputeNodeHashValues(CustomTestCase): + def setUp(self): + def mock_hash_page(start, end, parent_hash): + parts = [f"p{start}-{end}"] + if parent_hash is not None: + parts.append(parent_hash) + return "-".join(parts) + + self.mock_hash_page = mock_hash_page + + def _make_node(self, key_len, parent=None, parent_hash_values=None): + node = MagicMock() + node.key.__len__.return_value = key_len + node.key.hash_page = self.mock_hash_page + node.parent = parent + if parent is not None: + parent.hash_value = parent_hash_values + return node + + def test_single_page_root(self): + node = self._make_node(key_len=3) + result = compute_node_hash_values(node, page_size=16) + self.assertEqual(len(result), 1) + self.assertIn("p0-3", result[0]) + + def test_multiple_pages(self): + node = self._make_node(key_len=30) + result = compute_node_hash_values(node, page_size=16) + self.assertEqual(len(result), 2) + self.assertIn("p0-16", result[0]) + self.assertIn("p16-30", result[1]) + + def test_page_aligned_boundary(self): + node = self._make_node(key_len=32) + result = compute_node_hash_values(node, page_size=8) + self.assertEqual(len(result), 4) + self.assertIn("p24-32", result[3]) + + def test_key_shorter_than_page_size(self): + node = self._make_node(key_len=5) + result = compute_node_hash_values(node, page_size=16) + self.assertEqual(result, ["p0-5"]) + + def test_chained_parent_hash(self): + parent = MagicMock() + parent.key.__len__.return_value = 8 + parent.hash_value = ["parent_hash_0", "parent_hash_1"] + parent.key.hash_page = self.mock_hash_page + + child = self._make_node( + key_len=16, parent=parent, parent_hash_values=parent.hash_value + ) + result = compute_node_hash_values(child, page_size=8) + self.assertEqual(result, ["p0-8-parent_hash_1", "p8-16-p0-8-parent_hash_1"]) + + def test_parent_with_empty_key(self): + parent = MagicMock() + parent.key.__len__.return_value = 0 + parent.hash_value = ["some_hash"] + parent.key.hash_page = self.mock_hash_page + + child = self._make_node( + key_len=8, parent=parent, parent_hash_values=parent.hash_value + ) + result = compute_node_hash_values(child, page_size=8) + self.assertEqual(len(result), 1) + self.assertNotIn("some_hash", result[0]) + + def test_parent_without_hash_value(self): + parent = MagicMock() + parent.key.__len__.return_value = 8 + parent.hash_value = [] + parent.key.hash_page = self.mock_hash_page + + child = self._make_node(key_len=8, parent=parent, parent_hash_values=[]) + result = compute_node_hash_values(child, page_size=8) + self.assertEqual(result, ["p0-8"]) + + def test_parent_with_none_hash_value(self): + parent = MagicMock() + parent.key.__len__.return_value = 8 + parent.hash_value = None + parent.key.hash_page = self.mock_hash_page + + child = self._make_node(key_len=8, parent=parent, parent_hash_values=None) + result = compute_node_hash_values(child, page_size=8) + self.assertEqual(result, ["p0-8"]) + + +class TestSplitNodeHashValue(CustomTestCase): + def test_none_input_returns_none_tuple(self): + result = split_node_hash_value(None, 10, 4) + self.assertEqual(result, (None, None)) + + def test_page_size_one_split(self): + hash_values = [f"hash_{i}" for i in range(4)] + new_hash, child_hash = split_node_hash_value(hash_values, 2, 1) + self.assertEqual(new_hash, ["hash_0", "hash_1"]) + self.assertEqual(child_hash, ["hash_2", "hash_3"]) + + def test_page_size_one_split_at_zero(self): + hash_values = ["a", "b", "c"] + new_hash, child_hash = split_node_hash_value(hash_values, 0, 1) + self.assertEqual(new_hash, []) + self.assertEqual(child_hash, ["a", "b", "c"]) + + def test_page_size_one_split_at_len(self): + hash_values = ["a", "b", "c"] + new_hash, child_hash = split_node_hash_value(hash_values, 3, 1) + self.assertEqual(new_hash, ["a", "b", "c"]) + self.assertEqual(child_hash, []) + + def test_page_size_larger_than_one(self): + hash_values = ["p0", "p1", "p2", "p3", "p4", "p5"] + new_hash, child_hash = split_node_hash_value(hash_values, 4, 2) + self.assertEqual(new_hash, ["p0", "p1"]) + self.assertEqual(child_hash, ["p2", "p3", "p4", "p5"]) + + def test_page_size_split_exact_page_boundary(self): + hash_values = ["a", "b", "c", "d"] + new_hash, child_hash = split_node_hash_value(hash_values, 4, 4) + self.assertEqual(new_hash, ["a"]) + self.assertEqual(child_hash, ["b", "c", "d"]) + + def test_page_size_split_within_page(self): + hash_values = ["a", "b", "c"] + new_hash, child_hash = split_node_hash_value(hash_values, 6, 4) + self.assertEqual(new_hash, ["a"]) + self.assertEqual(child_hash, ["b", "c"]) + + def test_total_length_preserved(self): + hash_values = [f"h_{i}" for i in range(10)] + for split_len in [0, 1, 3, 5, 7, 9, 10]: + for page_size in [1, 2, 4]: + new_hash, child_hash = split_node_hash_value( + hash_values, split_len, page_size + ) + self.assertEqual( + len(new_hash) + len(child_hash), + len(hash_values), + f"split_len={split_len}, page_size={page_size}", + ) + + +if __name__ == "__main__": + unittest.main()