sglang 0.4.10.post2__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

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

@@ -0,0 +1,443 @@
+import argparse
+import atexit
+import json
+import logging
+import threading
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import requests
+from fastapi import FastAPI, HTTPException, Request, status
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs import Hf3fsMetadataInterface
+
+# --- Configuration ---
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+
+# --- Data Models ---
+class RankMetadata:
+    """Holds all metadata for a single rank."""
+
+    def __init__(self, num_pages: int):
+        self.lock = threading.RLock()
+        self.num_pages = num_pages
+        self.free_pages: List[int] = list(range(num_pages))
+        self.key_to_index: Dict[str, int] = {}
+        # Todo: Support multi files for HF3FS
+
+    def exists_keys(self, keys: List[str]) -> List[bool]:
+        """Check if keys exist in metadata."""
+        with self.lock:
+            return [key in self.key_to_index for key in keys]
+
+    def reserve_and_allocate_page_indices(
+        self, keys: List[Tuple[str, str]]
+    ) -> List[Tuple[bool, int]]:
+        """Reserve and allocate page indices for keys."""
+        with self.lock:
+            results = [None] * len(keys)
+            new_keys_to_process = []
+
+            for i, (key, prefix_key) in enumerate(keys):
+                if key in self.key_to_index:
+                    results[i] = (True, self.key_to_index[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)
+                else:
+                    results[i] = (False, -1)
+
+            return results
+
+    def confirm_write(
+        self,
+        written_keys_to_confirm: List[Tuple[str, int]],
+        pages_to_release: List[int],
+    ) -> None:
+        """Confirm write operations and release pages."""
+        with self.lock:
+            for key, page_index in written_keys_to_confirm:
+                self.key_to_index[key] = page_index
+
+            for page_index in pages_to_release:
+                if page_index not in self.free_pages:
+                    self.free_pages.append(page_index)
+
+    def delete_keys(self, keys: List[str]) -> int:
+        """Delete keys and return count of deleted keys."""
+        with self.lock:
+            count = 0
+            for key in keys:
+                if key in self.key_to_index:
+                    page_index = self.key_to_index.pop(key)
+                    if page_index not in self.free_pages:
+                        self.free_pages.append(page_index)
+                    count += 1
+            return count
+
+    def clear_all(self) -> None:
+        """Clear all metadata."""
+        with self.lock:
+            self.free_pages = list(range(self.num_pages))
+            self.key_to_index.clear()
+
+    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]
+
+
+class GlobalMetadataState:
+    """Manages the state for all ranks and persistence."""
+
+    def __init__(self, persistence_path: Optional[str], save_interval: int):
+        self.global_lock = threading.RLock()
+        self.ranks: Dict[int, RankMetadata] = {}
+        self.persistence_path = Path(persistence_path) if persistence_path else None
+        self.save_interval = save_interval
+        self.save_timer: Optional[threading.Timer] = None
+        self.is_shutting_down = False
+
+    def load_from_disk(self):
+        if not self.persistence_path or not self.persistence_path.exists():
+            logging.info("Persistence file not found. Starting with a clean state.")
+            return
+
+        logging.info(f"Loading state from {self.persistence_path}")
+        try:
+            with open(self.persistence_path, "r") as f:
+                persisted_data = json.load(f)
+
+            with self.global_lock:
+                for rank_id_str, data in persisted_data.items():
+                    rank_id = int(rank_id_str)
+                    num_pages = data["num_pages"]
+                    rank_meta = RankMetadata(num_pages)
+                    rank_meta.free_pages = data["free_pages"]
+                    rank_meta.key_to_index = dict(data["key_to_index"])
+                    self.ranks[rank_id] = rank_meta
+                logging.info(
+                    f"Successfully loaded metadata for {len(self.ranks)} ranks."
+                )
+        except (json.JSONDecodeError, KeyError, TypeError) as e:
+            logging.error(
+                f"Failed to load or parse persistence file: {e}. Starting fresh.",
+                exc_info=True,
+            )
+            self.ranks.clear()
+
+    def save_to_disk(self):
+        if not self.persistence_path:
+            return
+
+        logging.info("Persisting metadata to disk...")
+        with self.global_lock:
+            serializable_state = {}
+            for rank_id, rank_meta in self.ranks.items():
+                with rank_meta.lock:
+                    serializable_state[rank_id] = {
+                        "num_pages": rank_meta.num_pages,
+                        "free_pages": rank_meta.free_pages,
+                        "key_to_index": list(rank_meta.key_to_index.items()),
+                    }
+
+        try:
+            temp_path = self.persistence_path.with_suffix(".tmp")
+            with open(temp_path, "w") as f:
+                json.dump(serializable_state, f, indent=4)
+            temp_path.rename(self.persistence_path)
+            logging.info(f"Metadata successfully persisted to {self.persistence_path}")
+        except Exception as e:
+            logging.error(f"Failed to save metadata to disk: {e}", exc_info=True)
+
+    def schedule_save(self):
+        if self.is_shutting_down or not self.persistence_path:
+            return
+        self.save_to_disk()
+        self.save_timer = threading.Timer(self.save_interval, self.schedule_save)
+        self.save_timer.start()
+
+    def shutdown(self):
+        logging.info("Shutting down metadata server...")
+        self.is_shutting_down = True
+        if self.save_timer:
+            self.save_timer.cancel()
+        self.save_to_disk()
+        logging.info("Shutdown complete.")
+
+
+# --- Global MetadataServer implementation ---
+class Hf3fsMetadataServer:
+    """HF3FS Metadata Server that manages metadata for multiple ranks."""
+
+    def __init__(self, persistence_path: Optional[str] = None, save_interval: int = 60):
+        self.state = GlobalMetadataState(persistence_path, save_interval)
+        self.app = FastAPI()
+        self._setup_routes()
+
+    def _setup_routes(self):
+        """Setup FastAPI routes."""
+        self.app.post("/{rank}/initialize")(self.initialize)
+        self.app.post("/{rank}/exists")(self.exists)
+        self.app.post("/{rank}/reserve_and_allocate_page_indices")(
+            self.reserve_and_allocate_page_indices
+        )
+        self.app.post("/{rank}/confirm_write")(self.confirm_write)
+        self.app.post("/{rank}/delete_keys")(self.delete_keys)
+        self.app.post("/{rank}/clear")(self.clear)
+        self.app.post("/{rank}/get_page_indices")(self.get_page_indices)
+
+    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]
+
+    async def initialize(self, rank: int, request: Request):
+        """Initialize a rank with specified number of pages."""
+        data = await request.json()
+        num_pages = data["num_pages"]
+        with self.state.global_lock:
+            if rank in self.state.ranks:
+                logging.info(
+                    f"Rank {rank} already exists. Initialization request ignored."
+                )
+                if self.state.ranks[rank].num_pages != num_pages:
+                    logging.warning(
+                        f"Rank {rank} initialized with different num_pages. Existing: {self.state.ranks[rank].num_pages}, New: {num_pages}"
+                    )
+            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."}
+
+    async def exists(self, rank: int, request: Request):
+        """Check if keys exist in metadata."""
+        data = await request.json()
+        keys = data["keys"]
+        metadata = self.get_rank_metadata(rank)
+        results = metadata.exists_keys(keys)
+        return {"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()
+        metadata = self.get_rank_metadata(rank)
+        keys = data["keys"]
+        results = metadata.reserve_and_allocate_page_indices(keys)
+        return {"indices": results}
+
+    async def confirm_write(self, rank: int, request: Request):
+        """Confirm write operations and release pages."""
+        data = await request.json()
+        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."
+        }
+
+    async def delete_keys(self, rank: int, request: Request):
+        """Delete keys from metadata."""
+        data = await request.json()
+        metadata = self.get_rank_metadata(rank)
+        count = metadata.delete_keys(data["keys"])
+        return {"message": f"Rank {rank}: {count} keys deleted."}
+
+    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."}
+
+    async def get_page_indices(self, rank: int, request: Request):
+        """Get page indices for keys."""
+        data = await request.json()
+        metadata = self.get_rank_metadata(rank)
+        keys = data["keys"]
+        results = metadata.get_page_indices(keys)
+        return {"indices": results}
+
+    def run(self, host: str = "0.0.0.0", port: int = 18000):
+        """Run the metadata server."""
+        self.state.load_from_disk()
+        if self.state.persistence_path:
+            self.state.schedule_save()
+            atexit.register(self.state.shutdown)
+
+        import uvicorn
+
+        logging.info(f"Starting metadata server on http://{host}:{port}")
+        if self.state.persistence_path:
+            logging.info(
+                f"Persistence is ENABLED. Saving to '{self.state.persistence_path}' every {self.state.save_interval} seconds."
+            )
+        else:
+            logging.info("Persistence is DISABLED.")
+
+        uvicorn.run(self.app, host=host, port=port)
+
+
+# --- Client implementation ---
+class Hf3fsGlobalMetadataClient(Hf3fsMetadataInterface):
+    """Global http metadata client for HF3FS."""
+
+    def __init__(self, base_url: str, max_retries: int = 3):
+        self.base_url = base_url.rstrip("/")
+        self._session = requests.Session()
+
+        retry_strategy = Retry(
+            total=max_retries,
+            backoff_factor=0.3,
+            status_forcelist=[500, 502, 503, 504],
+            allowed_methods=["GET", "POST"],
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        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)
+            response.raise_for_status()
+            return response.json()
+        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
+
+    def initialize(self, rank: int, num_pages: int) -> None:
+        self._post(f"{rank}/initialize", {"num_pages": num_pages})
+
+    def reserve_and_allocate_page_indices(
+        self, rank: int, keys: List[Tuple[str, str]]
+    ) -> List[Tuple[bool, int]]:
+        response = self._post(
+            f"{rank}/reserve_and_allocate_page_indices", {"keys": keys}
+        )
+        return [tuple(item) for item in response.get("indices")]
+
+    def confirm_write(
+        self,
+        rank: int,
+        written_keys_to_confirm: List[Tuple[str, int]],
+        pages_to_release: List[int],
+    ) -> None:
+        self._post(
+            f"{rank}/confirm_write",
+            {
+                "written_keys_to_confirm": written_keys_to_confirm,
+                "pages_to_release": pages_to_release,
+            },
+        )
+
+    def delete_keys(self, rank: int, keys: List[str]) -> None:
+        self._post(f"{rank}/delete_keys", {"keys": keys})
+
+    def exists(self, rank: int, keys: List[str]) -> List[bool]:
+        response = self._post(f"{rank}/exists", {"keys": keys})
+        return response.get("exists", [])
+
+    def clear(self, rank: int) -> None:
+        self._post(f"{rank}/clear", {})
+
+    def get_page_indices(self, rank: int, keys: List[str]) -> List[Optional[int]]:
+        response = self._post(f"{rank}/get_page_indices", {"keys": keys})
+        return response.get("indices")
+
+
+class Hf3fsLocalMetadataClient(Hf3fsMetadataInterface):
+    """Local metadata client that directly operates on single RankMetadata in memory without metadata server."""
+
+    def __init__(self):
+        self.rank_metadata = None
+
+    def initialize(self, rank: int, num_pages: int) -> None:
+        self.rank_metadata = RankMetadata(num_pages)
+
+    def reserve_and_allocate_page_indices(
+        self, rank: int, keys: List[Tuple[str, str]]
+    ) -> List[Tuple[bool, int]]:
+        """Reserve and allocate page indices for keys."""
+        return self.rank_metadata.reserve_and_allocate_page_indices(keys)
+
+    def confirm_write(
+        self,
+        rank: int,
+        written_keys_to_confirm: List[Tuple[str, int]],
+        pages_to_release: List[int],
+    ) -> None:
+        """Confirm write operations."""
+        self.rank_metadata.confirm_write(written_keys_to_confirm, pages_to_release)
+
+    def delete_keys(self, rank: int, keys: List[str]) -> None:
+        """Delete keys."""
+        self.rank_metadata.delete_keys(keys)
+
+    def exists(self, rank: int, keys: List[str]) -> List[bool]:
+        """Check if keys exist."""
+        return self.rank_metadata.exists_keys(keys)
+
+    def clear(self, rank: int) -> None:
+        """Clear all metadata for rank."""
+        self.rank_metadata.clear_all()
+
+    def get_page_indices(self, rank: int, keys: List[str]) -> List[Optional[int]]:
+        """Get page indices for keys."""
+        return self.rank_metadata.get_page_indices(keys)
+
+
+def run_metadata_server(
+    host: str = "0.0.0.0",
+    port: int = 18000,
+    persistence_path: Optional[str] = None,
+    save_interval: int = 60,
+):
+    """Run the HF3FS metadata server."""
+    global server
+    server = Hf3fsMetadataServer(
+        persistence_path=persistence_path, save_interval=save_interval
+    )
+
+    server.run(host=host, port=port)
+
+
+# --- Main Execution ---
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="HF3FS Metadata Server")
+    parser.add_argument(
+        "--host", type=str, default="0.0.0.0", help="Host to bind the server to."
+    )
+    parser.add_argument(
+        "--port", type=int, default=18000, help="Port to run the server on."
+    )
+    parser.add_argument(
+        "--persistence-path",
+        type=str,
+        default=None,
+        help="Path to the file for persisting metadata. If not provided, persistence is disabled.",
+    )
+    parser.add_argument(
+        "--save-interval",
+        type=int,
+        default=60,
+        help="Interval in seconds for periodically saving metadata to disk.",
+    )
+    args = parser.parse_args()
+
+    run_metadata_server(args.host, args.port, args.persistence_path, args.save_interval)