utils.py 9.67 KB
Newer Older
Robert Shaw's avatar
Robert Shaw committed
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import tempfile
from collections import defaultdict
5
from collections.abc import Callable
6
from itertools import count
7
from typing import Any
Robert Shaw's avatar
Robert Shaw committed
8
9
10
11

import torch

from vllm import SamplingParams
12
13
14
15
16
17
18
19
20
from vllm.config import (
    CacheConfig,
    DeviceConfig,
    KVTransferConfig,
    ModelConfig,
    SchedulerConfig,
    VllmConfig,
)
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
21
from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import (  # noqa
22
23
    SharedStorageConnector,
)
24
from vllm.utils import sha256
25
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
26
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
Robert Shaw's avatar
Robert Shaw committed
27
from vllm.v1.core.sched.scheduler import Scheduler
28
29
30
31
32
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
)
33
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
Robert Shaw's avatar
Robert Shaw committed
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from vllm.v1.request import Request
from vllm.v1.structured_output import StructuredOutputManager

EOS_TOKEN_ID = 50256


def assert_scheduler_empty(scheduler: Scheduler):
    """Confirm the scheduler is "empty" - i.e. no leaks."""
    # Scheduler Metadata.
    assert len(scheduler.requests) == 0
    assert len(scheduler.waiting) == 0
    assert len(scheduler.running) == 0
    assert len(scheduler.finished_req_ids) == 0
    assert len(scheduler.finished_recving_kv_req_ids) == 0

    # EncoderCacheManager.
    assert len(scheduler.encoder_cache_manager.freed) == 0
    assert len(scheduler.encoder_cache_manager.cached) == 0

    # KVCache Manager.
54
55
56
57
58
59
60
61
62
63
64
65
66
67
    assert (
        len(
            scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks
        )
        == 0
    )
    assert (
        len(
            scheduler.kv_cache_manager.coordinator.single_type_managers[
                0
            ].num_cached_block
        )
        == 0
    )
Robert Shaw's avatar
Robert Shaw committed
68
    num_free_blocks = (
69
70
71
        scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks
    )
    assert num_free_blocks == (scheduler.kv_cache_manager.block_pool.num_gpu_blocks - 1)
Robert Shaw's avatar
Robert Shaw committed
72
73
74
75
76
77
78
79
80
81
82
83

    # NOTE(rob): just the ref count on blocks will be 0. The hash
    # value, etc will remain since we lazily evict for prefix cache.
    for block in scheduler.kv_cache_manager.block_pool.blocks:
        assert block.ref_cnt == 0


def create_vllm_config(
    model: str = "facebook/opt-125m",
    max_num_seqs: int = 16,
    max_num_batched_tokens: int = 64,
    block_size: int = 16,
84
85
    max_model_len: int = 10000,
    enable_chunked_prefill: bool = True,
Robert Shaw's avatar
Robert Shaw committed
86
87
88
89
90
) -> VllmConfig:
    """Initialize VllmConfig For Testing."""
    scheduler_config = SchedulerConfig(
        max_num_seqs=max_num_seqs,
        max_num_batched_tokens=max_num_batched_tokens,
91
92
        max_model_len=max_model_len,
        enable_chunked_prefill=enable_chunked_prefill,
Robert Shaw's avatar
Robert Shaw committed
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
    )
    model_config = ModelConfig(
        model=model,
        trust_remote_code=True,
        dtype="float16",
        seed=42,
    )
    # Cache config, optionally force APC
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
        swap_space=0,
        cache_dtype="auto",
        enable_prefix_caching=True,
    )
    kv_transfer_config = KVTransferConfig(
        kv_connector="NixlConnector",
        kv_role="kv_both",
    )
112
113
114
115
116
117
118
    return VllmConfig(
        scheduler_config=scheduler_config,
        model_config=model_config,
        cache_config=cache_config,
        kv_transfer_config=kv_transfer_config,
        device_config=DeviceConfig("cpu"),
    )
Robert Shaw's avatar
Robert Shaw committed
119
120
121
122
123
124
125
126
127
128


def create_scheduler(
    vllm_config: VllmConfig,
    num_blocks: int = 10000,
) -> Scheduler:
    """Initialize Scheduler For Testing."""
    block_size = vllm_config.cache_config.block_size
    kv_cache_config = KVCacheConfig(
        num_blocks=num_blocks,  # A large number of blocks to hold all requests
129
        kv_cache_tensors=[],
Robert Shaw's avatar
Robert Shaw committed
130
        kv_cache_groups=[
131
132
133
            KVCacheGroupSpec(
                ["layer"], FullAttentionSpec(block_size, 1, 1, torch.float32, False)
            )
Robert Shaw's avatar
Robert Shaw committed
134
135
136
137
138
139
140
141
        ],
    )
    vllm_config.cache_config.num_gpu_blocks = num_blocks
    return Scheduler(
        vllm_config=vllm_config,
        kv_cache_config=kv_cache_config,
        log_stats=True,
        structured_output_manager=StructuredOutputManager(vllm_config),
142
        block_size=block_size,
Robert Shaw's avatar
Robert Shaw committed
143
144
145
    )


146
_request_count = count(1)
147
148
149
_none_hash_initialized = False


150
def create_request(
151
    request_id: int | None = None,
152
153
154
155
156
157
158
159
160
    num_tokens: int = 10,
    common_prefix_len=0,
    max_tokens: int = 16,
    do_remote_decode: bool = False,
    do_remote_prefill: bool = False,
    num_remote_blocks: int = 3,
    block_size: int = 16,
    hash_fn: Callable = sha256,
) -> Request:
Robert Shaw's avatar
Robert Shaw committed
161
    """Make dummy request for testing."""
162
163
164
165
166
    assert num_tokens >= common_prefix_len >= 0

    if request_id is None:
        request_id = next(_request_count)

167
168
    global _none_hash_initialized
    if not _none_hash_initialized:
169
        init_none_hash(hash_fn)
170
        _none_hash_initialized = True
Robert Shaw's avatar
Robert Shaw committed
171

172
    kv_transfer_params: dict[str, Any] | None = None
173

Robert Shaw's avatar
Robert Shaw committed
174
175
    if do_remote_decode:
        assert not do_remote_prefill
176
        kv_transfer_params = dict(do_remote_prefill=False, do_remote_decode=True)
Robert Shaw's avatar
Robert Shaw committed
177
    elif do_remote_prefill:
178
179
180
181
182
183
184
185
        kv_transfer_params = dict(
            do_remote_prefill=True,
            do_remote_decode=False,
            remote_engine_id="my-engine-id",
            remote_block_ids=list(range(num_remote_blocks)),
            remote_host="my-host",
            remote_port=1234,
        )
Robert Shaw's avatar
Robert Shaw committed
186
187
188
189

    max_tokens = 1 if do_remote_decode else max_tokens
    sampling_params = SamplingParams(max_tokens=max_tokens)

190
191
192
    common_prefix = [1] * common_prefix_len if common_prefix_len > 0 else []
    suffix = [i * request_id for i in range(num_tokens - common_prefix_len)]
    prompt_token_ids = common_prefix + suffix
Robert Shaw's avatar
Robert Shaw committed
193
194
195
196
197

    req = Request(
        request_id=f"id-{request_id}",
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
198
        pooling_params=None,
199
        mm_features=None,
Robert Shaw's avatar
Robert Shaw committed
200
        eos_token_id=EOS_TOKEN_ID,
201
        block_hasher=get_request_block_hasher(block_size, hash_fn),
Robert Shaw's avatar
Robert Shaw committed
202
203
204
205
206
207
208
    )
    req.kv_transfer_params = kv_transfer_params
    return req


def create_model_runner_output(
    reqs: list[Request],
209
210
211
    finished_sending: set[str] | None = None,
    finished_recving: set[str] | None = None,
    invalid_block_ids: set[int] | None = None,
Robert Shaw's avatar
Robert Shaw committed
212
    use_eos: bool = False,
213
    token_id: int = 0,
Robert Shaw's avatar
Robert Shaw committed
214
215
216
217
218
219
220
221
) -> ModelRunnerOutput:
    """Make dummy model runner output for testing."""

    # Make request data.
    req_ids = [req.request_id for req in reqs]
    req_id_to_index = {req_id: idx for idx, req_id in enumerate(req_ids)}

    # Make sampled tokens.
222
    sampled_token = EOS_TOKEN_ID if use_eos else token_id
Robert Shaw's avatar
Robert Shaw committed
223
224
    sampled_token_ids = [[sampled_token] for _ in req_ids]

225
226
227
228
229
230
231
232
    kv_connector_output = (
        None
        if (
            finished_sending is None
            and finished_recving is None
            and invalid_block_ids is None
        )
        else KVConnectorOutput(
233
234
            finished_sending=finished_sending,
            finished_recving=finished_recving,
235
            invalid_block_ids=invalid_block_ids or set(),
236
        )
237
    )
238

Robert Shaw's avatar
Robert Shaw committed
239
240
241
242
243
244
245
    # Make output data structure.
    return ModelRunnerOutput(
        req_ids=req_ids,
        req_id_to_index=req_id_to_index,
        sampled_token_ids=sampled_token_ids,
        logprobs=None,
        prompt_logprobs_dict={},
246
        pooler_output=None,
247
        kv_connector_output=kv_connector_output,
Robert Shaw's avatar
Robert Shaw committed
248
    )
249
250
251
252
253
254
255
256


class TestSharedStorageConnector(SharedStorageConnector):
    def __init__(self, config: VllmConfig, role):
        self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
        self._connector = SharedStorageConnector(config, role)
        self.call_record: dict[str, int] = defaultdict(int)
        # Use a unique temp file per connector
257
258
259
260
        self._event_file = (
            tempfile.gettempdir()
            + f"/connector_{self.name}-{self.role.name}_events.log"
        )
261
262
263
264
265
        # Start with an empty file
        with open(self._event_file, "w") as _:
            pass

    def __getattribute__(self, name):
266
267
268
269
270
271
272
273
274
275
        if name in (
            "_connector",
            "call_record",
            "name",
            "_event_file",
            "__class__",
            "__dict__",
            "__getattribute__",
            "__init__",
        ):  # avoid recursion
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
            return object.__getattribute__(self, name)
        if not hasattr(self._connector, name):
            return object.__getattribute__(self, name)
        attr = getattr(self._connector, name)

        # Intercept calls to the connector interface and write an event
        # for each one to a file, which can be read back in the main test proc.
        if callable(attr):

            def wrapper(*args, **kwargs):
                self.call_record[name] += 1

                # Include args that we're interested in
                to_log = [name]
                for arg in args:
                    if isinstance(arg, int):
                        to_log.append(str(arg))
                    elif isinstance(arg, KVCacheBlocks):
294
                        to_log.append(f"num_blocks={[len(b) for b in arg.blocks]}")
295
296
297
298

                # Log the event as a line to the file
                try:
                    with open(self._event_file, "a") as f:
299
                        f.write(" ".join(to_log) + "\n")
300
                except Exception as e:
301
                    print(f"[ERROR] Could not log event {name} for {self.name}: {e}")
302
303
304
305
306
307
                return attr(*args, **kwargs)

            return wrapper
        return attr


308
309
310
KVConnectorFactory.register_connector(
    "TestSharedStorageConnector", __name__, TestSharedStorageConnector.__name__
)