vllm-omni/vllm_omni/distributed/omni_connectors/utils/initialization.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""Utilities for OmniConnector configuration and validation."""

import json
import os
import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal

from ..factory import OmniConnectorFactory
from .config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec, OmniTransferConfig
from .env import expand_env_int
from .local_rank import get_connector_local_rank
from .logging import get_connector_logger

if TYPE_CHECKING:
    from ..connectors.base import OmniConnectorBase
else:
    OmniConnectorBase = Any

logger = get_connector_logger(__name__)

# Reserve a separate port range for KV-transfer sockets so they do not
# collide with request-forwarding endpoints that share the same base port.
KV_TRANSFER_PORT_OFFSET = 100

# Port stride between TP ranks so each worker binds a unique ZMQ port
# when TP > 1.  Must be larger than the maximum number of pipeline stages.
# Formula:
#   zmq_port = base + KV_TRANSFER_PORT_OFFSET
#            + replica * KV_REPLICA_PORT_STRIDE
#            + rank * KV_RANK_PORT_STRIDE
#            + stage
KV_RANK_PORT_STRIDE = 16

# Port stride between Omni replicas of the same stage.  This reserves a
# comfortably sized block per replica for TP-rank and stage offsets.
KV_REPLICA_PORT_STRIDE = 1024

REQUEST_FORWARDING_PORT_OFFSET = 0
ORCHESTRATOR_PORT_OFFSET = 200

ConnectorPurpose = Literal["request_forwarding", "kv_transfer", "orchestrator"]


def compute_connector_zmq_port(
    base_port: int,
    *,
    purpose: ConnectorPurpose,
    from_stage: int,
    local_rank: int = 0,
    replica_id: int = 0,
) -> int:
    """Return a collision-free connector endpoint for an edge and worker."""
    purpose_offsets = {
        "request_forwarding": REQUEST_FORWARDING_PORT_OFFSET,
        "kv_transfer": KV_TRANSFER_PORT_OFFSET,
        "orchestrator": ORCHESTRATOR_PORT_OFFSET,
    }
    if purpose not in purpose_offsets:
        raise ValueError(f"Unsupported connector purpose: {purpose}")
    if not 0 <= from_stage < KV_RANK_PORT_STRIDE:
        raise ValueError(f"from_stage must be in [0, {KV_RANK_PORT_STRIDE}), got {from_stage}")
    if local_rank < 0:
        raise ValueError(f"local_rank must be non-negative, got {local_rank}")
    if replica_id < 0:
        raise ValueError(f"replica_id must be non-negative, got {replica_id}")

    port = (
        base_port
        + purpose_offsets[purpose]
        + replica_id * KV_REPLICA_PORT_STRIDE
        + local_rank * KV_RANK_PORT_STRIDE
        + from_stage
    )
    if not 1 <= port <= 65535:
        raise ValueError(f"Resolved connector port must be in [1, 65535], got {port}")
    return port


def resolve_connector_spec(
    spec: ConnectorSpec,
    *,
    stage_id: int,
    role: str | None,
    purpose: ConnectorPurpose = "request_forwarding",
    local_rank: int = 0,
    replica_id: int = 0,
) -> ConnectorSpec:
    """Copy and resolve a connector spec for one stage worker."""
    extra = dict(spec.extra or {})
    extra["stage_id"] = stage_id
    if role is not None:
        extra["role"] = role
    if spec.name not in TRANSFER_ENGINE_CONNECTOR_NAMES:
        return ConnectorSpec(name=spec.name, extra=extra)

    resolved_port = None
    if role != "receiver" or "sender_zmq_port" not in extra:
        resolved_port = compute_connector_zmq_port(
            expand_env_int(extra.get("zmq_port", 50051), "zmq_port"),
            purpose=purpose,
            from_stage=int(extra.get("from_stage", stage_id)),
            local_rank=local_rank,
            replica_id=replica_id,
        )
    if role == "receiver":
        extra.setdefault("sender_host", extra.get("host", "127.0.0.1"))
        extra.setdefault("sender_zmq_port", resolved_port)
        extra.pop("zmq_port", None)
        outgoing = extra.get("outgoing")
        if spec.name == "NixlConnector" and isinstance(outgoing, dict):
            extra["zmq_port"] = compute_connector_zmq_port(
                expand_env_int(outgoing.get("zmq_port", 50051), "zmq_port"),
                purpose=purpose,
                from_stage=int(outgoing.get("from_stage", stage_id)),
                local_rank=local_rank,
                replica_id=replica_id,
            )
            extra["host"] = outgoing.get("host", "auto")
    else:
        extra["zmq_port"] = resolved_port
    return ConnectorSpec(name=spec.name, extra=extra)


def initialize_connectors_from_config(
    config_path: str | Path | None = None,
    purpose: str = "request_forwarding",
    caller_stage_id: int | str | None = None,
    is_sender: bool | None = None,
) -> tuple[OmniTransferConfig | None, dict[tuple[str, str], OmniConnectorBase]]:
    """
    Initialize connectors from configuration file.

    Returns:
        tuple: (OmniTransferConfig, dict of {(from, to): connector_instance})
    """
    transfer_config = load_omni_transfer_config(config_path)

    if not transfer_config:
        logger.info("No OmniTransferConfig provided")
        return None, {}

    # create connectors from config
    connectors = create_connectors_from_config(
        transfer_config.connectors,
        purpose=purpose,
        caller_stage_id=caller_stage_id,
        is_sender=is_sender,
    )
    return transfer_config, connectors


def create_connectors_from_config(
    connectors_config: dict[tuple[str, str], ConnectorSpec],
    purpose: str = "request_forwarding",
    caller_stage_id: int | str | None = None,
    is_sender: bool | None = None,
) -> dict[tuple[str, str], OmniConnectorBase]:
    """
    Create connectors from config.

    Args:
        connectors_config: A dictionary of connector configurations.

    Returns:
        A dictionary of connectors.
    """
    connectors = {}
    for edge_key, connector_spec in connectors_config.items():
        from_stage, to_stage = edge_key
        try:
            if connector_spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES:
                extra = dict(connector_spec.extra) if connector_spec.extra else {}
                try:
                    from_stage_id = int(from_stage)
                except (TypeError, ValueError) as exc:
                    raise ValueError(f"Transfer-engine stage id must be an integer, got {from_stage!r}") from exc

                is_orchestrator = str(caller_stage_id) == "orchestrator"
                local_rank = 0 if is_orchestrator else get_connector_local_rank()
                resolved_purpose: ConnectorPurpose = "orchestrator" if is_orchestrator else purpose  # type: ignore[assignment]
                if is_sender is not None:
                    role = "sender" if is_sender else "receiver"
                elif caller_stage_id is not None and str(caller_stage_id) == to_stage:
                    role = "receiver"
                elif caller_stage_id is not None and str(caller_stage_id) == from_stage:
                    role = "sender"
                else:
                    role = str(extra.get("role", "auto"))
                extra["from_stage"] = from_stage_id
                extra["to_stage"] = int(to_stage)
                replica_id = max(int(os.environ.get("VLLM_OMNI_REPLICA_ID", "0")), 0)
                resolved_spec = resolve_connector_spec(
                    ConnectorSpec(name=connector_spec.name, extra=extra),
                    stage_id=int(caller_stage_id) if str(caller_stage_id).isdigit() else from_stage_id,
                    role=role,
                    purpose=resolved_purpose,
                    local_rank=local_rank,
                    replica_id=replica_id,
                )
                connector = OmniConnectorFactory.create_connector(resolved_spec)
            else:
                connector = OmniConnectorFactory.create_connector(connector_spec)
            connectors[edge_key] = connector
            logger.info(
                "Created connector for %s -> %s: %s",
                from_stage,
                to_stage,
                type(connector).__name__,
            )
        except Exception as e:
            raise RuntimeError(f"Failed to initialize connector for edge {edge_key}: {e}") from e

    return connectors


def get_connectors_config_for_stage(transfer_config: OmniTransferConfig | None, stage_id: str | int) -> dict[str, Any]:
    """
    Extract connector configurations relevant for a specific stage worker.

    Returns a dict compatible with worker initialization:
    {
        "from_stage_X": {
            "spec": {
                "name": "ConnectorName",
                "extra": {...}
            }
        },
        ...
    }
    """
    if not transfer_config:
        return {}

    stage_connectors_config = {}
    target_stage = str(stage_id)

    # Iterate through all configured edges and inject direction-specific role.
    # The shared edge-level ConnectorSpec is role-neutral; each stage gets
    # the correct role ("sender" or "receiver") based on its position in
    # the edge so that MooncakeTransferEngineConnector (and any future
    # role-aware connector) initializes correctly.
    for (from_stage, to_stage), spec in transfer_config.connectors.items():
        if to_stage == target_stage:
            # Incoming edge → this stage is the receiver
            extra = dict(spec.extra) if spec.extra else {}
            extra.setdefault("role", "receiver")
            extra["from_stage"] = int(from_stage)
            extra["to_stage"] = int(to_stage)
            stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": extra}}
        elif from_stage == target_stage and target_stage == "0":
            # Outgoing edge for stage 0 — included for async_chunk spec
            # extraction (omni_stage.py), NOT for connector instantiation.
            extra = dict(spec.extra) if spec.extra else {}
            extra.setdefault("role", "sender")
            extra["from_stage"] = int(from_stage)
            extra["to_stage"] = int(to_stage)
            stage_connectors_config[f"to_stage_{to_stage}"] = {"spec": {"name": spec.name, "extra": extra}}

    return stage_connectors_config


def load_omni_transfer_config(
    config_path: str | Path | None = None,
    config_dict: dict[str, Any] | None = None,
) -> OmniTransferConfig | None:
    """Load OmniTransferConfig from file or dict."""
    if config_path is None and config_dict is None:
        # Even if no config provided, we might want to return a default config with SHM connectors
        # But without stage info we can't do much.
        return None

    if config_path is not None:
        config_path = Path(config_path)
        if not config_path.exists():
            raise FileNotFoundError(f"Config file not found: {config_path}")

        with open(config_path, encoding="utf-8") as f:
            if config_path.suffix.lower() == ".json":
                config_dict = json.load(f)
            elif config_path.suffix.lower() in [".yaml", ".yml"]:
                try:
                    import yaml

                    config_dict = yaml.safe_load(f)
                except ImportError:
                    raise ImportError("PyYAML required for YAML config files")
            else:
                raise ValueError(f"Unsupported config file format: {config_path.suffix}")

    if config_dict is None:
        return None

    # Normalize new-schema (top-level ``connectors`` + ``stages``) into the
    # legacy ``runtime.connectors`` + ``stage_args`` shape the parser reads.
    if "stages" in config_dict and "stage_args" not in config_dict:
        normalized: dict[str, Any] = dict(config_dict)
        runtime = dict(normalized.get("runtime") or {})
        if "connectors" in normalized and "connectors" not in runtime:
            runtime["connectors"] = normalized["connectors"]
        if "edges" in normalized and "edges" not in runtime:
            runtime["edges"] = normalized["edges"]
        normalized["runtime"] = runtime
        normalized["stage_args"] = normalized["stages"]
        config_dict = normalized

    # Parse connectors
    connectors = {}
    runtime_config = config_dict.get("runtime", {})

    # Parse global connectors (from runtime.connectors)
    global_connectors = runtime_config.get("connectors", {})

    # Parse stage-level connectors
    stage_args = config_dict.get("stage_args", [])
    expected_edges: set[tuple[str, str]] = set()
    for stage_config in stage_args:
        stage_id = str(stage_config["stage_id"])

        # Input connectors (this stage is the receiver)
        # NOTE: role is NOT injected here — the shared edge-level ConnectorSpec
        # must remain role-neutral.  Role is injected per-stage in
        # get_connectors_config_for_stage() / resolve_omni_kv_config_for_stage().
        for input_key, conn_ref in stage_config.get("input_connectors", {}).items():
            if isinstance(conn_ref, str):
                # Reference to global connector
                if conn_ref in global_connectors:
                    conn_config = global_connectors[conn_ref]
                    extra = dict(conn_config.get("extra", {}))
                else:
                    raise ValueError(f"Undefined connector reference: {conn_ref}")
                connector = ConnectorSpec(name=conn_config["name"], extra=extra)
            else:
                # Inline connector definition
                extra = dict(conn_ref.get("extra", {}))
                connector = ConnectorSpec(name=conn_ref["name"], extra=extra)

            # Parse from_stage from key (e.g., "from_stage_0" -> "0")
            from_stage = input_key.replace("from_stage_", "")
            edge_key = (from_stage, stage_id)
            # Both sides of an edge may define the same connector reference;
            # verify consistency if already registered.
            if edge_key in connectors:
                existing = connectors[edge_key]
                if existing.name != connector.name:
                    raise ValueError(
                        f"Connector type mismatch for edge {edge_key[0]}->{edge_key[1]}: "
                        f"previously registered as '{existing.name}', "
                        f"but input_connectors of stage {stage_id} specifies '{connector.name}'"
                    )
            else:
                connectors[edge_key] = connector
            expected_edges.add(edge_key)

        # Output connectors (this stage is the sender)
        for output_key, conn_ref in stage_config.get("output_connectors", {}).items():
            if isinstance(conn_ref, str):
                # Reference to global connector
                if conn_ref in global_connectors:
                    conn_config = global_connectors[conn_ref]
                    extra = dict(conn_config.get("extra", {}))
                else:
                    raise ValueError(f"Undefined connector reference: {conn_ref}")
                connector = ConnectorSpec(name=conn_config["name"], extra=extra)
            else:
                # Inline connector definition
                extra = dict(conn_ref.get("extra", {}))
                connector = ConnectorSpec(name=conn_ref["name"], extra=extra)

            # Parse to_stage from key (e.g., "to_stage_1" -> "1")
            to_stage = output_key.replace("to_stage_", "")
            edge_key = (stage_id, to_stage)
            if edge_key in connectors:
                existing = connectors[edge_key]
                if existing.name != connector.name:
                    raise ValueError(
                        f"Connector type mismatch for edge {edge_key[0]}->{edge_key[1]}: "
                        f"previously registered as '{existing.name}', "
                        f"but output_connectors of stage {stage_id} specifies '{connector.name}'"
                    )
            else:
                connectors[edge_key] = connector
            expected_edges.add(edge_key)

    # Auto-configure SharedMemoryConnector for missing edges based on runtime edges / engine_input_source
    if stage_args:
        try:
            # Prefer explicit runtime edges if provided
            runtime_edges = runtime_config.get("edges", [])
            if isinstance(runtime_edges, list) and runtime_edges:
                for edge in runtime_edges:
                    from_stage = edge.get("from")
                    to_stage = edge.get("to")
                    if from_stage is None or to_stage is None:
                        continue
                    edge_key = (str(from_stage), str(to_stage))
                    expected_edges.add(edge_key)
                    if edge_key not in connectors:
                        logger.info(f"Auto-configuring SharedMemoryConnector for edge {edge_key}")
                        connectors[edge_key] = ConnectorSpec(name="SharedMemoryConnector")

            # Fallback: infer edges from engine_input_source for each stage
            for stage_config in stage_args:
                to_stage = str(stage_config["stage_id"])
                # Check explicit input sources
                sources = stage_config.get("engine_input_source", [])

                for from_stage in sources:
                    from_stage_str = str(from_stage)
                    edge_key = (from_stage_str, to_stage)
                    expected_edges.add(edge_key)

                    if edge_key not in connectors:
                        logger.info(f"Auto-configuring SharedMemoryConnector for edge {edge_key}")
                        connectors[edge_key] = ConnectorSpec(name="SharedMemoryConnector")

        except Exception as e:
            logger.warning(f"Failed to auto-configure SHM connectors: {e}")

    # Fail fast if any expected edge is still missing a connector
    missing_edges = [edge for edge in expected_edges if edge not in connectors]
    if missing_edges:
        missing_str = ", ".join([f"{f}->{t}" for f, t in missing_edges])
        raise ValueError(
            "Connector configuration missing for edges: "
            f"{missing_str}. Define connectors or allow auto SHM creation for these edges."
        )

    config = OmniTransferConfig(connectors=connectors)

    logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations")
    return config


# High-level management functions


def initialize_orchestrator_connectors(
    config_path: str | None,
) -> tuple[OmniTransferConfig | None, dict[tuple[str, str], OmniConnectorBase]]:
    """Initialize connectors shared at orchestrator level.
    Args:
        config_path: The path to the configuration file.
    Returns:
        A tuple containing the OmniTransferConfig and a dictionary of connectors.
    """
    transfer_config, connectors = initialize_connectors_from_config(
        config_path,
        purpose="request_forwarding",
        caller_stage_id="orchestrator",
        is_sender=True,
    )
    return transfer_config, connectors


def get_stage_connector_config(
    transfer_config: OmniTransferConfig | None,
    stage_id: int,
) -> dict[str, Any]:
    """Return the serialized connector config payload for a specific stage."""
    if transfer_config is None:
        return {}

    try:
        return get_connectors_config_for_stage(transfer_config, stage_id)
    except Exception as exc:  # pragma: no cover - defensive logging
        logger.warning(
            "Failed to build connector config for stage %s: %s. Using IPC fallback.",
            stage_id,
            exc,
        )
        return {}


def build_stage_connectors(
    stage_id: int,
    connectors_config: dict[str, Any],
    purpose: str = "request_forwarding",
) -> dict[tuple[str, str], Any] | None:
    """Instantiate OmniConnectors for a stage based on config.

    Deprecated: prefer ``get_stage_connector_config`` plus the unified
    connector factory. Kept as a thin shim so legacy callers keep working.
    """
    warnings.warn(
        "build_stage_connectors is deprecated; use get_stage_connector_config and OmniConnectorFactory instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    if not connectors_config:
        return {}

    logger.info(
        "[Stage-%s] Initializing OmniConnectors with config keys: %s",
        stage_id,
        list(connectors_config.keys()),
    )

    from .config import ConnectorSpec

    connectors: dict[tuple[str, str], Any] = {}
    # Convert dictionary-formatted config to ConnectorSpec objects.
    # Only instantiate INPUT connectors ("from_stage_X") — the stage worker
    # only receives via connectors.  Output connectors are handled at the
    # orchestrator level (try_send_via_connector uses orchestrator connectors).
    stage_connector_specs = {}
    for key, config in connectors_config.items():
        if not key.startswith("from_stage_"):
            continue

        from_stage = key.replace("from_stage_", "")
        spec_dict = config.get("spec", {})
        if not spec_dict:
            continue

        connector_spec = ConnectorSpec(
            name=spec_dict.get("name", "SharedMemoryConnector"),
            extra=spec_dict.get("extra", {}),
        )
        stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec

    try:
        # Use unified connector creation logic
        connectors = create_connectors_from_config(
            stage_connector_specs,
            purpose=purpose,
            caller_stage_id=stage_id,
            is_sender=False,
        )
    except Exception as exc:  # pragma: no cover - defensive logging
        # Fail fast so the stage does not start with missing connectors.
        logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc)
        raise

    return connectors


def resolve_omni_kv_config_for_stage(
    transfer_cfg: OmniTransferConfig | None, stage_id: int | str
) -> tuple[dict[str, Any] | None, str | None, str | None]:
    """Resolve connector configuration for a specific stage (Sender/Receiver).

    This determines the primary connector configuration to be injected into the
    engine arguments, prioritizing outgoing edges (Sender role).
    """
    if not transfer_cfg or not getattr(transfer_cfg, "connectors", None):
        return None, None, None

    stage_id_str = str(stage_id)

    # Find outgoing edges (Sender logic)
    outgoing = [
        (to_stage, spec)
        for (from_stage, to_stage), spec in transfer_cfg.connectors.items()
        if from_stage == stage_id_str
    ]

    # Find incoming edges (Receiver logic)
    incoming = [
        (from_stage, spec)
        for (from_stage, to_stage), spec in transfer_cfg.connectors.items()
        if to_stage == stage_id_str
    ]

    omni_conn_cfg = None
    omni_from = None
    omni_to = None

    # Prioritize outgoing (Sender) if exists, else check incoming (Receiver).
    # Inject direction-specific role so the connector initializes correctly.
    if outgoing:
        if len(outgoing) > 1:
            logger.debug(
                "Stage-%s has %d outgoing edges; using the smallest to_stage",
                stage_id,
                len(outgoing),
            )
        outgoing.sort(key=lambda x: int(x[0]) if str(x[0]).isdigit() else str(x[0]))
        to_s, spec = outgoing[0]
        omni_conn_cfg = {"type": spec.name, **(spec.extra or {})}
        omni_conn_cfg.setdefault("role", "sender")
        omni_from = stage_id_str
        omni_to = str(to_s)
    elif incoming:
        # For receiver, pick one incoming edge to configure the connector
        incoming.sort(key=lambda x: int(x[0]) if str(x[0]).isdigit() else str(x[0]))
        from_s, spec = incoming[0]
        omni_conn_cfg = {"type": spec.name, **(spec.extra or {})}
        omni_conn_cfg.setdefault("role", "receiver")
        omni_from = str(from_s)
        omni_to = stage_id_str

    return omni_conn_cfg, omni_from, omni_to