rein yangandSy03
ad3a162eb8
Signed-off-by: R2-Y <ruiruyang2@gmail.com> Co-authored-by: Sy03 <1370724210@qq.com>
77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from .logging import get_connector_logger
|
|
|
|
logger = get_connector_logger(__name__)
|
|
|
|
TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset(
|
|
{
|
|
"MooncakeTransferEngineConnector",
|
|
"MoriTransferEngineConnector",
|
|
"YuanrongTransferEngineConnector",
|
|
}
|
|
)
|
|
|
|
|
|
def get_stage_connector_role(model_config: Any) -> str | None:
|
|
"""Return the configured stage connector direction, if explicit."""
|
|
connector_config = getattr(model_config, "stage_connector_config", None)
|
|
if isinstance(connector_config, dict):
|
|
extra = connector_config.get("extra")
|
|
else:
|
|
extra = getattr(connector_config, "extra", None)
|
|
if isinstance(extra, dict):
|
|
role = extra.get("role")
|
|
return role if isinstance(role, str) else None
|
|
return None
|
|
|
|
|
|
def stage_receives_chunks(model_config: Any) -> bool:
|
|
"""Whether connector chunks, rather than the orchestrator, feed a stage."""
|
|
return get_stage_connector_role(model_config) != "sender"
|
|
|
|
|
|
def stage_sends_async_output(model_config: Any) -> bool:
|
|
"""Whether async output should be partitioned for connector transport."""
|
|
role = get_stage_connector_role(model_config)
|
|
if role is not None:
|
|
return role == "sender"
|
|
# Preserve legacy partitioning while keeping stage-0 orchestrator bridges
|
|
# on the normal RequestOutput path.
|
|
return getattr(model_config, "stage_id", None) != 0
|
|
|
|
|
|
@dataclass
|
|
class ConnectorSpec:
|
|
"""Specification for a connector instance."""
|
|
|
|
name: str # e.g., "MooncakeStoreConnector", "SharedMemoryConnector", "YuanrongConnector"
|
|
extra: dict[str, Any] = field(default_factory=dict) # backend-specific config
|
|
|
|
|
|
@dataclass
|
|
class OmniTransferConfig:
|
|
"""
|
|
Top-level configuration for OmniConnector system.
|
|
Members:
|
|
connectors: A dictionary of connectors, keyed by (from_stage, to_stage).
|
|
default_connector: The default connector to use if no connector is specified for an edge.
|
|
"""
|
|
|
|
# Direct mapping: (from_stage, to_stage) -> connector
|
|
connectors: dict[tuple[str, str], ConnectorSpec] = field(default_factory=dict)
|
|
default_connector: ConnectorSpec | None = None
|
|
|
|
def get_connector_for_edge(self, from_stage: str, to_stage: str) -> ConnectorSpec | None:
|
|
"""Get connector spec for a specific edge."""
|
|
edge_key = (from_stage, to_stage)
|
|
return self.connectors.get(edge_key, self.default_connector)
|
|
|
|
def has_connector_for_edge(self, from_stage: str, to_stage: str) -> bool:
|
|
"""Check if there's a connector configured for the edge."""
|
|
return self.get_connector_for_edge(from_stage, to_stage) is not None
|