utils.py 12.6 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
7
from dataclasses import dataclass
from itertools import chain, count
8
from typing import Any
Robert Shaw's avatar
Robert Shaw committed
9
10
11
12

import torch

from vllm import SamplingParams
13
from vllm.config import (
14
    AttentionConfig,
15
16
17
18
19
20
21
22
    CacheConfig,
    DeviceConfig,
    KVTransferConfig,
    ModelConfig,
    SchedulerConfig,
    VllmConfig,
)
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
23
24
25
26
27
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
    KVConnectorBase_V1,
    KVConnectorMetadata,
    KVConnectorRole,
)
28
29
from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (  # noqa
    ExampleConnector,
30
)
31
from vllm.utils.hashing import sha256
32
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
33
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
34
from vllm.v1.core.sched.scheduler import Scheduler, SchedulerOutput
35
36
37
38
39
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
)
40
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
Robert Shaw's avatar
Robert Shaw committed
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
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.
61
62
63
64
65
66
67
68
69
70
71
72
73
74
    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
75
    num_free_blocks = (
76
77
78
        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
79
80
81
82
83
84
85
86
87
88
89
90

    # 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,
91
92
    max_model_len: int = 10000,
    enable_chunked_prefill: bool = True,
93
    enable_permute_local_kv: bool = False,
94
95
96
97
    kv_connector_extra_config: dict[str, Any] | None = None,
    dtype: str = "float16",
    cache_dtype: str = "auto",
    hf_overrides: dict[str, Any] | None = None,
98
    attention_backend: str | None = None,
Robert Shaw's avatar
Robert Shaw committed
99
100
101
102
103
) -> VllmConfig:
    """Initialize VllmConfig For Testing."""
    model_config = ModelConfig(
        model=model,
        trust_remote_code=True,
104
        dtype=dtype,
Robert Shaw's avatar
Robert Shaw committed
105
        seed=42,
106
        hf_overrides=hf_overrides or {},
Robert Shaw's avatar
Robert Shaw committed
107
    )
108
109
110
111
112
113
114
    scheduler_config = SchedulerConfig(
        max_num_seqs=max_num_seqs,
        max_num_batched_tokens=max_num_batched_tokens,
        max_model_len=max_model_len,
        enable_chunked_prefill=enable_chunked_prefill,
        is_encoder_decoder=model_config.is_encoder_decoder,
    )
Robert Shaw's avatar
Robert Shaw committed
115
116
117
118
119
    # Cache config, optionally force APC
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
        swap_space=0,
120
        cache_dtype=cache_dtype,
Robert Shaw's avatar
Robert Shaw committed
121
122
123
124
125
        enable_prefix_caching=True,
    )
    kv_transfer_config = KVTransferConfig(
        kv_connector="NixlConnector",
        kv_role="kv_both",
126
        enable_permute_local_kv=enable_permute_local_kv,
127
        kv_connector_extra_config=kv_connector_extra_config or {},
Robert Shaw's avatar
Robert Shaw committed
128
    )
129
    attention_config = AttentionConfig(backend=attention_backend)
130
131
132
133
134
135
    return VllmConfig(
        scheduler_config=scheduler_config,
        model_config=model_config,
        cache_config=cache_config,
        kv_transfer_config=kv_transfer_config,
        device_config=DeviceConfig("cpu"),
136
        attention_config=attention_config,
137
    )
Robert Shaw's avatar
Robert Shaw committed
138
139
140
141
142
143
144
145
146
147


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
148
        kv_cache_tensors=[],
Robert Shaw's avatar
Robert Shaw committed
149
        kv_cache_groups=[
150
151
152
            KVCacheGroupSpec(
                ["layer"], FullAttentionSpec(block_size, 1, 1, torch.float32, False)
            )
Robert Shaw's avatar
Robert Shaw committed
153
154
155
156
157
158
159
160
        ],
    )
    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),
161
        block_size=block_size,
Robert Shaw's avatar
Robert Shaw committed
162
163
164
    )


165
_request_count = count(1)
166
167
168
_none_hash_initialized = False


169
def create_request(
170
    request_id: int | None = None,
171
172
173
174
175
176
177
178
179
    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
180
    """Make dummy request for testing."""
181
182
183
184
185
    assert num_tokens >= common_prefix_len >= 0

    if request_id is None:
        request_id = next(_request_count)

186
187
    global _none_hash_initialized
    if not _none_hash_initialized:
188
        init_none_hash(hash_fn)
189
        _none_hash_initialized = True
Robert Shaw's avatar
Robert Shaw committed
190

191
    kv_transfer_params: dict[str, Any] | None = None
192

Robert Shaw's avatar
Robert Shaw committed
193
194
    if do_remote_decode:
        assert not do_remote_prefill
195
        kv_transfer_params = dict(do_remote_prefill=False, do_remote_decode=True)
Robert Shaw's avatar
Robert Shaw committed
196
    elif do_remote_prefill:
197
198
199
200
        kv_transfer_params = dict(
            do_remote_prefill=True,
            do_remote_decode=False,
            remote_engine_id="my-engine-id",
201
            remote_request_id=f"prefill-{request_id}",
202
203
204
205
            remote_block_ids=list(range(num_remote_blocks)),
            remote_host="my-host",
            remote_port=1234,
        )
Robert Shaw's avatar
Robert Shaw committed
206
207
208
209

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

210
211
212
    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
213
214
215
216
217

    req = Request(
        request_id=f"id-{request_id}",
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
218
        pooling_params=None,
219
        mm_features=None,
Robert Shaw's avatar
Robert Shaw committed
220
        eos_token_id=EOS_TOKEN_ID,
221
        block_hasher=get_request_block_hasher(block_size, hash_fn),
Robert Shaw's avatar
Robert Shaw committed
222
223
224
225
226
227
228
    )
    req.kv_transfer_params = kv_transfer_params
    return req


def create_model_runner_output(
    reqs: list[Request],
229
230
231
    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
232
    use_eos: bool = False,
233
    token_id: int = 0,
Robert Shaw's avatar
Robert Shaw committed
234
235
236
237
238
239
240
241
) -> 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.
242
    sampled_token = EOS_TOKEN_ID if use_eos else token_id
243
    sampled_token_ids = [[sampled_token] for _ in req_ids]
Robert Shaw's avatar
Robert Shaw committed
244

245
246
247
248
249
250
251
252
    kv_connector_output = (
        None
        if (
            finished_sending is None
            and finished_recving is None
            and invalid_block_ids is None
        )
        else KVConnectorOutput(
253
254
            finished_sending=finished_sending,
            finished_recving=finished_recving,
255
            invalid_block_ids=invalid_block_ids or set(),
256
        )
257
    )
258

Robert Shaw's avatar
Robert Shaw committed
259
260
261
262
263
264
265
    # 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={},
266
        pooler_output=None,
267
        kv_connector_output=kv_connector_output,
Robert Shaw's avatar
Robert Shaw committed
268
    )
269
270


271
class TestExampleConnector(ExampleConnector):
272
    def __init__(self, config: VllmConfig, role, kv_cache_config):
273
        self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
274
        self._connector = ExampleConnector(config, role)
275
276
        self.call_record: dict[str, int] = defaultdict(int)
        # Use a unique temp file per connector
277
278
279
280
        self._event_file = (
            tempfile.gettempdir()
            + f"/connector_{self.name}-{self.role.name}_events.log"
        )
281
282
283
284
285
        # Start with an empty file
        with open(self._event_file, "w") as _:
            pass

    def __getattribute__(self, name):
286
287
288
289
290
291
292
293
294
295
        if name in (
            "_connector",
            "call_record",
            "name",
            "_event_file",
            "__class__",
            "__dict__",
            "__getattribute__",
            "__init__",
        ):  # avoid recursion
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
            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):
314
                        to_log.append(f"num_blocks={[len(b) for b in arg.blocks]}")
315
316
317
318

                # Log the event as a line to the file
                try:
                    with open(self._event_file, "a") as f:
319
                        f.write(" ".join(to_log) + "\n")
320
                except Exception as e:
321
                    print(f"[ERROR] Could not log event {name} for {self.name}: {e}")
322
323
324
325
326
327
                return attr(*args, **kwargs)

            return wrapper
        return attr


328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
@dataclass(frozen=True)
class MockKVConfig:
    matched_tokens: int = 0
    is_async: bool = False


class MockKVConnectorMetadata(KVConnectorMetadata):
    def __init__(self):
        # Scheduler tests check metadata.requests
        self.requests: list = []


class MockKVConnector(KVConnectorBase_V1):
    """Mock KV connector for scheduler tests, supporting both sync and async mode."""

    def __init__(
        self,
        vllm_config: VllmConfig,
        role: KVConnectorRole,
        kv_cache_config: KVCacheConfig | None = None,
    ):
        super().__init__(vllm_config, role, kv_cache_config)
        extra_config = self._kv_transfer_config.kv_connector_extra_config
        self.config = MockKVConfig(
            matched_tokens=extra_config["matched_tokens"],
            is_async=extra_config["is_async"],
        )

    def get_num_new_matched_tokens(
        self,
        request: Request,
        num_computed_tokens: int,
    ) -> tuple[int | None, bool]:
        return (self.config.matched_tokens, self.config.is_async)

    def update_state_after_alloc(
        self,
        request: Request,
        blocks: KVCacheBlocks,
        num_external_tokens: int,
    ):
        pass

    def build_connector_meta(
        self, scheduler_output: SchedulerOutput
    ) -> KVConnectorMetadata:
        metadata = MockKVConnectorMetadata()
        cached_reqs = scheduler_output.scheduled_cached_reqs
        for req_id in chain(
            (req.req_id for req in scheduler_output.scheduled_new_reqs),
            (
                req_id
                for req_id in cached_reqs.req_ids
                if req_id in cached_reqs.resumed_req_ids
            ),
        ):
            metadata.requests.append({"req_id": req_id})
        return metadata

    def start_load_kv(self, kv_caches, finished_req_ids):
        pass

    def wait_for_layer_load(self, layer_name):
        pass

    def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
        pass

    def wait_for_save(self):
        pass


400
KVConnectorFactory.register_connector(
401
    "TestExampleConnector", __name__, TestExampleConnector.__name__
402
)
403
404
405
406

KVConnectorFactory.register_connector(
    "MockKVConnector", __name__, MockKVConnector.__name__
)