From fd859d2e0689ce254bfb0e67483fba76405ca9df Mon Sep 17 00:00:00 2001 From: Schatten <3192396192@qq.com> Date: Wed, 5 Aug 2026 16:25:38 +0800 Subject: [PATCH] [KV Connector][Mooncake] Add store group semantics (#44956) Signed-off-by: Schatten Co-authored-by: OpenAI Codex Co-authored-by: Yifan Qiao --- .../unit/test_mooncake_store_worker.py | 465 +++++++++++++++++- .../kv_connector/v1/mooncake/store/worker.py | 69 ++- 2 files changed, 531 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 7d7a2f277d..db64e79701 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -76,6 +76,8 @@ def _make_store_sending_thread( tp_rank: int = 0, put_step: int = 1, replicate_config: object | None = None, + enable_group_semantics: bool = False, + supports_group_ids: bool = False, ) -> mooncake_store_worker.KVCacheStoreSendingThread: if coord is None: coord = _default_send_coord() @@ -94,6 +96,8 @@ def _make_store_sending_thread( kv_role="kv_producer", ready_event=threading.Event(), replicate_config=replicate_config, + enable_group_semantics=enable_group_semantics, + supports_group_ids=supports_group_ids, ) thread.request_queue.task_done = MagicMock() return thread @@ -162,6 +166,16 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta: ) +def _make_multi_group_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta: + return ReqMeta( + req_id=req_id, + token_len_chunk=32, + block_ids=([0, 1], [2, 3]), + block_hashes=block_hashes, + can_save=True, + ) + + _DISK_OFFLOAD_SINGLE_KEY_BYTES = worker._estimate_disk_offload_staging_bytes([256]) _DISK_OFFLOAD_USABLE_BUDGET_RATIO = 0.9 _DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS = 4 * _DISK_OFFLOAD_SINGLE_KEY_BYTES @@ -276,6 +290,7 @@ def _patch_worker_runtime( monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group) monkeypatch.setattr(worker, "get_dcp_group", lambda: dcp_group) monkeypatch.setattr(worker, "get_ip", lambda: local_ip) + monkeypatch.setattr(worker, "LookupKeyServer", MagicMock()) def test_pool_key_to_string_without_prefix_is_unchanged(): @@ -599,7 +614,13 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): assert store.batch_put_from_multi_buffers.call_args.args[0] == keys -def _make_partial_tail_send_thread(store): +def _make_partial_tail_send_thread( + store, + *, + replicate_config=None, + enable_group_semantics=False, + supports_group_ids=False, +): coord = SimpleNamespace( enable_partial_hash_hits=True, hash_block_size=4, @@ -616,6 +637,9 @@ def _make_partial_tail_send_thread(store): store, coord=coord, token_databases=[db], + replicate_config=replicate_config, + enable_group_semantics=enable_group_semantics, + supports_group_ids=supports_group_ids, ) @@ -648,6 +672,32 @@ def test_partial_tail_offload_skips_null_source_blocks(): assert addrs == [[0x1000 + 2 * 256], [0x1000 + 3 * 256]] +def test_partial_tail_offload_replaces_stale_group_ids_after_filtering(): + store = MagicMock() + store.batch_is_exist.return_value = [1, 0, 0] + store.batch_put_from_multi_buffers.return_value = [256, 256] + replicate_config = SimpleNamespace(group_ids=["stale"]) + thread = _make_partial_tail_send_thread( + store, + replicate_config=replicate_config, + enable_group_semantics=True, + supports_group_ids=True, + ) + + assert thread._maybe_offload_partial_tail(_make_partial_tail_req([1, 2, 3])) + + keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args + assert keys == [ + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131", + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132", + ] + assert config is replicate_config + assert config.group_ids == [ + "vllm-mooncake-store:test-model@6131", + "vllm-mooncake-store:test-model@6132", + ] + + def test_partial_tail_offload_honors_active_pressure_gate(): store = MagicMock() thread = _make_partial_tail_send_thread(store) @@ -969,6 +1019,291 @@ def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_ assert call_args[3] is replicate_config +def test_group_id_detection_uses_class_attribute_without_instantiating(): + class ReplicateConfigWithClassGroupIdsAndRequiredInit: + group_ids = None + + def __init__(self, required): + self.group_ids = required + + assert worker._replicate_config_supports_group_ids( + ReplicateConfigWithClassGroupIdsAndRequiredInit, + object(), + ) + + +def test_group_id_detection_uses_instance_attribute_when_class_lacks_attribute(): + class ReplicateConfigWithInstanceGroupIds: + def __init__(self) -> None: + self.group_ids = None + + replicate_config = ReplicateConfigWithInstanceGroupIds() + + assert worker._replicate_config_supports_group_ids( + ReplicateConfigWithInstanceGroupIds, + replicate_config, + ) + + +def test_group_id_detection_rejects_old_replicate_config(): + class ReplicateConfigWithoutGroupIds: + __slots__ = () + + assert not worker._replicate_config_supports_group_ids( + ReplicateConfigWithoutGroupIds, + ReplicateConfigWithoutGroupIds(), + ) + + +def test_store_sending_thread_sets_group_ids_when_enabled(): + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.return_value = [256, 256] + replicate_config = SimpleNamespace(group_ids=None) + thread = _make_store_sending_thread( + store, + replicate_config=replicate_config, + enable_group_semantics=True, + supports_group_ids=True, + ) + + thread.add_stored_request("req-a") + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + + assert store.batch_put_from_multi_buffers.call_count == 1 + keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args + assert config is replicate_config + assert config.group_ids == [ + "vllm-mooncake-store:test-model@6130", + "vllm-mooncake-store:test-model@6131", + ] + assert len(config.group_ids) == len(keys) + + +def test_store_sending_thread_leaves_group_ids_unchanged_when_flag_disabled(): + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.return_value = [256, 256] + replicate_config = SimpleNamespace(group_ids=["existing"]) + thread = _make_store_sending_thread( + store, + replicate_config=replicate_config, + enable_group_semantics=False, + supports_group_ids=True, + ) + + thread.add_stored_request("req-a") + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + + assert store.batch_put_from_multi_buffers.call_count == 1 + assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config + assert replicate_config.group_ids == ["existing"] + + +def test_store_sending_thread_leaves_group_ids_unchanged_when_unsupported(): + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.return_value = [256, 256] + replicate_config = SimpleNamespace(group_ids=["existing"]) + thread = _make_store_sending_thread( + store, + replicate_config=replicate_config, + enable_group_semantics=True, + supports_group_ids=False, + ) + + thread.add_stored_request("req-a") + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + + assert store.batch_put_from_multi_buffers.call_count == 1 + assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config + assert replicate_config.group_ids == ["existing"] + + +def test_mooncake_group_id_uses_logical_chunk_and_excludes_physical_ranks(): + group_id = worker._make_mooncake_group_id( + KeyMetadata( + "test-model", + tp_rank=3, + pcp_rank=5, + dcp_rank=7, + pp_rank=11, + group_id=13, + ), + "abcdef", + ) + + assert group_id == "vllm-mooncake-store:test-model@abcdef" + assert "tp_rank" not in group_id + assert "pp_rank" not in group_id + assert "pcp" not in group_id + assert "dcp" not in group_id + assert "group:" not in group_id + + +def test_mooncake_group_id_uses_cache_prefix_namespace(): + metadata_a = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="deployment-a") + metadata_b = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="deployment-b") + + group_id_a = worker._make_mooncake_group_id(metadata_a, "abcdef") + group_id_b = worker._make_mooncake_group_id(metadata_b, "abcdef") + + assert group_id_a == "vllm-mooncake-store:deployment-a@test-model@abcdef" + assert group_id_b == "vllm-mooncake-store:deployment-b@test-model@abcdef" + assert group_id_a != group_id_b + + +def test_store_sending_thread_group_id_excludes_physical_sharding(): + store = MagicMock() + store.batch_is_exist.return_value = [0, 0] + store.batch_put_from_multi_buffers.return_value = [256, 256] + replicate_config = SimpleNamespace(group_ids=None) + db = ChunkedTokenDatabase( + KeyMetadata( + "test-model", + tp_rank=2, + pcp_rank=0, + dcp_rank=0, + pp_rank=0, + group_id=0, + ), + block_size=16, + ) + db.set_kv_caches_base_addr([0x1000]) + db.set_block_len([256]) + thread = _make_store_sending_thread( + store, + token_databases=[db], + tp_rank=2, + replicate_config=replicate_config, + enable_group_semantics=True, + supports_group_ids=True, + ) + + thread.add_stored_request("req-a") + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + + keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args + assert keys == [ + "test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:0@6130", + "test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:0@6131", + ] + assert config.group_ids == [ + "vllm-mooncake-store:test-model@6130", + "vllm-mooncake-store:test-model@6131", + ] + + +def test_store_sending_thread_multiple_segments_share_logical_group_id(): + store = MagicMock() + store.batch_is_exist.return_value = [0, 0] + store.batch_put_from_multi_buffers.return_value = [512, 512] + replicate_config = SimpleNamespace(group_ids=None) + db = ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16) + db.set_kv_caches_base_addr([0x1000, 0x2000]) + db.set_block_len([256, 256]) + thread = _make_store_sending_thread( + store, + token_databases=[db], + replicate_config=replicate_config, + enable_group_semantics=True, + supports_group_ids=True, + ) + + thread.add_stored_request("req-a") + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + + keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args + assert keys == [ + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130", + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131", + ] + assert addrs == [[0x1000, 0x2000], [0x1100, 0x2100]] + assert sizes == [[256, 256], [256, 256]] + assert config.group_ids == [ + "vllm-mooncake-store:test-model@6130", + "vllm-mooncake-store:test-model@6131", + ] + + +def test_store_sending_thread_group_ids_share_across_kv_cache_groups(): + from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec + + store = MagicMock() + store.batch_is_exist.return_value = [0, 0, 0, 0] + store.batch_put_from_multi_buffers.return_value = [256, 256, 256, 256] + replicate_config = SimpleNamespace(group_ids=None) + spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + coord = mooncake_store_worker.MooncakeStoreCoordinator( + [ + KVCacheGroupSpec(["layer0"], spec), + KVCacheGroupSpec(["layer1"], spec), + ], + scheduler_block_size=16, + hash_block_size=16, + ) + token_databases = [] + for group_id, base_addr in enumerate([0x1000, 0x3000]): + db = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id), + block_size=16, + ) + db.set_kv_caches_base_addr([base_addr]) + db.set_block_len([256]) + token_databases.append(db) + thread = _make_store_sending_thread( + store, + coord=coord, + token_databases=token_databases, + replicate_config=replicate_config, + enable_group_semantics=True, + supports_group_ids=True, + ) + + thread.add_stored_request("req-a") + thread._handle_request(_make_multi_group_store_req("req-a", [b"a0", b"a1"])) + + keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args + assert keys == [ + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130", + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131", + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6130", + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6131", + ] + assert addrs == [[0x1000], [0x1100], [0x3200], [0x3300]] + assert sizes == [[256], [256], [256], [256]] + # Different vLLM KV cache groups for the same prefix chunk share the + # same Mooncake lifecycle group id. + assert config.group_ids == [ + "vllm-mooncake-store:test-model@6130", + "vllm-mooncake-store:test-model@6131", + "vllm-mooncake-store:test-model@6130", + "vllm-mooncake-store:test-model@6131", + ] + assert config.group_ids[0] == config.group_ids[2] + assert config.group_ids[1] == config.group_ids[3] + + +def test_store_sending_thread_group_ids_follow_missing_key_filter(): + store = MagicMock() + store.batch_is_exist.return_value = [1, 0] + store.batch_put_from_multi_buffers.return_value = [256] + replicate_config = SimpleNamespace(group_ids=None) + thread = _make_store_sending_thread( + store, + replicate_config=replicate_config, + enable_group_semantics=True, + supports_group_ids=True, + ) + + thread.add_stored_request("req-a") + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + + keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args + assert keys == ["test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131"] + assert config.group_ids == ["vllm-mooncake-store:test-model@6131"] + + def test_estimate_disk_offload_staging_bytes_sums_multi_segment_sizes(): assert worker._estimate_disk_offload_staging_bytes([256, 512]) == 12288 @@ -1392,6 +1727,130 @@ def test_worker_put_striding_covers_every_rank_get_namespace( ) +def test_requester_worker_group_semantics_falls_back_without_group_ids( + tmp_path, + monkeypatch, +): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch) + warning = MagicMock() + monkeypatch.setattr(worker.logger, "warning", warning) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "device_name": "", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + w = worker.MooncakeStoreWorker( + _make_vllm_config( + extra_config={ + "enable_group_semantics": True, + } + ), + _make_kv_cache_config(), + ) + + assert w.enable_group_semantics is True + assert w._supports_group_ids is False + assert any( + "does not support ReplicateConfig.group_ids" in call.args[0] + for call in warning.call_args_list + ) + + +def test_requester_worker_group_semantics_string_false_does_not_enable( + tmp_path, + monkeypatch, +): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch) + warning = MagicMock() + monkeypatch.setattr(worker.logger, "warning", warning) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "device_name": "", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + w = worker.MooncakeStoreWorker( + _make_vllm_config( + extra_config={ + "enable_group_semantics": "false", + } + ), + _make_kv_cache_config(), + ) + + assert w.enable_group_semantics is False + assert not any( + "does not support ReplicateConfig.group_ids" in call.args[0] + for call in warning.call_args_list + ) + + +def test_requester_worker_group_semantics_string_true_enables( + tmp_path, + monkeypatch, +): + store = MagicMock() + store.setup.return_value = 0 + + class FakeReplicateConfig: + def __init__(self) -> None: + self.group_ids = None + + fake_store_module = types.ModuleType("mooncake.store") + fake_store_module.MooncakeDistributedStore = lambda: store # type: ignore[attr-defined] + fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined] + fake_mooncake_module = types.ModuleType("mooncake") + fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module) + monkeypatch.setitem(sys.modules, "mooncake.store", fake_store_module) + _patch_worker_runtime(monkeypatch) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "device_name": "", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + w = worker.MooncakeStoreWorker( + _make_vllm_config( + extra_config={ + "enable_group_semantics": "true", + } + ), + _make_kv_cache_config(), + ) + + assert w.enable_group_semantics is True + assert w._supports_group_ids is True + + # --------------------------------------------------------------------------- # Helpers for register_kv_caches tests # --------------------------------------------------------------------------- @@ -1481,7 +1940,7 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = ( - lambda keys, addrs, sizes, replicate_config: [256] * len(keys) + lambda keys, addrs, sizes, *_args: [256] * len(keys) ) full_spec = FullAttentionSpec( @@ -1767,6 +2226,8 @@ def _make_bare_worker( worker.disk_offload_buffer_budget_bytes = None worker.store_replicate_config = SimpleNamespace() + worker.enable_group_semantics = False + worker._supports_group_ids = False worker._kv_connector_stats_lock = threading.Lock() worker.kv_connector_stats = MooncakeStoreConnectorStats() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 769aaae1a4..97bb89d2ec 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -89,6 +89,22 @@ def _rotate_list(values: list[_T], offset: int) -> list[_T]: return values[offset:] + values[:offset] +def _replicate_config_supports_group_ids( + replicate_config_cls: type[Any], + replicate_config: Any, +) -> bool: + if hasattr(replicate_config_cls, "group_ids"): + return True + return hasattr(replicate_config, "group_ids") + + +def _make_mooncake_group_id(metadata: KeyMetadata, chunk_hash: str) -> str: + # Mooncake group ids describe the lifecycle unit. For vLLM, that unit is + # a prefix chunk, so shard dimensions stay only in the object key. + prefix = f"{metadata.cache_prefix}@" if metadata.cache_prefix else "" + return f"vllm-mooncake-store:{prefix}{metadata.model_name}@{chunk_hash}" + + # Mirrors FileStorageConfig::local_buffer_size in Mooncake C++. DEFAULT_MOONCAKE_DISK_STAGING_BUFFER_BYTES = 1280 * 1024 * 1024 @@ -459,6 +475,8 @@ class KVCacheStoreSendingThread(KVTransferThread): ready_event: threading.Event, enable_kv_event: bool = False, replicate_config: Any = None, + enable_group_semantics: bool = False, + supports_group_ids: bool = False, record_operation: Callable[..., None] | None = None, ): super().__init__( @@ -479,6 +497,8 @@ class KVCacheStoreSendingThread(KVTransferThread): # Caller always passes a non-None ReplicateConfig — see # MooncakeStoreWorker.__init__ where store_replicate_config is built. self.replicate_config = replicate_config + self.enable_group_semantics = enable_group_semantics + self.supports_group_ids = supports_group_ids # Pause store requests when CPU/disk offloading is under pressure. self._store_pressure_active = False @@ -567,6 +587,9 @@ class KVCacheStoreSendingThread(KVTransferThread): keys: list[str] = [] addrs: list[list[int]] = [] sizes: list[list[int]] = [] + group_ids: list[str] | None = ( + [] if self.enable_group_semantics and self.supports_group_ids else None + ) saved = self._saved_offset.get(req_meta.req_id, 0) for g_idx, db in enumerate(self.token_databases): group_blocks = req_meta.block_ids[g_idx] @@ -603,9 +626,17 @@ class KVCacheStoreSendingThread(KVTransferThread): ) continue addr, size = db.prepare_value_for_block(block_id) - keys.append(db.key_for(key_hash)) + key = db.key_for(key_hash) + keys.append(key) addrs.append(addr) sizes.append(size) + if group_ids is not None: + group_ids.append( + _make_mooncake_group_id( + db.metadata, + key.rsplit("@", 1)[-1], + ) + ) if not keys: return True @@ -633,9 +664,14 @@ class KVCacheStoreSendingThread(KVTransferThread): keys = [keys[i] for i in missing] addrs = [addrs[i] for i in missing] sizes = [sizes[i] for i in missing] + if group_ids is not None: + group_ids = [group_ids[i] for i in missing] if req_meta.current_event is not None: # Fence the CoW block copy enqueued earlier this step. req_meta.current_event.synchronize() + if group_ids is not None: + assert len(group_ids) == len(keys) + self.replicate_config.group_ids = group_ids batch_bytes = _sum_batch_bytes(sizes) put_start = time.perf_counter() try: @@ -797,6 +833,18 @@ class KVCacheStoreSendingThread(KVTransferThread): ] group_indices = [group_indices[i] for i in missing_indices] + group_ids = ( + [ + _make_mooncake_group_id( + self.token_databases[g_idx].metadata, + key.rsplit("@", 1)[-1], + ) + for key, g_idx in zip(keys, group_indices, strict=True) + ] + if self.enable_group_semantics and self.supports_group_ids + else None + ) + logger.debug( "Storing KV cache for %d blocks (groups=%s) for request %s", len(keys), @@ -855,6 +903,10 @@ class KVCacheStoreSendingThread(KVTransferThread): if current_event is not None: current_event.synchronize() + if group_ids is not None: + assert len(group_ids) == len(keys) + self.replicate_config.group_ids = group_ids + batch_bytes = _sum_batch_bytes(sizes) put_start = time.perf_counter() try: @@ -1210,6 +1262,19 @@ class MooncakeStoreWorker: preferred_segment = rdma_utils.get_configured_preferred_segment(extra_config) self.preferred_segment = preferred_segment self.store_replicate_config = ReplicateConfig() + self.enable_group_semantics = ( + str(extra_config.get("enable_group_semantics", "False")).strip().lower() + == "true" + ) + self._supports_group_ids = _replicate_config_supports_group_ids( + ReplicateConfig, self.store_replicate_config + ) + if self.enable_group_semantics and not self._supports_group_ids: + logger.warning( + "Mooncake group semantics is enabled, but the installed " + "Mooncake package does not support ReplicateConfig.group_ids. " + "Falling back to the existing batch_put_from_multi_buffers path." + ) if preferred_segment is not None: self.store_replicate_config.preferred_segment = preferred_segment @@ -1503,6 +1568,8 @@ class MooncakeStoreWorker: ready_event_sending, self.enable_kv_events, self.store_replicate_config, + enable_group_semantics=self.enable_group_semantics, + supports_group_ids=self._supports_group_ids, record_operation=self._record_kv_connector_operation, ) self.kv_send_thread.start()