From 32bfb3dc41e8888da6457f2034cdb2dc86ecf555 Mon Sep 17 00:00:00 2001 From: 0oshowero0 Date: Thu, 15 Oct 2026 17:17:00 +0800 Subject: [PATCH 1/2] [fix] put: raise when the controller rejects the written data notify_data_update returned False both when no ACK arrived and when the controller rejected the update (e.g. a dtype conflict), and put_data and restore ignored it. A rejected kv_put returned normally while none of its fields became ready: the samples silently never reached kv_batch_get or the dataloader, and the later error pointed at kv_update(finalize=False). The controller ACK now says why it rejected an update and whether the partition still exists, and notify_data_update raises on a rejection. Unchanged: a missing ACK still returns False, since the controller may still apply it, and a partition cleared while its data was in flight only logs a warning, so a late writer does not fail. Signed-off-by: 0oshowero0 Co-authored-by: Cursor --- tests/e2e/test_kv_interface_e2e.py | 17 +++++++ tests/test_controller_data_partitions.py | 5 +- tests/test_storage_request_retry.py | 41 ++++++++++++++++- transfer_queue/controller.py | 46 +++++++++++++------ transfer_queue/interface.py | 8 ++++ transfer_queue/storage/managers/base.py | 31 ++++++++++--- .../managers/simple_storage_manager.py | 15 +++--- 7 files changed, 133 insertions(+), 30 deletions(-) diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index e1e52dcf..07e9c51b 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -470,6 +470,23 @@ 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_raises_when_the_controller_rejects_its_dtype(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_clear(keys=["float_x", "int_x"], partition_id=partition_id) + + def test_notify_for_a_cleared_partition_warns_instead_of_raising(self, tq_system): + # 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}} + assert asyncio.run(manager.notify_data_update("no_such_partition", [0], schema)) is False + 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..3574b064 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 @@ -70,7 +71,7 @@ async def __aexit__(self, *_args): @pytest.mark.asyncio @pytest.mark.parametrize("accepted", [False, True]) -async def test_controller_notification_returns_ack_status(accepted): +async def test_controller_notification_returns_ack_body(accepted): manager = _manager() manager.controller_info = MagicMock(id="controller") socket = MagicMock() @@ -85,7 +86,39 @@ async def test_controller_notification_returns_ack_status(accepted): manager.notify_pool = MagicMock() manager.notify_pool.alease.return_value = _SocketLease(socket) - assert await manager._notify_and_wait([b"request"]) is accepted + assert await manager._notify_and_wait([b"request"]) == {"success": accepted} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "ack, published", + [ + ({"success": True, "partition_exists": True, "error": None}, True), + (None, False), # no ACK in time: the outcome is unknown, so the put is not failed + ({"success": False, "partition_exists": False, "error": None}, False), + ], +) +async def test_notify_publishes_or_tolerates_ack(ack, published, 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 (ack is not None and not ack["partition_exists"]) + + +@pytest.mark.asyncio +async def test_notify_raises_when_the_controller_rejects_the_data(): + manager = _manager() + manager.controller_info = MagicMock(id="controller") + manager._notify_loop = asyncio.get_running_loop() + rejected = {"success": False, "partition_exists": True, "error": "Field 'x' dtype mismatch"} + manager._notify_and_wait = AsyncMock(return_value=rejected) + + with pytest.raises(RuntimeError, match=r"rejected the data for partition 'p': Field 'x' dtype mismatch"): + await manager.notify_data_update("p", [0], {}) @pytest.mark.asyncio @@ -142,6 +175,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("the controller rejected the data: dtype")) + with pytest.raises(RuntimeError, match="Storage update committed, but the controller rejected.*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..31b5b156 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -521,15 +521,18 @@ 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; nothing is changed. """ - 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}") + # 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}") + try: # Derive field_names from field_schema to guarantee consistency field_names = list(field_schema.keys()) @@ -1155,6 +1158,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; nothing is changed. """ field_names = list(field_schema.keys()) partition = self._get_partition(partition_id) @@ -2139,13 +2145,23 @@ def _handle_notify_data_update_request(self, request_msg: ZMQMessage) -> ZMQMess params = request_msg.body partition_id = params.get("partition_id") global_indexes = params.get("global_indexes", []) - success = self.update_production_status( - partition_id=cast(str, 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), - ) + # A rejected write makes the writer raise. A partition cleared while its data was in flight + # only makes the writer warn: that data is no longer wanted. + partition_exists = self._get_partition(cast(str, partition_id)) is not None + success, error = False, None + if partition_exists: + try: + success = self.update_production_status( + partition_id=cast(str, 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), + ) + except ValueError as e: + error = str(e) + if not success and error is None: + error = "the controller failed to record it; see the controller log" if success: if self._metrics is not None: self._metrics.record_samples("NOTIFY_DATA_UPDATE", len(global_indexes)) @@ -2158,6 +2174,8 @@ def _handle_notify_data_update_request(self, request_msg: ZMQMessage) -> ZMQMess "controller_id": self.controller_id, "partition_id": partition_id, "success": success, + "partition_exists": partition_exists, + "error": error, }, ) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 7fd23610..aee5c1f6 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 rejects the data, e.g. a field's dtype conflicts with + earlier writes to it. The value may already be stored but is never made ready. 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 rejects the data, e.g. a field's dtype conflicts with + earlier writes to it. The values may already be stored but are never made ready. 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 rejects the data, e.g. a field's dtype conflicts with + earlier writes to it. The value may already be stored but is never made ready. 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 rejects the data, e.g. a field's dtype conflicts with + earlier writes to it. The values may already be stored but are never made ready. Example: >>> import transfer_queue as tq diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 88f7b37a..b9942b75 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 rejected the data, e.g. on a dtype conflict. """ if not self.controller_info: @@ -274,10 +281,21 @@ async def notify_data_update( self._notify_and_wait(request_msg), self._notify_loop, ) - return await asyncio.wrap_future(thread_future) + ack = await asyncio.wrap_future(thread_future) + if ack is None: + return False + if ack.get("success"): + return True + if not ack.get("partition_exists", True): + logger.warning( + f"[{self.storage_manager_id}]: Partition {partition_id!r} no longer exists; " + f"the controller dropped {len(global_indexes)} written samples" + ) + return False + raise RuntimeError(f"the controller rejected the data for partition {partition_id!r}: {ack.get('error')}") - 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) -> dict[str, Any] | None: + """Send a data status notification to the controller and return its ACK body, or None if none arrived.""" # 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: @@ -301,18 +319,17 @@ async def _notify_and_wait(self, request_msg: list) -> bool: 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}." + f"from controller id #{response_msg.sender_id}: {response_msg.body}." ) - return success + return response_msg.body 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 @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..3508d614 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -605,12 +605,15 @@ 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 {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; " From 8d3230764d0ef3ac355dc6201ceadde0834a08d8 Mon Sep 17 00:00:00 2001 From: 0oshowero0 Date: Thu, 15 Oct 2026 20:25:50 +0800 Subject: [PATCH 2/2] [fix] Reuse controller errors for write notifications Signed-off-by: 0oshowero0 --- tests/e2e/test_kv_interface_e2e.py | 32 ++++++++++- tests/test_storage_request_retry.py | 38 ++++++++----- transfer_queue/controller.py | 56 +++++++------------ transfer_queue/interface.py | 16 +++--- transfer_queue/storage/managers/base.py | 36 ++++++------ .../managers/simple_storage_manager.py | 4 +- 6 files changed, 103 insertions(+), 79 deletions(-) diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index 07e9c51b..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,7 +471,7 @@ 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_raises_when_the_controller_rejects_its_dtype(self, tq_api): + 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])}) @@ -479,13 +480,38 @@ def test_kv_put_raises_when_the_controller_rejects_its_dtype(self, tq_api): 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_notify_for_a_cleared_partition_warns_instead_of_raising(self, tq_system): + 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}} - assert asyncio.run(manager.notify_data_update("no_such_partition", [0], schema)) is 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: diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py index 3574b064..4d49e6a2 100644 --- a/tests/test_storage_request_retry.py +++ b/tests/test_storage_request_retry.py @@ -71,7 +71,7 @@ async def __aexit__(self, *_args): @pytest.mark.asyncio @pytest.mark.parametrize("accepted", [False, True]) -async def test_controller_notification_returns_ack_body(accepted): +async def test_controller_notification_returns_ack_status(accepted): manager = _manager() manager.controller_info = MagicMock(id="controller") socket = MagicMock() @@ -86,19 +86,19 @@ async def test_controller_notification_returns_ack_body(accepted): manager.notify_pool = MagicMock() manager.notify_pool.alease.return_value = _SocketLease(socket) - assert await manager._notify_and_wait([b"request"]) == {"success": accepted} + assert await manager._notify_and_wait([b"request"]) is accepted @pytest.mark.asyncio @pytest.mark.parametrize( - "ack, published", + "ack, published, warns", [ - ({"success": True, "partition_exists": True, "error": None}, True), - (None, False), # no ACK in time: the outcome is unknown, so the put is not failed - ({"success": False, "partition_exists": False, "error": None}, False), + (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, caplog): +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() @@ -106,19 +106,29 @@ async def test_notify_publishes_or_tolerates_ack(ack, published, caplog): 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 (ack is not None and not ack["partition_exists"]) + assert ("no longer exists" in caplog.text) is warns @pytest.mark.asyncio -async def test_notify_raises_when_the_controller_rejects_the_data(): +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() - rejected = {"success": False, "partition_exists": True, "error": "Field 'x' dtype mismatch"} - manager._notify_and_wait = AsyncMock(return_value=rejected) + 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=r"rejected the data for partition 'p': Field 'x' dtype mismatch"): + 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 @@ -175,8 +185,8 @@ 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("the controller rejected the data: dtype")) - with pytest.raises(RuntimeError, match="Storage update committed, but the controller rejected.*do not retry"): + 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) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 31b5b156..c44ac236 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -523,9 +523,10 @@ def update_production_status( True if update was successful, False on error Raises: - ValueError: If a field's dtype conflicts with the partition's; nothing is changed. + ValueError: If a field's dtype conflicts with the partition's; metadata and readiness are unchanged. """ - # Same dtype check as FieldMeta.update, run first so a conflicting write changes nothing. + # 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") @@ -1160,7 +1161,7 @@ def update_production_status( True if update was successful, False otherwise Raises: - ValueError: If a field's dtype conflicts with the partition's; nothing is changed. + 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) @@ -2143,41 +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", []) - # A rejected write makes the writer raise. A partition cleared while its data was in flight - # only makes the writer warn: that data is no longer wanted. - partition_exists = self._get_partition(cast(str, partition_id)) is not None - success, error = False, None - if partition_exists: - try: - success = self.update_production_status( - partition_id=cast(str, 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), - ) - except ValueError as e: - error = str(e) - if not success and error is None: - error = "the controller failed to record it; see the controller log" - 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}") + 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), + ): + 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, - "partition_exists": partition_exists, - "error": error, - }, - ) + 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 aee5c1f6..0b4ab04b 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -399,8 +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 rejects the data, e.g. a field's dtype conflicts with - earlier writes to it. The value may already be stored but is never made ready. + 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 @@ -461,8 +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 rejects the data, e.g. a field's dtype conflicts with - earlier writes to it. The values may already be stored but are never made ready. + 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 @@ -805,8 +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 rejects the data, e.g. a field's dtype conflicts with - earlier writes to it. The value may already be stored but is never made ready. + 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 @@ -916,8 +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 rejects the data, e.g. a field's dtype conflicts with - earlier writes to it. The values may already be stored but are never made ready. + 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 b9942b75..9f6eb739 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -243,7 +243,7 @@ async def notify_data_update( unknown) or the partition no longer exists (the data is dropped, with a warning). Raises: - RuntimeError: If the controller rejected the data, e.g. on a dtype conflict. + RuntimeError: If the controller failed to record the data, e.g. on a dtype conflict. """ if not self.controller_info: @@ -281,21 +281,16 @@ async def notify_data_update( self._notify_and_wait(request_msg), self._notify_loop, ) - ack = await asyncio.wrap_future(thread_future) - if ack is None: - return False - if ack.get("success"): - return True - if not ack.get("partition_exists", True): + 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 False - raise RuntimeError(f"the controller rejected the data for partition {partition_id!r}: {ack.get('error')}") + return accepted is True - async def _notify_and_wait(self, request_msg: list) -> dict[str, Any] | None: - """Send a data status notification to the controller and return its ACK body, or None if none arrived.""" + 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: @@ -318,12 +313,11 @@ async def _notify_and_wait(self, request_msg: list) -> dict[str, Any] | None: 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] - logger.debug( - f"[{self.storage_manager_id}]: Received data status update ACK " - f"from controller id #{response_msg.sender_id}: {response_msg.body}." - ) - return response_msg.body + 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. @@ -331,6 +325,14 @@ async def _notify_and_wait(self, request_msg: list) -> dict[str, Any] | None: sock.close(linger=0) 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( self, data: TensorDict, metadata: BatchMeta, data_parser: Callable[[Any], Any] | None = None diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 3508d614..25b0eeed 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -613,7 +613,9 @@ async def update_data( ready=finalize, ) except RuntimeError as e: - raise RuntimeError(f"Storage update committed, but {e}; do not retry this merge blindly") from 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; "