From 4c82bb3252c65288f12f09ce5d0dce68ea910e54 Mon Sep 17 00:00:00 2001 From: LZW <99333079+Lin-z-w@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:17:41 +0800 Subject: [PATCH] Add Mooncake tenant id support (#30256) --- python/sglang/srt/environ.py | 1 + .../storage/mooncake_store/README.md | 14 +- .../mooncake_embedding_store.py | 36 ++- .../storage/mooncake_store/mooncake_store.py | 25 +++ .../mem_cache/test_mooncake_tenant_config.py | 212 ++++++++++++++++++ 5 files changed, 276 insertions(+), 12 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_mooncake_tenant_config.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index b4ff61057..b125d9add 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -524,6 +524,7 @@ class Envs: MOONCAKE_STANDALONE_STORAGE = EnvBool(False) MOONCAKE_ENABLE_SSD_OFFLOAD = EnvBool(False) MOONCAKE_OFFLOAD_FILE_STORAGE_PATH = EnvStr(None) + MOONCAKE_TENANT_ID = EnvStr("default") # MoRI KV Transfer # Send CPU-resident AUX data via RDMA instead of ZMQ TCP (default: TCP). diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/README.md b/python/sglang/srt/mem_cache/storage/mooncake_store/README.md index 419ee5c5f..b9078f387 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/README.md +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/README.md @@ -203,7 +203,7 @@ python -m sglang.launch_server \ --enable-hierarchical-cache \ --hicache-storage-backend mooncake \ --model-path [model_path] \ - --hicache-storage-backend-extra-config '{"master_server_address": "127.0.0.1:50051", "local_hostname": "localhost", "metadata_server": "http://127.0.0.1:8080/metadata", "global_segment_size": "4gb", "protocol": "rdma", "device_name": ""}' + --hicache-storage-backend-extra-config '{"master_server_address": "127.0.0.1:50051", "local_hostname": "localhost", "metadata_server": "http://127.0.0.1:8080/metadata", "global_segment_size": "4gb", "protocol": "rdma", "device_name": "", "tenant_id": "tenant-a"}' ``` **Using JSON file to configure Mooncake** @@ -219,7 +219,8 @@ echo '{ "master_server_address": "127.0.0.1:50051", "protocol": "rdma", "device_name": "", - "global_segment_size": "4gb" + "global_segment_size": "4gb", + "tenant_id": "tenant-a" }' > ${SGLANG_HICACHE_MOONCAKE_CONFIG_PATH} python -m sglang.launch_server \ @@ -236,6 +237,7 @@ MOONCAKE_MASTER="127.0.0.1:50051" \ MOONCAKE_PROTOCOL="rdma" \ MOONCAKE_DEVICE="" \ MOONCAKE_GLOBAL_SEGMENT_SIZE="4gb" \ +MOONCAKE_TENANT_ID="tenant-a" \ python -m sglang.launch_server \ --enable-hierarchical-cache \ --hicache-storage-backend mooncake\ @@ -250,6 +252,14 @@ In particular, for the `global segment size`, if at least one `store service` in **Important:** when `tp > 1`, each Tensor Parallel (TP) rank launches its own Mooncake backend instance and contributes `1/global_segment_size` memory. Therefore, the total memory consumption equals `global segment size`. +**Tenant Isolation (`tenant_id`):** + +When `tenant_id` is set, SGLang forwards it to `MooncakeDistributedStore.setup(..., tenant_id=...)`. Producers and consumers that should share HiCache data must use the same `tenant_id`. + +You can configure it through `tenant_id` in `--hicache-storage-backend-extra-config`, `tenant_id` in the JSON config file, or `MOONCAKE_TENANT_ID`. + +> **Note:** strict isolation between tenants requires a Mooncake master started with `--enable_multi_tenants=true` and a tenant quota policy that explicitly registers each tenant. When strict multi-tenant mode is disabled, Mooncake ignores request tenant IDs for object placement and all objects use the `default` namespace. Non-default `tenant_id` also requires a Mooncake version that supports the `tenant_id` parameter in `MooncakeDistributedStore.setup()`. In `standalone_storage` mode, start the external `mooncake_client` with the matching `--tenant_id` because that process owns the real Mooncake client. + **SSD Offload (`enable_ssd_offload`):** When `enable_ssd_offload` is set to `true`, SGLang will request that Mooncake enable SSD offloading for the KV cache. This allows Mooncake to spill overflow data from DRAM to local SSDs, effectively expanding the available L3 cache capacity. diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_embedding_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_embedding_store.py index 578259f7c..79d007760 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_embedding_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_embedding_store.py @@ -1,7 +1,10 @@ import logging from typing import Any, List -from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import MooncakeBaseStore +from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( + DEFAULT_TENANT_ID, + MooncakeBaseStore, +) logger = logging.getLogger(__name__) @@ -16,15 +19,28 @@ class MooncakeEmbeddingStore(MooncakeBaseStore): MooncakeDistributedStore = self._import_mooncake_store() self.store = MooncakeDistributedStore() self.config = self._load_config(storage_config) - ret_code = self.store.setup( - self.config.local_hostname, - self.config.metadata_server, - self.config.global_segment_size, - 16 * 1024 * 1024, # Internal local buffer size - self.config.protocol, - self.config.device_name, - self.config.master_server_address, - ) + setup_kwargs = {} + if self.config.tenant_id != DEFAULT_TENANT_ID: + setup_kwargs["tenant_id"] = self.config.tenant_id + try: + ret_code = self.store.setup( + self.config.local_hostname, + self.config.metadata_server, + self.config.global_segment_size, + 16 * 1024 * 1024, # Internal local buffer size + self.config.protocol, + self.config.device_name, + self.config.master_server_address, + **setup_kwargs, + ) + except TypeError as e: + if "tenant_id" in setup_kwargs and "tenant_id" in str(e): + raise RuntimeError( + "The installed Mooncake version does not support tenant_id " + "in MooncakeDistributedStore.setup(). Please upgrade " + "Mooncake to use non-default Mooncake tenants with SGLang." + ) from e + raise if ret_code != 0: raise RuntimeError(f"Failed to setup Mooncake Embedding Store: {ret_code}") diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index dd42e2cac..0051f1241 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -27,6 +27,7 @@ from sglang.srt.observability.metrics_collector import StorageMetrics DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MB SETUP_TIMEOUT = 600 # 10min +DEFAULT_TENANT_ID = "default" logger = logging.getLogger(__name__) @@ -82,6 +83,13 @@ def _parse_global_segment_size(value) -> int: return int(value) +def _normalize_tenant_id(value) -> str: + if value is None: + return DEFAULT_TENANT_ID + tenant_id = str(value).strip() + return tenant_id if tenant_id else DEFAULT_TENANT_ID + + @dataclass class MooncakeStoreConfig: local_hostname: str @@ -96,6 +104,7 @@ class MooncakeStoreConfig: client_server_address: str enable_ssd_offload: bool = False ssd_offload_path: Optional[str] = None + tenant_id: str = DEFAULT_TENANT_ID @staticmethod def from_file() -> "MooncakeStoreConfig": @@ -152,6 +161,9 @@ class MooncakeStoreConfig: ssd_offload_path=config.get( "ssd_offload_path", envs.MOONCAKE_OFFLOAD_FILE_STORAGE_PATH.default ), + tenant_id=_normalize_tenant_id( + config.get("tenant_id", envs.MOONCAKE_TENANT_ID.default) + ), ) @staticmethod @@ -193,6 +205,7 @@ class MooncakeStoreConfig: client_server_address=envs.MOONCAKE_CLIENT.get(), enable_ssd_offload=envs.MOONCAKE_ENABLE_SSD_OFFLOAD.get(), ssd_offload_path=envs.MOONCAKE_OFFLOAD_FILE_STORAGE_PATH.get(), + tenant_id=_normalize_tenant_id(envs.MOONCAKE_TENANT_ID.get()), ) @staticmethod @@ -241,6 +254,9 @@ class MooncakeStoreConfig: ssd_offload_path=extra_config.get( "ssd_offload_path", envs.MOONCAKE_OFFLOAD_FILE_STORAGE_PATH.default ), + tenant_id=_normalize_tenant_id( + extra_config.get("tenant_id", envs.MOONCAKE_TENANT_ID.default) + ), ) @@ -486,6 +502,8 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): setup_kwargs["enable_ssd_offload"] = True if self.config.ssd_offload_path is not None: setup_kwargs["ssd_offload_path"] = self.config.ssd_offload_path + if self.config.tenant_id != DEFAULT_TENANT_ID: + setup_kwargs["tenant_id"] = self.config.tenant_id while True: try: @@ -507,6 +525,13 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): ] if not unsupported_kwargs: raise + if "tenant_id" in unsupported_kwargs: + raise RuntimeError( + "The installed Mooncake version does not support " + "tenant_id in MooncakeDistributedStore.setup(). " + "Please upgrade Mooncake to use non-default " + "Mooncake tenants with SGLang." + ) from e logger.warning( "The installed Mooncake version does not support the " f"{', '.join(unsupported_kwargs)} parameter(s) in setup(). " diff --git a/test/registered/unit/mem_cache/test_mooncake_tenant_config.py b/test/registered/unit/mem_cache/test_mooncake_tenant_config.py new file mode 100644 index 000000000..339e93726 --- /dev/null +++ b/test/registered/unit/mem_cache/test_mooncake_tenant_config.py @@ -0,0 +1,212 @@ +import json +import tempfile +import types +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.environ import envs +from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( + DEFAULT_TENANT_ID, + MooncakeStoreConfig, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _fake_mooncake_modules(fake_store_cls): + mooncake = types.ModuleType("mooncake") + mooncake_store = types.ModuleType("mooncake.store") + mooncake_store.MooncakeDistributedStore = fake_store_cls + + class ReplicateConfig: + pass + + mooncake_store.ReplicateConfig = ReplicateConfig + return { + "mooncake": mooncake, + "mooncake.store": mooncake_store, + } + + +class FakeMooncakeDistributedStore: + instances = [] + + def __init__(self): + self.setup_calls = [] + self.objects = {} + type(self).instances.append(self) + + def setup(self, *args, **kwargs): + self.setup_calls.append((args, kwargs)) + return 0 + + def register_buffer(self, *args, **kwargs): + return 0 + + def put(self, key, value, *args): + self.objects[key] = value + return 0 + + def is_exist(self, key): + return 1 if key in self.objects else 0 + + def get(self, key): + return self.objects.get(key) + + +class OldMooncakeDistributedStore(FakeMooncakeDistributedStore): + instances = [] + + def setup(self, *args, **kwargs): + if "tenant_id" in kwargs: + raise TypeError("tenant_id is an invalid keyword argument") + return super().setup(*args, **kwargs) + + +def _make_storage_config(tenant_id=DEFAULT_TENANT_ID): + extra_config = { + "master_server_address": "127.0.0.1:50051", + "check_server": False, + "global_segment_size": 1024 * 1024, + } + if tenant_id is not None: + extra_config["tenant_id"] = tenant_id + + return SimpleNamespace( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=False, + enable_storage_metrics=False, + is_page_first_layout=True, + model_name="test", + tp_lcm_size=None, + should_split_heads=False, + extra_config=extra_config, + ) + + +class TestMooncakeTenantConfig(unittest.TestCase): + def test_load_from_extra_config_normalizes_tenant(self): + cfg = MooncakeStoreConfig.load_from_extra_config( + { + "master_server_address": "127.0.0.1:50051", + "tenant_id": " tenant-extra ", + } + ) + self.assertEqual(cfg.tenant_id, "tenant-extra") + + cfg = MooncakeStoreConfig.load_from_extra_config( + { + "master_server_address": "127.0.0.1:50051", + "tenant_id": " ", + } + ) + self.assertEqual(cfg.tenant_id, DEFAULT_TENANT_ID) + + cfg = MooncakeStoreConfig.load_from_extra_config( + { + "master_server_address": "127.0.0.1:50051", + "tenant_id": None, + } + ) + self.assertEqual(cfg.tenant_id, DEFAULT_TENANT_ID) + + def test_load_from_env_reads_mooncake_tenant_id(self): + with ( + envs.MOONCAKE_MASTER.override("127.0.0.1:50051"), + envs.MOONCAKE_TENANT_ID.override("tenant-env"), + ): + cfg = MooncakeStoreConfig.load_from_env() + + self.assertEqual(cfg.tenant_id, "tenant-env") + + def test_load_from_file_reads_tenant_id(self): + with tempfile.NamedTemporaryFile("w", suffix=".json") as config_file: + json.dump( + { + "master_server_address": "127.0.0.1:50051", + "tenant_id": "tenant-file", + }, + config_file, + ) + config_file.flush() + + with envs.SGLANG_HICACHE_MOONCAKE_CONFIG_PATH.override(config_file.name): + cfg = MooncakeStoreConfig.from_file() + + self.assertEqual(cfg.tenant_id, "tenant-file") + + def test_mooncake_store_forwards_non_default_tenant_id(self): + FakeMooncakeDistributedStore.instances = [] + with patch.dict( + "sys.modules", + _fake_mooncake_modules(FakeMooncakeDistributedStore), + ): + from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( + MooncakeStore, + ) + + MooncakeStore(_make_storage_config("tenant-a")) + + fake_store = FakeMooncakeDistributedStore.instances[-1] + self.assertEqual(fake_store.setup_calls[0][1]["tenant_id"], "tenant-a") + + def test_mooncake_store_keeps_default_tenant_backward_compatible(self): + FakeMooncakeDistributedStore.instances = [] + with patch.dict( + "sys.modules", + _fake_mooncake_modules(FakeMooncakeDistributedStore), + ): + from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( + MooncakeStore, + ) + + MooncakeStore(_make_storage_config()) + + fake_store = FakeMooncakeDistributedStore.instances[-1] + self.assertNotIn("tenant_id", fake_store.setup_calls[0][1]) + + def test_non_default_tenant_requires_new_mooncake(self): + OldMooncakeDistributedStore.instances = [] + with patch.dict( + "sys.modules", + _fake_mooncake_modules(OldMooncakeDistributedStore), + ): + from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( + MooncakeStore, + ) + + with self.assertRaisesRegex(RuntimeError, "tenant_id"): + MooncakeStore(_make_storage_config("tenant-a")) + + def test_embedding_store_forwards_tenant_id(self): + FakeMooncakeDistributedStore.instances = [] + storage_config = SimpleNamespace( + extra_config={ + "master_server_address": "127.0.0.1:50051", + "tenant_id": "tenant-embedding", + } + ) + + with patch.dict( + "sys.modules", + _fake_mooncake_modules(FakeMooncakeDistributedStore), + ): + from sglang.srt.mem_cache.storage.mooncake_store.mooncake_embedding_store import ( + MooncakeEmbeddingStore, + ) + + MooncakeEmbeddingStore(storage_config) + + fake_store = FakeMooncakeDistributedStore.instances[-1] + self.assertEqual(fake_store.setup_calls[0][1]["tenant_id"], "tenant-embedding") + + +if __name__ == "__main__": + unittest.main(verbosity=2)