Unverified Commit b2370a97 authored by Olga Andreeva's avatar Olga Andreeva Committed by GitHub
Browse files

fix: Make lmcache dependency optional for PdConnector (#4319)


Signed-off-by: default avatarOlga Andreeva <oandreeva@nvidia.com>
Signed-off-by: default avatarOlga Andreeva <124622579+oandreeva-nv@users.noreply.github.com>
Co-authored-by: default avatarCopilot <175728472+Copilot@users.noreply.github.com>
parent d6cda975
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional, Type
from kvbm.vllm_integration.connector.dynamo_connector import DynamoConnector
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector import (
LMCacheConnectorV1,
)
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
MultiConnector,
MultiKVConnectorMetadata,
......@@ -15,8 +12,22 @@ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import NixlConnector
from vllm.v1.core.sched.output import SchedulerOutput
# Optional import for LMCache support
_LMCacheConnectorV1: Optional[Type] = None
try:
from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector import (
LMCacheConnectorV1,
)
_LMCacheConnectorV1 = LMCacheConnectorV1
except ImportError:
pass
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector import (
LMCacheConnectorV1,
)
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
from vllm.v1.request import Request
......@@ -41,9 +52,18 @@ class PdConnector(MultiConnector):
raise ValueError(
f"PdConnector requires exactly two connectors (got {len(self._connectors)})"
)
if not isinstance(self._connectors[0], (DynamoConnector, LMCacheConnectorV1)):
# Build allowed types for first connector
allowed_first_types: list[Type] = [DynamoConnector]
if _LMCacheConnectorV1 is not None:
allowed_first_types.append(_LMCacheConnectorV1)
if not isinstance(self._connectors[0], tuple(allowed_first_types)):
allowed_names = ["DynamoConnector"]
if _LMCacheConnectorV1 is not None:
allowed_names.append("LMCacheConnectorV1")
raise TypeError(
f"Expected first connector to be DynamoConnector or LMCacheConnector, "
f"Expected first connector to be {' or '.join(allowed_names)}, "
f"got {type(self._connectors[0]).__name__}"
)
if not isinstance(self._connectors[1], NixlConnector):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment