Add Mooncake tenant id support (#30256)

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)