sglang 0.5.1.post3__py3-none-any.whl → 0.5.2__py3-none-any.whl

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(

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

@@ -4,10 +4,12 @@ import json
 import logging
 import threading
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, OrderedDict, Tuple
 
+import orjson
 import requests
-from fastapi import FastAPI, HTTPException, Request, status
+from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi.responses import ORJSONResponse
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 
@@ -24,10 +26,10 @@ class RankMetadata:
     """Holds all metadata for a single rank."""
 
     def __init__(self, num_pages: int):
-        self.lock = threading.RLock()
+        self.lock = threading.Lock()
         self.num_pages = num_pages
         self.free_pages: List[int] = list(range(num_pages))
-        self.key_to_index: Dict[str, int] = {}
+        self.key_to_index: OrderedDict[str, int] = OrderedDict()
         # Todo: Support multi files for HF3FS
 
     def exists_keys(self, keys: List[str]) -> List[bool]:
@@ -46,16 +48,18 @@ class RankMetadata:
             for i, (key, prefix_key) in enumerate(keys):
                 if key in self.key_to_index:
                     results[i] = (True, self.key_to_index[key])
+                    self.key_to_index.move_to_end(key)
                 else:
                     new_keys_to_process.append((i, key, prefix_key))
 
             # Todo: Implementing data eviction logic after HiCache supports prefix information pass-through
             for i, key, prefix_key in new_keys_to_process:
                 if len(self.free_pages) > 0:
-                    page_idx = self.free_pages.pop()
-                    results[i] = (False, page_idx)
+                    page_index = self.free_pages.pop()
                 else:
-                    results[i] = (False, -1)
+                    page_index = self.key_to_index.popitem(last=False)[1]
+
+                results[i] = (False, page_index)
 
             return results
 
@@ -68,6 +72,7 @@ class RankMetadata:
         with self.lock:
             for key, page_index in written_keys_to_confirm:
                 self.key_to_index[key] = page_index
+                self.key_to_index.move_to_end(key)
 
             for page_index in pages_to_release:
                 if page_index not in self.free_pages:
@@ -94,7 +99,14 @@ class RankMetadata:
     def get_page_indices(self, keys: List[str]) -> List[Optional[int]]:
         """Get page indices for keys."""
         with self.lock:
-            return [self.key_to_index.get(key) for key in keys]
+            results = []
+            for key in keys:
+                if key in self.key_to_index:
+                    results.append(self.key_to_index[key])
+                    self.key_to_index.move_to_end(key)
+                else:
+                    results.append(None)
+            return results
 
 
 class GlobalMetadataState:
@@ -182,7 +194,8 @@ class Hf3fsMetadataServer:
 
     def __init__(self, persistence_path: Optional[str] = None, save_interval: int = 60):
         self.state = GlobalMetadataState(persistence_path, save_interval)
-        self.app = FastAPI()
+        self.app = FastAPI(default_response_class=ORJSONResponse)
+
         self._setup_routes()
 
     def _setup_routes(self):
@@ -199,17 +212,25 @@ class Hf3fsMetadataServer:
 
     def get_rank_metadata(self, rank: int) -> RankMetadata:
         """Get rank metadata with proper error handling."""
-        with self.state.global_lock:
-            if rank not in self.state.ranks:
-                raise HTTPException(
-                    status_code=404,
-                    detail=f"Rank {rank} not initialized. Please call /{{rank}}/initialize first.",
-                )
-            return self.state.ranks[rank]
+        if rank not in self.state.ranks:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Rank {rank} not initialized. Please call /{rank}/initialize first.",
+            )
+        return self.state.ranks[rank]
+
+    async def _read_json(self, request: Request) -> dict:
+        """Parse request JSON using orjson if available."""
+        body = await request.body()
+        return orjson.loads(body)
+
+    def _json_response(self, content: dict):
+        """Return ORJSONResponse when available to bypass jsonable_encoder."""
+        return ORJSONResponse(content)
 
     async def initialize(self, rank: int, request: Request):
         """Initialize a rank with specified number of pages."""
-        data = await request.json()
+        data = await self._read_json(request)
         num_pages = data["num_pages"]
         with self.state.global_lock:
             if rank in self.state.ranks:
@@ -223,57 +244,55 @@ class Hf3fsMetadataServer:
             else:
                 logging.info(f"Initializing new Rank {rank} with {num_pages} pages.")
                 self.state.ranks[rank] = RankMetadata(num_pages)
-        return {"message": f"Rank {rank} is ready."}
+        return Response(status_code=204)
 
     async def exists(self, rank: int, request: Request):
         """Check if keys exist in metadata."""
-        data = await request.json()
+        data = await self._read_json(request)
         keys = data["keys"]
         metadata = self.get_rank_metadata(rank)
         results = metadata.exists_keys(keys)
-        return {"exists": results}
+        return self._json_response({"exists": results})
 
     async def reserve_and_allocate_page_indices(self, rank: int, request: Request):
         """Reserve and allocate page indices for keys."""
-        data = await request.json()
+        data = await self._read_json(request)
         metadata = self.get_rank_metadata(rank)
         keys = data["keys"]
         results = metadata.reserve_and_allocate_page_indices(keys)
-        return {"indices": results}
+        return self._json_response({"indices": results})
 
     async def confirm_write(self, rank: int, request: Request):
         """Confirm write operations and release pages."""
-        data = await request.json()
+        data = await self._read_json(request)
         metadata = self.get_rank_metadata(rank)
         success_written_keys = data.get("written_keys_to_confirm", [])
         released_pages = data.get("pages_to_release", [])
 
         metadata.confirm_write(success_written_keys, released_pages)
 
-        return {
-            "message": f"Rank {rank}: Write confirmed for {len(success_written_keys)} keys. {len(released_pages)} pages released."
-        }
+        return Response(status_code=204)
 
     async def delete_keys(self, rank: int, request: Request):
         """Delete keys from metadata."""
-        data = await request.json()
+        data = await self._read_json(request)
         metadata = self.get_rank_metadata(rank)
         count = metadata.delete_keys(data["keys"])
-        return {"message": f"Rank {rank}: {count} keys deleted."}
+        return Response(status_code=204)
 
     async def clear(self, rank: int):
         """Clear all metadata for a rank."""
         metadata = self.get_rank_metadata(rank)
         metadata.clear_all()
-        return {"message": f"Rank {rank}: Metadata cleared."}
+        return Response(status_code=204)
 
     async def get_page_indices(self, rank: int, request: Request):
         """Get page indices for keys."""
-        data = await request.json()
+        data = await self._read_json(request)
         metadata = self.get_rank_metadata(rank)
         keys = data["keys"]
         results = metadata.get_page_indices(keys)
-        return {"indices": results}
+        return self._json_response({"indices": results})
 
     def run(self, host: str = "0.0.0.0", port: int = 18000):
         """Run the metadata server."""
@@ -309,14 +328,22 @@ class Hf3fsGlobalMetadataClient(Hf3fsMetadataInterface):
             status_forcelist=[500, 502, 503, 504],
             allowed_methods=["GET", "POST"],
         )
-        adapter = HTTPAdapter(max_retries=retry_strategy)
+        adapter = HTTPAdapter(
+            max_retries=retry_strategy, pool_connections=256, pool_maxsize=256
+        )
         self._session.mount("http://", adapter)
 
     def _post(self, endpoint: str, json_data: dict) -> dict:
         try:
-            response = self._session.post(f"{self.base_url}/{endpoint}", json=json_data)
+            url = f"{self.base_url}/{endpoint}"
+            headers = {"Content-Type": "application/json"}
+            payload = orjson.dumps(json_data)  # type: ignore[union-attr]
+            response = self._session.post(url, data=payload, headers=headers)
             response.raise_for_status()
-            return response.json()
+
+            if response.status_code == 204 or not response.content:
+                return {}
+            return orjson.loads(response.content)  # type: ignore[union-attr]
         except requests.exceptions.RequestException as e:
             logging.error(f"Failed to POST to {endpoint} after retries: {e}")
             raise RuntimeError(f"Failed to connect to metadata server: {e}") from e