@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d36f24b66a

[KV Connector][Mooncake] Add tenant ID support to MooncakeStoreConnector (#48069)

Signed-off-by: lzw <2994587066@qq.com> Co-authored-by: Roger Wang <hey@rogerw.io>

LZW committed Aug 6, 2026 at 08:31 UTC d36f24b66a9ce71cf15f1cb9249f5d4092f32d11

3 files changed +149 -2

docs/features/mooncake_store_connector_usage.md

+23

@@ -61,6 +61,9 @@ Create a JSON configuration file (e.g., `mooncake_config.json`):

61 large prefills do not exceed the owner's SSD-write budget. Set this together

62 with the matching `--enable_offload=true` flag on `mooncake_master` and on

63 the external `mooncake_client` (if any).

64 +- `tenant_id`: Optional Mooncake tenant namespace. Producers and consumers

65 + that should share store data must use the same tenant id. Default:

66 + `"default"`.

68 Set the config path via environment variable:

@@ -181,6 +184,26 @@ Mooncake environment variables (`MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`,

184 `MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES`, etc.). Those are independent of

185 the vLLM JSON config.

186

187 +### Tenant Isolation

188 +

189 +Set `tenant_id` in the Mooncake JSON config when different vLLM deployments should use separate Mooncake tenant namespaces:

190 +

191 +```json

192 +{

193 + "mode": "embedded",

194 + "metadata_server": "P2PHANDSHAKE",

195 + "master_server_address": "127.0.0.1:50051",

196 + "global_segment_size": "80GB",

197 + "local_buffer_size": "4GB",

198 + "protocol": "rdma",

199 + "device_name": "",

200 + "enable_offload": false,

201 + "tenant_id": "tenant-a"

202 +}

203 +```

204 +

205 +Strict isolation requires a Mooncake master started with `--enable_multi_tenants=true` and a tenant quota policy that registers each tenant. Non-default `tenant_id` also requires a Mooncake version whose `MooncakeDistributedStore.setup()` accepts the `tenant_id` parameter. In `standalone-store` mode, start the external `mooncake_client` with the matching tenant id because that process owns the real store client.

206 +

207 ## Environment Variables

208

209 | Variable | Description | Default |

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+106 -1

@@ -1599,7 +1599,7 @@ def test_requester_worker_init_prefers_local_hostname_override(

1599 },

1600 ),

1601 )

1602 - worker.MooncakeStoreWorker(_make_vllm_config(), _make_kv_cache_config())

1602 + worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config())

1603

1604 assert store.setup.call_args.args[0] == "worker-a:50053"

1605

@@ -2857,6 +2857,7 @@ def test_config_defaults_to_embedded():

2857 assert cfg.global_segment_size == worker.DEFAULT_GLOBAL_SEGMENT_SIZE

2858 assert cfg.local_buffer_size == worker.DEFAULT_LOCAL_BUFFER_SIZE

2859 assert cfg.enable_offload is False

2860 + assert cfg.tenant_id == worker.DEFAULT_TENANT_ID

2861

2862

2863 def test_config_pr40900_unchanged(tmp_path):

@@ -2878,6 +2879,51 @@ def test_config_pr40900_unchanged(tmp_path):

2879 assert cfg.global_segment_size == 4 * 1024**3

2880 assert cfg.local_buffer_size == 4 * 1024**3

2881 assert cfg.enable_offload is False

2882 + assert cfg.tenant_id == worker.DEFAULT_TENANT_ID

2883 +

2884 +

2885 +def test_config_from_file_normalizes_tenant_id(tmp_path):

2886 + config_path = _write_mooncake_config(

2887 + tmp_path,

2888 + {

2889 + "metadata_server": "http://metadata/endpoint",

2890 + "master_server_address": "10.0.0.7:50051",

2891 + "tenant_id": " tenant-a ",

2892 + },

2893 + )

2894 +

2895 + cfg = worker.MooncakeStoreConfig.from_file(config_path)

2896 +

2897 + assert cfg.tenant_id == "tenant-a"

2898 +

2899 +

2900 +def test_config_from_file_normalizes_empty_tenant_id_to_default(tmp_path):

2901 + config_path = _write_mooncake_config(

2902 + tmp_path,

2903 + {

2904 + "metadata_server": "http://metadata/endpoint",

2905 + "master_server_address": "10.0.0.7:50051",

2906 + "tenant_id": " ",

2907 + },

2908 + )

2909 +

2910 + cfg = worker.MooncakeStoreConfig.from_file(config_path)

2911 +

2912 + assert cfg.tenant_id == worker.DEFAULT_TENANT_ID

2913 +

2914 +

2915 +def test_config_from_file_rejects_non_string_tenant_id(tmp_path):

2916 + config_path = _write_mooncake_config(

2917 + tmp_path,

2918 + {

2919 + "metadata_server": "http://metadata/endpoint",

2920 + "master_server_address": "10.0.0.7:50051",

2921 + "tenant_id": False,

2922 + },

2923 + )

2924 +

2925 + with pytest.raises(TypeError, match="tenant_id must be a string or null"):

2926 + worker.MooncakeStoreConfig.from_file(config_path)

2927

2928

2929 def test_config_embedded_rejects_zero_segment():

@@ -2961,6 +3007,7 @@ def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch):

3007 "mlx5_0",

3008 "10.0.0.7:50051",

3009 )

3010 + assert store.setup.call_args.kwargs == {}

3011 # ReplicateConfig is built and carries the preferred_segment.

3012 assert isinstance(w.store_replicate_config, fake_replicate_config_cls)

3013 assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"

@@ -3004,6 +3051,7 @@ def test_topology_embedded_cpu_only(tmp_path, monkeypatch):

3051 "mlx5_0",

3052 "10.0.0.7:50051",

3053 )

3054 + assert store.setup.call_args.kwargs == {}

3055 # No preferred_segment — ReplicateConfig is default-constructed (so the

3056 # preferred_segment field keeps its default value).

3057 assert w.preferred_segment is None

@@ -3013,6 +3061,63 @@ def test_topology_embedded_cpu_only(tmp_path, monkeypatch):

3061 assert w.disk_offload_buffer_budget_bytes is None

3062

3063

3064 +def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch):

3065 + store = MagicMock()

3066 + store.setup.return_value = 0

3067 + _install_fake_mooncake(monkeypatch, store)

3068 + _patch_worker_runtime(monkeypatch)

3069 + monkeypatch.setenv(

3070 + "MOONCAKE_CONFIG_PATH",

3071 + _write_mooncake_config(

3072 + tmp_path,

3073 + {

3074 + "metadata_server": "http://metadata/endpoint",

3075 + "global_segment_size": "4GB",

3076 + "local_buffer_size": "4GB",

3077 + "protocol": "rdma",

3078 + "device_name": "mlx5_0",

3079 + "master_server_address": "10.0.0.7:50051",

3080 + "tenant_id": " tenant-a ",

3081 + },

3082 + ),

3083 + )

3084 +

3085 + worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config())

3086 +

3087 + assert store.setup.call_args.kwargs == {"tenant_id": "tenant-a"}

3088 +

3089 +

3090 +def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch):

3091 + store = MagicMock()

3092 + setup_error = TypeError(

3093 + "setup(): incompatible function arguments; "

3094 + "supported signature includes tenant_id"

3095 + )

3096 + store.setup.side_effect = setup_error

3097 + _install_fake_mooncake(monkeypatch, store)

3098 + _patch_worker_runtime(monkeypatch)

3099 + monkeypatch.setenv(

3100 + "MOONCAKE_CONFIG_PATH",

3101 + _write_mooncake_config(

3102 + tmp_path,

3103 + {

3104 + "metadata_server": "http://metadata/endpoint",

3105 + "global_segment_size": "4GB",

3106 + "local_buffer_size": "4GB",

3107 + "protocol": "rdma",

3108 + "device_name": "mlx5_0",

3109 + "master_server_address": "10.0.0.7:50051",

3110 + "tenant_id": "tenant-a",

3111 + },

3112 + ),

3113 + )

3114 +

3115 + with pytest.raises(TypeError) as exc_info:

3116 + worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config())

3117 +

3118 + assert exc_info.value is setup_error

3119 +

3120 +

3121 # ---------------------------------------------------------------------------

3122 # Stats/metrics tests (PR-35 port)

3123 # ---------------------------------------------------------------------------

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+20 -1

@@ -80,6 +80,7 @@ logger = init_logger(__name__)

81 DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB

82 DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB

83 +DEFAULT_TENANT_ID = "default"

85 MOONCAKE_NO_AVAILABLE_HANDLE = -200

86 _T = TypeVar("_T")

@@ -134,6 +135,7 @@ class MooncakeStoreConfig:

135 global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE

136 local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE

137 enable_offload: bool = False

138 + tenant_id: str = DEFAULT_TENANT_ID

139

140 def __post_init__(self) -> None:

141 if self.mode not in ("embedded", "standalone-store"):

@@ -162,6 +164,7 @@ class MooncakeStoreConfig:

164 config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)

165 ),

166 enable_offload=bool(config.get("enable_offload", False)),

167 + tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)),

168 )

169

170 @staticmethod

@@ -174,6 +177,17 @@ class MooncakeStoreConfig:

177 return MooncakeStoreConfig.from_file(config_path)

178

179

180 +def _normalize_tenant_id(value: Any) -> str:

181 + if value is None:

182 + return DEFAULT_TENANT_ID

183 + if not isinstance(value, str):

184 + raise TypeError(

185 + f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}"

186 + )

187 + tenant_id = value.strip()

188 + return tenant_id if tenant_id else DEFAULT_TENANT_ID

189 +

190 +

191 def _parse_size(value: Any) -> int:

192 """Parse storage size strings with units: GB, MB, KB, B."""

193 if isinstance(value, int):

@@ -1245,6 +1259,9 @@ class MooncakeStoreWorker:

1259 self.store = MooncakeDistributedStore()

1260 local_ip = get_ip()

1261 local_hostname = rdma_utils.get_requester_local_hostname(local_ip)

1262 + setup_kwargs: dict[str, str] = {}

1263 + if store_config.tenant_id != DEFAULT_TENANT_ID:

1264 + setup_kwargs["tenant_id"] = store_config.tenant_id

1265 ret = self.store.setup(

1266 local_hostname,

1267 store_config.metadata_server,

@@ -1253,6 +1270,7 @@ class MooncakeStoreWorker:

1270 store_config.protocol,

1271 store_config.device_name,

1272 store_config.master_server_address,

1273 + **setup_kwargs,

1274 )

1275 if ret != 0:

1276 msg = "Initialize MooncakeDistributedStore failed."

@@ -1280,12 +1298,13 @@ class MooncakeStoreWorker:

1298

1299 logger.info(

1300 "Mooncake mode=%s (global_segment_size=%d, local_buffer_size=%d, "

1283 - "preferred_segment=%s, enable_offload=%s)",

1301 + "preferred_segment=%s, enable_offload=%s, tenant_id=%s)",

1302 store_config.mode,

1303 store_config.global_segment_size,

1304 store_config.local_buffer_size,

1305 preferred_segment or "<none>",

1306 store_config.enable_offload,

1307 + store_config.tenant_id,

1308 )

1309 if store_config.mode == "embedded":

1310 if store_config.enable_offload and preferred_segment is None: