diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index e1e52dcf..572336a8 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/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.""" diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index 3d92e175..fd458ed1 100644 --- a/tests/test_controller_data_partitions.py +++ b/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" diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py index 854c35ac..4d49e6a2 100644 --- a/tests/test_storage_request_retry.py +++ b/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(): diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index d19152cd..c44ac236 100644 --- a/transfer_queue/controller.py +++ b/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"] diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 7fd23610..0b4ab04b 100644 --- a/transfer_queue/interface.py +++ b/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 diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 88f7b37a..9f6eb739 100644 --- a/transfer_queue/storage/managers/base.py +++ b/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( diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index b17fa087..25b0eeed 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/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; "