| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import tempfile |
| 4 | from collections import defaultdict |
| 5 | from collections.abc import Callable |
| 6 | from dataclasses import dataclass |
| 7 | from itertools import chain, count |
| 8 | from typing import Any, Literal |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm import SamplingParams |
| 13 | from vllm.config import ( |
| 14 | AttentionConfig, |
| 15 | CacheConfig, |
| 16 | DeviceConfig, |
| 17 | KVTransferConfig, |
| 18 | ModelConfig, |
| 19 | SchedulerConfig, |
| 20 | SpeculativeConfig, |
| 21 | VllmConfig, |
| 22 | ) |
| 23 | from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory |
| 24 | from vllm.distributed.kv_transfer.kv_connector.v1.base import ( |
| 25 | KVConnectorBase_V1, |
| 26 | KVConnectorMetadata, |
| 27 | KVConnectorRole, |
| 28 | KVConnectorWorkerMetadata, |
| 29 | ) |
| 30 | from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( # noqa |
| 31 | ExampleConnector, |
| 32 | ) |
| 33 | from vllm.utils.hashing import sha256 |
| 34 | from vllm.v1.core.kv_cache_manager import KVCacheBlocks |
| 35 | from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash |
| 36 | from vllm.v1.core.sched.async_scheduler import AsyncScheduler |
| 37 | from vllm.v1.core.sched.scheduler import Scheduler, SchedulerOutput |
| 38 | from vllm.v1.kv_cache_interface import ( |
| 39 | FullAttentionSpec, |
| 40 | KVCacheConfig, |
| 41 | KVCacheGroupSpec, |
| 42 | MambaSpec, |
| 43 | SlidingWindowSpec, |
| 44 | ) |
| 45 | from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput |
| 46 | from vllm.v1.request import Request |
| 47 | from vllm.v1.structured_output import StructuredOutputManager |
| 48 | |
| 49 | EOS_TOKEN_ID = 50256 |
| 50 | |
| 51 | |
| 52 | def assert_scheduler_empty(scheduler: Scheduler): |
| 53 | """Confirm the scheduler is "empty" - i.e. no leaks.""" |
| 54 | # Scheduler Metadata. |
| 55 | assert len(scheduler.requests) == 0 |
| 56 | assert len(scheduler.waiting) == 0 |
| 57 | assert len(scheduler.kv_holding_waiting) == 0 |
| 58 | assert not scheduler.deferred_waiting |
| 59 | assert len(scheduler.running) == 0 |
| 60 | assert len(scheduler.finished_req_ids) == 0 |
| 61 | assert len(scheduler.finished_recving_kv_req_ids) == 0 |
| 62 | assert len(scheduler._inflight_prefills) == 0 |
| 63 | assert not scheduler._kv_fetch_stages |
| 64 | assert not any(scheduler._kv_fetch_counts.values()) |
| 65 | |
| 66 | # EncoderCacheManager. |
| 67 | assert len(scheduler.encoder_cache_manager.freed) == 0 |
| 68 | assert len(scheduler.encoder_cache_manager.cached) == 0 |
| 69 | |
| 70 | # KVCache Manager. |
| 71 | assert ( |
| 72 | len( |
| 73 | scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks |
| 74 | ) |
| 75 | == 0 |
| 76 | ) |
| 77 | assert ( |
| 78 | len( |
| 79 | scheduler.kv_cache_manager.coordinator.single_type_managers[ |
| 80 | 0 |
| 81 | ].num_cached_block |
| 82 | ) |
| 83 | == 0 |
| 84 | ) |
| 85 | num_free_blocks = ( |
| 86 | scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks |
| 87 | ) |
| 88 | assert num_free_blocks == (scheduler.kv_cache_manager.block_pool.num_gpu_blocks - 1) |
| 89 | |
| 90 | # NOTE(rob): just the ref count on blocks will be 0. The hash |
| 91 | # value, etc will remain since we lazily evict for prefix cache. |
| 92 | for block in scheduler.kv_cache_manager.block_pool.blocks: |
| 93 | assert block.ref_cnt == 0 |
| 94 | |
| 95 | |
| 96 | def create_vllm_config( |
| 97 | model: str = "facebook/opt-125m", |
| 98 | max_num_seqs: int = 16, |
| 99 | max_num_batched_tokens: int = 64, |
| 100 | block_size: int = 16, |
| 101 | max_model_len: int = 10000, |
| 102 | enable_chunked_prefill: bool = True, |
| 103 | enable_permute_local_kv: bool = False, |
| 104 | kv_connector_extra_config: dict[str, Any] | None = None, |
| 105 | dtype: str = "float16", |
| 106 | cache_dtype: str = "auto", |
| 107 | hf_overrides: dict[str, Any] | None = None, |
| 108 | attention_backend: str | None = None, |
| 109 | kv_load_failure_policy: Literal["recompute", "fail"] = "fail", |
| 110 | kv_connector: str = "NixlConnector", |
| 111 | kv_connector_module_path: str | None = None, |
| 112 | kv_role: str = "kv_consumer", |
| 113 | disable_hybrid_kv_cache_manager: bool | None = None, |
| 114 | num_speculative_tokens: int | None = None, |
| 115 | ) -> VllmConfig: |
| 116 | """Initialize VllmConfig For Testing.""" |
| 117 | model_config = ModelConfig( |
| 118 | model=model, |
| 119 | trust_remote_code=True, |
| 120 | dtype=dtype, |
| 121 | seed=42, |
| 122 | hf_overrides=hf_overrides or {}, |
| 123 | ) |
| 124 | scheduler_config = SchedulerConfig( |
| 125 | max_num_seqs=max_num_seqs, |
| 126 | max_num_batched_tokens=max_num_batched_tokens, |
| 127 | max_model_len=max_model_len, |
| 128 | enable_chunked_prefill=enable_chunked_prefill, |
| 129 | is_encoder_decoder=model_config.is_encoder_decoder, |
| 130 | disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager, |
| 131 | ) |
| 132 | # Cache config, optionally force APC |
| 133 | cache_config = CacheConfig( |
| 134 | block_size=block_size, |
| 135 | gpu_memory_utilization=0.9, |
| 136 | cache_dtype=cache_dtype, |
| 137 | enable_prefix_caching=True, |
| 138 | ) |
| 139 | # Connectors are constructed after layout resolution; mirror that here. |
| 140 | cache_config.kv_cache_layout = "LBNHC" |
| 141 | kv_transfer_config = KVTransferConfig( |
| 142 | kv_connector=kv_connector, |
| 143 | kv_connector_module_path=kv_connector_module_path, |
| 144 | kv_role=kv_role, |
| 145 | enable_permute_local_kv=enable_permute_local_kv, |
| 146 | kv_connector_extra_config=kv_connector_extra_config or {}, |
| 147 | kv_load_failure_policy=kv_load_failure_policy, |
| 148 | ) |
| 149 | attention_config = AttentionConfig(backend=attention_backend) |
| 150 | speculative_config = ( |
| 151 | SpeculativeConfig(model="ngram", num_speculative_tokens=num_speculative_tokens) |
| 152 | if num_speculative_tokens is not None |
| 153 | else None |
| 154 | ) |
| 155 | return VllmConfig( |
| 156 | scheduler_config=scheduler_config, |
| 157 | model_config=model_config, |
| 158 | cache_config=cache_config, |
| 159 | kv_transfer_config=kv_transfer_config, |
| 160 | device_config=DeviceConfig("cpu"), |
| 161 | attention_config=attention_config, |
| 162 | speculative_config=speculative_config, |
| 163 | ) |
| 164 | |
| 165 | |
| 166 | def create_scheduler( |
| 167 | vllm_config: VllmConfig, |
| 168 | num_blocks: int = 10000, |
| 169 | kv_cache_config: KVCacheConfig | None = None, |
| 170 | ) -> Scheduler | AsyncScheduler: |
| 171 | """Initialize Scheduler For Testing.""" |
| 172 | block_size = vllm_config.cache_config.block_size |
| 173 | if kv_cache_config is None: |
| 174 | kv_cache_config = KVCacheConfig( |
| 175 | num_blocks=num_blocks, # A large number of blocks to hold all requests |
| 176 | kv_cache_tensors=[], |
| 177 | kv_cache_groups=[ |
| 178 | KVCacheGroupSpec( |
| 179 | ["layer"], |
| 180 | FullAttentionSpec( |
| 181 | block_size=block_size, |
| 182 | num_kv_heads=1, |
| 183 | head_size=1, |
| 184 | dtype=torch.float32, |
| 185 | ), |
| 186 | ) |
| 187 | ], |
| 188 | ) |
| 189 | vllm_config.cache_config.num_gpu_blocks = num_blocks |
| 190 | |
| 191 | scheduler_cls = ( |
| 192 | AsyncScheduler if vllm_config.scheduler_config.async_scheduling else Scheduler |
| 193 | ) |
| 194 | return scheduler_cls( |
| 195 | vllm_config=vllm_config, |
| 196 | kv_cache_config=kv_cache_config, |
| 197 | log_stats=True, |
| 198 | structured_output_manager=StructuredOutputManager(vllm_config), |
| 199 | block_size=block_size, |
| 200 | ) |
| 201 | |
| 202 | |
| 203 | _request_count = count(1) |
| 204 | _none_hash_initialized = False |
| 205 | |
| 206 | |
| 207 | def create_request( |
| 208 | request_id: int | None = None, |
| 209 | num_tokens: int = 10, |
| 210 | common_prefix_len=0, |
| 211 | max_tokens: int = 16, |
| 212 | do_remote_decode: bool = False, |
| 213 | do_remote_prefill: bool = False, |
| 214 | num_remote_blocks: int = 3, |
| 215 | block_size: int = 16, |
| 216 | hash_fn: Callable = sha256, |
| 217 | ) -> Request: |
| 218 | """Make dummy request for testing.""" |
| 219 | assert num_tokens >= common_prefix_len >= 0 |
| 220 | |
| 221 | if request_id is None: |
| 222 | request_id = next(_request_count) |
| 223 | |
| 224 | global _none_hash_initialized |
| 225 | if not _none_hash_initialized: |
| 226 | init_none_hash(hash_fn) |
| 227 | _none_hash_initialized = True |
| 228 | |
| 229 | kv_transfer_params: dict[str, Any] | None = None |
| 230 | |
| 231 | if do_remote_decode: |
| 232 | assert not do_remote_prefill |
| 233 | kv_transfer_params = dict(do_remote_prefill=False, do_remote_decode=True) |
| 234 | elif do_remote_prefill: |
| 235 | kv_transfer_params = dict( |
| 236 | do_remote_prefill=True, |
| 237 | do_remote_decode=False, |
| 238 | remote_engine_id="my-engine-id", |
| 239 | remote_request_id=f"prefill-{request_id}", |
| 240 | remote_block_ids=list(range(num_remote_blocks)), |
| 241 | remote_host="my-host", |
| 242 | remote_port=1234, |
| 243 | tp_size=1, |
| 244 | ) |
| 245 | |
| 246 | max_tokens = 1 if do_remote_decode else max_tokens |
| 247 | sampling_params = SamplingParams(max_tokens=max_tokens) |
| 248 | sampling_params.update_from_generation_config({}, EOS_TOKEN_ID) |
| 249 | |
| 250 | common_prefix = [1] * common_prefix_len if common_prefix_len > 0 else [] |
| 251 | suffix = [i * request_id for i in range(num_tokens - common_prefix_len)] |
| 252 | prompt_token_ids = common_prefix + suffix |
| 253 | |
| 254 | req = Request( |
| 255 | request_id=f"id-{request_id}", |
| 256 | prompt_token_ids=prompt_token_ids, |
| 257 | sampling_params=sampling_params, |
| 258 | pooling_params=None, |
| 259 | mm_features=None, |
| 260 | block_hasher=get_request_block_hasher(block_size, hash_fn), |
| 261 | ) |
| 262 | req.kv_transfer_params = kv_transfer_params |
| 263 | return req |
| 264 | |
| 265 | |
| 266 | def create_model_runner_output( |
| 267 | reqs: list[Request], |
| 268 | finished_sending: set[str] | None = None, |
| 269 | finished_recving: set[str] | None = None, |
| 270 | invalid_block_ids: set[int] | None = None, |
| 271 | use_eos: bool = False, |
| 272 | token_id: int = 0, |
| 273 | kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None, |
| 274 | ) -> ModelRunnerOutput: |
| 275 | """Make dummy model runner output for testing.""" |
| 276 | # Make request data. |
| 277 | req_ids = [req.request_id for req in reqs] |
| 278 | req_id_to_index = {req_id: idx for idx, req_id in enumerate(req_ids)} |
| 279 | |
| 280 | # Make sampled tokens. |
| 281 | sampled_token = EOS_TOKEN_ID if use_eos else token_id |
| 282 | sampled_token_ids = [[sampled_token] for _ in req_ids] |
| 283 | |
| 284 | kv_connector_output = ( |
| 285 | None |
| 286 | if ( |
| 287 | finished_sending is None |
| 288 | and finished_recving is None |
| 289 | and invalid_block_ids is None |
| 290 | and kv_connector_worker_meta is None |
| 291 | ) |
| 292 | else KVConnectorOutput( |
| 293 | finished_sending=finished_sending, |
| 294 | finished_recving=finished_recving, |
| 295 | invalid_block_ids=invalid_block_ids or set(), |
| 296 | kv_connector_worker_meta=kv_connector_worker_meta, |
| 297 | ) |
| 298 | ) |
| 299 | |
| 300 | # Make output data structure. |
| 301 | return ModelRunnerOutput( |
| 302 | req_ids=req_ids, |
| 303 | req_id_to_index=req_id_to_index, |
| 304 | sampled_token_ids=sampled_token_ids, |
| 305 | logprobs=None, |
| 306 | prompt_logprobs_dict={}, |
| 307 | pooler_output=None, |
| 308 | kv_connector_output=kv_connector_output, |
| 309 | ) |
| 310 | |
| 311 | |
| 312 | class TestExampleConnector(ExampleConnector): |
| 313 | def __init__( |
| 314 | self, |
| 315 | config: VllmConfig, |
| 316 | role: KVConnectorRole, |
| 317 | kv_cache_config: KVCacheConfig, |
| 318 | ): |
| 319 | self.name = config.kv_transfer_config.kv_connector_extra_config["name"] |
| 320 | self._connector = ExampleConnector(config, role, kv_cache_config) |
| 321 | self.call_record: dict[str, int] = defaultdict(int) |
| 322 | # Use a unique temp file per connector |
| 323 | self._event_file = ( |
| 324 | tempfile.gettempdir() |
| 325 | + f"/connector_{self.name}-{self.role.name}_events.log" |
| 326 | ) |
| 327 | # Start with an empty file |
| 328 | with open(self._event_file, "w") as _: |
| 329 | pass |
| 330 | |
| 331 | def __getattribute__(self, name): |
| 332 | if name in ( |
| 333 | "_connector", |
| 334 | "call_record", |
| 335 | "name", |
| 336 | "_event_file", |
| 337 | "__class__", |
| 338 | "__dict__", |
| 339 | "__getattribute__", |
| 340 | "__init__", |
| 341 | ): # avoid recursion |
| 342 | return object.__getattribute__(self, name) |
| 343 | if not hasattr(self._connector, name): |
| 344 | return object.__getattribute__(self, name) |
| 345 | attr = getattr(self._connector, name) |
| 346 | |
| 347 | # Intercept calls to the connector interface and write an event |
| 348 | # for each one to a file, which can be read back in the main test proc. |
| 349 | if callable(attr): |
| 350 | |
| 351 | def wrapper(*args, **kwargs): |
| 352 | self.call_record[name] += 1 |
| 353 | |
| 354 | # Include args that we're interested in |
| 355 | to_log = [name] |
| 356 | for arg in args: |
| 357 | if isinstance(arg, int): |
| 358 | to_log.append(str(arg)) |
| 359 | elif isinstance(arg, KVCacheBlocks): |
| 360 | to_log.append(f"num_blocks={[len(b) for b in arg.blocks]}") |
| 361 | |
| 362 | # Log the event as a line to the file |
| 363 | try: |
| 364 | with open(self._event_file, "a") as f: |
| 365 | f.write(" ".join(to_log) + "\n") |
| 366 | except Exception as e: |
| 367 | print(f"[ERROR] Could not log event {name} for {self.name}: {e}") |
| 368 | return attr(*args, **kwargs) |
| 369 | |
| 370 | return wrapper |
| 371 | return attr |
| 372 | |
| 373 | |
| 374 | @dataclass(frozen=True) |
| 375 | class MockKVConfig: |
| 376 | matched_tokens: int = 0 |
| 377 | is_async: bool = False |
| 378 | num_defers_before_matching: int = 0 |
| 379 | supports_divergent_local_hybrid_hits: bool = False |
| 380 | |
| 381 | |
| 382 | class MockKVConnectorMetadata(KVConnectorMetadata): |
| 383 | def __init__(self): |
| 384 | # Scheduler tests check metadata.requests |
| 385 | self.requests: list = [] |
| 386 | |
| 387 | |
| 388 | class MockKVConnector(KVConnectorBase_V1): |
| 389 | """Mock KV connector for scheduler tests, supporting both sync and async mode.""" |
| 390 | |
| 391 | def __init__( |
| 392 | self, |
| 393 | vllm_config: VllmConfig, |
| 394 | role: KVConnectorRole, |
| 395 | kv_cache_config: KVCacheConfig, |
| 396 | ): |
| 397 | super().__init__(vllm_config, role, kv_cache_config) |
| 398 | extra_config = self._kv_transfer_config.kv_connector_extra_config |
| 399 | self.config = MockKVConfig( |
| 400 | matched_tokens=extra_config["matched_tokens"], |
| 401 | is_async=extra_config["is_async"], |
| 402 | num_defers_before_matching=extra_config.get( |
| 403 | "num_defers_before_matching", 0 |
| 404 | ), |
| 405 | supports_divergent_local_hybrid_hits=extra_config.get( |
| 406 | "supports_divergent_local_hybrid_hits", False |
| 407 | ), |
| 408 | ) |
| 409 | self._defers_left: defaultdict[str, int] = defaultdict( |
| 410 | lambda: self.config.num_defers_before_matching |
| 411 | ) |
| 412 | |
| 413 | @property |
| 414 | def supports_divergent_local_hybrid_hits(self) -> bool: |
| 415 | return self.config.supports_divergent_local_hybrid_hits |
| 416 | |
| 417 | def get_num_new_matched_tokens( |
| 418 | self, |
| 419 | request: Request, |
| 420 | num_computed_tokens: int, |
| 421 | ) -> tuple[int | None, bool]: |
| 422 | if self._defers_left[request.request_id] > 0: |
| 423 | self._defers_left[request.request_id] -= 1 |
| 424 | return (None, False) |
| 425 | return (self.config.matched_tokens, self.config.is_async) |
| 426 | |
| 427 | def update_state_after_alloc( |
| 428 | self, |
| 429 | request: Request, |
| 430 | blocks: KVCacheBlocks, |
| 431 | num_external_tokens: int, |
| 432 | ): |
| 433 | pass |
| 434 | |
| 435 | def build_connector_meta( |
| 436 | self, scheduler_output: SchedulerOutput |
| 437 | ) -> KVConnectorMetadata: |
| 438 | metadata = MockKVConnectorMetadata() |
| 439 | cached_reqs = scheduler_output.scheduled_cached_reqs |
| 440 | for req_id in chain( |
| 441 | (req.req_id for req in scheduler_output.scheduled_new_reqs), |
| 442 | ( |
| 443 | req_id |
| 444 | for req_id in cached_reqs.req_ids |
| 445 | if req_id in cached_reqs.resumed_req_ids |
| 446 | ), |
| 447 | ): |
| 448 | metadata.requests.append({"req_id": req_id}) |
| 449 | return metadata |
| 450 | |
| 451 | def start_load_kv(self, kv_caches, finished_req_ids): |
| 452 | pass |
| 453 | |
| 454 | def wait_for_layer_load(self, layer_name): |
| 455 | pass |
| 456 | |
| 457 | def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs): |
| 458 | pass |
| 459 | |
| 460 | def wait_for_save(self): |
| 461 | pass |
| 462 | |
| 463 | |
| 464 | KVConnectorFactory.register_connector( |
| 465 | "TestExampleConnector", __name__, TestExampleConnector.__name__ |
| 466 | ) |
| 467 | |
| 468 | KVConnectorFactory.register_connector( |
| 469 | "MockKVConnector", __name__, MockKVConnector.__name__ |
| 470 | ) |
| 471 | |
| 472 | |
| 473 | def make_kv_cache_config( |
| 474 | block_size: int, |
| 475 | swa_enabled: bool = False, |
| 476 | mamba_enabled: bool = False, |
| 477 | sw_size: int = 128, |
| 478 | num_blocks: int = 100, |
| 479 | mamba_cache_mode: Literal["align", "none"] = "none", |
| 480 | ) -> KVCacheConfig: |
| 481 | kv_cache_groups = [ |
| 482 | KVCacheGroupSpec( |
| 483 | ["layer0", "layer2"], |
| 484 | FullAttentionSpec( |
| 485 | block_size=block_size, |
| 486 | num_kv_heads=4, |
| 487 | head_size=16, |
| 488 | dtype=torch.float16, |
| 489 | ), |
| 490 | ) |
| 491 | ] |
| 492 | if swa_enabled: |
| 493 | kv_cache_groups.append( |
| 494 | KVCacheGroupSpec( |
| 495 | ["layer1", "layer3"], |
| 496 | SlidingWindowSpec( |
| 497 | block_size=block_size, |
| 498 | num_kv_heads=4, |
| 499 | head_size=16, |
| 500 | dtype=torch.float16, |
| 501 | sliding_window=sw_size, |
| 502 | ), |
| 503 | ) |
| 504 | ) |
| 505 | if mamba_enabled: |
| 506 | kv_cache_groups.append( |
| 507 | KVCacheGroupSpec( |
| 508 | ["mamba0", "mamba1"], |
| 509 | MambaSpec( |
| 510 | block_size=block_size, |
| 511 | shapes=((16,), (16,)), |
| 512 | dtypes=(torch.float16,), |
| 513 | mamba_cache_mode=mamba_cache_mode, |
| 514 | ), |
| 515 | ) |
| 516 | ) |
| 517 | return KVCacheConfig( |
| 518 | num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups |
| 519 | ) |
| 520 | |
| 521 | |
| 522 | def make_nixl_scheduler( |
| 523 | has_mamba: bool = False, |
| 524 | is_hma_required: bool = False, |
| 525 | heartbeat: bool = False, |
| 526 | kv_lease_duration: int = 30, |
| 527 | ): |
| 528 | """Create a NixlConnectorScheduler via __new__ (skipping __init__). |
| 529 | |
| 530 | Only sets the flags needed by the tests. When *heartbeat=True* the |
| 531 | scheduler-side heartbeat bookkeeping fields are also initialised. |
| 532 | """ |
| 533 | from types import SimpleNamespace |
| 534 | |
| 535 | from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( |
| 536 | NixlConnectorScheduler, |
| 537 | ) |
| 538 | |
| 539 | sched = object.__new__(NixlConnectorScheduler) |
| 540 | sched._has_mamba = has_mamba |
| 541 | sched._is_hma_required = is_hma_required |
| 542 | sched.kv_cache_config = make_kv_cache_config( |
| 543 | block_size=16, |
| 544 | mamba_enabled=has_mamba, |
| 545 | ) |
| 546 | sched.vllm_config = SimpleNamespace(num_prefill_lookahead_tokens=0) |
| 547 | |
| 548 | if heartbeat: |
| 549 | sched._heartbeat_by_engine = {} |
| 550 | sched._heartbeat_req_engine = {} |
| 551 | sched._last_heartbeat_time = 0.0 |
| 552 | sched._kv_lease_duration = kv_lease_duration |
| 553 | sched._heartbeat_interval = kv_lease_duration // 6 |
| 554 | # Fields touched by build_connector_meta / request_finished: |
| 555 | sched._reqs_need_recv = {} |
| 556 | sched._hisparse_host_blocks_to_recv = {} |
| 557 | sched._reqs_need_send = {} |
| 558 | sched._reqs_in_batch = set() |
| 559 | sched._reqs_not_processed = set() |
| 560 | sched._reqs_need_save = {} |
| 561 | sched.use_host_buffer = False |
| 562 | sched.engine_id = "test-engine" |
| 563 | sched.transfer_tp_size = 1 |
| 564 | sched.side_channel_host = "localhost" |
| 565 | sched.side_channel_port = 5555 |
| 566 | sched.blocks_per_sw = [] |
| 567 | sched.is_bidirectional_kv_xfer_enabled = False |
| 568 | return sched |
| 569 | |
| 570 | |
| 571 | def make_nixl_push_scheduler( |
| 572 | *, |
| 573 | decoder_kv_blocks_ttl: float = 30.0, |
| 574 | push_registration_timeout: float | None = None, |
| 575 | is_bidirectional_kv_xfer_enabled: bool = False, |
| 576 | has_mamba: bool = False, |
| 577 | ): |
| 578 | """Create a NixlPushConnectorScheduler via __new__ (skipping __init__). |
| 579 | |
| 580 | The push scheduler can't reuse :func:`make_nixl_scheduler` because it |
| 581 | is a different class (``NixlPushConnectorScheduler`` vs |
| 582 | ``NixlConnectorScheduler``) and carries push-specific state. Only the |
| 583 | fields touched by the unit tests are populated. |
| 584 | """ |
| 585 | from unittest.mock import MagicMock |
| 586 | |
| 587 | from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import ( |
| 588 | NixlPushConnectorScheduler, |
| 589 | ) |
| 590 | |
| 591 | sched = object.__new__(NixlPushConnectorScheduler) |
| 592 | |
| 593 | # Base scheduler fields (shared with pull / heartbeat path). |
| 594 | sched._reqs_need_recv = {} |
| 595 | sched._reqs_need_send = {} |
| 596 | sched._reqs_in_batch = set() |
| 597 | sched._reqs_not_processed = set() |
| 598 | sched._reqs_need_save = {} |
| 599 | sched._kv_lease_duration = 30 |
| 600 | sched.decoder_kv_blocks_ttl = decoder_kv_blocks_ttl |
| 601 | sched.use_host_buffer = False |
| 602 | sched.engine_id = "decode-engine" |
| 603 | sched.transfer_tp_size = 1 |
| 604 | sched.side_channel_host = "127.0.0.1" |
| 605 | sched.side_channel_port = 5600 |
| 606 | sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled |
| 607 | sched._has_mamba = has_mamba |
| 608 | sched.kv_cache_config = make_kv_cache_config( |
| 609 | block_size=16, |
| 610 | mamba_enabled=has_mamba, |
| 611 | ) |
| 612 | |
| 613 | # vllm_config is consulted for parallel_config.tensor_parallel_size, and by |
| 614 | # `_prefill_backoff` on both the P and D prefill paths. |
| 615 | vllm_config = MagicMock() |
| 616 | vllm_config.parallel_config.tensor_parallel_size = 1 |
| 617 | vllm_config.num_prefill_lookahead_tokens = 0 |
| 618 | sched.vllm_config = vllm_config |
| 619 | |
| 620 | # Push-specific state. |
| 621 | sched._push_pending_registrations = {} |
| 622 | sched._push_registration_deadlines = {} |
| 623 | sched._finished_request_blocks = {} |
| 624 | sched._newly_finished_push_blocks = {} |
| 625 | sched._push_registration_timeout = ( |
| 626 | push_registration_timeout |
| 627 | if push_registration_timeout is not None |
| 628 | else decoder_kv_blocks_ttl |
| 629 | ) |
| 630 | |
| 631 | # Heartbeat fields touched by base request_finished / |
| 632 | # update_connector_output. |
| 633 | sched._heartbeat_by_engine = {} |
| 634 | sched._heartbeat_req_engine = {} |
| 635 | sched._last_heartbeat_time = 0.0 |
| 636 | sched.blocks_per_sw = [] |
| 637 | |
| 638 | return sched |
| 639 | |
| 640 | |
| 641 | def make_moriio_writer(fake_worker: Any) -> Any: |
| 642 | """Build a MoRIIOWriter with internals stubbed for unit tests. |
| 643 | |
| 644 | Bypasses ``__init__`` and wires only the write/finalize state the tests |
| 645 | touch, including the deferred-task fields used by the routing suite. |
| 646 | """ |
| 647 | import threading |
| 648 | from queue import Queue |
| 649 | |
| 650 | from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( |
| 651 | MoRIIOWriter, |
| 652 | ) |
| 653 | |
| 654 | writer = MoRIIOWriter.__new__(MoRIIOWriter) |
| 655 | writer._worker_ref = lambda: fake_worker |
| 656 | writer._write_task_q = Queue() |
| 657 | writer._write_state_lock = threading.Lock() |
| 658 | writer._scheduled_writes = defaultdict(int) |
| 659 | writer._scheduled_layers = defaultdict(set) |
| 660 | writer._sealed_writes = {} |
| 661 | writer._deferred_tasks = [] |
| 662 | writer._defer_timeout = 60.0 |
| 663 | writer.ensure_worker_started = lambda: None |
| 664 | return writer |