sglang 0.4.9.post6__py3-none-any.whl → 0.4.10.post1__py3-none-any.whl

sglang/srt/mem_cache/nixl/nixl_utils.py

@@ -0,0 +1,238 @@
+import logging
+import os
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class NixlBackendSelection:
+    """Handles NIXL backend selection and creation."""
+
+    # Priority order for File-based plugins in case of auto selection
+    FILE_PLUGINS = ["3FS", "POSIX", "GDS_MT", "GDS"]
+    # Priority order for File-based plugins in case of auto selection (add more as needed)
+    OBJ_PLUGINS = ["OBJ"]  # Based on Amazon S3 SDK
+
+    def __init__(self, plugin: str = "auto"):
+        """Initialize backend selection.
+        Args:
+            plugin: Plugin to use (default "auto" selects best available).
+                   Can be a file plugin (3FS, POSIX, GDS, GDS_MT) or
+                   an object plugin (OBJ).
+        """
+        self.plugin = plugin
+        self.backend_name = None
+        self.mem_type = None
+
+    def set_bucket(self, bucket_name: str) -> None:
+        """Set AWS bucket name in environment variable."""
+        os.environ["AWS_DEFAULT_BUCKET"] = bucket_name
+        logger.debug(f"Set AWS bucket name to: {bucket_name}")
+
+    def create_backend(self, agent) -> bool:
+        """Create the appropriate NIXL backend based on configuration."""
+        try:
+            plugin_list = agent.get_plugin_list()
+            logger.debug(f"Available NIXL plugins: {plugin_list}")
+
+            # Handle explicit plugin selection or auto priority
+            if self.plugin == "auto":
+                # Try all file plugins first
+                for plugin in self.FILE_PLUGINS:
+                    if plugin in plugin_list:
+                        self.backend_name = plugin
+                        break
+                # If no file plugin found, try object plugins
+                if not self.backend_name:
+                    for plugin in self.OBJ_PLUGINS:
+                        if plugin in plugin_list:
+                            self.backend_name = plugin
+                            break
+            else:
+                # Use explicitly requested plugin
+                self.backend_name = self.plugin
+
+            if self.backend_name not in plugin_list:
+                logger.error(
+                    f"Backend {self.backend_name} not available in plugins: {plugin_list}"
+                )
+                return False
+
+            # Create backend and set memory type
+            if self.backend_name in self.OBJ_PLUGINS:
+                bucket = os.environ.get("AWS_DEFAULT_BUCKET")
+                if not bucket:
+                    logger.error(
+                        "AWS_DEFAULT_BUCKET environment variable must be set for object storage"
+                    )
+                    return False
+                agent.create_backend(self.backend_name, {"bucket": bucket})
+            else:
+                agent.create_backend(self.backend_name)
+
+            self.mem_type = "OBJ" if self.backend_name in self.OBJ_PLUGINS else "FILE"
+            logger.debug(
+                f"Created NIXL backend: {self.backend_name} with memory type: {self.mem_type}"
+            )
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to create NIXL backend: {e}")
+            return False
+
+
+class NixlRegistration:
+    """Handles NIXL memory registration."""
+
+    def __init__(self, agent):
+        self.agent = agent
+
+    def create_query_tuples(
+        self, key: str, mem_type: str, file_manager=None
+    ) -> List[Tuple]:
+        """Create NIXL tuples for querying memory.
+        Args:
+            key: Key to query (file path for FILE or object key for OBJ)
+            mem_type: Memory type ("FILE" or "OBJ")
+            file_manager: Optional NixlFileManager for FILE memory type
+        Returns:
+            List of NIXL tuples for querying
+        """
+        if mem_type == "FILE":
+            if file_manager is None:
+                logger.error("file_manager required for FILE memory type")
+                return []
+            return [(0, 0, 0, file_manager.get_file_path(key))]
+        else:  # OBJ
+            return [(0, 0, key)]
+
+    def _register_memory(
+        self, items: Union[List[tuple], List[torch.Tensor]], mem_type: str, desc: str
+    ) -> Optional[Any]:
+        """Common registration logic for files, objects, and buffers.
+        Args:
+            items: List of tuples or tensors to register
+            mem_type: Memory type ("FILE", "OBJ", "DRAM", "VRAM")
+            desc: Description for logging
+        """
+        try:
+            if not items:
+                return None
+
+            reg_descs = self.agent.get_reg_descs(items, mem_type)
+            if reg_descs is None:
+                logger.error("Failed to create registration descriptors")
+                return None
+
+            registered_memory = self.agent.register_memory(reg_descs)
+            if registered_memory:
+                return registered_memory
+            else:
+                logger.error("Failed to register with NIXL")
+                return None
+
+        except Exception as e:
+            logger.error(f"Failed to register {desc}: {e}")
+            return None
+
+    def register_buffers(
+        self, buffers: Union[torch.Tensor, List[torch.Tensor]]
+    ) -> Optional[Any]:
+        """Register tensors/buffers with NIXL."""
+        if isinstance(buffers, torch.Tensor):
+            buffers = [buffers]
+
+        if not buffers:
+            return None
+
+        # Determine memory type based on tensor device
+        mem_type = "VRAM" if buffers[0].device.type == "cuda" else "DRAM"
+        return self._register_memory(buffers, mem_type, "buffers")
+
+    def register_files(self, tuples: List[tuple]) -> Optional[Any]:
+        """Register files with NIXL using (0, 0, fd, file_path) tuples."""
+        return self._register_memory(tuples, "FILE", "files")
+
+    def register_objects(
+        self, keys: List[str], tensors: Optional[List[torch.Tensor]] = None
+    ) -> Optional[Any]:
+        """Register objects with NIXL."""
+        if not keys:
+            return None
+
+        # Create object tuples with proper sizes
+        tuples = [
+            (0, tensor.element_size() * tensor.numel() if tensor else 0, key)
+            for key, tensor in zip(keys, tensors or [None] * len(keys))
+        ]
+        return self._register_memory(tuples, "OBJ", "objects")
+
+
+class NixlFileManager:
+    """Handles file system operations for NIXL."""
+
+    def __init__(self, base_dir: str):
+        """
+        Initialize file manager.
+        Args:
+            base_dir: Base directory for storing tensor files
+        """
+        self.base_dir = base_dir
+        if base_dir == "":
+            logger.debug(f"Initialized file manager without a base directory")
+        else:
+            os.makedirs(base_dir, exist_ok=True)
+            logger.debug(f"Initialized file manager with base directory: {base_dir}")
+
+    def get_file_path(self, key: str) -> str:
+        """Get full file path for a given key."""
+        return os.path.join(self.base_dir, key)
+
+    def create_file(self, file_path: str) -> bool:
+        """Create a file if it doesn't exist."""
+        try:
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            if not os.path.exists(file_path):
+                with open(file_path, "wb") as f:
+                    pass  # Create empty file
+            return True
+        except Exception as e:
+            logger.error(f"Failed to create file {file_path}: {e}")
+            return False
+
+    def open_file(self, file_path: str) -> Optional[int]:
+        """Open a file and return its file descriptor."""
+        try:
+            fd = os.open(file_path, os.O_RDWR)
+            return fd
+        except Exception as e:
+            logger.error(f"Failed to open file {file_path}: {e}")
+            return None
+
+    def close_file(self, fd: int) -> bool:
+        """Close a file descriptor."""
+        try:
+            os.close(fd)
+            return True
+        except Exception as e:
+            logger.error(f"Failed to close file descriptor {fd}: {e}")
+            return False
+
+    def files_to_nixl_tuples(
+        self, file_paths: List[str], open_file: bool = True
+    ) -> List[Tuple[int, int, int, str]]:
+        """Create NIXL tuples (offset, length, fd, file_path) for given files."""
+        if not open_file:
+            return [(0, 0, 0, path) for path in file_paths]
+
+        tuples = []
+        for path in file_paths:
+            if (fd := self.open_file(path)) is None:
+                # Clean up on failure
+                for t in tuples:
+                    self.close_file(t[2])
+                return []
+            tuples.append((0, 0, fd, path))
+        return tuples

sglang/srt/mem_cache/nixl/test_hicache_nixl_storage.py

@@ -0,0 +1,216 @@
+#!/usr/bin/env python3
+
+import os
+import unittest
+from typing import List, Optional
+from unittest.mock import MagicMock
+
+import torch
+
+from sglang.srt.mem_cache.nixl.hicache_nixl import HiCacheNixl
+from sglang.srt.mem_cache.nixl.nixl_utils import NixlFileManager, NixlRegistration
+
+
+class TestNixlUnified(unittest.TestCase):
+    """Unified test suite for all NIXL components."""
+
+    def setUp(self):
+        """Set up test environment."""
+        # Create test directories
+        self.test_dir = "/tmp/test_nixl_unified"
+        os.makedirs(self.test_dir, exist_ok=True)
+
+        # Mock NIXL agent for registration tests
+        self.mock_agent = MagicMock()
+        self.mock_agent.get_reg_descs.return_value = "mock_reg_descs"
+        self.mock_agent.register_memory.return_value = "mock_registered_memory"
+
+        # Create instances
+        self.file_manager = NixlFileManager(self.test_dir)
+        self.registration = NixlRegistration(self.mock_agent)
+        try:
+            self.hicache = HiCacheNixl(file_path=self.test_dir, plugin="POSIX")
+        except ImportError:
+            self.skipTest("NIXL not available, skipping NIXL storage tests")
+
+    def tearDown(self):
+        """Clean up test directories."""
+        if os.path.exists(self.test_dir):
+            import shutil
+
+            shutil.rmtree(self.test_dir)
+
+    def delete_test_file(self, file_path: str) -> bool:
+        """Helper method to delete a test file.
+
+        Args:
+            file_path: Path to the file to delete
+
+        Returns:
+            bool: True if file was deleted or didn't exist, False on error
+        """
+        try:
+            if os.path.exists(file_path):
+                os.remove(file_path)
+            return True
+        except Exception as e:
+            return False
+
+    def verify_tensors_equal(self, expected: torch.Tensor, actual: torch.Tensor):
+        """Helper to verify tensor equality."""
+        self.assertIsNotNone(actual, "Retrieved tensor is None")
+        self.assertTrue(
+            torch.allclose(expected, actual, atol=1e-6),
+            f"Tensors not equal:\nExpected: {expected}\nActual: {actual}",
+        )
+
+    def verify_tensor_lists_equal(
+        self, expected: List[torch.Tensor], actual: List[torch.Tensor]
+    ):
+        """Helper to verify lists of tensors are equal."""
+        self.assertEqual(len(expected), len(actual), "Lists have different lengths")
+        for exp, act in zip(expected, actual):
+            self.verify_tensors_equal(exp, act)
+
+    # ============================================================================
+    # HiCache Integration Tests
+    # ============================================================================
+
+    def test_single_set_get(self):
+        """Test single tensor set/get operations."""
+        key = "test_key"
+        value = torch.randn(10, 10, device="cpu")
+        dst_tensor = torch.zeros_like(value, device="cpu")
+
+        # Test set
+        self.assertTrue(self.hicache.set(key, value))
+        self.assertTrue(self.hicache.exists(key))
+
+        # Test get
+        retrieved = self.hicache.get(key, dst_tensor)
+        self.verify_tensors_equal(value, retrieved)
+
+    def test_batch_set_get(self):
+        """Test batch tensor set/get operations."""
+        keys = ["key1", "key2", "key3"]
+        values = [
+            torch.randn(5, 5, device="cpu"),
+            torch.randn(3, 3, device="cpu"),
+            torch.randn(7, 7, device="cpu"),
+        ]
+        dst_tensors = [torch.zeros_like(v, device="cpu") for v in values]
+
+        # Test batch set
+        self.assertTrue(self.hicache.batch_set(keys, values))
+        self.assertTrue(all(self.hicache.exists(key) for key in keys))
+
+        # Test batch get
+        retrieved = self.hicache.batch_get(keys, dst_tensors)
+        self.verify_tensor_lists_equal(values, retrieved)
+
+    def test_mixed_operations(self):
+        """Test mixing single and batch operations."""
+        # Test interleaved set/get operations
+        key1, key2 = "key1", "key2"
+        value1 = torch.randn(4, 4, device="cpu")
+        value2 = torch.randn(6, 6, device="cpu")
+        dst1 = torch.zeros_like(value1)
+        dst2 = torch.zeros_like(value2)
+
+        # Single set/get
+        self.assertTrue(self.hicache.set(key1, value1))
+        retrieved1 = self.hicache.get(key1, dst1)
+        self.verify_tensors_equal(value1, retrieved1)
+
+        # Batch set/get
+        self.assertTrue(self.hicache.batch_set([key2], [value2]))
+        retrieved2 = self.hicache.batch_get([key2], [dst2])
+        self.verify_tensors_equal(value2, retrieved2[0])
+
+    def test_data_integrity(self):
+        """Test data integrity across operations."""
+        # Test with various tensor types and sizes
+        test_cases = [
+            ("float32", torch.randn(10, 10, dtype=torch.float32)),
+            ("float64", torch.randn(5, 5, dtype=torch.float64)),
+            ("int32", torch.randint(-100, 100, (8, 8), dtype=torch.int32)),
+            ("int64", torch.randint(-100, 100, (6, 6), dtype=torch.int64)),
+            ("bool", torch.randint(0, 2, (4, 4)).bool()),
+        ]
+
+        for name, tensor in test_cases:
+            with self.subTest(tensor_type=name):
+                key = f"test_{name}"
+                dst_tensor = torch.zeros_like(tensor)
+
+                # Set and immediately get
+                self.assertTrue(self.hicache.set(key, tensor))
+                retrieved1 = self.hicache.get(key, dst_tensor)
+                self.verify_tensors_equal(tensor, retrieved1)
+
+                # Get again to verify persistence
+                dst_tensor.zero_()
+                retrieved2 = self.hicache.get(key, dst_tensor)
+                self.verify_tensors_equal(tensor, retrieved2)
+
+    def test_basic_file_operations(self):
+        """Test basic file operations."""
+        test_file = os.path.join(self.test_dir, "test_file.bin")
+        self.file_manager.create_file(test_file)
+        self.assertTrue(os.path.exists(test_file))
+        self.assertEqual(os.path.getsize(test_file), 0)  # Empty file
+
+        # Test file deletion
+        self.assertTrue(self.delete_test_file(test_file))
+        self.assertFalse(os.path.exists(test_file))
+
+    def test_create_nixl_tuples(self):
+        """Test creation of NIXL tuples."""
+        test_file = os.path.join(self.test_dir, "test_file.bin")
+        self.file_manager.create_file(test_file)
+
+        # Test tuple creation
+        tuples = self.file_manager.files_to_nixl_tuples([test_file], False)
+        self.assertIsNotNone(tuples)
+        self.assertTrue(len(tuples) > 0)
+
+    def test_error_handling(self):
+        """Test error handling in file operations."""
+        # Test non-existent file
+        self.assertTrue(
+            self.delete_test_file("nonexistent_file.bin")
+        )  # Returns True if file doesn't exist
+
+        # Test invalid file path
+        self.assertFalse(self.file_manager.create_file(""))  # Empty path should fail
+
+    def test_register_buffers(self):
+        """Test registration of memory buffers."""
+        # Create test tensor
+        tensor = torch.randn(10, 10)
+
+        # Test buffer registration
+        self.assertIsNotNone(self.registration.register_buffers(tensor))
+
+        # Test batch registration
+        tensors = [torch.randn(5, 5) for _ in range(3)]
+        self.assertIsNotNone(self.registration.register_buffers(tensors))
+
+    def test_register_files_with_tuples(self):
+        """Test registration of files using NIXL tuples."""
+        files = [os.path.join(self.test_dir, f"test_file_{i}.bin") for i in range(3)]
+        for file in files:
+            self.file_manager.create_file(file)
+
+        # Create tuples and register
+        tuples = self.file_manager.files_to_nixl_tuples(files, False)
+        self.registration.register_files(tuples)
+
+        # Verify tuples
+        self.assertEqual(len(tuples), len(files))
+        for t, f in zip(tuples, files):
+            self.assertEqual(t[3], f)  # Check file path
+
+
+if __name__ == "__main__":
+    unittest.main()

sglang/srt/mem_cache/storage/hf3fs/client_hf3fs.py

@@ -0,0 +1,183 @@
+import logging
+import multiprocessing
+import os
+import threading
+from functools import wraps
+from pathlib import Path
+from typing import List
+
+import torch
+from torch.utils.cpp_extension import load
+
+root = Path(__file__).parent.resolve()
+hf3fs_utils = load(name="hf3fs_utils", sources=[f"{root}/hf3fs_utils.cpp"])
+
+logger = logging.getLogger(__name__)
+
+HF3FS_AVAILABLE = True
+try:
+    from hf3fs_fuse.io import (
+        deregister_fd,
+        extract_mount_point,
+        make_ioring,
+        make_iovec,
+        register_fd,
+    )
+except ImportError:
+    HF3FS_AVAILABLE = False
+
+
+def rsynchronized():
+    def _decorator(func):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            with self.rlock:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return _decorator
+
+
+def wsynchronized():
+    def _decorator(func):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            with self.wlock:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return _decorator
+
+
+class Hf3fsClient:
+    def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
+        if not HF3FS_AVAILABLE:
+            raise ImportError(
+                "hf3fs_fuse.io is not available. Please install the hf3fs_fuse package."
+            )
+
+        self.path = path
+        self.size = size
+        self.bytes_per_page = bytes_per_page
+        self.entries = entries
+
+        self.file = os.open(self.path, os.O_RDWR | os.O_CREAT)
+        os.ftruncate(self.file, size)
+        register_fd(self.file)
+
+        self.hf3fs_mount_point = extract_mount_point(path)
+        self.bs = self.bytes_per_page
+        self.shm_r = multiprocessing.shared_memory.SharedMemory(
+            size=self.bs * self.entries, create=True
+        )
+        self.shm_w = multiprocessing.shared_memory.SharedMemory(
+            size=self.bs * self.entries, create=True
+        )
+
+        self.shm_r_tensor = torch.frombuffer(self.shm_r.buf, dtype=torch.uint8)
+        self.shm_w_tensor = torch.frombuffer(self.shm_w.buf, dtype=torch.uint8)
+
+        self.numa = -1
+        self.ior_r = make_ioring(
+            self.hf3fs_mount_point,
+            self.entries,
+            for_read=True,
+            timeout=1,
+            numa=self.numa,
+        )
+        self.ior_w = make_ioring(
+            self.hf3fs_mount_point,
+            self.entries,
+            for_read=False,
+            timeout=1,
+            numa=self.numa,
+        )
+        self.iov_r = make_iovec(self.shm_r, self.hf3fs_mount_point)
+        self.iov_w = make_iovec(self.shm_w, self.hf3fs_mount_point)
+
+        self.rlock = threading.RLock()
+        self.wlock = threading.RLock()
+
+    @rsynchronized()
+    def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        self.check(offsets, tensors)
+
+        # prepare
+        current = 0
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+            self.ior_r.prepare(
+                self.iov_r[current : current + size], True, self.file, offset
+            )
+            current += size
+
+        # submit
+        ionum = len(offsets)
+        resv = self.ior_r.submit().wait(min_results=ionum)
+
+        # results
+        hf3fs_utils.read_shm(self.shm_r_tensor, tensors)
+        results = [res.result for res in resv]
+
+        return results
+
+    @wsynchronized()
+    def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        self.check(offsets, tensors)
+
+        # prepare
+        hf3fs_utils.write_shm(tensors, self.shm_w_tensor)
+        current = 0
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+            self.ior_w.prepare(
+                self.iov_w[current : current + size], False, self.file, offset
+            )
+            current += size
+
+        # submit
+        ionum = len(offsets)
+        resv = self.ior_w.submit().wait(min_results=ionum)
+
+        # results
+        results = [res.result for res in resv]
+
+        return results
+
+    def check(self, offsets: List[int], tensors: List[torch.Tensor]) -> None:
+        sizes = [t.numel() * t.itemsize for t in tensors]
+        if any(
+            [
+                len(offsets) > self.entries,
+                len(offsets) != len(sizes),
+                all(
+                    [
+                        offset < 0 or offset + size > self.size
+                        for offset, size in zip(offsets, sizes)
+                    ]
+                ),
+                all([size > self.bytes_per_page for size in sizes]),
+            ]
+        ):
+            self.close()
+            raise ValueError(f"Hf3fsClient.check: {offsets=}, {sizes=}")
+
+    def get_size(self) -> int:
+        return self.size
+
+    def close(self) -> None:
+        deregister_fd(self.file)
+        os.close(self.file)
+        del self.ior_r
+        del self.ior_w
+        del self.iov_r
+        del self.iov_w
+        self.shm_r.close()
+        self.shm_w.close()
+        self.shm_r.unlink()
+        self.shm_w.unlink()
+
+    def flush(self) -> None:
+        os.fsync(self.file)