From 6d3338d2453e133b014fe3f7ff71769400215b8c Mon Sep 17 00:00:00 2001
From: jathonzhang <jathonzhang@tencent.com>
Date: Wed, 9 Sep 2026 20:47:03 +0800
Subject: [PATCH 1/8] [fix] Retry a lost storage-unit request instead of
 failing the job

A storage unit stops answering and every client routed to it fails when its own
recv timeout expires. Raising that timeout (400s to 1800s in our deployment)
changed nothing.

Evidence from one such failure: the unit's node was alive and serving other
traffic throughout, a TCP connect to its put_get_socket succeeded, and the unit's
own counters showed it fully healthy (1025 GET_DATA served, 9.6ms p99, 1.33GB
RSS). It had served exactly one GET_DATA fewer than its cohort. So the unit never
saw the request that timed out; it was lost between the two ends, not queued
behind slow work. #168 protected the worker thread from dying, which is a
different cause of the same symptom; here the thread was intact.

That loss is invisible by construction. ZMQ connect is asynchronous and SNDHWM is
0, so a DEALER accepts send() into an unbounded local queue for a peer it has not
reached yet. The message sits there and the caller only learns anything when its
own RCVTIMEO expires, which is why no timeout value can distinguish a lost
request from a slow one. Lowering SNDHWM would not help either, it only trades
silent queuing for silent dropping.

Retry the request on a new socket and TCP connection, which is the part that
matters, up to TQ_SIMPLE_STORAGE_MAX_ATTEMPTS (default 3). Only a missing answer
is retried: zmq.error.Again now raises StorageUnitTimeout, while an error the
unit actually reported still surfaces on the first attempt. Replaying an attempt
is safe: put is keyed by global index and overwrites, get is read-only.

Make the residual failure self-diagnosing, so a next occurrence does not need
another round of manual probing. After the last attempt, ask the unit for its own
counters over a fresh socket with a short timeout. That probe is served by the
same worker thread as put and get, so an answer proves the unit is serving and
the request was lost in flight, while silence means the unit itself stopped. The
failure log now carries that verdict plus tcp reachability, the unit's op counts
and RSS, and the shape of the request that failed.

Log volume is unchanged in the steady state. A recovered request logs one line
and skips the diagnosis entirely; storage units log a request only above
TQ_STORAGE_SLOW_REQUEST_SECONDS (5s) or TQ_STORAGE_LARGE_PAYLOAD_MB (256MB), both
far above the single-digit-millisecond norm, so tripping either one is itself the
finding. The put_data failure log no longer dumps every routed unit id, which on
a large job was thousands of them per line, matching what get_data already does.

Tests cover recovery on retry, the bounded attempt count, that reported errors
are not retried, and each diagnosis verdict.

Signed-off-by: jathonzhang <jathonzhang@tencent.com>
---
 tests/test_storage_request_retry.py           | 215 ++++++++++++++++++
 .../managers/simple_storage_manager.py        | 189 +++++++++++++--
 transfer_queue/storage/simple_storage.py      |  24 +-
 transfer_queue/utils/common.py                |  21 ++
 4 files changed, 423 insertions(+), 26 deletions(-)
 create mode 100644 tests/test_storage_request_retry.py

diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py
new file mode 100644
index 00000000..b15c42a7
--- /dev/null
+++ b/tests/test_storage_request_retry.py
@@ -0,0 +1,215 @@
+# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for storage-unit request retry and the timeout diagnosis that classifies a failure."""
+
+import logging
+from unittest.mock import patch
+
+import pytest
+import torch
+import zmq
+
+from transfer_queue.storage.managers import simple_storage_manager as ssm
+from transfer_queue.storage.managers.simple_storage_manager import (
+    AsyncSimpleStorageManager,
+    StorageUnitTimeout,
+)
+from transfer_queue.utils.common import estimate_payload_bytes
+from transfer_queue.utils.enum_utils import Role
+from transfer_queue.utils.zmq_utils import ZMQServerInfo
+
+
+def _manager(**units: ZMQServerInfo) -> AsyncSimpleStorageManager:
+    """Build a manager carrying only the state the retry and diagnosis helpers read."""
+    manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager)
+    manager.storage_manager_id = "TQ_STORAGE_test"
+    manager.storage_unit_infos = dict(units)
+    return manager
+
+
+def _server_info(unit_id: str, ip: str, port: int) -> ZMQServerInfo:
+    return ZMQServerInfo(role=Role.STORAGE, id=unit_id, ip=ip, ports={"put_get_socket": port})
+
+
+@pytest.mark.asyncio
+async def test_lost_request_recovers_on_retry():
+    """A request lost in flight must be recovered by a second attempt, not kill the job."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+    attempts = []
+
+    async def flaky():
+        attempts.append(1)
+        if len(attempts) == 1:
+            raise StorageUnitTimeout("no answer")
+        return "payload"
+
+    with patch.object(manager, "_diagnose_storage_unit") as diagnose:
+        result = await manager._request_with_retry("get", "unit_a", "samples=4", flaky)
+
+    assert result == "payload"
+    assert len(attempts) == 2, "the retry must issue a second attempt on a new connection"
+    assert diagnose.call_count == 0, "a recovered request must not pay for a diagnosis"
+
+
+@pytest.mark.asyncio
+async def test_retry_gives_up_after_configured_attempts():
+    """Attempts are bounded, and the final failure still raises for the caller to handle."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+    attempts = []
+
+    async def always_timeout():
+        attempts.append(1)
+        raise StorageUnitTimeout("no answer")
+
+    with (
+        patch.object(ssm, "TQ_SIMPLE_STORAGE_MAX_ATTEMPTS", 3),
+        patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis") as diagnose,
+        pytest.raises(StorageUnitTimeout),
+    ):
+        await manager._request_with_retry("get", "unit_a", "samples=4", always_timeout)
+
+    assert len(attempts) == 3
+    diagnose.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_errors_reported_by_the_unit_are_not_retried():
+    """Only a missing answer is worth another connection; a real error must surface at once."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+    attempts = []
+
+    async def hard_failure():
+        attempts.append(1)
+        raise RuntimeError("storage unit rejected the request")
+
+    with pytest.raises(RuntimeError, match="rejected"):
+        await manager._request_with_retry("get", "unit_a", "samples=4", hard_failure)
+
+    assert len(attempts) == 1
+
+
+@pytest.mark.asyncio
+async def test_retry_logs_endpoint_and_request_shape(caplog):
+    """The retry warning must name the endpoint and the request, to correlate both ends."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+    attempts = []
+
+    async def flaky():
+        attempts.append(1)
+        if len(attempts) == 1:
+            raise StorageUnitTimeout("no answer")
+        return "payload"
+
+    with caplog.at_level(logging.WARNING):
+        await manager._request_with_retry("get", "unit_a", "samples=4 fields=['input_ids']", flaky)
+
+    warning = next(r for r in caplog.records if "retry" in r.message)
+    assert "10.0.0.7:5555" in warning.message
+    assert "samples=4" in warning.message
+
+
+@pytest.mark.asyncio
+async def test_diagnosis_blames_the_link_when_the_unit_still_answers():
+    """A unit that answers a fresh probe was not the one that stalled."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+    probe_body = {
+        "active_keys": 4,
+        "process_rss_bytes": 1_330_159_616,
+        "op_stats": {"GET_DATA": {"request_count": 1025}},
+    }
+
+    async def reachable(*args, **kwargs):
+        return None, _Writer()
+
+    with (
+        patch.object(ssm.asyncio, "open_connection", reachable),
+        patch.object(manager, "_probe_storage_unit", return_value=probe_body),
+    ):
+        diagnosis = await manager._diagnose_storage_unit("unit_a")
+
+    assert "tcp=up" in diagnosis
+    assert "verdict=request_lost_in_flight" in diagnosis
+    assert "1025" in diagnosis
+
+
+@pytest.mark.asyncio
+async def test_diagnosis_blames_the_unit_when_the_probe_times_out():
+    """A probe timeout means the worker thread itself stopped serving."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+
+    async def reachable(*args, **kwargs):
+        return None, _Writer()
+
+    async def probe_timeout(**kwargs):
+        raise zmq.error.Again()
+
+    with (
+        patch.object(ssm.asyncio, "open_connection", reachable),
+        patch.object(manager, "_probe_storage_unit", probe_timeout),
+    ):
+        diagnosis = await manager._diagnose_storage_unit("unit_a")
+
+    assert "verdict=unit_not_serving" in diagnosis
+
+
+@pytest.mark.asyncio
+async def test_diagnosis_reports_an_unreachable_endpoint():
+    """When the port itself is gone the diagnosis must say so rather than blame the unit."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+
+    async def refused(*args, **kwargs):
+        raise ConnectionRefusedError()
+
+    async def probe_timeout(**kwargs):
+        raise zmq.error.Again()
+
+    with (
+        patch.object(ssm.asyncio, "open_connection", refused),
+        patch.object(manager, "_probe_storage_unit", probe_timeout),
+    ):
+        diagnosis = await manager._diagnose_storage_unit("unit_a")
+
+    assert "tcp=down(ConnectionRefusedError)" in diagnosis
+
+
+@pytest.mark.asyncio
+async def test_diagnosis_never_raises_on_an_unknown_unit():
+    """Diagnosis runs while another failure is being reported and must not mask it."""
+    manager = _manager()
+
+    assert "unit_not_registered" in await manager._diagnose_storage_unit("missing_unit")
+
+
+class _Writer:
+    """Minimal stand-in for the writer half returned by ``asyncio.open_connection``."""
+
+    def close(self):
+        pass
+
+
+def test_estimate_payload_bytes_handles_put_and_get_shapes():
+    """Both the batched put shape and the per-sample get shape must be measurable."""
+    batched = {"input_ids": torch.zeros(4, 8, dtype=torch.int64)}
+    per_sample = {"input_ids": [torch.zeros(8, dtype=torch.int64) for _ in range(4)]}
+
+    assert estimate_payload_bytes(batched) == 4 * 8 * 8
+    assert estimate_payload_bytes(per_sample) == 4 * 8 * 8
+
+
+def test_estimate_payload_bytes_degrades_to_zero_on_unmeasurable_input():
+    """Size estimation is diagnostic only and must never break the path it reports on."""
+    assert estimate_payload_bytes(object()) == 0
+    assert estimate_payload_bytes({"meta": "not a tensor"}) == 0
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index 88340006..2ffca7ae 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -18,6 +18,7 @@
 import warnings
 from collections import defaultdict
 from collections.abc import Mapping
+from functools import partial
 from operator import itemgetter
 from pathlib import Path
 from typing import Any, Callable, NamedTuple
@@ -30,6 +31,7 @@
 from transfer_queue.metadata import BatchMeta, extract_field_schema
 from transfer_queue.storage.managers.base import StorageManager, StorageManagerFactory
 from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER, StorageKeyNotFoundError
+from transfer_queue.utils.common import estimate_payload_bytes
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
@@ -42,6 +44,21 @@
 
 TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 200))  # seconds
 
+# Attempts per storage-unit request, including the first. Each one must use a new connection: a
+# DEALER silently queues messages for a peer it never reached, so no timeout can catch that.
+TQ_SIMPLE_STORAGE_MAX_ATTEMPTS = int(os.environ.get("TQ_SIMPLE_STORAGE_MAX_ATTEMPTS", 3))
+
+# Timeout for the post-failure probe, which only has to answer whether the unit still serves.
+TQ_SIMPLE_STORAGE_PROBE_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_PROBE_TIMEOUT", 10))
+
+
+class StorageUnitTimeout(RuntimeError):
+    """A storage unit did not answer within the send/recv timeout.
+
+    Distinct from an error the unit reported: only a missing answer is worth a new connection.
+    """
+
+
 _SU_SUBDIR = "simple_storage"
 _SU_INFO_FILE = "storage_unit_info.json"
 
@@ -57,6 +74,16 @@
     timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT,
 )
 
+# Same endpoint as above but with the short diagnostic timeout, used only after a failure.
+with_storage_unit_probe_socket = with_zmq_socket(
+    "put_get_socket",
+    get_identity=lambda self: f"{self.storage_manager_id}_probe",
+    get_peer=lambda self, target: self.storage_unit_infos[target],
+    get_context=lambda self: self.zmq_context,
+    resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"),
+    timeout=TQ_SIMPLE_STORAGE_PROBE_TIMEOUT,
+)
+
 
 class RoutingGroup(NamedTuple):
     """Routing result for a single storage unit."""
@@ -144,6 +171,103 @@ def _group_by_hash(self, global_indexes: list[int]) -> dict[str, RoutingGroup]:
             pos_lists[key].append(pos)
         return {key: RoutingGroup(gi_lists[key], pos_lists[key]) for key in gi_lists}
 
+    def _describe_storage_unit(self, storage_unit_id: str) -> str:
+        """Return ``ip:port`` for a storage unit, for use in diagnostics.
+
+        The unit id is a random uuid4 fragment that carries no location, so a bare id in an
+        error message cannot be traced back to a node. Never raises: it is only ever called
+        while reporting another failure, and must not mask it.
+        """
+        info = self.storage_unit_infos.get(storage_unit_id)
+        if info is None:
+            return "endpoint unknown (unit not registered with this manager)"
+        return f"{info.ip}:{info.ports.get('put_get_socket')}"
+
+    @with_storage_unit_probe_socket
+    async def _probe_storage_unit(self, target_storage_unit: str, socket: zmq.Socket = None) -> dict[str, Any]:
+        """Ask a storage unit for its own counters over a brand-new socket.
+
+        Served by the same worker thread as put and get, so an answer proves the unit is serving.
+        """
+        request_msg = ZMQMessage.create(
+            request_type=ZMQRequestType.GET_METRICS,  # type: ignore[arg-type]
+            sender_id=f"{self.storage_manager_id}_probe",
+            receiver_id=target_storage_unit,
+            body={},
+        )
+        await socket.send_multipart(request_msg.serialize())
+        messages = await socket.recv_multipart(copy=False)
+        response_msg = ZMQMessage.deserialize(messages)
+        if response_msg.request_type != ZMQRequestType.METRICS_RESPONSE:
+            raise RuntimeError(f"unexpected probe response type {response_msg.request_type}")
+        return response_msg.body
+
+    async def _diagnose_storage_unit(self, target_storage_unit: str) -> str:
+        """Classify a timeout as a lost request, a stuck unit, or an unreachable node.
+
+        Returns one log line and never raises: it runs while another failure is being reported.
+        """
+        info = self.storage_unit_infos.get(target_storage_unit)
+        if info is None:
+            return "verdict=unknown(unit_not_registered)"
+
+        try:
+            _, writer = await asyncio.wait_for(
+                asyncio.open_connection(info.ip, info.ports.get("put_get_socket")), timeout=5
+            )
+            writer.close()
+            tcp = "tcp=up"
+        except Exception as e:
+            tcp = f"tcp=down({type(e).__name__})"
+
+        try:
+            body = await self._probe_storage_unit(target_storage_unit=target_storage_unit)
+            op_counts = {op: stats.get("request_count") for op, stats in (body.get("op_stats") or {}).items()}
+            return (
+                f"{tcp} verdict=request_lost_in_flight (unit answered a fresh probe: "
+                f"ops={op_counts} active_keys={body.get('active_keys')} "
+                f"rss_gb={body.get('process_rss_bytes', 0) / 2**30:.2f})"
+            )
+        except zmq.error.Again:
+            return f"{tcp} verdict=unit_not_serving (no probe answer in {TQ_SIMPLE_STORAGE_PROBE_TIMEOUT}s)"
+        except Exception as e:
+            return f"{tcp} verdict=unknown (probe failed: {type(e).__name__}: {e})"
+
+    async def _request_with_retry(
+        self,
+        operation: str,
+        target_storage_unit: str,
+        request_context: str,
+        make_request: Callable[[], Any],
+    ):
+        """Run one storage-unit request, retrying a missing answer on a fresh connection.
+
+        Args:
+            operation: Operation name for logs, e.g. ``get`` or ``put``.
+            target_storage_unit: Unit this request is routed to.
+            request_context: Request shape, so both ends of a failure can be correlated.
+            make_request: Zero-arg callable returning a coroutine for one attempt. Must build a
+                new socket per call, which ``with_storage_unit_socket`` does.
+        """
+        endpoint = self._describe_storage_unit(target_storage_unit)
+        for attempt in range(1, TQ_SIMPLE_STORAGE_MAX_ATTEMPTS + 1):
+            try:
+                return await make_request()
+            except StorageUnitTimeout:
+                if attempt < TQ_SIMPLE_STORAGE_MAX_ATTEMPTS:
+                    logger.warning(
+                        f"[{self.storage_manager_id}]: no answer from {target_storage_unit} at {endpoint} in "
+                        f"{TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s, {operation} retry "
+                        f"{attempt + 1}/{TQ_SIMPLE_STORAGE_MAX_ATTEMPTS}. {request_context}"
+                    )
+                    continue
+                logger.error(
+                    f"[{self.storage_manager_id}]: {operation} to {target_storage_unit} at {endpoint} failed "
+                    f"after {attempt}x{TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s. {request_context} "
+                    f"{await self._diagnose_storage_unit(target_storage_unit)}"
+                )
+                raise
+
     @staticmethod
     def _select_by_positions(field_data, positions: list[int]):
         """Slice a single field's data by non-contiguous batch positions.
@@ -257,24 +381,37 @@ async def put_data(
         field_schema = extract_field_schema(data)
 
         routing = self._group_by_hash(metadata.global_indexes)
-        tasks = [
-            self._put_to_single_storage_unit(
-                group.global_indexes,
-                {f: self._select_by_positions(data[f], group.batch_positions) for f in data.keys()},
-                target_storage_unit=su_id,
-                data_parser=data_parser,
+        # A retried put is replayed as-is: it overwrites the same global indexes, and re-runs
+        # data_parser on the unit, so a parser must be free of external side effects.
+        tasks = []
+        for su_id, group in routing.items():
+            storage_data = {f: self._select_by_positions(data[f], group.batch_positions) for f in data.keys()}
+            tasks.append(
+                self._request_with_retry(
+                    "put",
+                    su_id,
+                    f"samples={len(group.global_indexes)} fields={list(storage_data.keys())} "
+                    f"payload_mb={estimate_payload_bytes(storage_data) / 2**20:.1f}",
+                    partial(
+                        self._put_to_single_storage_unit,
+                        group.global_indexes,
+                        storage_data,
+                        target_storage_unit=su_id,
+                        data_parser=data_parser,
+                    ),
+                )
             )
-            for su_id, group in routing.items()
-        ]
 
         try:
             await asyncio.gather(*tasks)
         except Exception as e:
+            # The offending unit is named in the error itself; the full routing list can run to
+            # hundreds of ids, which buries it.
             logger.error(
                 f"[{self.storage_manager_id}]: put_data failed. "
                 f"partition_id={metadata.partition_ids[0]}, "
                 f"num_samples={metadata.size}, "
-                f"storage_units={list(routing.keys())}, "
+                f"num_storage_units={len(routing)}, "
                 f"error={type(e).__name__}: {e}"
             )
             raise
@@ -318,14 +455,9 @@ async def _put_to_single_storage_unit(
                     f"{response_msg.body.get('message', 'Unknown error')}"
                 )
         except zmq.error.Again as e:
-            timeout_sec = TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT
-            logger.error(
-                f"[{self.storage_manager_id}]: ZMQ recv timeout ({timeout_sec}s) "
-                f"during put to storage unit {target_storage_unit}. "
-                f"The storage unit may be overloaded or crashed."
-            )
-            raise RuntimeError(
-                f"ZMQ recv timeout ({timeout_sec}s) during put to storage unit {target_storage_unit}"
+            raise StorageUnitTimeout(
+                f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during put to storage unit "
+                f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}"
             ) from e
         except Exception as e:
             logger.error(
@@ -405,7 +537,17 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict:
         routing = self._group_by_hash(metadata.global_indexes)
 
         tasks = [
-            self._get_from_single_storage_unit(group.global_indexes, metadata.field_names, target_storage_unit=su_id)
+            self._request_with_retry(
+                "get",
+                su_id,
+                f"samples={len(group.global_indexes)} fields={list(metadata.field_names)}",
+                partial(
+                    self._get_from_single_storage_unit,
+                    group.global_indexes,
+                    metadata.field_names,
+                    target_storage_unit=su_id,
+                ),
+            )
             for su_id, group in routing.items()
         ]
         try:
@@ -464,13 +606,10 @@ async def _get_from_single_storage_unit(
                 error_type = StorageKeyNotFoundError if KEY_NOT_FOUND_MARKER in message else RuntimeError
                 raise error_type(f"Failed to get data from storage unit {target_storage_unit}: {message}")
         except zmq.error.Again as e:
-            timeout_sec = TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT
-            logger.error(
-                f"[{self.storage_manager_id}]: ZMQ recv timeout ({timeout_sec}s) "
-                f"from storage unit {target_storage_unit}. "
-                f"The storage unit may be overloaded or crashed."
-            )
-            raise RuntimeError(f"ZMQ recv timeout ({timeout_sec}s) from storage unit {target_storage_unit}") from e
+            raise StorageUnitTimeout(
+                f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s from storage unit "
+                f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}"
+            ) from e
         except StorageKeyNotFoundError:
             # Already logged at debug by the storage unit; propagate for the caller to classify.
             raise
diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index bf9d616f..b1a8ce57 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -25,7 +25,7 @@
 import ray
 import zmq
 
-from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads
+from transfer_queue.utils.common import estimate_payload_bytes, limit_pytorch_auto_parallel_threads
 from transfer_queue.utils.enum_utils import Role
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.perf_utils import IntervalPerfMonitor
@@ -48,6 +48,11 @@
 TQ_STORAGE_POLLER_TIMEOUT = int(os.environ.get("TQ_STORAGE_POLLER_TIMEOUT", 5))  # in seconds
 TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8))
 
+# Thresholds above which a single request earns a log line. Both sit far above the normal range
+# of single-digit milliseconds, so tripping either one is itself the finding.
+TQ_STORAGE_SLOW_REQUEST_SECONDS = float(os.environ.get("TQ_STORAGE_SLOW_REQUEST_SECONDS", 5.0))
+TQ_STORAGE_LARGE_PAYLOAD_MB = float(os.environ.get("TQ_STORAGE_LARGE_PAYLOAD_MB", 256))
+
 # Marks a GET_ERROR reply as "the key is gone" so the caller can tell it apart from a real fault.
 KEY_NOT_FOUND_MARKER = "TQKeyNotFound"
 
@@ -394,6 +399,17 @@ def _worker_routine(self) -> None:
         poller.unregister(worker_socket)
         worker_socket.close(linger=0)
 
+    def _log_if_heavy(self, operation: str, started: float, field_data: Any, num_samples: int, fields: Any) -> None:
+        """Log a request that was slow or unusually large, and stay silent otherwise."""
+        elapsed = time.perf_counter() - started
+        payload_mb = estimate_payload_bytes(field_data) / 2**20
+        if elapsed < TQ_STORAGE_SLOW_REQUEST_SECONDS and payload_mb < TQ_STORAGE_LARGE_PAYLOAD_MB:
+            return
+        logger.warning(
+            f"[{self.storage_unit_id}]: heavy {operation} samples={num_samples} "
+            f"payload_mb={payload_mb:.1f} elapsed={elapsed:.2f}s fields={list(fields)}"
+        )
+
     def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage:
         """
         Handle put request, add or update data into storage unit.
@@ -409,6 +425,7 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage:
             field_data = data_parts.body["data"]  # field_data should be a dict.
             data_parser = data_parts.body.get("data_parser", None)
 
+            started = time.perf_counter()
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_put"
             ):
@@ -456,6 +473,8 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage:
                             )
                 self.storage_data.put_data(field_data, global_indexes)
 
+            self._log_if_heavy("PUT_DATA", started, field_data, len(global_indexes), field_data.keys())
+
             # After put operation finish, send a message to the client
             response_msg = ZMQMessage.create(
                 request_type=ZMQRequestType.PUT_DATA_RESPONSE,  # type: ignore[arg-type]
@@ -488,11 +507,14 @@ def _handle_get(self, data_parts: ZMQMessage) -> ZMQMessage:
             fields = data_parts.body["fields"]
             global_indexes = data_parts.body["global_indexes"]
 
+            started = time.perf_counter()
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_get"
             ):
                 result_data = self.storage_data.get_data(fields, global_indexes)
 
+            self._log_if_heavy("GET_DATA", started, result_data, len(global_indexes), fields)
+
             response_msg = ZMQMessage.create(
                 request_type=ZMQRequestType.GET_DATA_RESPONSE,  # type: ignore[arg-type]
                 sender_id=self.storage_unit_id,
diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py
index 999e8183..d92b628f 100644
--- a/transfer_queue/utils/common.py
+++ b/transfer_queue/utils/common.py
@@ -14,7 +14,9 @@
 # limitations under the License.
 
 import os
+from collections.abc import Mapping
 from contextlib import contextmanager
+from typing import Any
 
 import psutil
 import ray
@@ -133,3 +135,22 @@ def get_env_bool(env_key: str, default: bool = False) -> bool:
 
     true_values = {"true", "1", "yes", "y", "on"}
     return env_value_lower in true_values
+
+
+def estimate_payload_bytes(field_data: Any) -> int:
+    """Best-effort size of a request payload in bytes; 0 when it cannot be measured.
+
+    Walks two levels: covers both put's ``field -> value`` and get's ``field -> per-sample list``.
+    """
+    total = 0
+    try:
+        values = field_data.values() if isinstance(field_data, Mapping) else field_data
+        for value in values:
+            items = value if isinstance(value, list | tuple) else [value]
+            for item in items:
+                nbytes = getattr(item, "nbytes", None)
+                if isinstance(nbytes, int):
+                    total += nbytes
+    except Exception:
+        return 0
+    return total

From 4185ac1c3a958478354ea729372434939a060cd0 Mon Sep 17 00:00:00 2001
From: jathonzhang <jathonzhang@tencent.com>
Date: Wed, 9 Sep 2026 21:33:02 +0800
Subject: [PATCH 2/8] [fix] Always issue at least one storage-unit request
 attempt

Codex review on #171: with TQ_SIMPLE_STORAGE_MAX_ATTEMPTS set to 0 or a negative
value the retry loop never ran, so _request_with_retry returned None without
raising. For a put that reads as success, and put_data then notified the
controller, publishing metadata that points at data which was never sent.

Floor the attempt count where it is consumed, so every caller gets one attempt
regardless of how the value was configured.

Signed-off-by: jathonzhang <jathonzhang@tencent.com>
---
 tests/test_storage_request_retry.py             | 17 +++++++++++++++++
 .../storage/managers/simple_storage_manager.py  |  9 ++++++---
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py
index b15c42a7..beb423bb 100644
--- a/tests/test_storage_request_retry.py
+++ b/tests/test_storage_request_retry.py
@@ -85,6 +85,23 @@ async def always_timeout():
     diagnose.assert_called_once()
 
 
+@pytest.mark.asyncio
+async def test_a_nonpositive_attempt_count_still_issues_one_request():
+    """Skipping the request would report success, and let put_data publish absent data."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+    attempts = []
+
+    async def succeed():
+        attempts.append(1)
+        return "payload"
+
+    with patch.object(ssm, "TQ_SIMPLE_STORAGE_MAX_ATTEMPTS", 0):
+        result = await manager._request_with_retry("get", "unit_a", "samples=4", succeed)
+
+    assert result == "payload"
+    assert len(attempts) == 1
+
+
 @pytest.mark.asyncio
 async def test_errors_reported_by_the_unit_are_not_retried():
     """Only a missing answer is worth another connection; a real error must surface at once."""
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index 2ffca7ae..fb3d11b1 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -250,15 +250,18 @@ async def _request_with_retry(
                 new socket per call, which ``with_storage_unit_socket`` does.
         """
         endpoint = self._describe_storage_unit(target_storage_unit)
-        for attempt in range(1, TQ_SIMPLE_STORAGE_MAX_ATTEMPTS + 1):
+        # Floored at one: a nonpositive count would skip the request and report success, which for
+        # put would publish metadata for data that was never sent.
+        attempts_allowed = max(1, TQ_SIMPLE_STORAGE_MAX_ATTEMPTS)
+        for attempt in range(1, attempts_allowed + 1):
             try:
                 return await make_request()
             except StorageUnitTimeout:
-                if attempt < TQ_SIMPLE_STORAGE_MAX_ATTEMPTS:
+                if attempt < attempts_allowed:
                     logger.warning(
                         f"[{self.storage_manager_id}]: no answer from {target_storage_unit} at {endpoint} in "
                         f"{TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s, {operation} retry "
-                        f"{attempt + 1}/{TQ_SIMPLE_STORAGE_MAX_ATTEMPTS}. {request_context}"
+                        f"{attempt + 1}/{attempts_allowed}. {request_context}"
                     )
                     continue
                 logger.error(

From 28d7a130b2d336a4bf779e1389b36888c0654503 Mon Sep 17 00:00:00 2001
From: jathonzhang <jathonzhang@tencent.com>
Date: Wed, 9 Sep 2026 21:40:30 +0800
Subject: [PATCH 3/8] [fix] Do not replay a parser-backed put on retry

Codex review on #171: a put whose reply was lost has already been committed by
the unit, so a retry re-runs data_parser there a second time. Replaying the write
itself is harmless, it overwrites the same global indexes, but the parser is not:
kv_put, kv_batch_put and put accept an arbitrary callable and constrain only its
keys, element count and ordering, never its side effects. A parser that consumes
references or writes externally would see those effects duplicated.

Send a parser-backed put once, so it fails exactly as it did before this series,
and keep the retry for the ordinary put, which is a plain overwrite.

Signed-off-by: jathonzhang <jathonzhang@tencent.com>
---
 tests/test_storage_request_retry.py           | 46 ++++++++++++++++++-
 .../managers/simple_storage_manager.py        | 12 +++--
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py
index beb423bb..44701ed2 100644
--- a/tests/test_storage_request_retry.py
+++ b/tests/test_storage_request_retry.py
@@ -16,12 +16,15 @@
 """Tests for storage-unit request retry and the timeout diagnosis that classifies a failure."""
 
 import logging
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
+import numpy as np
 import pytest
 import torch
 import zmq
+from tensordict import TensorDict
 
+from transfer_queue.metadata import BatchMeta
 from transfer_queue.storage.managers import simple_storage_manager as ssm
 from transfer_queue.storage.managers.simple_storage_manager import (
     AsyncSimpleStorageManager,
@@ -37,6 +40,7 @@ def _manager(**units: ZMQServerInfo) -> AsyncSimpleStorageManager:
     manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager)
     manager.storage_manager_id = "TQ_STORAGE_test"
     manager.storage_unit_infos = dict(units)
+    manager.close = lambda: None  # __init__ is skipped, so there is no socket or thread to close
     return manager
 
 
@@ -44,6 +48,46 @@ def _server_info(unit_id: str, ip: str, port: int) -> ZMQServerInfo:
     return ZMQServerInfo(role=Role.STORAGE, id=unit_id, ip=ip, ports={"put_get_socket": port})
 
 
+def _single_sample_batch() -> tuple[TensorDict, BatchMeta]:
+    """One sample routed to one unit, the smallest input put_data accepts."""
+    metadata = BatchMeta(
+        global_indexes=[0],
+        partition_ids=["0"],
+        field_schema={"input_ids": {"dtype": torch.int64, "shape": (2,), "is_nested": False, "is_non_tensor": False}},
+        production_status=np.ones(1, dtype=np.int8),
+    )
+    return TensorDict({"input_ids": torch.zeros(1, 2, dtype=torch.int64)}, batch_size=1), metadata
+
+
+async def _failing_put_attempts(data_parser) -> int:
+    """Count the attempts put_data makes when the unit never answers."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+    manager.notify_data_update = AsyncMock()
+    manager._put_to_single_storage_unit = AsyncMock(side_effect=StorageUnitTimeout("no answer"))
+    data, metadata = _single_sample_batch()
+
+    with (
+        patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"),
+        pytest.raises(StorageUnitTimeout),
+    ):
+        await manager.put_data(data, metadata, data_parser=data_parser)
+
+    manager.notify_data_update.assert_not_awaited()
+    return manager._put_to_single_storage_unit.await_count
+
+
+@pytest.mark.asyncio
+async def test_parser_backed_put_is_not_replayed():
+    """A parser re-runs on the unit, and the public API does not constrain its side effects."""
+    assert await _failing_put_attempts(data_parser=lambda field_data: field_data) == 1
+
+
+@pytest.mark.asyncio
+async def test_put_without_a_parser_is_still_retried():
+    """The retry must stay in force for the ordinary put, which is a plain overwrite."""
+    assert await _failing_put_attempts(data_parser=None) == ssm.TQ_SIMPLE_STORAGE_MAX_ATTEMPTS
+
+
 @pytest.mark.asyncio
 async def test_lost_request_recovers_on_retry():
     """A request lost in flight must be recovered by a second attempt, not kill the job."""
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index fb3d11b1..c49f4265 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -239,6 +239,7 @@ async def _request_with_retry(
         target_storage_unit: str,
         request_context: str,
         make_request: Callable[[], Any],
+        max_attempts: int | None = None,
     ):
         """Run one storage-unit request, retrying a missing answer on a fresh connection.
 
@@ -248,11 +249,13 @@ async def _request_with_retry(
             request_context: Request shape, so both ends of a failure can be correlated.
             make_request: Zero-arg callable returning a coroutine for one attempt. Must build a
                 new socket per call, which ``with_storage_unit_socket`` does.
+            max_attempts: Attempts allowed. Defaults to ``TQ_SIMPLE_STORAGE_MAX_ATTEMPTS``; pass 1
+                for a request that must not be replayed.
         """
         endpoint = self._describe_storage_unit(target_storage_unit)
         # Floored at one: a nonpositive count would skip the request and report success, which for
         # put would publish metadata for data that was never sent.
-        attempts_allowed = max(1, TQ_SIMPLE_STORAGE_MAX_ATTEMPTS)
+        attempts_allowed = max(1, TQ_SIMPLE_STORAGE_MAX_ATTEMPTS if max_attempts is None else max_attempts)
         for attempt in range(1, attempts_allowed + 1):
             try:
                 return await make_request()
@@ -384,8 +387,10 @@ async def put_data(
         field_schema = extract_field_schema(data)
 
         routing = self._group_by_hash(metadata.global_indexes)
-        # A retried put is replayed as-is: it overwrites the same global indexes, and re-runs
-        # data_parser on the unit, so a parser must be free of external side effects.
+        # Replaying a put is safe because it overwrites the same global indexes, but a retry also
+        # re-runs data_parser on the unit, and the public API does not require a parser to be free
+        # of side effects. So a parser-backed put is sent once and fails as it did before.
+        max_attempts = 1 if data_parser is not None else None
         tasks = []
         for su_id, group in routing.items():
             storage_data = {f: self._select_by_positions(data[f], group.batch_positions) for f in data.keys()}
@@ -402,6 +407,7 @@ async def put_data(
                         target_storage_unit=su_id,
                         data_parser=data_parser,
                     ),
+                    max_attempts=max_attempts,
                 )
             )
 

From 44b7573d9bfa2d60f63ae42c7d95f5445c290e78 Mon Sep 17 00:00:00 2001
From: jathonzhang <jathonzhang@tencent.com>
Date: Thu, 10 Sep 2026 11:36:00 +0800
Subject: [PATCH 4/8] [chore] Shorten the parser-backed put retry comment

Signed-off-by: jathonzhang <jathonzhang@tencent.com>
---
 transfer_queue/storage/managers/simple_storage_manager.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index c49f4265..9140c0f4 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -387,9 +387,7 @@ async def put_data(
         field_schema = extract_field_schema(data)
 
         routing = self._group_by_hash(metadata.global_indexes)
-        # Replaying a put is safe because it overwrites the same global indexes, but a retry also
-        # re-runs data_parser on the unit, and the public API does not require a parser to be free
-        # of side effects. So a parser-backed put is sent once and fails as it did before.
+        # Parser-backed puts are not replayed: the public API does not constrain parser side effects.
         max_attempts = 1 if data_parser is not None else None
         tasks = []
         for su_id, group in routing.items():

From 2bcbcf8cc9b54e7c7f30078c67e7ad0efc08ede7 Mon Sep 17 00:00:00 2001
From: jathonzhang <jathonzhang@tencent.com>
Date: Thu, 10 Sep 2026 11:37:07 +0800
Subject: [PATCH 5/8] [fix] Report wire size from serialized frames on put
 timeout

estimate_payload_bytes only walked tensor nbytes, so NonTensorStack fields,
msgpack/pickle overhead, and multipart framing were all invisible. Measure the
frames after serialize in _put_to_single_storage_unit and put that size on
StorageUnitTimeout; the retry layer now logs the exception text so the true
wire size shows up on both intermediate retries and the final failure.

Signed-off-by: jathonzhang <jathonzhang@tencent.com>
---
 tests/test_storage_request_retry.py           | 19 ++++++++++++++++++
 .../managers/simple_storage_manager.py        | 20 ++++++++++---------
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py
index 44701ed2..157abf3e 100644
--- a/tests/test_storage_request_retry.py
+++ b/tests/test_storage_request_retry.py
@@ -182,6 +182,25 @@ async def flaky():
     assert "samples=4" in warning.message
 
 
+@pytest.mark.asyncio
+async def test_retry_logs_timeout_detail(caplog):
+    """The retry warning must carry the size reported by the failed attempt."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+    attempts = []
+
+    async def flaky():
+        attempts.append(1)
+        if len(attempts) == 1:
+            raise StorageUnitTimeout("serialized_mb=12.5")
+        return "payload"
+
+    with caplog.at_level(logging.WARNING):
+        await manager._request_with_retry("put", "unit_a", "samples=4", flaky)
+
+    warning = next(r for r in caplog.records if "retry" in r.message)
+    assert "serialized_mb=12.5" in warning.message
+
+
 @pytest.mark.asyncio
 async def test_diagnosis_blames_the_link_when_the_unit_still_answers():
     """A unit that answers a fresh probe was not the one that stalled."""
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index 9140c0f4..d6e781bf 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -31,12 +31,12 @@
 from transfer_queue.metadata import BatchMeta, extract_field_schema
 from transfer_queue.storage.managers.base import StorageManager, StorageManagerFactory
 from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER, StorageKeyNotFoundError
-from transfer_queue.utils.common import estimate_payload_bytes
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
     ZMQRequestType,
     ZMQServerInfo,
+    frame_nbytes,
     with_zmq_socket,
 )
 
@@ -259,17 +259,17 @@ async def _request_with_retry(
         for attempt in range(1, attempts_allowed + 1):
             try:
                 return await make_request()
-            except StorageUnitTimeout:
+            except StorageUnitTimeout as e:
                 if attempt < attempts_allowed:
                     logger.warning(
                         f"[{self.storage_manager_id}]: no answer from {target_storage_unit} at {endpoint} in "
                         f"{TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s, {operation} retry "
-                        f"{attempt + 1}/{attempts_allowed}. {request_context}"
+                        f"{attempt + 1}/{attempts_allowed}. {request_context} {e}"
                     )
                     continue
                 logger.error(
                     f"[{self.storage_manager_id}]: {operation} to {target_storage_unit} at {endpoint} failed "
-                    f"after {attempt}x{TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s. {request_context} "
+                    f"after {attempt}x{TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s. {request_context} {e} "
                     f"{await self._diagnose_storage_unit(target_storage_unit)}"
                 )
                 raise
@@ -396,8 +396,7 @@ async def put_data(
                 self._request_with_retry(
                     "put",
                     su_id,
-                    f"samples={len(group.global_indexes)} fields={list(storage_data.keys())} "
-                    f"payload_mb={estimate_payload_bytes(storage_data) / 2**20:.1f}",
+                    f"samples={len(group.global_indexes)} fields={list(storage_data.keys())}",
                     partial(
                         self._put_to_single_storage_unit,
                         group.global_indexes,
@@ -450,9 +449,11 @@ async def _put_to_single_storage_unit(
             body={"global_indexes": global_indexes, "data": storage_data, "data_parser": data_parser},
         )
 
+        serialized_bytes = 0
         try:
-            data = request_msg.serialize()
-            await socket.send_multipart(data, copy=False)
+            frames = request_msg.serialize()
+            serialized_bytes = sum(frame_nbytes(frame) or 0 for frame in frames)
+            await socket.send_multipart(frames, copy=False)
             messages = await socket.recv_multipart(copy=False)
             response_msg = ZMQMessage.deserialize(messages)
 
@@ -464,7 +465,8 @@ async def _put_to_single_storage_unit(
         except zmq.error.Again as e:
             raise StorageUnitTimeout(
                 f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during put to storage unit "
-                f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}"
+                f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}; "
+                f"samples={len(global_indexes)} serialized_mb={serialized_bytes / 2**20:.1f}"
             ) from e
         except Exception as e:
             logger.error(

From 91a9e46f01ac6d304b814061c69dbf96cd39d9b2 Mon Sep 17 00:00:00 2001
From: jathonzhang <jathonzhang@tencent.com>
Date: Thu, 10 Sep 2026 11:39:18 +0800
Subject: [PATCH 6/8] [fix] Log heavy puts on the manager and heavy gets on the
 unit

The manager sees the wire size and end-to-end RTT of a put; the unit only sees
local processing after deserialize. Move put heavy logging to the manager using
the serialized frame size already measured for timeouts, and keep get heavy
logging on the unit where the response payload is built.

Signed-off-by: jathonzhang <jathonzhang@tencent.com>
---
 tests/test_storage_request_retry.py           | 17 ++++++++++++
 .../managers/simple_storage_manager.py        | 26 ++++++++++++++++++-
 transfer_queue/storage/simple_storage.py      |  8 +++---
 3 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py
index 157abf3e..80ee453d 100644
--- a/tests/test_storage_request_retry.py
+++ b/tests/test_storage_request_retry.py
@@ -16,6 +16,7 @@
 """Tests for storage-unit request retry and the timeout diagnosis that classifies a failure."""
 
 import logging
+import time
 from unittest.mock import AsyncMock, patch
 
 import numpy as np
@@ -201,6 +202,22 @@ async def flaky():
     assert "serialized_mb=12.5" in warning.message
 
 
+def test_manager_logs_heavy_put(caplog):
+    """A put that is large on the wire is logged on the manager, not the unit."""
+    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
+
+    with (
+        patch.object(ssm, "TQ_STORAGE_LARGE_PAYLOAD_MB", 0.0),
+        patch.object(ssm, "TQ_STORAGE_SLOW_REQUEST_SECONDS", 1e9),
+        caplog.at_level(logging.WARNING),
+    ):
+        manager._log_if_heavy_put(time.perf_counter(), 512 * 2**20, 4, ["input_ids"], "unit_a")
+
+    warning = next(r for r in caplog.records if "heavy put" in r.message)
+    assert "serialized_mb=512.0" in warning.message
+    assert "10.0.0.7:5555" in warning.message
+
+
 @pytest.mark.asyncio
 async def test_diagnosis_blames_the_link_when_the_unit_still_answers():
     """A unit that answers a fresh probe was not the one that stalled."""
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index d6e781bf..f964ad3b 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -15,6 +15,7 @@
 
 import asyncio
 import os
+import time
 import warnings
 from collections import defaultdict
 from collections.abc import Mapping
@@ -30,7 +31,12 @@
 
 from transfer_queue.metadata import BatchMeta, extract_field_schema
 from transfer_queue.storage.managers.base import StorageManager, StorageManagerFactory
-from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER, StorageKeyNotFoundError
+from transfer_queue.storage.simple_storage import (
+    KEY_NOT_FOUND_MARKER,
+    TQ_STORAGE_LARGE_PAYLOAD_MB,
+    TQ_STORAGE_SLOW_REQUEST_SECONDS,
+    StorageKeyNotFoundError,
+)
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
@@ -429,6 +435,20 @@ async def put_data(
             field_schema,
         )
 
+    def _log_if_heavy_put(
+        self, started: float, serialized_bytes: int, num_samples: int, fields: Any, target_storage_unit: str
+    ) -> None:
+        """Log a put that was slow or unusually large on the wire; stay silent otherwise."""
+        elapsed = time.perf_counter() - started
+        payload_mb = serialized_bytes / 2**20
+        if elapsed < TQ_STORAGE_SLOW_REQUEST_SECONDS and payload_mb < TQ_STORAGE_LARGE_PAYLOAD_MB:
+            return
+        logger.warning(
+            f"[{self.storage_manager_id}]: heavy put to {target_storage_unit} "
+            f"at {self._describe_storage_unit(target_storage_unit)} samples={num_samples} "
+            f"serialized_mb={payload_mb:.1f} elapsed={elapsed:.2f}s fields={list(fields)}"
+        )
+
     @with_storage_unit_socket
     async def _put_to_single_storage_unit(
         self,
@@ -450,6 +470,7 @@ async def _put_to_single_storage_unit(
         )
 
         serialized_bytes = 0
+        started = time.perf_counter()
         try:
             frames = request_msg.serialize()
             serialized_bytes = sum(frame_nbytes(frame) or 0 for frame in frames)
@@ -462,6 +483,9 @@ async def _put_to_single_storage_unit(
                     f"Failed to put data to storage unit {target_storage_unit}: "
                     f"{response_msg.body.get('message', 'Unknown error')}"
                 )
+            self._log_if_heavy_put(
+                started, serialized_bytes, len(global_indexes), storage_data.keys(), target_storage_unit
+            )
         except zmq.error.Again as e:
             raise StorageUnitTimeout(
                 f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during put to storage unit "
diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index b1a8ce57..50dccfbf 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -400,7 +400,10 @@ def _worker_routine(self) -> None:
         worker_socket.close(linger=0)
 
     def _log_if_heavy(self, operation: str, started: float, field_data: Any, num_samples: int, fields: Any) -> None:
-        """Log a request that was slow or unusually large, and stay silent otherwise."""
+        """Log a get that was slow or unusually large on this unit; stay silent otherwise.
+
+        Put heavy logging lives on the manager, which sees the wire size and end-to-end RTT.
+        """
         elapsed = time.perf_counter() - started
         payload_mb = estimate_payload_bytes(field_data) / 2**20
         if elapsed < TQ_STORAGE_SLOW_REQUEST_SECONDS and payload_mb < TQ_STORAGE_LARGE_PAYLOAD_MB:
@@ -425,7 +428,6 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage:
             field_data = data_parts.body["data"]  # field_data should be a dict.
             data_parser = data_parts.body.get("data_parser", None)
 
-            started = time.perf_counter()
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_put"
             ):
@@ -473,8 +475,6 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage:
                             )
                 self.storage_data.put_data(field_data, global_indexes)
 
-            self._log_if_heavy("PUT_DATA", started, field_data, len(global_indexes), field_data.keys())
-
             # After put operation finish, send a message to the client
             response_msg = ZMQMessage.create(
                 request_type=ZMQRequestType.PUT_DATA_RESPONSE,  # type: ignore[arg-type]

From f3d31af9938abf55f444fd1fb26aff0d3cd997b7 Mon Sep 17 00:00:00 2001
From: jathonzhang <jathonzhang@tencent.com>
Date: Thu, 10 Sep 2026 18:04:35 +0800
Subject: [PATCH 7/8] [fix] Measure get wire size and share heavy-request
 logging

Both ends already serialize the payload they care about, so report that size
instead of estimating tensor nbytes. Fold the duplicated threshold check into
log_heavy_operation and drop estimate_payload_bytes.

Signed-off-by: jathonzhang <jathonzhang@tencent.com>
---
 tests/test_storage_request_retry.py           | 44 ++++++++-----------
 .../managers/simple_storage_manager.py        | 28 +++++-------
 transfer_queue/storage/simple_storage.py      | 39 +++++++---------
 transfer_queue/utils/common.py                | 37 +++++++++-------
 4 files changed, 64 insertions(+), 84 deletions(-)

diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py
index 80ee453d..08bfdf0e 100644
--- a/tests/test_storage_request_retry.py
+++ b/tests/test_storage_request_retry.py
@@ -16,7 +16,6 @@
 """Tests for storage-unit request retry and the timeout diagnosis that classifies a failure."""
 
 import logging
-import time
 from unittest.mock import AsyncMock, patch
 
 import numpy as np
@@ -31,7 +30,7 @@
     AsyncSimpleStorageManager,
     StorageUnitTimeout,
 )
-from transfer_queue.utils.common import estimate_payload_bytes
+from transfer_queue.utils import common
 from transfer_queue.utils.enum_utils import Role
 from transfer_queue.utils.zmq_utils import ZMQServerInfo
 
@@ -202,20 +201,30 @@ async def flaky():
     assert "serialized_mb=12.5" in warning.message
 
 
-def test_manager_logs_heavy_put(caplog):
-    """A put that is large on the wire is logged on the manager, not the unit."""
-    manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
-
+def test_log_heavy_operation_reports_the_measured_wire_size(caplog):
+    """Both call sites hand over serialized bytes, so the message reports them as measured."""
     with (
-        patch.object(ssm, "TQ_STORAGE_LARGE_PAYLOAD_MB", 0.0),
-        patch.object(ssm, "TQ_STORAGE_SLOW_REQUEST_SECONDS", 1e9),
+        patch.object(common, "TQ_STORAGE_LARGE_PAYLOAD_MB", 0.0),
+        patch.object(common, "TQ_STORAGE_SLOW_REQUEST_SECONDS", 1e9),
         caplog.at_level(logging.WARNING),
     ):
-        manager._log_if_heavy_put(time.perf_counter(), 512 * 2**20, 4, ["input_ids"], "unit_a")
+        common.log_heavy_operation("TQ_STORAGE_test", "put", 0.25, 512 * 2**20, "to unit_a samples=4")
 
     warning = next(r for r in caplog.records if "heavy put" in r.message)
     assert "serialized_mb=512.0" in warning.message
-    assert "10.0.0.7:5555" in warning.message
+    assert "to unit_a samples=4" in warning.message
+
+
+def test_log_heavy_operation_stays_silent_below_both_thresholds(caplog):
+    """Only requests past a threshold are worth a line; the normal path must not log."""
+    with (
+        patch.object(common, "TQ_STORAGE_LARGE_PAYLOAD_MB", 256.0),
+        patch.object(common, "TQ_STORAGE_SLOW_REQUEST_SECONDS", 5.0),
+        caplog.at_level(logging.WARNING),
+    ):
+        common.log_heavy_operation("TQ_STORAGE_test", "get", 0.01, 2**20, "samples=4")
+
+    assert not caplog.records
 
 
 @pytest.mark.asyncio
@@ -295,18 +304,3 @@ class _Writer:
 
     def close(self):
         pass
-
-
-def test_estimate_payload_bytes_handles_put_and_get_shapes():
-    """Both the batched put shape and the per-sample get shape must be measurable."""
-    batched = {"input_ids": torch.zeros(4, 8, dtype=torch.int64)}
-    per_sample = {"input_ids": [torch.zeros(8, dtype=torch.int64) for _ in range(4)]}
-
-    assert estimate_payload_bytes(batched) == 4 * 8 * 8
-    assert estimate_payload_bytes(per_sample) == 4 * 8 * 8
-
-
-def test_estimate_payload_bytes_degrades_to_zero_on_unmeasurable_input():
-    """Size estimation is diagnostic only and must never break the path it reports on."""
-    assert estimate_payload_bytes(object()) == 0
-    assert estimate_payload_bytes({"meta": "not a tensor"}) == 0
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index f964ad3b..64cb38d6 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -33,10 +33,9 @@
 from transfer_queue.storage.managers.base import StorageManager, StorageManagerFactory
 from transfer_queue.storage.simple_storage import (
     KEY_NOT_FOUND_MARKER,
-    TQ_STORAGE_LARGE_PAYLOAD_MB,
-    TQ_STORAGE_SLOW_REQUEST_SECONDS,
     StorageKeyNotFoundError,
 )
+from transfer_queue.utils.common import log_heavy_operation
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
@@ -435,20 +434,6 @@ async def put_data(
             field_schema,
         )
 
-    def _log_if_heavy_put(
-        self, started: float, serialized_bytes: int, num_samples: int, fields: Any, target_storage_unit: str
-    ) -> None:
-        """Log a put that was slow or unusually large on the wire; stay silent otherwise."""
-        elapsed = time.perf_counter() - started
-        payload_mb = serialized_bytes / 2**20
-        if elapsed < TQ_STORAGE_SLOW_REQUEST_SECONDS and payload_mb < TQ_STORAGE_LARGE_PAYLOAD_MB:
-            return
-        logger.warning(
-            f"[{self.storage_manager_id}]: heavy put to {target_storage_unit} "
-            f"at {self._describe_storage_unit(target_storage_unit)} samples={num_samples} "
-            f"serialized_mb={payload_mb:.1f} elapsed={elapsed:.2f}s fields={list(fields)}"
-        )
-
     @with_storage_unit_socket
     async def _put_to_single_storage_unit(
         self,
@@ -483,8 +468,15 @@ async def _put_to_single_storage_unit(
                     f"Failed to put data to storage unit {target_storage_unit}: "
                     f"{response_msg.body.get('message', 'Unknown error')}"
                 )
-            self._log_if_heavy_put(
-                started, serialized_bytes, len(global_indexes), storage_data.keys(), target_storage_unit
+            # This end serializes the put payload and waits out the round trip, so it holds both
+            # the true wire size and the end-to-end latency.
+            log_heavy_operation(
+                self.storage_manager_id,
+                "put",
+                time.perf_counter() - started,
+                serialized_bytes,
+                f"to {target_storage_unit} at {self._describe_storage_unit(target_storage_unit)} "
+                f"samples={len(global_indexes)} fields={list(storage_data.keys())}",
             )
         except zmq.error.Again as e:
             raise StorageUnitTimeout(
diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index 50dccfbf..cd0d2b4b 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -25,7 +25,7 @@
 import ray
 import zmq
 
-from transfer_queue.utils.common import estimate_payload_bytes, limit_pytorch_auto_parallel_threads
+from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads, log_heavy_operation
 from transfer_queue.utils.enum_utils import Role
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.perf_utils import IntervalPerfMonitor
@@ -36,6 +36,7 @@
     ZMQServerInfo,
     create_zmq_socket,
     format_zmq_address,
+    frame_nbytes,
     get_free_port,
     get_node_ip_address,
 )
@@ -48,11 +49,6 @@
 TQ_STORAGE_POLLER_TIMEOUT = int(os.environ.get("TQ_STORAGE_POLLER_TIMEOUT", 5))  # in seconds
 TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8))
 
-# Thresholds above which a single request earns a log line. Both sit far above the normal range
-# of single-digit milliseconds, so tripping either one is itself the finding.
-TQ_STORAGE_SLOW_REQUEST_SECONDS = float(os.environ.get("TQ_STORAGE_SLOW_REQUEST_SECONDS", 5.0))
-TQ_STORAGE_LARGE_PAYLOAD_MB = float(os.environ.get("TQ_STORAGE_LARGE_PAYLOAD_MB", 256))
-
 # Marks a GET_ERROR reply as "the key is gone" so the caller can tell it apart from a real fault.
 KEY_NOT_FOUND_MARKER = "TQKeyNotFound"
 
@@ -349,6 +345,7 @@ def _worker_routine(self) -> None:
                     worker_socket.send_multipart([identity] + error_msg.serialize(), copy=False)
                     continue
                 operation = request_msg.request_type
+                started = time.perf_counter()
 
                 try:
                     logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}")
@@ -393,26 +390,23 @@ def _worker_routine(self) -> None:
                     )
 
                 # Send response back with identity for routing
-                worker_socket.send_multipart([identity] + response_msg.serialize(), copy=False)
+                response_frames = response_msg.serialize()
+                if operation == ZMQRequestType.GET_DATA:  # type: ignore[arg-type]
+                    # This end serializes the get response, so its frames give the true wire size.
+                    log_heavy_operation(
+                        self.storage_unit_id,
+                        "get",
+                        time.perf_counter() - started,
+                        sum(frame_nbytes(frame) or 0 for frame in response_frames),
+                        f"samples={len(request_msg.body.get('global_indexes', []))} "
+                        f"fields={list(request_msg.body.get('fields', []))}",
+                    )
+                worker_socket.send_multipart([identity] + response_frames, copy=False)
 
         logger.info(f"[{self.storage_unit_id}]: worker stopped.")
         poller.unregister(worker_socket)
         worker_socket.close(linger=0)
 
-    def _log_if_heavy(self, operation: str, started: float, field_data: Any, num_samples: int, fields: Any) -> None:
-        """Log a get that was slow or unusually large on this unit; stay silent otherwise.
-
-        Put heavy logging lives on the manager, which sees the wire size and end-to-end RTT.
-        """
-        elapsed = time.perf_counter() - started
-        payload_mb = estimate_payload_bytes(field_data) / 2**20
-        if elapsed < TQ_STORAGE_SLOW_REQUEST_SECONDS and payload_mb < TQ_STORAGE_LARGE_PAYLOAD_MB:
-            return
-        logger.warning(
-            f"[{self.storage_unit_id}]: heavy {operation} samples={num_samples} "
-            f"payload_mb={payload_mb:.1f} elapsed={elapsed:.2f}s fields={list(fields)}"
-        )
-
     def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage:
         """
         Handle put request, add or update data into storage unit.
@@ -507,14 +501,11 @@ def _handle_get(self, data_parts: ZMQMessage) -> ZMQMessage:
             fields = data_parts.body["fields"]
             global_indexes = data_parts.body["global_indexes"]
 
-            started = time.perf_counter()
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_get"
             ):
                 result_data = self.storage_data.get_data(fields, global_indexes)
 
-            self._log_if_heavy("GET_DATA", started, result_data, len(global_indexes), fields)
-
             response_msg = ZMQMessage.create(
                 request_type=ZMQRequestType.GET_DATA_RESPONSE,  # type: ignore[arg-type]
                 sender_id=self.storage_unit_id,
diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py
index d92b628f..1278dfd4 100644
--- a/transfer_queue/utils/common.py
+++ b/transfer_queue/utils/common.py
@@ -14,9 +14,7 @@
 # limitations under the License.
 
 import os
-from collections.abc import Mapping
 from contextlib import contextmanager
-from typing import Any
 
 import psutil
 import ray
@@ -137,20 +135,25 @@ def get_env_bool(env_key: str, default: bool = False) -> bool:
     return env_value_lower in true_values
 
 
-def estimate_payload_bytes(field_data: Any) -> int:
-    """Best-effort size of a request payload in bytes; 0 when it cannot be measured.
+# Thresholds above which a single request earns a log line. Both sit far above the normal range
+# of single-digit milliseconds, so tripping either one is itself the finding.
+TQ_STORAGE_SLOW_REQUEST_SECONDS = float(os.environ.get("TQ_STORAGE_SLOW_REQUEST_SECONDS", 5.0))
+TQ_STORAGE_LARGE_PAYLOAD_MB = float(os.environ.get("TQ_STORAGE_LARGE_PAYLOAD_MB", 256))
 
-    Walks two levels: covers both put's ``field -> value`` and get's ``field -> per-sample list``.
+
+def log_heavy_operation(component_id: str, operation: str, elapsed: float, payload_bytes: int, detail: str) -> None:
+    """Warn about one slow or unusually large request; stay silent otherwise.
+
+    Args:
+        component_id: Storage manager or storage unit reporting the request.
+        operation: Operation name, e.g. ``put`` or ``get``.
+        elapsed: Wall time spent on the request, in seconds.
+        payload_bytes: Serialized size of the payload on the wire.
+        detail: Request shape, appended to the message verbatim.
     """
-    total = 0
-    try:
-        values = field_data.values() if isinstance(field_data, Mapping) else field_data
-        for value in values:
-            items = value if isinstance(value, list | tuple) else [value]
-            for item in items:
-                nbytes = getattr(item, "nbytes", None)
-                if isinstance(nbytes, int):
-                    total += nbytes
-    except Exception:
-        return 0
-    return total
+    payload_mb = payload_bytes / 2**20
+    if elapsed < TQ_STORAGE_SLOW_REQUEST_SECONDS and payload_mb < TQ_STORAGE_LARGE_PAYLOAD_MB:
+        return
+    logger.warning(
+        f"[{component_id}]: heavy {operation} {detail} serialized_mb={payload_mb:.1f} elapsed={elapsed:.2f}s"
+    )

From 7f28a136b56109875082730cda9ae05e2434e6d4 Mon Sep 17 00:00:00 2001
From: jathonzhang <jathonzhang@tencent.com>
Date: Thu, 10 Sep 2026 21:43:37 +0800
Subject: [PATCH 8/8] [fix] Stop repeating the failure detail in retry logs

The timeout exception is what put_data and get_data callers see, so it names the
unit, its endpoint and the timeout; the retry logs now add only the attempt and
the request shape. Both operations report the same fields.

Signed-off-by: jathonzhang <jathonzhang@tencent.com>
---
 tests/test_storage_request_retry.py           | 11 ++++++-----
 .../managers/simple_storage_manager.py        | 19 ++++++++++---------
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py
index 08bfdf0e..75a41d36 100644
--- a/tests/test_storage_request_retry.py
+++ b/tests/test_storage_request_retry.py
@@ -163,23 +163,24 @@ async def hard_failure():
 
 
 @pytest.mark.asyncio
-async def test_retry_logs_endpoint_and_request_shape(caplog):
-    """The retry warning must name the endpoint and the request, to correlate both ends."""
+async def test_retry_log_adds_request_shape_without_echoing_the_failure(caplog):
+    """The failure already names the unit and endpoint; the retry line adds shape, not an echo."""
     manager = _manager(unit_a=_server_info("unit_a", "10.0.0.7", 5555))
     attempts = []
 
     async def flaky():
         attempts.append(1)
         if len(attempts) == 1:
-            raise StorageUnitTimeout("no answer")
+            raise StorageUnitTimeout("no answer in 200s during get from storage unit unit_a at 10.0.0.7:5555")
         return "payload"
 
     with caplog.at_level(logging.WARNING):
         await manager._request_with_retry("get", "unit_a", "samples=4 fields=['input_ids']", flaky)
 
     warning = next(r for r in caplog.records if "retry" in r.message)
-    assert "10.0.0.7:5555" in warning.message
-    assert "samples=4" in warning.message
+    assert "samples=4 fields=['input_ids']" in warning.message
+    assert warning.message.count("10.0.0.7:5555") == 1
+    assert warning.message.count("unit_a") == 1
 
 
 @pytest.mark.asyncio
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index 64cb38d6..43b7a242 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -61,6 +61,8 @@ class StorageUnitTimeout(RuntimeError):
     """A storage unit did not answer within the send/recv timeout.
 
     Distinct from an error the unit reported: only a missing answer is worth a new connection.
+    The message must name the unit, its endpoint and the timeout, because it is what the callers
+    of ``put_data`` and ``get_data`` see, and the retry logs rely on it instead of repeating them.
     """
 
 
@@ -257,7 +259,6 @@ async def _request_with_retry(
             max_attempts: Attempts allowed. Defaults to ``TQ_SIMPLE_STORAGE_MAX_ATTEMPTS``; pass 1
                 for a request that must not be replayed.
         """
-        endpoint = self._describe_storage_unit(target_storage_unit)
         # Floored at one: a nonpositive count would skip the request and report success, which for
         # put would publish metadata for data that was never sent.
         attempts_allowed = max(1, TQ_SIMPLE_STORAGE_MAX_ATTEMPTS if max_attempts is None else max_attempts)
@@ -265,17 +266,17 @@ async def _request_with_retry(
             try:
                 return await make_request()
             except StorageUnitTimeout as e:
+                # The exception already names the unit, its endpoint and the timeout, so these
+                # lines only add what it cannot know: which attempt this was, and the shape.
                 if attempt < attempts_allowed:
                     logger.warning(
-                        f"[{self.storage_manager_id}]: no answer from {target_storage_unit} at {endpoint} in "
-                        f"{TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s, {operation} retry "
-                        f"{attempt + 1}/{attempts_allowed}. {request_context} {e}"
+                        f"[{self.storage_manager_id}]: {operation} retry {attempt + 1}/{attempts_allowed} "
+                        f"on a new connection. {request_context} {e}"
                     )
                     continue
                 logger.error(
-                    f"[{self.storage_manager_id}]: {operation} to {target_storage_unit} at {endpoint} failed "
-                    f"after {attempt}x{TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s. {request_context} {e} "
-                    f"{await self._diagnose_storage_unit(target_storage_unit)}"
+                    f"[{self.storage_manager_id}]: {operation} failed after {attempt} attempts. "
+                    f"{request_context} {e} {await self._diagnose_storage_unit(target_storage_unit)}"
                 )
                 raise
 
@@ -482,7 +483,7 @@ async def _put_to_single_storage_unit(
             raise StorageUnitTimeout(
                 f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during put to storage unit "
                 f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}; "
-                f"samples={len(global_indexes)} serialized_mb={serialized_bytes / 2**20:.1f}"
+                f"serialized_mb={serialized_bytes / 2**20:.1f}"
             ) from e
         except Exception as e:
             logger.error(
@@ -632,7 +633,7 @@ async def _get_from_single_storage_unit(
                 raise error_type(f"Failed to get data from storage unit {target_storage_unit}: {message}")
         except zmq.error.Again as e:
             raise StorageUnitTimeout(
-                f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s from storage unit "
+                f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during get from storage unit "
                 f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}"
             ) from e
         except StorageKeyNotFoundError: