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.