[fix] Raise in put when the controller rejects the written data

tests/e2e/test_kv_interface_e2e.py

@@ -21,6 +21,7 @@
 """
 
 import asyncio
+import logging
 import os
 
 import pytest
@@ -470,6 +471,48 @@ def test_kv_put_returns_cumulative_fields(self, controller, tq_api):
 
         tq_api.kv_clear(keys=key, partition_id=partition_id)
 
+    def test_kv_put_can_retry_after_dtype_rejection(self, tq_api):
+        partition_id = "test_partition"
+        tq_api.kv_put(key="float_x", partition_id=partition_id, fields={"x": torch.tensor([1.0, 2.0])})
+
+        with pytest.raises(RuntimeError, match=r"Field 'x' dtype mismatch: torch.float32 != torch.int64"):
+            tq_api.kv_put(key="int_x", partition_id=partition_id, fields={"x": torch.tensor([3, 4])})
+
+        stored = tq_api.kv_batch_get(keys="float_x", partition_id=partition_id, select_fields="x")
+        assert_tensor_equal(stored["x"][0], torch.tensor([1.0, 2.0]))
+
+        tq_api.kv_put(key="int_x", partition_id=partition_id, fields={"x": torch.tensor([3.0, 4.0])})
+        stored = tq_api.kv_batch_get(keys="int_x", partition_id=partition_id, select_fields="x")
+        assert_tensor_equal(stored["x"][0], torch.tensor([3.0, 4.0]))
+        tq_api.kv_clear(keys=["float_x", "int_x"], partition_id=partition_id)
+
+    def test_rejected_overwrite_stays_readable_until_corrected(self, tq_api, backend_name):
+        if backend_name != "SimpleStorage":
+            pytest.skip("Reading the overwritten value with stale dtype metadata depends on the backend")
+
+        partition_id = "test_partition"
+        tq_api.kv_put(key="a", partition_id=partition_id, fields={"x": torch.tensor([1.0, 2.0])})
+        with pytest.raises(RuntimeError, match="dtype mismatch"):
+            tq_api.kv_put(key="a", partition_id=partition_id, fields={"x": torch.tensor([3, 4])})
+
+        stored = tq_api.kv_batch_get(keys="a", partition_id=partition_id, select_fields="x")
+        assert stored["x"][0].dtype == torch.int64
+        assert_tensor_equal(stored["x"][0], torch.tensor([3, 4]))
+
+        tq_api.kv_put(key="a", partition_id=partition_id, fields={"x": torch.tensor([5.0, 6.0])})
+        stored = tq_api.kv_batch_get(keys="a", partition_id=partition_id, select_fields="x")
+        assert stored["x"][0].dtype == torch.float32
+        assert_tensor_equal(stored["x"][0], torch.tensor([5.0, 6.0]))
+        tq_api.kv_clear(keys="a", partition_id=partition_id)
+
+    def test_notify_for_a_cleared_partition_warns_instead_of_raising(self, tq_system, caplog):
+        # A writer that finishes after its partition was cleared must not fail.
+        manager = tq.get_client().storage_manager
+        schema = {"x": {"dtype": torch.float32, "shape": (2,), "is_nested": False, "is_non_tensor": False}}
+        with caplog.at_level(logging.WARNING):
+            assert asyncio.run(manager.notify_data_update("no_such_partition", [0], schema)) is False
+        assert "no longer exists" in caplog.text
+
 
 class TestKVBatchPutE2E:
     """End-to-end tests for kv_batch_put functionality."""

tests/test_controller_data_partitions.py

@@ -1363,12 +1363,15 @@ def test_unfinalized_write_records_metadata_but_stays_hidden(self):
         assert partition.scan_data_status(["pv"], task_name="trainer") == [0, 1]
 
     def test_dtype_conflict_leaves_readiness_unchanged(self):
+        import pytest
+
         partition = self._make_partition()
         f32 = {"dtype": "torch.float32", "shape": (4,), "is_nested": False, "is_non_tensor": False}
         i64 = {**f32, "dtype": "torch.int64"}
         partition.update_production_status([0, 1], ["pv"], {"pv": f32})
 
-        assert not partition.update_production_status([1], ["pv"], {"pv": i64}, ready=False)
+        with pytest.raises(ValueError, match="dtype mismatch"):
+            partition.update_production_status([1], ["pv"], {"pv": i64}, ready=False)
         assert partition.scan_data_status(["pv"], task_name="trainer") == [0, 1]
         assert partition.field_metadata["pv"].dtype == "torch.float32"
 

tests/test_storage_request_retry.py

@@ -15,6 +15,7 @@
 
 """Storage-unit retry and timeout-diagnosis tests."""
 
+import asyncio
 import logging
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -88,6 +89,48 @@ async def test_controller_notification_returns_ack_status(accepted):
     assert await manager._notify_and_wait([b"request"]) is accepted
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "ack, published, warns",
+    [
+        (True, True, False),
+        (None, False, False),  # no ACK in time: the outcome is unknown, so the put is not failed
+        (False, False, True),
+    ],
+)
+async def test_notify_publishes_or_tolerates_ack(ack, published, warns, caplog):
+    manager = _manager()
+    manager.controller_info = MagicMock(id="controller")
+    manager._notify_loop = asyncio.get_running_loop()
+    manager._notify_and_wait = AsyncMock(return_value=ack)
+
+    with caplog.at_level(logging.WARNING):
+        assert await manager.notify_data_update("p", [0, 1], {}) is published
+    assert ("no longer exists" in caplog.text) is warns
+
+
+@pytest.mark.asyncio
+async def test_notify_raises_on_request_error_without_waiting_for_timeout():
+    manager = _manager()
+    manager.controller_info = MagicMock(id="controller")
+    manager._notify_loop = asyncio.get_running_loop()
+    socket = MagicMock()
+    socket.send_multipart = AsyncMock()
+    error = ZMQMessage.create(
+        request_type=ZMQRequestType.REQUEST_ERROR,
+        sender_id="controller",
+        body={"success": False, "message": "Field 'x' dtype mismatch"},
+    )
+    socket.recv_multipart = AsyncMock(side_effect=[error.serialize(), TimeoutError("no further reply")])
+    manager.notify_pool = MagicMock()
+    manager.notify_pool.alease.return_value = _SocketLease(socket)
+
+    with pytest.raises(RuntimeError, match="Field 'x' dtype mismatch"):
+        await manager.notify_data_update("p", [0], {})
+    socket.recv_multipart.assert_awaited_once()
+    socket.close.assert_not_called()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("data_parser, expected_attempts", [(None, 3), (lambda data: data, 1)])
 async def test_put_retries_only_without_a_parser(data_parser, expected_attempts):
@@ -142,6 +185,10 @@ async def test_update_reports_controller_rejection_after_storage_commit():
     manager._update_to_single_storage_unit.assert_awaited_once()
     manager.notify_data_update.assert_awaited_once()
 
+    manager.notify_data_update = AsyncMock(side_effect=RuntimeError("Controller failed to record written data: dtype"))
+    with pytest.raises(RuntimeError, match="Storage update committed.*notification failed.*dtype.*do not retry"):
+        await manager.update_data(metadata, values, lambda old, new: torch.cat([old, new]), _HELD_INPUT_IDS)
+
 
 @pytest.mark.asyncio
 async def test_update_timeout_is_not_retried_or_published():

transfer_queue/controller.py

@@ -521,15 +521,19 @@ def update_production_status(
 
         Returns:
             True if update was successful, False on error
+
+        Raises:
+            ValueError: If a field's dtype conflicts with the partition's; metadata and readiness are unchanged.
         """
-        try:
-            # Same dtype check as FieldMeta.update, run first so a conflicting write changes nothing.
-            for name, incoming in field_schema.items():
-                existing = self.field_metadata.get(name)
-                dtype = incoming.get("dtype")
-                if existing is not None and None not in (existing.dtype, dtype) and existing.dtype != dtype:
-                    raise ValueError(f"Field {name!r} dtype mismatch: {existing.dtype} != {dtype}")
+        # A dtype conflict must leave controller metadata and readiness unchanged.
+        # Payloads are already stored, so this check cannot roll back the write.
+        for name, incoming in field_schema.items():
+            existing = self.field_metadata.get(name)
+            dtype = incoming.get("dtype")
+            if existing is not None and None not in (existing.dtype, dtype) and existing.dtype != dtype:
+                raise ValueError(f"Field {name!r} dtype mismatch: {existing.dtype} != {dtype}")
 
+        try:
             # Derive field_names from field_schema to guarantee consistency
             field_names = list(field_schema.keys())
 
@@ -1155,6 +1159,9 @@ def update_production_status(
 
         Returns:
             True if update was successful, False otherwise
+
+        Raises:
+            ValueError: If a field's dtype conflicts with the partition's; metadata and readiness are unchanged.
         """
         field_names = list(field_schema.keys())
         partition = self._get_partition(partition_id)
@@ -2137,29 +2144,24 @@ def _handle_get_meta_request(self, request_msg: ZMQMessage) -> ZMQMessage:
 
     def _handle_notify_data_update_request(self, request_msg: ZMQMessage) -> ZMQMessage:
         params = request_msg.body
-        partition_id = params.get("partition_id")
+        partition_id = params["partition_id"]
         global_indexes = params.get("global_indexes", [])
-        success = self.update_production_status(
-            partition_id=cast(str, partition_id),
+        if self._get_partition(partition_id) is None:
+            return self._make_response(request_msg, ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, {"success": False})
+
+        if not self.update_production_status(
+            partition_id=partition_id,
             global_indexes=global_indexes,
             field_schema=params.get("field_schema", {}),
             custom_backend_meta=params.get("custom_backend_meta", {}),
             ready=params.get("ready", True),
-        )
-        if success:
-            if self._metrics is not None:
-                self._metrics.record_samples("NOTIFY_DATA_UPDATE", len(global_indexes))
-            logger.debug(f"[{self.controller_id}]: Updated production status for partition {partition_id}")
+        ):
+            raise RuntimeError(f"Failed to record data for partition {partition_id!r}; see the controller log")
+        if self._metrics is not None:
+            self._metrics.record_samples("NOTIFY_DATA_UPDATE", len(global_indexes))
+        logger.debug(f"[{self.controller_id}]: Updated production status for partition {partition_id}")
 
-        return self._make_response(
-            request_msg,
-            ZMQRequestType.NOTIFY_DATA_UPDATE_ACK,
-            {
-                "controller_id": self.controller_id,
-                "partition_id": partition_id,
-                "success": success,
-            },
-        )
+        return self._make_response(request_msg, ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, {"success": True})
 
     def _handle_get_partition_meta_request(self, request_msg: ZMQMessage) -> ZMQMessage:
         partition_id = request_msg.body["partition_id"]

transfer_queue/interface.py

@@ -399,6 +399,8 @@ def kv_put(
         ValueError: If neither fields nor tag is provided
         ValueError: If nested tensors are provided (use kv_batch_put instead)
         RuntimeError: If retrieved BatchMeta size doesn't match length of `keys`
+        RuntimeError: If the controller fails to record the written data, e.g. on a dtype conflict.
+            Storage changes are not rolled back; overwrites of ready fields may remain visible.
 
     Example:
         >>> import transfer_queue as tq
@@ -459,6 +461,8 @@ def kv_batch_put(
         ValueError: When ``tags`` length mismatches key count.
         ValueError: When ``keys`` contains duplicates.
         RuntimeError: When retrieved metadata size mismatches input key count.
+        RuntimeError: When the controller fails to record the written data, e.g. on a dtype conflict.
+            Storage changes are not rolled back; overwrites of ready fields may remain visible.
 
     Example:
         >>> import transfer_queue as tq
@@ -801,6 +805,8 @@ async def async_kv_put(
         ValueError: If neither fields nor tag is provided
         ValueError: If nested tensors are provided (use kv_batch_put instead)
         RuntimeError: If retrieved BatchMeta size doesn't match length of `keys`
+        RuntimeError: If the controller fails to record the written data, e.g. on a dtype conflict.
+            Storage changes are not rolled back; overwrites of ready fields may remain visible.
 
     Example:
         >>> import transfer_queue as tq
@@ -910,6 +916,8 @@ async def async_kv_batch_put(
         ValueError: If length of `keys` doesn't match length of `tags` or the batch_size of `fields` TensorDict
         ValueError: If `keys` contains duplicates
         RuntimeError: If retrieved BatchMeta size doesn't match length of `keys`
+        RuntimeError: If the controller fails to record the written data, e.g. on a dtype conflict.
+            Storage changes are not rolled back; overwrites of ready fields may remain visible.
 
     Example:
         >>> import transfer_queue as tq

transfer_queue/storage/managers/base.py

@@ -237,6 +237,13 @@ async def notify_data_update(
             field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}}.
             custom_backend_meta: Per-field custom_meta for each sample, in {global_index: {field: custom_meta}} format.
             ready: Whether consumers may see the data now; False for a kv_update with finalize=False.
+
+        Returns:
+            True if the controller recorded the data. False if no ACK arrived (the outcome is
+            unknown) or the partition no longer exists (the data is dropped, with a warning).
+
+        Raises:
+            RuntimeError: If the controller failed to record the data, e.g. on a dtype conflict.
         """
 
         if not self.controller_info:
@@ -274,10 +281,16 @@ async def notify_data_update(
             self._notify_and_wait(request_msg),
             self._notify_loop,
         )
-        return await asyncio.wrap_future(thread_future)
+        accepted = await asyncio.wrap_future(thread_future)
+        if accepted is False:
+            logger.warning(
+                f"[{self.storage_manager_id}]: Partition {partition_id!r} no longer exists; "
+                f"the controller dropped {len(global_indexes)} written samples"
+            )
+        return accepted is True
 
-    async def _notify_and_wait(self, request_msg: list) -> bool:
-        """Send a data status notification to the controller and block until ACK is received."""
+    async def _notify_and_wait(self, request_msg: list) -> bool | None:
+        """Return ACK success, or None if no reply arrived; raise on REQUEST_ERROR."""
         # Acquiring the lease sits outside the handler below: a missing socket name or a dead
         # context is a configuration/lifecycle fault the caller must see, not a slow ACK.
         async with self.notify_pool.alease(self.controller_info) as sock:
@@ -300,19 +313,25 @@ async def _notify_and_wait(self, request_msg: list) -> bool:
                     messages = await asyncio.wait_for(sock.recv_multipart(copy=False), timeout=remaining)
                     response_msg = ZMQMessage.deserialize(messages)
 
-                    if response_msg.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE_ACK:  # type: ignore[arg-type]
-                        success = bool(response_msg.body.get("success"))
-                        logger.debug(
-                            f"[{self.storage_manager_id}]: Received data status update ACK "
-                            f"from controller id #{response_msg.sender_id}: success={success}."
-                        )
-                        return success
+                    if response_msg.request_type in (
+                        ZMQRequestType.NOTIFY_DATA_UPDATE_ACK,
+                        ZMQRequestType.REQUEST_ERROR,
+                    ):
+                        break
             except Exception as e:
                 # Logged rather than raised, so a slow controller does not fail the put. Close
                 # the socket: a late ACK would otherwise be read as the next lessee's reply.
                 logger.error(f"[{self.storage_manager_id}]: Data status update failed: {type(e).__name__}: {e}")
                 sock.close(linger=0)
-                return False
+                return None
+
+        logger.debug(
+            f"[{self.storage_manager_id}]: Received data status update response "
+            f"from controller id #{response_msg.sender_id}: {response_msg.body}."
+        )
+        if response_msg.request_type == ZMQRequestType.REQUEST_ERROR:
+            raise RuntimeError(f"Controller failed to record written data: {response_msg.body['message']}")
+        return bool(response_msg.body["success"])
 
     @abstractmethod
     async def put_data(

transfer_queue/storage/managers/simple_storage_manager.py

@@ -605,12 +605,17 @@ async def update_data(
             metadata.global_indexes, list(zip([g.global_indexes for g in routing.values()], described, strict=True))
         )
 
-        published = await self.notify_data_update(
-            metadata.partition_ids[0],
-            metadata.global_indexes,
-            field_schema,
-            ready=finalize,
-        )
+        try:
+            published = await self.notify_data_update(
+                metadata.partition_ids[0],
+                metadata.global_indexes,
+                field_schema,
+                ready=finalize,
+            )
+        except RuntimeError as e:
+            raise RuntimeError(
+                f"Storage update committed, but controller notification failed: {e}; do not retry this merge blindly"
+            ) from e
         if not published:
             raise RuntimeError(
                 "Storage update committed, but the controller did not accept its metadata; "