From d1ccf4141b6b6b8646af9abf5731b4cdd376f137 Mon Sep 17 00:00:00 2001 From: 0oshowero0 Date: Thu, 15 Oct 2026 18:22:17 +0800 Subject: [PATCH] [chore] clear_data: log each failed storage unit once, with what it kept A failed clear logged two lines per unit: one from _clear_single_storage_unit and one naming only the gather task index. Log one line per failed unit instead, with the unit id, the number of samples and fields it still holds, and the error. Behavior is unchanged: the failure is still logged rather than raised, so a cleanup failure does not stop the caller's job. Signed-off-by: 0oshowero0 Co-authored-by: Cursor --- tests/test_async_simple_storage_manager.py | 20 ++++++++- .../managers/simple_storage_manager.py | 42 +++++++++---------- 2 files changed, 40 insertions(+), 22 deletions(-) diff --git a/tests/test_async_simple_storage_manager.py b/tests/test_async_simple_storage_manager.py index 609be511..b4dbfbf5 100644 --- a/tests/test_async_simple_storage_manager.py +++ b/tests/test_async_simple_storage_manager.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging from unittest.mock import AsyncMock, Mock, patch import numpy as np @@ -317,7 +318,7 @@ def parser(data, args): @pytest.mark.asyncio -async def test_clear_data_routes_from_hash(): +async def test_clear_data_routes_from_hash(caplog): """clear_data should route using global_idx % num_su (hash routing).""" storage_unit_infos = { "storage_0": ZMQServerInfo( @@ -361,6 +362,23 @@ async def fake_clear(global_indexes, target_storage_unit=None, **kwargs): assert set(called_with.get("storage_0", [])) == {0, 2} assert set(called_with.get("storage_1", [])) == {1, 3} + # A failed unit is logged, not raised, and the other units are still cleared. + called_with.clear() + + async def fail_on_storage_1(global_indexes, target_storage_unit=None, **kwargs): + called_with[target_storage_unit] = list(global_indexes) + if target_storage_unit == "storage_1": + raise RuntimeError("unit timeout") + + manager._clear_single_storage_unit = fail_on_storage_1 + with caplog.at_level(logging.ERROR): + await manager.clear_data(batch_meta) + + assert set(called_with) == {"storage_0", "storage_1"} + failures = [record.message for record in caplog.records if "Failed to clear" in record.message] + assert len(failures) == 1 + assert "2 samples" in failures[0] and "storage_1" in failures[0] and "unit timeout" in failures[0] + @pytest.mark.asyncio async def test_hash_routing_stable_across_batch_sizes(): diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index b17fa087..8d29cfa2 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -937,33 +937,33 @@ async def clear_data(self, metadata: BatchMeta, fields: list[str] | None = None) results = await asyncio.gather(*tasks, return_exceptions=True) - for i, result in enumerate(results): + # Logged rather than raised: the caller no longer wants this data, and failing its cleanup + # would stop the job. The unit keeps the payloads until their indexes are reused and cleared. + for (su_id, group), result in zip(routing.items(), results, strict=True): if isinstance(result, Exception): - logger.error(f"[{self.storage_manager_id}]: Error in clear operation task {i}: {result}") + logger.error( + f"[{self.storage_manager_id}]: Failed to clear {len(group.global_indexes)} samples " + f"(fields={fields or 'all'}) on storage unit {su_id}; they stay in its memory: " + f"{type(result).__name__}: {result}" + ) @with_storage_unit_socket async def _clear_single_storage_unit(self, global_indexes, fields=None, target_storage_unit=None, socket=None): - try: - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.CLEAR_DATA, - sender_id=self.storage_manager_id, - receiver_id=target_storage_unit, - body={"global_indexes": global_indexes, "fields": fields}, - ) - - await socket.send_multipart(request_msg.serialize()) - messages = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(messages) + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.CLEAR_DATA, + sender_id=self.storage_manager_id, + receiver_id=target_storage_unit, + body={"global_indexes": global_indexes, "fields": fields}, + ) - if response_msg.request_type != ZMQRequestType.CLEAR_DATA_RESPONSE: - raise RuntimeError( - f"Failed to clear storage {target_storage_unit}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) + await socket.send_multipart(request_msg.serialize()) + messages = await socket.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(messages) - except Exception as e: - logger.error(f"[{self.storage_manager_id}]: Error clearing storage unit {target_storage_unit}: {str(e)}") - raise + if response_msg.request_type != ZMQRequestType.CLEAR_DATA_RESPONSE: + raise RuntimeError( + f"Failed to clear storage {target_storage_unit}: {response_msg.body.get('message', 'Unknown error')}" + ) def get_zmq_server_info(self) -> dict[str, ZMQServerInfo]: """Get ZMQ server information for all storage units.