Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 0 additions & 5 deletions tests/e2e/test_e2e_lifecycle_consistency.py
Original file line number Diff line number Diff line change
Expand Up @@ -424,11 +424,6 @@ def test_core_consistency(e2e_client):
def test_cross_shard_complex_update(e2e_client):
"""Cross-shard update: put A + put B, update overlapping region, verify all regions."""

# FIXME: Add data update test to MooncakeStore after Upsert function is ready
# https://github.com/kvcache-ai/Mooncake/issues/1645
if os.environ.get("TQ_TEST_BACKEND", "SimpleStorage") == "MooncakeStore":
return

client = e2e_client
partition_id = "test_cross_shard_update"
task_name = "cross_shard_task"
Expand Down
17 changes: 7 additions & 10 deletions transfer_queue/storage/managers/mooncake_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,22 +16,19 @@
from typing import Any

from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory
from transfer_queue.utils.logging_utils import get_logger
from transfer_queue.utils.zmq_utils import ZMQServerInfo

logger = get_logger(__name__)


@StorageManagerFactory.register("MooncakeStore")
class MooncakeStorageManager(KVStorageManager):
"""Storage manager for MooncakeStorage backend."""
"""Storage manager for MooncakeStorage backend.

def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]):
logger.warning(
"MooncakeStore backend doesn't support key update (upsert) for now. "
"You must delete the key before updating it. "
"Refer to https://github.com/kvcache-ai/Mooncake/issues/1645 for details."
)
Key update (upsert) is supported natively via the MooncakeStore client's
``batch_upsert_from`` (zero-copy tensor path) and ``upsert_batch`` (raw bytes
path). See ``mooncake-integration/store/store_py.cpp`` upstream for the
pybind bindings.
"""

def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]):
config["client_name"] = "MooncakeStoreClient"
super().__init__(controller_info, config)
Loading