fix(storage): enforce Mooncake correctness contracts

tests/test_mooncake_utils.py

@@ -34,6 +34,13 @@
 _has_cuda_python = importlib.util.find_spec("cuda") is not None
 
 
+def test_mooncake_correctness_contract_version_is_public():
+    import transfer_queue as tq
+
+    assert tq.MOONCAKE_CORRECTNESS_CONTRACT_VERSION == 1
+    assert "MOONCAKE_CORRECTNESS_CONTRACT_VERSION" in tq.__all__
+
+
 def _aligned(n: int) -> int:
     return (n + _DEFAULT_ALIGN - 1) // _DEFAULT_ALIGN * _DEFAULT_ALIGN
 
@@ -376,12 +383,18 @@ def test_no_gdr_meta_none_no_warning(self, caplog):
             client.clear(["k0"], custom_backend_meta=None)
         assert "custom_backend_meta" not in caplog.text
 
-    def test_error_code_triggers_log(self, caplog):
+    def test_non_idempotent_failure_is_raised(self):
         client = _make_clear_client(use_gdr=False)
         client._store.batch_remove.side_effect = lambda keys, force: [-1] * len(keys)
-        with caplog.at_level(logging.ERROR, logger="transfer_queue.storage.clients.mooncake_client"):
+        with pytest.raises(RuntimeError, match=r"batch_remove failed: k0=-1"):
             client.clear(["k0"])
-        assert "remove failed" in caplog.text
+
+    def test_short_result_is_raised(self):
+        client = _make_clear_client(use_gdr=False)
+        client._store.batch_remove.return_value = [0]
+        client._store.batch_remove.side_effect = None
+        with pytest.raises(RuntimeError, match="returned 1 results, expected 2"):
+            client.clear(["k0", "k1"])
 
     def test_already_removed_code_704_is_silent(self, caplog):
         client = _make_clear_client(use_gdr=False)
@@ -395,3 +408,51 @@ def test_success_code_zero_is_silent(self, caplog):
         with caplog.at_level(logging.ERROR):
             client.clear(["k0"])
         assert "remove failed" not in caplog.text
+
+
+class _SequenceStore:
+    """Return configured results from each low-level Mooncake batch call."""
+
+    def __init__(self, results):
+        self.results = iter(results)
+
+    def batch_upsert_from(self, keys, ptrs, sizes, config=None):
+        return next(self.results)
+
+    def batch_get_into(self, keys, ptrs, sizes):
+        return next(self.results)
+
+
+def _make_retry_client(store):
+    from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient
+
+    client = object.__new__(MooncakeStoreClient)
+    client._store = store
+    client.replica_config = None
+    return client
+
+
+class TestBatchResultValidation:
+    def test_upsert_retry_short_result_is_raised(self, monkeypatch):
+        monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0)
+        client = _make_retry_client(_SequenceStore([[-1, -1], [0]]))
+
+        with pytest.raises(RuntimeError, match="batch_upsert_from returned 1 results, expected 2"):
+            client._batch_upsert_with_retry(["k0", "k1"], [1, 2], [8, 8])
+
+    def test_get_retry_short_result_is_raised(self, monkeypatch):
+        monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0)
+        client = _make_retry_client(_SequenceStore([[-1, -1], [0]]))
+
+        with pytest.raises(RuntimeError, match="batch_get_into returned 1 results, expected 2"):
+            client._batch_get_into_with_retry(["k0", "k1"], [1, 2], [8, 8])
+
+    @pytest.mark.parametrize("operation", ["upsert", "get"])
+    def test_non_sized_result_is_raised(self, operation):
+        client = _make_retry_client(_SequenceStore([None]))
+
+        with pytest.raises(RuntimeError, match="returned a non-sized result, expected 2 codes"):
+            if operation == "upsert":
+                client._batch_upsert_with_retry(["k0", "k1"], [1, 2], [8, 8])
+            else:
+                client._batch_get_into_with_retry(["k0", "k1"], [1, 2], [8, 8])

tests/test_storage_manager_notifications.py

@@ -0,0 +1,115 @@
+# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import pytest
+
+from transfer_queue.storage.managers.base import StorageManager
+from transfer_queue.utils.zmq_utils import ZMQRequestType
+
+
+class _FakeNotifySocket:
+    def __init__(self, connect_error: Exception | None = None) -> None:
+        self.closed = False
+        self.connect_error = connect_error
+
+    def setsockopt(self, *args, **kwargs) -> None:
+        pass
+
+    def connect(self, *args, **kwargs) -> None:
+        if self.connect_error is not None:
+            raise self.connect_error
+
+    async def send_multipart(self, request) -> None:
+        pass
+
+    async def recv_multipart(self, copy=False):
+        return [b"ack"]
+
+    def close(self, linger=0) -> None:
+        self.closed = True
+
+
+def _manager(controller_info=None):
+    if controller_info is None:
+        controller_info = SimpleNamespace(
+            id="controller",
+            ip="127.0.0.1",
+            to_addr=lambda name: "inproc://controller",
+        )
+    return SimpleNamespace(
+        storage_manager_id="notification-test",
+        zmq_context=object(),
+        controller_info=controller_info,
+    )
+
+
+def _ack(success: bool, partition_id: str = "p0"):
+    return SimpleNamespace(
+        request_type=ZMQRequestType.NOTIFY_DATA_UPDATE_ACK,
+        sender_id="controller",
+        body={"success": success, "partition_id": partition_id},
+    )
+
+
+@pytest.mark.asyncio
+async def test_notify_data_update_rejects_missing_controller():
+    manager = _manager(controller_info=False)
+
+    with pytest.raises(RuntimeError, match="has no controller"):
+        await StorageManager.notify_data_update(manager, "p0", [], {}, {})
+
+
+@pytest.mark.asyncio
+async def test_notify_and_wait_requires_positive_ack(monkeypatch):
+    socket = _FakeNotifySocket()
+    monkeypatch.setattr("transfer_queue.storage.managers.base.create_zmq_socket", lambda **kwargs: socket)
+    monkeypatch.setattr("transfer_queue.storage.managers.base.ZMQMessage.deserialize", lambda messages: _ack(False))
+
+    with pytest.raises(RuntimeError, match="rejected the production-status update"):
+        await StorageManager._notify_and_wait(_manager(), [b"request"])
+    assert socket.closed is True
+
+
+@pytest.mark.asyncio
+async def test_notify_and_wait_accepts_positive_ack(monkeypatch):
+    socket = _FakeNotifySocket()
+    monkeypatch.setattr("transfer_queue.storage.managers.base.create_zmq_socket", lambda **kwargs: socket)
+    monkeypatch.setattr("transfer_queue.storage.managers.base.ZMQMessage.deserialize", lambda messages: _ack(True))
+
+    await StorageManager._notify_and_wait(_manager(), [b"request"])
+    assert socket.closed is True
+
+
+@pytest.mark.asyncio
+async def test_notify_and_wait_times_out_without_ack(monkeypatch):
+    socket = _FakeNotifySocket()
+    monkeypatch.setattr("transfer_queue.storage.managers.base.create_zmq_socket", lambda **kwargs: socket)
+    monkeypatch.setattr("transfer_queue.storage.managers.base.TQ_DATA_UPDATE_RESPONSE_TIMEOUT", 0)
+
+    with pytest.raises(TimeoutError, match="production-status ACK"):
+        await StorageManager._notify_and_wait(_manager(), [b"request"])
+    assert socket.closed is True
+
+
+@pytest.mark.asyncio
+async def test_notify_and_wait_closes_socket_when_connect_fails(monkeypatch):
+    socket = _FakeNotifySocket(connect_error=ConnectionError("controller unavailable"))
+    monkeypatch.setattr("transfer_queue.storage.managers.base.create_zmq_socket", lambda **kwargs: socket)
+
+    with pytest.raises(ConnectionError, match="controller unavailable"):
+        await StorageManager._notify_and_wait(_manager(), [b"request"])
+    assert socket.closed is True

transfer_queue/__init__.py

@@ -45,6 +45,10 @@
 from .sampler.sequential_sampler import SequentialSampler
 from .sampler.streaming_token_budget_sampler import StreamingTokenBudgetSampler
 
+# Version 1 guarantees Mooncake batch/retry result-count validation,
+# batch_remove failure propagation, and fail-closed production-status ACKs.
+MOONCAKE_CORRECTNESS_CONTRACT_VERSION = 1
+
 __all__ = (
     [
         # High-Level KV Interface
@@ -80,6 +84,7 @@
         "get_client",
         "BatchMeta",
         "TransferQueueClient",
+        "MOONCAKE_CORRECTNESS_CONTRACT_VERSION",
     ]
     + [
         # Sampler

transfer_queue/storage/clients/mooncake_client.py

@@ -45,6 +45,17 @@
 MAX_SERIAL_WORKER_THREADS = 4
 MAX_RETRIES = 3
 RETRY_DELAY_SECONDS = 1.0
+_MOONCAKE_OBJECT_NOT_FOUND = -704
+
+
+def _validate_batch_result_count(operation: str, keys: list[str], results: Any) -> None:
+    """Require one Mooncake result code for every requested key."""
+    try:
+        actual = len(results)
+    except TypeError as error:
+        raise RuntimeError(f"{operation} returned a non-sized result, expected {len(keys)} codes") from error
+    if actual != len(keys):
+        raise RuntimeError(f"{operation} returned {actual} results, expected {len(keys)}")
 
 
 @StorageClientFactory.register("MooncakeStoreClient")
@@ -529,9 +540,15 @@ def clear(self, keys: list[str], custom_backend_meta: list[Any] | None = None) -
             actual_keys = keys
 
         ret_codes = self._store.batch_remove(actual_keys, force=True)
-        for i, ret in enumerate(ret_codes):
-            if not (ret == 0 or ret == -704):
-                logger.error(f"remove failed for key `{actual_keys[i]}` with error code: {ret}")
+        _validate_batch_result_count("batch_remove", actual_keys, ret_codes)
+        failures = [
+            (key, code)
+            for key, code in zip(actual_keys, ret_codes, strict=True)
+            if code not in (0, _MOONCAKE_OBJECT_NOT_FOUND)
+        ]
+        if failures:
+            detail = ", ".join(f"{key}={code}" for key, code in failures)
+            raise RuntimeError(f"batch_remove failed: {detail}")
 
     def close(self):
         """Closes MooncakeStore."""
@@ -549,8 +566,7 @@ def _batch_upsert_with_retry(self, batch_keys: list[str], batch_ptrs: list[int],
         backing tensors/buffers).
         """
         results = self._store.batch_upsert_from(batch_keys, batch_ptrs, batch_sizes, config=self.replica_config)
-        if len(results) != len(batch_keys):
-            raise RuntimeError(f"batch_upsert_from returned {len(results)} results, expected {len(batch_keys)}")
+        _validate_batch_result_count("batch_upsert_from", batch_keys, results)
 
         failed_indices = [j for j, r in enumerate(results) if r != 0]
         if not failed_indices:
@@ -572,6 +588,7 @@ def _batch_upsert_with_retry(self, batch_keys: list[str], batch_ptrs: list[int],
             retry_results = self._store.batch_upsert_from(
                 current_failed_keys, retry_ptrs, retry_sizes, config=self.replica_config
             )
+            _validate_batch_result_count("batch_upsert_from", current_failed_keys, retry_results)
 
             next_failed_indices = []
             next_failed_keys = []
@@ -612,8 +629,7 @@ def _batch_get_into_with_retry(
         Caller owns the receive buffers (allocate/register/unregister).
         """
         ret_codes = self._store.batch_get_into(batch_keys, batch_buffer_ptrs, batch_nbytes)
-        if len(ret_codes) != len(batch_keys):
-            raise RuntimeError(f"batch_get_into returned {len(ret_codes)} results, expected {len(batch_keys)}")
+        _validate_batch_result_count("batch_get_into", batch_keys, ret_codes)
 
         failed_indices = [i for i, ret in enumerate(ret_codes) if ret < 0]
         if not failed_indices:
@@ -634,6 +650,7 @@ def _batch_get_into_with_retry(
             retry_nbytes = [batch_nbytes[i] for i in current_failed_indices]
 
             retry_codes = self._store.batch_get_into(current_failed_keys, retry_ptrs, retry_nbytes)
+            _validate_batch_result_count("batch_get_into", current_failed_keys, retry_codes)
 
             next_failed_indices = []
             next_failed_keys = []