Add Mooncake tenant id support (#30256)
This commit is contained in:
LZW
@@ -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)
|
||||
Reference in New Issue
Block a user