pegaflow-llm 0.0.2__cp310-cp310-manylinux_2_34_x86_64.whl

pegaflow/ipc_wrapper.py

@@ -0,0 +1,183 @@
+"""CUDA IPC Wrapper for cross-process GPU memory sharing.
+
+This module provides a wrapper class for PyTorch tensors that enables
+cross-process GPU memory sharing via CUDA IPC handles. The wrapper can
+be serialized (via pickle) and sent across process boundaries.
+
+This implementation handles CUDA_VISIBLE_DEVICES correctly by using GPU UUIDs
+instead of device indices for device identification.
+"""
+
+import threading
+import torch
+from typing import Tuple
+
+
+class CudaIPCWrapper:
+    """Wrapper for CUDA IPC handle with tensor metadata.
+    
+    This class wraps a PyTorch CUDA tensor and extracts its IPC handle,
+    allowing the tensor to be reconstructed in another process. It correctly
+    handles CUDA_VISIBLE_DEVICES by using GPU UUIDs for device identification.
+    
+    Attributes:
+        handle: CUDA IPC handle tuple (device, ipc_handle, size, offset, ...)
+        dtype: PyTorch dtype of the tensor
+        shape: Shape tuple of the tensor
+        device_uuid: UUID string of the GPU device
+    
+    Example:
+        # Process 1 (sender)
+        tensor = torch.randn(10, device='cuda:0')
+        wrapper = CudaIPCWrapper(tensor)
+        serialized = pickle.dumps(wrapper)
+        # ... send serialized bytes to another process ...
+        
+        # Process 2 (receiver)
+        wrapper = pickle.loads(serialized)
+        tensor = wrapper.to_tensor()  # Reconstruct tensor
+        ptr = tensor.data_ptr()  # Get GPU pointer
+    """
+    
+    # Class-level cache for device UUID to index mapping
+    _discovered_device_mapping: dict[str, int] = {}
+    _device_mapping_lock = threading.Lock()
+    
+    @staticmethod
+    def _get_device_uuid(device_index: int) -> str:
+        """Get the UUID of a GPU device given its index.
+        
+        Args:
+            device_index: CUDA device index (relative to CUDA_VISIBLE_DEVICES)
+        
+        Returns:
+            UUID string of the GPU device
+        """
+        return str(torch.cuda.get_device_properties(device_index).uuid)
+    
+    @staticmethod
+    def _discover_gpu_devices():
+        """Discover all available GPU devices and map their UUIDs to
+        the physical device ordinals (relative to CUDA_VISIBLE_DEVICES).
+        """
+        if not torch.cuda.is_available():
+            return
+        
+        num_devices = torch.cuda.device_count()
+        with CudaIPCWrapper._device_mapping_lock:
+            if CudaIPCWrapper._discovered_device_mapping:
+                return  # Already discovered
+            
+            for i in range(num_devices):
+                device_uuid = CudaIPCWrapper._get_device_uuid(i)
+                CudaIPCWrapper._discovered_device_mapping[device_uuid] = i
+    
+    @staticmethod
+    def _get_device_index_from_uuid(device_uuid: str) -> int:
+        """Get the physical device ordinal from its UUID.
+        
+        Args:
+            device_uuid: UUID string of the GPU device
+        
+        Returns:
+            Device index relative to CUDA_VISIBLE_DEVICES
+        
+        Raises:
+            RuntimeError: If the device UUID is not found
+        """
+        CudaIPCWrapper._discover_gpu_devices()
+        
+        with CudaIPCWrapper._device_mapping_lock:
+            device_index = CudaIPCWrapper._discovered_device_mapping.get(
+                device_uuid, None
+            )
+        
+        if device_index is None:
+            raise RuntimeError(
+                f"Device UUID {device_uuid} not found in the discovered devices. "
+                "Please make sure the process can see all the GPU devices."
+            )
+        return device_index
+    
+    def __init__(self, tensor: torch.Tensor):
+        """Create IPC wrapper from a CUDA tensor.
+        
+        Args:
+            tensor: PyTorch CUDA tensor to wrap. Must be contiguous and
+                   have zero storage offset.
+        
+        Raises:
+            AssertionError: If tensor is not contiguous or has non-zero offset.
+        """
+        assert tensor.storage_offset() == 0, "Tensor must have zero storage offset"
+        assert tensor.is_contiguous(), "Tensor must be contiguous"
+        
+        # Get the underlying storage and create IPC handle
+        storage = tensor.untyped_storage()
+        handle = storage._share_cuda_()
+        
+        # Store metadata needed to reconstruct the tensor
+        self.handle = handle
+        self.dtype = tensor.dtype
+        self.shape = tensor.shape
+        
+        # Store device UUID instead of device index to handle CUDA_VISIBLE_DEVICES
+        device_index = tensor.device.index
+        self.device_uuid = CudaIPCWrapper._get_device_uuid(device_index)
+    
+    def to_tensor(self) -> torch.Tensor:
+        """Reconstruct tensor from IPC handle.
+        
+        This method creates a new tensor in the current process that shares
+        the same GPU memory as the original tensor (via CUDA IPC).
+        
+        Returns:
+            PyTorch tensor that shares GPU memory with the original tensor.
+        
+        Note:
+            The reconstructed tensor shares memory with the original. Any
+            modifications to one will be visible in the other.
+            
+            This function may break if torch.cuda is not initialized.
+            Call torch.cuda.init() before using this function if needed.
+        """
+        # Get the correct device index in the current process based on UUID
+        device = CudaIPCWrapper._get_device_index_from_uuid(self.device_uuid)
+        
+        # Reconstruct storage from IPC handle
+        storage = torch.UntypedStorage._new_shared_cuda(device, *self.handle[1:])
+        
+        # Create empty tensor on the correct device
+        t = torch.tensor([], device=device, dtype=self.dtype)
+        
+        # Set the tensor to use the shared storage
+        t.set_(storage)
+        
+        # Reshape to original shape
+        return t.view(self.shape)
+    
+    def __eq__(self, other) -> bool:
+        """Check equality with another CudaIPCWrapper.
+        
+        Args:
+            other: Object to compare with
+        
+        Returns:
+            True if the wrappers refer to the same tensor, False otherwise
+        """
+        if not isinstance(other, CudaIPCWrapper):
+            return False
+        return (
+            self.handle == other.handle
+            and self.dtype == other.dtype
+            and self.shape == other.shape
+            and self.device_uuid == other.device_uuid
+        )
+    
+    def __repr__(self) -> str:
+        return (f"CudaIPCWrapper(shape={self.shape}, dtype={self.dtype}, "
+                f"device_uuid={self.device_uuid})")
+
+
+__all__ = ["CudaIPCWrapper"]
+

pegaflow/logging_utils.py

@@ -0,0 +1,61 @@
+"""Logging utilities for PegaFlow connector.
+
+This module provides timing decorators and logger configuration.
+"""
+
+import functools
+import logging
+import os
+import time
+
+# Environment variable to control timing logging
+ENABLE_TIMING = os.environ.get("PEGAFLOW_ENABLE_TIMING", "1") == "1"
+
+# Module logger
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+if not logger.hasHandlers():
+    _handler = logging.StreamHandler()
+    _handler.setLevel(logging.NOTSET)
+    _handler.setFormatter(logging.Formatter("%(message)s"))
+    logger.addHandler(_handler)
+    logger.propagate = False
+
+
+def get_connector_logger() -> logging.Logger:
+    """Get a logger for the connector module."""
+    connector_logger = logging.getLogger("pegaflow.connector")
+    connector_logger.setLevel(logging.INFO)
+    if not connector_logger.hasHandlers():
+        handler = logging.StreamHandler()
+        handler.setLevel(logging.NOTSET)
+        handler.setFormatter(logging.Formatter("%(message)s"))
+        connector_logger.addHandler(handler)
+        connector_logger.propagate = False
+    return connector_logger
+
+
+def timing_wrapper(func):
+    """Decorator to log function name and execution time when enabled.
+
+    Enable by setting environment variable: PEGAFLOW_ENABLE_TIMING=1
+    """
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if not ENABLE_TIMING:
+            return func(*args, **kwargs)
+
+        start = time.perf_counter()
+        result = func(*args, **kwargs)
+        elapsed_ms = (time.perf_counter() - start) * 1000
+        logger.debug(
+            "[PegaKVConnector] %s took %.2f ms",
+            func.__name__,
+            elapsed_ms,
+        )
+        return result
+    return wrapper
+
+
+__all__ = ["ENABLE_TIMING", "timing_wrapper", "get_connector_logger"]

pegaflow_llm-0.0.2.dist-info/METADATA

@@ -0,0 +1,100 @@
+Metadata-Version: 2.4
+Name: pegaflow-llm
+Version: 0.0.2
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Rust
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Summary: High-performance key-value storage engine with Python bindings
+Keywords: storage,kv-store,rust,vllm,inference
+Author: PegaFlow Contributors
+License: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+
+# PegaFlow Python Package
+
+High-performance key-value storage engine with Python bindings, built with Rust and PyO3.
+
+## Features
+
+- **PegaEngine**: Fast Rust-based key-value storage with Python bindings
+- **PegaKVConnector**: vLLM KV connector for distributed inference with KV cache transfer
+
+## Installation
+
+### From Source
+
+```bash
+# Install maturin if you haven't already
+pip install maturin
+
+# Build and install in development mode
+cd python
+maturin develop
+
+# Or build a wheel
+maturin build --release
+```
+
+### From PyPI (coming soon)
+
+```bash
+pip install pegaflow
+```
+
+## Usage
+
+### Basic KV Storage
+
+```python
+from pegaflow import PegaEngine
+
+# Create a new engine
+engine = PegaEngine()
+
+# Store key-value pairs
+engine.put("name", "PegaFlow")
+engine.put("version", "0.1.0")
+
+# Retrieve values
+name = engine.get("name")  # Returns "PegaFlow"
+missing = engine.get("nonexistent")  # Returns None
+
+# Remove keys
+removed = engine.remove("name")  # Returns "PegaFlow"
+```
+
+### vLLM KV Connector
+
+```python
+from vllm import LLM
+from vllm.distributed.kv_transfer.kv_transfer_agent import KVTransferConfig
+
+# Configure vLLM to use PegaKVConnector
+kv_transfer_config = KVTransferConfig(
+    kv_connector="PegaKVConnector",
+    kv_role="kv_both",
+    kv_connector_module_path="pegaflow.connector",
+)
+
+# Create LLM with KV transfer enabled
+llm = LLM(
+    model="gpt2",
+    kv_transfer_config=kv_transfer_config,
+)
+```
+
+## Development
+
+See the [examples](../examples/) directory for more usage examples.
+
+## License
+
+MIT
+
+

pegaflow_llm-0.0.2.dist-info/RECORD

@@ -0,0 +1,15 @@
+pegaflow/__init__.py,sha256=F6E-sxO1kxm1WyJZsY4Eba3muE9Z3IYxG1RxWbGFlyU,741
+pegaflow/_server.py,sha256=CL0HSfBAP9rx5Rb4oh77FwFhxfcwQr9pSrEX4IlPBYg,1224
+pegaflow/connector copy.py,sha256=ELvaMvMTOfA4g48IZ8rjmKzwQRN3g5tUPOH99_fGA4w,35994
+pegaflow/connector/__init__.py,sha256=6kr5KwOR5GpU06BMGdw0RSwTvHk4EtASsp3DZ8rR88U,7315
+pegaflow/connector/common.py,sha256=2Qj3-H4-XNd8o3mORtYONp7N3OjTuO8t1rUYSGc-trA,9769
+pegaflow/connector/scheduler.py,sha256=7JQ6bB2C9Titd6oFTsv3PCLuqPIc4ZctB3fAnFFjxVU,7490
+pegaflow/connector/worker.py,sha256=EWJfx5UmFT-RSYUL1-AV1dPQTFiGF2YvGpEXfVvA7jU,15782
+pegaflow/ipc_wrapper.py,sha256=r5OHvTLBpyyBgr_1lyt5yPCBDS6D6aeqNvOY5Kjm0Q4,6600
+pegaflow/logging_utils.py,sha256=D3el9wRKjLQhfmK6rJcr0laUy4ey0CZEuRrLDrQxdh4,1798
+pegaflow/pegaflow-server-py,sha256=VWfPVGfy5cOarvWgR9kPkQJxl249MUtj3vf2_sQPNFw,9505568
+pegaflow/pegaflow.cpython-310-x86_64-linux-gnu.so,sha256=Qp_jaEa7bUsHa9uthOnREgTzbDb35_QuImqdib7T5xU,5556608
+pegaflow_llm-0.0.2.dist-info/METADATA,sha256=EdFrIDevJdZuuASVdhZTP59Ch1JVymVpYRZWj3oZGDI,2232
+pegaflow_llm-0.0.2.dist-info/WHEEL,sha256=hwQ1X0enL4h--Y0T24JPMOc_YcS3ZMkjluo8hhmik0c,109
+pegaflow_llm-0.0.2.dist-info/entry_points.txt,sha256=ayii4bkhIRmNW0s6u0ys3ybZ27zGhrKVnUUWSY04Zi0,56
+pegaflow_llm-0.0.2.dist-info/RECORD,,

pegaflow_llm-0.0.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: maturin (1.10.2)
+Root-Is-Purelib: false
+Tag: cp310-cp310-manylinux_2_34_x86_64

pegaflow_llm-0.0.2.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+pegaflow-server=pegaflow._server:main