| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared configuration and setup for EC and KV Mooncake Store clients.""" |
| 4 | |
| 5 | import json |
| 6 | import os |
| 7 | from dataclasses import dataclass |
| 8 | from typing import Any, Literal |
| 9 | |
| 10 | import regex as re |
| 11 | |
| 12 | DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB |
| 13 | DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB |
| 14 | DEFAULT_TENANT_ID = "default" |
| 15 | |
| 16 | |
| 17 | MooncakeMode = Literal["embedded", "standalone-store"] |
| 18 | |
| 19 | |
| 20 | @dataclass |
| 21 | class MooncakeStoreConfig: |
| 22 | """Configuration for MooncakeDistributedStore. |
| 23 | |
| 24 | ``mode`` selects the topology: ``embedded`` (each rank contributes |
| 25 | ``global_segment_size`` in-process) or ``standalone-store`` (rank |
| 26 | contributes 0; an external ``mooncake_client`` process owns the pool |
| 27 | and the SSD tier). |
| 28 | """ |
| 29 | |
| 30 | metadata_server: str |
| 31 | master_server_address: str |
| 32 | protocol: str |
| 33 | device_name: str |
| 34 | mode: MooncakeMode = "embedded" |
| 35 | global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE |
| 36 | local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE |
| 37 | enable_offload: bool = False |
| 38 | tenant_id: str = DEFAULT_TENANT_ID |
| 39 | |
| 40 | def __post_init__(self) -> None: |
| 41 | if self.mode not in ("embedded", "standalone-store"): |
| 42 | raise ValueError(f"unknown Mooncake mode: {self.mode!r}") |
| 43 | if self.local_buffer_size <= 0: |
| 44 | raise ValueError("local_buffer_size must be > 0") |
| 45 | if self.mode == "embedded" and self.global_segment_size == 0: |
| 46 | raise ValueError("embedded mode requires global_segment_size > 0") |
| 47 | if self.mode == "standalone-store" and self.global_segment_size != 0: |
| 48 | raise ValueError("standalone-store mode requires global_segment_size == 0") |
| 49 | |
| 50 | @staticmethod |
| 51 | def from_file(file_path: str) -> "MooncakeStoreConfig": |
| 52 | with open(file_path) as file: |
| 53 | config = json.load(file) |
| 54 | return MooncakeStoreConfig( |
| 55 | metadata_server=config.get("metadata_server", ""), |
| 56 | master_server_address=config.get("master_server_address", ""), |
| 57 | protocol=config.get("protocol", "rdma"), |
| 58 | device_name=config.get("device_name", ""), |
| 59 | mode=config.get("mode", "embedded"), |
| 60 | global_segment_size=_parse_size( |
| 61 | config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE) |
| 62 | ), |
| 63 | local_buffer_size=_parse_size( |
| 64 | config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE) |
| 65 | ), |
| 66 | enable_offload=bool(config.get("enable_offload", False)), |
| 67 | tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)), |
| 68 | ) |
| 69 | |
| 70 | @staticmethod |
| 71 | def load_from_config() -> "MooncakeStoreConfig": |
| 72 | config_path = os.getenv("MOONCAKE_CONFIG_PATH") |
| 73 | if not config_path: |
| 74 | raise ValueError( |
| 75 | "The environment variable 'MOONCAKE_CONFIG_PATH' is not set." |
| 76 | ) |
| 77 | return MooncakeStoreConfig.from_file(config_path) |
| 78 | |
| 79 | |
| 80 | def _normalize_tenant_id(value: Any) -> str: |
| 81 | if value is None: |
| 82 | return DEFAULT_TENANT_ID |
| 83 | if not isinstance(value, str): |
| 84 | raise TypeError( |
| 85 | f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}" |
| 86 | ) |
| 87 | tenant_id = value.strip() |
| 88 | return tenant_id if tenant_id else DEFAULT_TENANT_ID |
| 89 | |
| 90 | |
| 91 | def _parse_size(value: Any) -> int: |
| 92 | """Parse storage size strings with units: GB, MB, KB, B.""" |
| 93 | if isinstance(value, int): |
| 94 | return value |
| 95 | if not isinstance(value, str): |
| 96 | try: |
| 97 | return int(value) |
| 98 | except (TypeError, ValueError) as e: |
| 99 | raise TypeError(f"Unsupported type for size: {type(value)}") from e |
| 100 | |
| 101 | cleaned = value.strip().lower() |
| 102 | if not cleaned: |
| 103 | raise ValueError("Size cannot be empty.") |
| 104 | |
| 105 | unit_multipliers = { |
| 106 | "gb": 1024**3, |
| 107 | "mb": 1024**2, |
| 108 | "kb": 1024, |
| 109 | "b": 1, |
| 110 | } |
| 111 | match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned) |
| 112 | if not match: |
| 113 | raise ValueError(f"Invalid format: '{value}'") |
| 114 | |
| 115 | number_str = match.group(1) |
| 116 | unit = match.group(2) or "b" |
| 117 | multiplier = unit_multipliers[unit] |
| 118 | |
| 119 | try: |
| 120 | numeric_value = float(number_str) |
| 121 | except ValueError as exc: |
| 122 | raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc |
| 123 | return int(numeric_value * multiplier) |
| 124 | |
| 125 | |
| 126 | def setup_mooncake_store( |
| 127 | store: Any, config: MooncakeStoreConfig, local_hostname: str |
| 128 | ) -> None: |
| 129 | """Initialize a native Store, preserving explicit tenant configuration.""" |
| 130 | setup_kwargs: dict[str, str] = {} |
| 131 | if config.tenant_id != DEFAULT_TENANT_ID: |
| 132 | setup_kwargs["tenant_id"] = config.tenant_id |
| 133 | ret = store.setup( |
| 134 | local_hostname, |
| 135 | config.metadata_server, |
| 136 | config.global_segment_size, |
| 137 | config.local_buffer_size, |
| 138 | config.protocol, |
| 139 | config.device_name, |
| 140 | config.master_server_address, |
| 141 | **setup_kwargs, |
| 142 | ) |
| 143 | if ret != 0: |
| 144 | raise RuntimeError("Initialize MooncakeDistributedStore failed.") |