pegaflow-llm 0.0.2__cp310-cp310-manylinux_2_34_x86_64.whl

pegaflow/__init__.py

@@ -0,0 +1,21 @@
+"""PegaFlow - High-performance key-value storage engine with Python bindings.
+
+This package provides:
+1. PegaEngine: Rust-based high-performance KV storage (via PyO3)
+2. PegaKVConnector: vLLM KV connector for distributed inference
+"""
+
+# Import Rust-based PegaEngine from the compiled extension
+try:
+    from .pegaflow import EngineRpcClient, PegaEngine, PyLoadState
+except ImportError:
+    # Fallback for development when the Rust extension is not built
+    EngineRpcClient = None
+    PegaEngine = None
+    PyLoadState = None
+
+# Import Python-based vLLM connector
+from .connector import PegaKVConnector, KVConnectorRole
+
+__version__ = "0.0.1"
+__all__ = ["PegaEngine", "EngineRpcClient", "PyLoadState", "PegaKVConnector", "KVConnectorRole"]

pegaflow/_server.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+"""
+Wrapper script to launch pegaflow-server binary from the installed package.
+"""
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+def get_server_binary():
+    """Locate the pegaflow-server-py binary in the installed package."""
+    # The binary is in the same directory as this Python module
+    module_dir = Path(__file__).parent
+    binary_path = module_dir / "pegaflow-server-py"
+
+    if binary_path.exists() and binary_path.is_file():
+        return str(binary_path)
+
+    # Fallback: try to find in PATH
+    return "pegaflow-server-py"
+
+
+def main():
+    """Launch pegaflow-server with command-line arguments."""
+    server_binary = get_server_binary()
+
+    try:
+        # Pass through all command-line arguments
+        result = subprocess.run(
+            [server_binary] + sys.argv[1:],
+            check=False,
+        )
+        sys.exit(result.returncode)
+    except FileNotFoundError:
+        print(f"Error: pegaflow-server-py binary not found at {server_binary}", file=sys.stderr)
+        print("Please ensure pegaflow is properly installed.", file=sys.stderr)
+        sys.exit(1)
+    except KeyboardInterrupt:
+        sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

pegaflow/connector/__init__.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+"""
+Facade for the PegaFlow vLLM connector, split into scheduler/worker implementations.
+"""
+
+import os
+import torch
+from typing import Any, Iterable, Optional, Tuple
+
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorRole,
+)
+from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
+
+from pegaflow.connector.common import (
+    ConnectorContext,
+    ENGINE_ENDPOINT,
+    PegaConnectorMetadata,
+    RequestPhase,
+    RequestTracker,
+    derive_namespace,
+    logger,
+    resolve_instance_id,
+)
+from pegaflow.connector.scheduler import SchedulerConnector
+from pegaflow.connector.worker import WorkerConnector
+from pegaflow.logging_utils import timing_wrapper
+from pegaflow.pegaflow import EngineRpcClient
+
+
+class PegaKVConnector(KVConnectorBase_V1):
+    """v1 KV connector for PegaFlow with separated scheduler/worker logic."""
+
+    def __init__(self, vllm_config, role: KVConnectorRole):
+        super().__init__(vllm_config, role)
+
+        instance_id = resolve_instance_id(vllm_config)
+        tp_size = vllm_config.parallel_config.tensor_parallel_size
+        namespace = derive_namespace(vllm_config, tp_size)
+        num_layers = getattr(vllm_config.model_config.hf_text_config,
+                             "num_hidden_layers", 0)
+        block_size = vllm_config.cache_config.block_size
+
+        tp_rank: Optional[int] = None
+        device_id: Optional[int] = None
+        if role == KVConnectorRole.WORKER:
+            tp_rank = get_tensor_model_parallel_rank()
+            if torch.cuda.is_available():
+                device_id = _resolve_device_id()
+
+        self._engine_endpoint = ENGINE_ENDPOINT
+        engine_client = EngineRpcClient(self._engine_endpoint)
+        logger.info("[PegaKVConnector] Connected to engine server at %s",
+                    self._engine_endpoint)
+
+        self._ctx = ConnectorContext(
+            instance_id=instance_id,
+            namespace=namespace,
+            block_size=block_size,
+            num_layers=num_layers,
+            tp_size=tp_size,
+            tp_rank=tp_rank,
+            device_id=device_id,
+            engine_client=engine_client,
+        )
+
+        self._scheduler: SchedulerConnector | None = None
+        self._worker: WorkerConnector | None = None
+        if role == KVConnectorRole.SCHEDULER:
+            self._scheduler = SchedulerConnector(self._ctx)
+        else:
+            self._worker = WorkerConnector(self._ctx)
+
+        logger.info(
+            "[PegaKVConnector] Initialized role=%s instance_id=%s device=%s tp_rank=%s tp_size=%d layers=%d namespace=%s",
+            role.name,
+            instance_id,
+            device_id if device_id is not None else "cpu",
+            tp_rank if tp_rank is not None else "N/A",
+            tp_size,
+            num_layers,
+            namespace,
+        )
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+    @timing_wrapper
+    def start_load_kv(self, forward_context, **kwargs: Any) -> None:
+        if not self._worker:
+            return
+        metadata = self._get_connector_metadata()
+        if metadata is None:
+            return
+        self._worker.start_load_kv(metadata, forward_context, **kwargs)
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        if not self._worker:
+            return
+        self._worker.wait_for_layer_load(layer_name)
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: "torch.Tensor",
+        attn_metadata,
+        **kwargs: Any,
+    ) -> None:
+        if not self._worker:
+            return
+        metadata = self._get_connector_metadata()
+        if metadata is None:
+            return
+        self._worker.save_kv_layer(metadata, layer_name, kv_layer,
+                                   attn_metadata, **kwargs)
+
+    @timing_wrapper
+    def wait_for_save(self) -> None:
+        if not self._worker:
+            return
+        self._worker.wait_for_save()
+
+    def get_finished(
+            self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        if not self._worker:
+            return (None, None)
+        return self._worker.get_finished(finished_req_ids)
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        if not self._worker:
+            return
+        self._worker.register_kv_caches(kv_caches)
+
+    def unregister_context(self) -> None:
+        if self._worker:
+            self._worker.unregister_context()
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+    def update_connector_output(self, connector_output) -> None:
+        if self._scheduler:
+            self._scheduler.update_connector_output(connector_output)
+
+    def request_finished(
+        self,
+        request,
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        if self._scheduler:
+            return self._scheduler.request_finished(request, block_ids)
+        return (False, None)
+
+    def take_events(self) -> Iterable:
+        return ()
+
+    @timing_wrapper
+    def get_num_new_matched_tokens(
+        self,
+        request,
+        num_computed_tokens: int,
+    ) -> Tuple[Optional[int], bool]:
+        if not self._scheduler:
+            return (0, False)
+        return self._scheduler.get_num_new_matched_tokens(
+            request, num_computed_tokens)
+
+    @timing_wrapper
+    def update_state_after_alloc(
+        self,
+        request,
+        blocks,
+        num_external_tokens: int,
+    ) -> None:
+        if self._scheduler:
+            self._scheduler.update_state_after_alloc(
+                request, blocks, num_external_tokens)
+
+    @timing_wrapper
+    def build_connector_meta(self, scheduler_output) -> PegaConnectorMetadata:
+        if not self._scheduler:
+            return PegaConnectorMetadata()
+        return self._scheduler.build_connector_meta(scheduler_output)
+
+    # ==============================
+    # Defaults and shutdown
+    # ==============================
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        return set()
+
+    def get_kv_connector_stats(self):
+        return None
+
+    def get_handshake_metadata(self):
+        return None
+
+    def set_host_xfer_buffer_ops(self, copy_operation):
+        return
+
+    def get_finished_count(self) -> int | None:
+        return None
+
+    def shutdown(self):
+        if self._worker:
+            self._worker.shutdown()
+
+
+def _resolve_device_id() -> int:
+    """
+    Return the global CUDA device id even when CUDA_VISIBLE_DEVICES masks GPUs.
+
+    torch.cuda.current_device() returns the local index within the visible set,
+    but we need the actual global device ID for operations like CUDA IPC.
+    This function maps the local index back to the global device ID.
+    """
+    local_id = torch.cuda.current_device()
+    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
+    if not visible:
+        return local_id
+
+    slots = [slot.strip() for slot in visible.split(",") if slot.strip()]
+    try:
+        mapped = slots[local_id]
+    except IndexError:
+        return local_id
+
+    try:
+        return int(mapped)
+    except ValueError:
+        return local_id
+
+
+__all__ = ["PegaKVConnector", "KVConnectorRole"]

pegaflow/connector/common.py

@@ -0,0 +1,308 @@
+"""
+Shared types and helpers for the PegaFlow vLLM connector.
+"""
+
+import enum
+import hashlib
+import os
+import uuid
+from dataclasses import dataclass, field
+
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
+
+from pegaflow.logging_utils import get_connector_logger
+from pegaflow.pegaflow import EngineRpcClient
+
+logger = get_connector_logger()
+
+# Engine server endpoint (gRPC URL)
+ENGINE_ENDPOINT = os.environ.get("PEGAFLOW_ENGINE_ENDPOINT",
+                                 "http://127.0.0.1:50055")
+
+
+@dataclass(frozen=True)
+class ConnectorContext:
+    """Shared configuration for scheduler/worker connectors."""
+    instance_id: str
+    namespace: str
+    block_size: int
+    num_layers: int
+    tp_size: int
+    tp_rank: int | None
+    device_id: int | None
+    engine_client: EngineRpcClient
+
+
+class RequestPhase(enum.Enum):
+    """Lifecycle phase of a request in the KV connector."""
+    LOOKUP = "lookup"  # Waiting for lookup result from external storage
+    LOADING = "loading"  # Need to load KV from external storage
+    ACTIVE = "active"  # Actively generating (may be saving concurrently)
+    DRAINING = "draining"  # Generation done, waiting for async save to complete
+    DONE = "done"  # Fully completed
+
+
+@dataclass(frozen=True)
+class LoadIntent:
+    """Intent for a KV load operation."""
+    block_ids: tuple[int, ...]
+    block_hashes: tuple[bytes, ...]
+    num_tokens: int
+
+
+@dataclass(slots=True)
+class LoadState:
+    """
+    Mutable state for an in-progress load operation.
+
+    Lifecycle:
+    - Created by on_lookup() when cache hit is detected
+    - Updated by on_alloc() with allocated block IDs
+    - Consumed by consume_load_intent() which returns LoadIntent and clears state
+    """
+    hit_blocks: int
+    computed_blocks: int
+    allocated_blocks: list[int] = field(default_factory=list)
+    external_tokens: int = 0
+
+    def to_intent(
+        self,
+        block_hashes: tuple[bytes, ...],
+        block_size: int,
+    ) -> LoadIntent | None:
+        """
+        Convert to LoadIntent if conditions are met.
+
+        Returns None if:
+        - No external tokens to load
+        - All hits are already computed (prefix cache)
+        - No blocks to load after accounting for computed blocks
+        """
+        if self.external_tokens <= 0 or self.hit_blocks <= self.computed_blocks:
+            return None
+
+        num_blocks = min(self.hit_blocks, len(self.allocated_blocks), len(block_hashes))
+        load_blocks = num_blocks - self.computed_blocks
+        if load_blocks <= 0:
+            return None
+
+        start = self.computed_blocks
+        return LoadIntent(
+            block_ids=tuple(self.allocated_blocks[start:start + load_blocks]),
+            block_hashes=block_hashes[start:start + load_blocks],
+            num_tokens=load_blocks * block_size,
+        )
+
+
+@dataclass(frozen=True)
+class SaveIntent:
+    """Intent for a KV save operation."""
+    block_ids: tuple[int, ...]
+    block_hashes: tuple[bytes, ...]
+
+
+class RequestTracker:
+    """
+    Tracks the KV cache state for a single request.
+
+    Load state lifecycle:
+    - on_lookup() creates LoadState (or None if no hit)
+    - on_alloc() updates LoadState with allocated blocks
+    - consume_load_intent() returns LoadIntent and clears LoadState
+    - Preemption: next on_lookup() replaces stale LoadState
+    """
+
+    __slots__ = (
+        'request_id',
+        '_block_hashes',
+        '_block_size',
+        '_load',
+        '_allocated_blocks',
+        '_scheduled_tokens',
+        '_stored_blocks',
+        '_total_layers',
+        '_saved_layers',
+        '_finished',
+    )
+
+    def __init__(
+        self,
+        request_id: str,
+        block_hashes: list[bytes],
+        block_size: int,
+        num_layers: int,
+    ):
+        self.request_id = request_id
+        self._block_hashes = tuple(block_hashes)
+        self._block_size = block_size
+        self._load: LoadState | None = None
+        self._allocated_blocks: list[int] = []
+        self._scheduled_tokens: int = 0
+        self._stored_blocks: int = 0
+        self._total_layers = num_layers
+        self._saved_layers: int = 0
+        self._finished: bool = False
+
+    @property
+    def phase(self) -> RequestPhase:
+        if self._load is not None:
+            return RequestPhase.LOADING
+        if not self._finished:
+            return RequestPhase.ACTIVE
+        if self._saved_layers < self._total_layers:
+            return RequestPhase.DRAINING
+        return RequestPhase.DONE
+
+    @property
+    def num_blocks(self) -> int:
+        return len(self._block_hashes)
+
+    def on_lookup(self, hit_blocks: int, computed_blocks: int) -> None:
+        """New lookup = fresh load state. Handles preemption implicitly."""
+        self._load = (
+            LoadState(hit_blocks=hit_blocks, computed_blocks=computed_blocks)
+            if hit_blocks > computed_blocks
+            else None
+        )
+        self._allocated_blocks = []
+
+    def on_alloc(self, block_ids: list[int], num_external_tokens: int) -> None:
+        self._allocated_blocks.extend(block_ids)
+        if self._load is not None:
+            self._load.allocated_blocks.extend(block_ids)
+            if num_external_tokens > 0:
+                self._load.external_tokens = num_external_tokens
+
+    def consume_load_intent(self) -> LoadIntent | None:
+        load, self._load = self._load, None
+        if load is None:
+            return None
+        return load.to_intent(self._block_hashes, self._block_size)
+
+    def on_scheduled(self, num_tokens: int) -> None:
+        self._scheduled_tokens += num_tokens
+
+    def on_layer_saved(self) -> None:
+        self._saved_layers += 1
+
+    def on_finished(self) -> None:
+        self._finished = True
+
+    def consume_save_intent(self) -> SaveIntent | None:
+        saveable = min(
+            len(self._block_hashes),
+            len(self._allocated_blocks),
+            self._scheduled_tokens // self._block_size,
+        )
+        new_blocks = saveable - self._stored_blocks
+        if new_blocks <= 0:
+            return None
+
+        start = self._stored_blocks
+        self._stored_blocks = start + new_blocks
+        return SaveIntent(
+            block_ids=tuple(self._allocated_blocks[start:self._stored_blocks]),
+            block_hashes=self._block_hashes[start:self._stored_blocks],
+        )
+
+    def should_hold_blocks(self) -> bool:
+        return (self._finished and self._stored_blocks > 0 and
+                self._saved_layers < self._total_layers)
+
+    def is_done(self) -> bool:
+        return self.phase == RequestPhase.DONE
+
+    def __repr__(self) -> str:
+        return (
+            f"RequestTracker({self.request_id}, {self.phase.value}, "
+            f"load={self._load}, alloc={len(self._allocated_blocks)}, "
+            f"stored={self._stored_blocks}, saved={self._saved_layers}/{self._total_layers})"
+        )
+
+
+class PegaConnectorMetadata(KVConnectorMetadata):
+    """Metadata passed from scheduler to worker for KV cache operations."""
+
+    def __init__(
+        self,
+        load_intents: dict[str, LoadIntent] | None = None,
+        save_intents: dict[str, SaveIntent] | None = None,
+    ):
+        super().__init__()
+        # Maps request_id -> intent
+        self.load_intents: dict[str, LoadIntent] = load_intents or {}
+        self.save_intents: dict[str, SaveIntent] = save_intents or {}
+
+    def __repr__(self) -> str:
+        return (f"PegaConnectorMetadata(loads={len(self.load_intents)}, "
+                f"saves={len(self.save_intents)})")
+
+
+def resolve_instance_id(vllm_config, dp_rank_suffix: bool = True) -> str:
+    """Resolve or generate connector instance_id with optional DP rank suffix."""
+    instance_id = vllm_config.kv_transfer_config.engine_id
+    if instance_id:
+        logger.info(
+            "[PegaKVConnector] Using kv_transfer_config.engine_id: %s",
+            instance_id)
+        return instance_id
+
+    instance_id = vllm_config.instance_id or os.environ.get(
+        "PEGAFLOW_INSTANCE_ID", "")
+    if not instance_id:
+        instance_id = uuid.uuid4().hex
+        logger.info(
+            "[PegaKVConnector] No instance_id from vLLM; generated fallback %s",
+            instance_id)
+
+    if dp_rank_suffix:
+        parallel_config = vllm_config.parallel_config
+        if parallel_config.data_parallel_size > 1:
+            local_dp_rank = parallel_config.data_parallel_rank_local
+            if local_dp_rank is not None:
+                instance_id = f"{instance_id}_dp{local_dp_rank}"
+                logger.info(
+                    "[PegaKVConnector] Appended DP rank to instance_id: %s (dp_size=%d, local_dp_rank=%d)",
+                    instance_id,
+                    parallel_config.data_parallel_size,
+                    local_dp_rank,
+                )
+
+    return instance_id
+
+
+def derive_namespace(vllm_config, tp_size: int) -> str:
+    """
+    Derive namespace for storage isolation.
+    """
+    model_config = vllm_config.model_config
+    cache_config = vllm_config.cache_config
+
+    factors = {
+        "model": model_config.model,
+        "dtype": str(model_config.dtype),
+        "tp_size": tp_size,
+        "num_kv_heads": model_config.get_total_num_kv_heads(),
+        "head_size": model_config.get_head_size(),
+        "num_hidden_layers": model_config.get_total_num_hidden_layers(),
+        "cache_dtype": str(cache_config.cache_dtype),
+    }
+
+    factor_str = str(sorted(factors.items()))
+    hash_suffix = hashlib.sha256(factor_str.encode()).hexdigest()[:8]
+    return f"{hash_suffix}"
+
+
+__all__ = [
+    "ConnectorContext",
+    "ENGINE_ENDPOINT",
+    "LoadIntent",
+    "LoadState",
+    "PegaConnectorMetadata",
+    "RequestPhase",
+    "RequestTracker",
+    "SaveIntent",
+    "derive_namespace",
+    "logger",
+    "resolve_instance_id",
+]