Unverified Commit 8a6efec6 authored by Kris Hung's avatar Kris Hung Committed by GitHub
Browse files

fix: Replace NixlAgentMetadata with NixlHandshakePayload for vllm 0.13.0 (#5351)

parent 16451f48
...@@ -13,8 +13,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( ...@@ -13,8 +13,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
MultiKVConnectorMetadata, MultiKVConnectorMetadata,
) )
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
NixlAgentMetadata,
NixlConnector, NixlConnector,
NixlHandshakePayload,
) )
from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.output import SchedulerOutput
...@@ -111,14 +111,14 @@ class PdConnector(MultiConnector): ...@@ -111,14 +111,14 @@ class PdConnector(MultiConnector):
metadata so decode workers can connect for KV transfer coordination. metadata so decode workers can connect for KV transfer coordination.
Returns: Returns:
NixlAgentMetadata from the NIXL connector, or None if not available. NixlHandshakePayload from the NIXL connector, or None if not available.
""" """
# Get handshake metadata from the NIXL connector (second connector) # Get handshake metadata from the NIXL connector (second connector)
nixl_connector = self._connectors[1] nixl_connector = self._connectors[1]
metadata = nixl_connector.get_handshake_metadata() metadata = nixl_connector.get_handshake_metadata()
if metadata is not None and not isinstance(metadata, NixlAgentMetadata): if metadata is not None and not isinstance(metadata, NixlHandshakePayload):
raise TypeError( raise TypeError(
f"Expected NixlAgentMetadata from NIXL connector, " f"Expected NixlHandshakePayload from NIXL connector, "
f"got {type(metadata).__name__}" f"got {type(metadata).__name__}"
) )
return metadata return metadata
......
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