zbtrs2
/
Mooncake
forked from mooncake-track/Mooncake
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Store] feat: support load config from env for mooncake store_service (#951) 

Co-authored-by: 玖宇 <guotongyu.gty@alibaba-inc.com>
This commit is contained in:
Syspretor 2025-10-23 14:47:40 +08:00 committed by GitHub
parent 545302ccaf
commit dde5c1d894
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 81 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md
View File

@ -157,6 +157,18 @@ Then start the `store service`:
python -m mooncake.mooncake_store_service --config=[config_path]
```
Mooncake `store service` configuration can also be provided via environment variables:
```bash
MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" \
MOONCAKE_GLOBAL_SEGMENT_SIZE=4294967296 \
MOONCAKE_PROTOCOL="rdma" \
MOONCAKE_DEVICE="erdma_0,erdma_1" \
MOONCAKE_MASTER=127.0.0.1:50051 \
python -m mooncake.mooncake_store_service
```
Note: If `MOONCAKE_GLOBAL_SEGMENT_SIZE` is set to a non-zero value when starting the `SGLang server`, launching the `store service` can be skipped. In this case, the `SGLang server` also takes on the role of the `store service`, which simplifies deployment but couples the two components together. Users can choose the deployment approach that best fits their needs.
**Start the `SGLang server` with Mooncake enabled:**

43
mooncake-wheel/mooncake/mooncake_config.py
View File

@ -10,6 +10,21 @@ from typing import Optional
DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200 # 3.125 GiB
DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 # 1.0 GiB
def _parse_global_segment_size(value) -> int:
if isinstance(value, int):
return value
if isinstance(value, str):
s = value.strip().lower()
if s.endswith("gb"):
num = s[:-2].strip()
if not num:
raise ValueError(
"Invalid global_segment_size: missing number before 'gb'"
)
return int(num) * 1024 * 1024 * 1024
return int(s)
return int(value)
@dataclass
class MooncakeConfig:
"""The configuration class for Mooncake.
@ -58,8 +73,9 @@ class MooncakeConfig:
return MooncakeConfig(
local_hostname=config.get("local_hostname"),
metadata_server=config.get("metadata_server"),
global_segment_size=config.get("global_segment_size",
DEFAULT_GLOBAL_SEGMENT_SIZE),
global_segment_size=_parse_global_segment_size(
config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
),
local_buffer_size=config.get("local_buffer_size",
DEFAULT_LOCAL_BUFFER_SIZE),
protocol=config.get("protocol", "tcp"),
@ -69,9 +85,26 @@ class MooncakeConfig:
@staticmethod
def load_from_env() -> 'MooncakeConfig':
"""Load config from a file specified in the environment variable."""
"""Load config from a file specified in the environment variable.
export MOONCAKE_MASTER=10.13.3.232:50051
export MOONCAKE_PROTOCOL="rdma"
export MOONCAKE_DEVICE=""
export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE"
"""
config_file_path = os.getenv('MOONCAKE_CONFIG_PATH')
if config_file_path is None:
raise ValueError(
"The environment variable 'MOONCAKE_CONFIG_PATH' is not set.")
if not os.getenv("MOONCAKE_MASTER"):
raise ValueError("Neither the environment variable 'MOONCAKE_CONFIG_PATH' nor 'MOONCAKE_MASTER' is set.")
return MooncakeConfig(
local_hostname=os.getenv("LOCAL_HOSTNAME", "localhost"),
metadata_server=os.getenv("MOONCAKE_TE_META_DATA_SERVER", "P2PHANDSHAKE"),
global_segment_size=_parse_global_segment_size(
os.getenv("MOONCAKE_GLOBAL_SEGMENT_SIZE", DEFAULT_GLOBAL_SEGMENT_SIZE)
),
# Zero copy interface does not need local buffer
local_buffer_size=DEFAULT_LOCAL_BUFFER_SIZE,
protocol=os.getenv("MOONCAKE_PROTOCOL", "tcp"),
device_name=os.getenv("MOONCAKE_DEVICE", ""),
master_server_address=os.getenv("MOONCAKE_MASTER"),
)
return MooncakeConfig.from_file(config_file_path)

34
mooncake-wheel/tests/test_mooncake_config.py
View File

@ -70,8 +70,8 @@ class TestMooncakeConfig(unittest.TestCase):
MooncakeConfig.from_file(self.config_file)
self.assertIn(f"Missing required config field: {field}", str(cm.exception))
def test_load_from_env(self):
"""Test loading configuration from environment variable"""
def test_load_from_config_path_env(self):
"""Test loading configuration from environment variable MOONCAKE_CONFIG_PATH"""
self.write_config(self.valid_config)
# Set environment variable
@ -84,11 +84,39 @@ class TestMooncakeConfig(unittest.TestCase):
# Clean up environment variable
del os.environ['MOONCAKE_CONFIG_PATH']
def test_load_from_config_env(self):
"""Test loading configuration from environment variable MOONCAKE_MASTER"""
# Set environment variable
os.environ['MOONCAKE_MASTER'] = self.valid_config["master_server_address"]
os.environ['LOCAL_HOSTNAME'] = self.valid_config["local_hostname"]
os.environ['MOONCAKE_TE_META_DATA_SERVER'] = self.valid_config["metadata_server"]
os.environ['MOONCAKE_GLOBAL_SEGMENT_SIZE'] = str(self.valid_config["global_segment_size"])
os.environ['MOONCAKE_PROTOCOL'] = self.valid_config["protocol"]
os.environ['MOONCAKE_DEVICE'] = self.valid_config["device_name"]
try:
config = MooncakeConfig.load_from_env()
self.assertEqual(config.master_server_address, self.valid_config["master_server_address"])
self.assertEqual(config.metadata_server, self.valid_config["metadata_server"])
self.assertEqual(config.local_hostname, self.valid_config["local_hostname"])
self.assertEqual(config.global_segment_size, self.valid_config["global_segment_size"])
self.assertEqual(config.protocol, self.valid_config["protocol"])
self.assertEqual(config.device_name, self.valid_config["device_name"])
finally:
# Clean up environment variable
del os.environ['MOONCAKE_MASTER']
del os.environ['LOCAL_HOSTNAME']
del os.environ['MOONCAKE_TE_META_DATA_SERVER']
del os.environ['MOONCAKE_GLOBAL_SEGMENT_SIZE']
del os.environ['MOONCAKE_PROTOCOL']
del os.environ['MOONCAKE_DEVICE']
def test_load_from_env_missing(self):
"""Test loading configuration from environment variable when not set"""
with self.assertRaises(ValueError) as cm:
MooncakeConfig.load_from_env()
self.assertIn("The environment variable 'MOONCAKE_CONFIG_PATH' is not set", str(cm.exception))
self.assertIn("Neither the environment variable 'MOONCAKE_CONFIG_PATH' nor 'MOONCAKE_MASTER' is set.", str(cm.exception))
if __name__ == '__main__':
unittest.main()