sglang 0.5.2rc1__py3-none-any.whl → 0.5.3__py3-none-any.whl

sglang/srt/mem_cache/storage/eic/test_unit.py

@@ -0,0 +1,115 @@
+import argparse
+import os
+
+import eic
+import torch
+import yaml
+
+
+def pase_args():
+    parser = argparse.ArgumentParser(description="EIC Storage Unit Test")
+    parser.add_argument(
+        "--config",
+        "-c",
+        type=str,
+        default="/sgl-workspace/config/remote-eic.yaml",
+        help="EIC yaml config",
+    )
+    args, _ = parser.parse_known_args()
+    return args
+
+
+def init_eic_client():
+    args = pase_args()
+    config_path = os.path.abspath(args.config)
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+    with open(config_path, "r") as fin:
+        config = yaml.safe_load(fin)
+
+    remote_url = config.get("remote_url", None)
+    if remote_url is None:
+        AssertionError("remote_url is None")
+    endpoint = remote_url[len("eic://") :]
+    eic_instance_id = config.get("eic_instance_id", None)
+    eic_log_dir = config.get("eic_log_dir", None)
+    eic_log_level = config.get("eic_log_level", 2)
+    eic_trans_type = config.get("eic_trans_type", 3)
+    eic_flag_file = config.get("eic_flag_file", None)
+
+    if not os.path.exists(eic_log_dir):
+        os.makedirs(eic_log_dir, exist_ok=True)
+    eic_client = eic.Client()
+    init_option = eic.InitOption()
+    init_option.log_dir = eic_log_dir
+    init_option.log_level = eic.LogLevel(eic_log_level)
+    init_option.transport_type = eic.TransportType(eic_trans_type)
+    init_option.flag_file = eic_flag_file
+    ret = eic_client.init(eic_instance_id, endpoint, init_option)
+    if ret != 0:
+        raise RuntimeError(f"EIC Client init failed with error code: {ret}")
+    return eic_client
+
+
+def test_set(eic_client):
+    test_key = ["test_key_" + str(i) for i in range(16)]
+    tensors = [
+        torch.ones([12, 6, 1, 512], dtype=torch.bfloat16, device="cpu")
+        for _ in range(16)
+    ]
+    data_keys = eic.StringVector()
+    data_vals = eic.IOBuffers()
+    for i in range(16):
+        data_keys.append(test_key[i])
+        data_vals.append(
+            tensors[i].data_ptr(), tensors[i].numel() * tensors[i].element_size(), False
+        )
+    set_opt = eic.SetOption()
+    set_opt.ttl_second = 3
+    status_code, set_outcome = eic_client.mset(data_keys, data_vals, set_opt)
+    assert (
+        status_code == eic.StatusCode.SUCCESS
+    ), f"Set failed with status code: {status_code}"
+
+
+def test_get(eic_client):
+    test_key = ["test_key_" + str(i) for i in range(16)]
+    tensors = [
+        torch.zeros([12, 6, 1, 512], dtype=torch.bfloat16, device="cpu")
+        for _ in range(16)
+    ]
+    data_keys = eic.StringVector()
+    data_vals = eic.IOBuffers()
+    for i in range(16):
+        data_keys.append(test_key[i])
+        data_vals.append(
+            tensors[i].data_ptr(), tensors[i].numel() * tensors[i].element_size(), False
+        )
+    get_opt = eic.GetOption()
+    status_code, data_vals, get_outcome = eic_client.mget(data_keys, get_opt, data_vals)
+    assert (
+        status_code == eic.StatusCode.SUCCESS
+    ), f"Get failed with status code: {status_code}"
+
+
+def test_exists(eic_client):
+    test_key = ["test_key_" + str(i) for i in range(16)]
+    data_keys = eic.StringVector()
+    for key in test_key:
+        data_keys.append(key)
+    exists_opt = eic.ExistOption()
+    status_code, exists_outcome = eic_client.mexist(data_keys, exists_opt)
+    assert (
+        status_code == eic.StatusCode.SUCCESS
+    ), f"Exists failed with status code: {status_code}"
+
+
+def main():
+    eic_client = init_eic_client()
+    test_set(eic_client)
+    test_exists(eic_client)
+    test_get(eic_client)
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,164 @@
+import logging
+import os
+import threading
+from abc import ABC, abstractmethod
+from typing import List
+
+import torch
+
+
+class Hf3fsClient(ABC):
+    """Abstract interface for HF3FS clients."""
+
+    @abstractmethod
+    def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
+        """Initialize the HF3FS client.
+
+        Args:
+            path: File path for storage
+            size: Total size of storage file
+            bytes_per_page: Bytes per page
+            entries: Number of entries for batch operations
+        """
+        pass
+
+    @abstractmethod
+    def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        """Batch read from storage."""
+        pass
+
+    @abstractmethod
+    def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        """Batch write to storage."""
+        pass
+
+    @abstractmethod
+    def check(self, offsets: List[int], tensors: List[torch.Tensor]) -> None:
+        """Validate batch operation parameters."""
+        pass
+
+    @abstractmethod
+    def get_size(self) -> int:
+        """Get total storage size."""
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the client and cleanup resources."""
+        pass
+
+    @abstractmethod
+    def flush(self) -> None:
+        """Flush data to disk."""
+        pass
+
+
+logger = logging.getLogger(__name__)
+
+
+class Hf3fsMockClient(Hf3fsClient):
+    """Mock implementation of Hf3fsClient for CI testing purposes."""
+
+    def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
+        """Initialize mock HF3FS client."""
+        self.path = path
+        self.size = size
+        self.bytes_per_page = bytes_per_page
+        self.entries = entries
+
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(self.path), exist_ok=True)
+
+        # Create and initialize the file
+        self.file = os.open(self.path, os.O_RDWR | os.O_CREAT)
+        os.ftruncate(self.file, size)
+
+        logger.info(
+            f"Hf3fsMockClient initialized: path={path}, size={size}, "
+            f"bytes_per_page={bytes_per_page}, entries={entries}"
+        )
+
+    def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        """Batch read from mock storage."""
+        self.check(offsets, tensors)
+
+        results = []
+
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+
+            try:
+                os.lseek(self.file, offset, os.SEEK_SET)
+                bytes_read = os.read(self.file, size)
+
+                if len(bytes_read) == size:
+                    # Convert bytes to tensor and copy to target
+                    bytes_tensor = torch.frombuffer(bytes_read, dtype=torch.uint8)
+                    typed_tensor = bytes_tensor.view(tensor.dtype).view(tensor.shape)
+                    tensor.copy_(typed_tensor)
+                    results.append(size)
+                else:
+                    logger.warning(
+                        f"Short read: expected {size}, got {len(bytes_read)}"
+                    )
+                    results.append(len(bytes_read))
+
+            except Exception as e:
+                logger.error(f"Error reading from offset {offset}: {e}")
+                results.append(0)
+
+        return results
+
+    def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        """Batch write to mock storage."""
+        self.check(offsets, tensors)
+
+        results = []
+
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+
+            try:
+                # Convert tensor to bytes and write directly to file
+                tensor_bytes = tensor.contiguous().view(torch.uint8).flatten()
+                data = tensor_bytes.numpy().tobytes()
+
+                os.lseek(self.file, offset, os.SEEK_SET)
+                bytes_written = os.write(self.file, data)
+
+                if bytes_written == size:
+                    results.append(size)
+                else:
+                    logger.warning(f"Short write: expected {size}, got {bytes_written}")
+                    results.append(bytes_written)
+
+            except Exception as e:
+                logger.error(f"Error writing to offset {offset}: {e}")
+                results.append(0)
+
+        return results
+
+    def check(self, offsets: List[int], tensors: List[torch.Tensor]) -> None:
+        """Validate batch operation parameters."""
+        pass
+
+    def get_size(self) -> int:
+        """Get total storage size."""
+        return self.size
+
+    def close(self) -> None:
+        """Close the mock client and cleanup resources."""
+        try:
+            if hasattr(self, "file") and self.file >= 0:
+                os.close(self.file)
+                self.file = -1  # Mark as closed
+            logger.info(f"MockHf3fsClient closed: {self.path}")
+        except Exception as e:
+            logger.error(f"Error closing MockHf3fsClient: {e}")
+
+    def flush(self) -> None:
+        """Flush data to disk."""
+        try:
+            os.fsync(self.file)
+        except Exception as e:
+            logger.error(f"Error flushing MockHf3fsClient: {e}")

sglang/srt/mem_cache/storage/hf3fs/{client_hf3fs.py → hf3fs_usrbio_client.py}

@@ -9,6 +9,8 @@ from typing import List
 import torch
 from torch.utils.cpp_extension import load
 
+from sglang.srt.mem_cache.storage.hf3fs.hf3fs_client import Hf3fsClient
+
 root = Path(__file__).parent.resolve()
 hf3fs_utils = load(name="hf3fs_utils", sources=[f"{root}/hf3fs_utils.cpp"])
 
@@ -51,7 +53,9 @@ def wsynchronized():
     return _decorator
 
 
-class Hf3fsClient:
+class Hf3fsUsrBioClient(Hf3fsClient):
+    """HF3FS client implementation using usrbio."""
+
     def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
         if not HF3FS_AVAILABLE:
             raise ImportError(