@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 192 lines 7.95 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for HF3FS metadata server data structures and allocation logic:
4 - RankFileMetadata : page allocation / release primitives
5 - KeyMetadata : per-key rank-page tracking and completion detection
6 - GlobalMetadataState : coordinated allocation with cache-hit semantics
7 """
8
9 import pytest
10
11 from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_metadata_server import (
12 GlobalMetadataState,
13 KeyMetadata,
14 RankFileMetadata,
15 )
16
17 # ===========================================================================
18 # TestRankFileMetadata
19 # ===========================================================================
20
21
22 class TestRankFileMetadata:
23 """Unit tests for RankFileMetadata page allocation primitives."""
24
25 @pytest.mark.parametrize(
26 "alloc_count, expected_pages",
27 [(3, 3), (5, 0)],
28 ids=["alloc_partial", "alloc_exceeds"],
29 )
30 def test_allocate_pages(self, alloc_count, expected_pages):
31 """allocate_pages returns correct pages or empty list when insufficient."""
32 rank_meta = RankFileMetadata(rank_id=0, num_pages=3, free_pages=list(range(3)))
33 pages = rank_meta.allocate_pages(alloc_count)
34 assert len(pages) == expected_pages
35 if expected_pages > 0:
36 rank_meta.release_pages(pages)
37 assert rank_meta.get_free_page_count() == 3
38
39 def test_release_pages_restores_count(self):
40 """Releasing allocated pages returns them to the free pool."""
41 rank_meta = RankFileMetadata(rank_id=0, num_pages=4, free_pages=list(range(4)))
42 pages = rank_meta.allocate_pages(2)
43 assert rank_meta.get_free_page_count() == 2
44 rank_meta.release_pages(pages)
45 assert rank_meta.get_free_page_count() == 4
46
47 def test_release_pages_no_duplicates(self):
48 """Releasing the same page twice must not create duplicates."""
49 rank_meta = RankFileMetadata(rank_id=0, num_pages=3, free_pages=list(range(3)))
50 rank_meta.allocate_pages(1) # takes page 0
51 rank_meta.release_pages([0])
52 rank_meta.release_pages([0]) # second release of the same page
53 assert rank_meta.get_free_page_count() == 3
54
55
56 # ===========================================================================
57 # TestKeyMetadata
58 # ===========================================================================
59
60
61 class TestKeyMetadata:
62 """Unit tests for KeyMetadata completion tracking."""
63
64 def test_is_complete_false_until_all_ranks(self):
65 """is_complete() returns True only when all ranks confirmed."""
66 key_meta = KeyMetadata(key="k", rank_to_page={}, tp_world_size=2)
67 assert key_meta.is_complete() is False
68 key_meta.add_rank_page(0, 5)
69 assert key_meta.is_complete() is False
70 key_meta.add_rank_page(1, 10)
71 assert key_meta.is_complete() is True
72
73 def test_get_rank_page_returns_none_for_missing_rank(self):
74 """get_rank_page() returns None when the rank has no entry."""
75 key_meta = KeyMetadata(key="k", rank_to_page={0: 3}, tp_world_size=2)
76 assert key_meta.get_rank_page(0) == 3
77 assert key_meta.get_rank_page(1) is None
78
79 def test_get_all_pages(self):
80 """get_all_pages() returns all (rank, page) pairs."""
81 key_meta = KeyMetadata(key="k", rank_to_page={0: 1, 1: 2}, tp_world_size=2)
82 pairs = key_meta.get_all_pages()
83 assert set(pairs) == {(0, 1), (1, 2)}
84
85
86 # ===========================================================================
87 # TestGlobalMetadataStateAllocation
88 # ===========================================================================
89
90
91 class TestGlobalMetadataStateAllocation:
92 """Tests for GlobalMetadataState allocation and cache-hit semantics."""
93
94 def test_uninitialized_rank_raises_on_allocate(self):
95 """allocate_pages_for_keys raises ValueError for unknown rank."""
96 state = GlobalMetadataState()
97 with pytest.raises((ValueError, Exception)):
98 state.allocate_pages_for_keys(99, [("key", "")])
99
100 def test_uninitialized_rank_raises_on_get_locations(self):
101 """get_key_locations raises ValueError for unknown rank."""
102 state = GlobalMetadataState()
103 with pytest.raises((ValueError, Exception)):
104 state.get_key_locations(99, ["any_key"])
105
106 def test_basic_allocation_and_confirm(self):
107 """Allocating a page and confirming it marks the key as complete."""
108 state = GlobalMetadataState()
109 state.initialize_rank(0, 4)
110
111 results = state.allocate_pages_for_keys(0, [("K", "")])
112 assert results["K"] >= 0
113
114 state.confirm_write_for_keys(0, [("K", results["K"])])
115 assert state.batch_key_exists(["K"]) == [True]
116 locations = state.get_key_locations(0, ["K"])
117 assert locations == [results["K"]]
118
119 def test_allocate_pages_cache_hit_does_not_leak_pages(self):
120 """Cache-hit key must not consume a page from the free pool;
121 the pre-allocated slot must be returned before reusing the existing page.
122 """
123 state = GlobalMetadataState()
124 state.initialize_rank(0, 5) # 5 free pages: [0,1,2,3,4]
125
126 # Simulate a key that has already been fully written and confirmed.
127 state.key_metadata["K_cached"] = KeyMetadata(
128 key="K_cached", rank_to_page={0: 2}, tp_world_size=1
129 )
130
131 free_before = state.rank_metadata[0].get_free_page_count() # 5
132
133 results = state.allocate_pages_for_keys(0, [("K_cached", ""), ("K_new", "")])
134
135 free_after = state.rank_metadata[0].get_free_page_count()
136
137 # Cache-hit key must reuse its existing page.
138 assert results["K_cached"] == 2, (
139 f"Cache-hit key should reuse page 2, got {results['K_cached']}"
140 )
141 # New key must receive a valid page.
142 assert results["K_new"] >= 0, (
143 f"New key should get a valid page, got {results['K_new']}"
144 )
145 # Exactly one page consumed from the free pool.
146 assert free_before - free_after == 1, (
147 f"Expected 1 page consumed, got delta={free_before - free_after}"
148 )
149
150 def test_allocate_pages_all_cache_hits_frees_all_slots(self):
151 """When every key in the batch is a cache hit, no pages are consumed."""
152 state = GlobalMetadataState()
153 state.initialize_rank(0, 5)
154
155 for key, page in (("K1", 0), ("K2", 1)):
156 state.key_metadata[key] = KeyMetadata(
157 key=key, rank_to_page={0: page}, tp_world_size=1
158 )
159
160 free_before = state.rank_metadata[0].get_free_page_count()
161 results = state.allocate_pages_for_keys(0, [("K1", ""), ("K2", "")])
162 free_after = state.rank_metadata[0].get_free_page_count()
163
164 assert results["K1"] == 0
165 assert results["K2"] == 1
166 assert free_after == free_before, (
167 f"All-cache-hit batch must not consume free pages; "
168 f"before={free_before}, after={free_after}"
169 )
170
171 def test_allocate_returns_minus_one_when_pool_exhausted(self):
172 """If the free pool is exhausted, all new keys receive -1."""
173 state = GlobalMetadataState()
174 state.initialize_rank(0, 1) # only 1 free page
175
176 results = state.allocate_pages_for_keys(0, [("K1", ""), ("K2", "")])
177 # allocate_pages uses all-or-nothing: 2 needed but only 1 available → []
178 assert all(v == -1 for v in results.values()), f"Expected all -1, got {results}"
179
180 def test_confirm_write_releases_pages(self):
181 """confirm_write_for_keys with pages_to_release returns them to pool."""
182 state = GlobalMetadataState()
183 state.initialize_rank(0, 3)
184
185 results = state.allocate_pages_for_keys(0, [("K", "")])
186 page = results["K"]
187 free_after_alloc = state.rank_metadata[0].get_free_page_count()
188
189 state.confirm_write_for_keys(0, [("K", page)], pages_to_release=[page])
190 free_after_release = state.rank_metadata[0].get_free_page_count()
191
192 assert free_after_release == free_after_alloc + 1