[bugfix] ODPS reader: pin pyodps==0.12.5.1 (retry connection resets) + log session/table context on errors (#537)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: tiankongdeguiji

requirements/runtime.txt

@@ -15,7 +15,7 @@ psutil
 pyfg @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/pyfg-1.0.5-cp312-cp312-linux_x86_64.whl ; python_version=="3.12"
 pyfg @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/pyfg-1.0.5-cp311-cp311-linux_x86_64.whl ; python_version=="3.11"
 pyfg @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/pyfg-1.0.5-cp310-cp310-linux_x86_64.whl ; python_version=="3.10"
-pyodps>=0.12.4
+pyodps==0.12.5.1
 safetensors
 scikit-learn
 tensorboard

tzrec/datasets/odps_dataset.py

@@ -206,21 +206,52 @@ def _parse_table_path(
     return str_list[2], table_name, table_partitions, schema
 
 
+def _storage_debug_info(client: StorageApiArrowClient, **req_fields: Any) -> str:
+    """Format MaxCompute storage-api context for error logging."""
+    parts = {
+        "table": client.table.full_table_name,
+        "quota": getattr(client, "_quota_name", None),
+        "endpoint": getattr(client, "_rest_endpoint", None),
+        **req_fields,
+    }
+    return ", ".join(f"{k}={v}" for k, v in parts.items())
+
+
 def _read_rows_arrow_with_retry(
     client: StorageApiArrowClient,
     read_req: ReadRowsRequest,
 ) -> ArrowReader:
+    def debug_info() -> str:
+        return _storage_debug_info(
+            client,
+            session_id=read_req.session_id,
+            row_index=read_req.row_index,
+            row_count=read_req.row_count,
+            max_batch_rows=read_req.max_batch_rows,
+        )
+
     max_retry_count = 3
     retry_cnt = 0
     while True:
         try:
             reader = client.read_rows_arrow(read_req)
         except ODPSError as e:
             if retry_cnt >= max_retry_count:
+                logger.error(
+                    f"read_rows_arrow failed after {retry_cnt} retries "
+                    f"({debug_info()}): {e!r}"
+                )
                 raise e
             retry_cnt += 1
+            logger.warning(
+                f"read_rows_arrow retry {retry_cnt}/{max_retry_count} "
+                f"({debug_info()}): {e!r}"
+            )
             time.sleep(random.choice([5, 9, 12]))
             continue
+        except Exception as e:
+            logger.error(f"read_rows_arrow failed ({debug_info()}): {e!r}")
+            raise
         break
     return reader
 
@@ -231,7 +262,12 @@ def _get_session_record_count(
 ) -> int:
     """Get record count from a session, waiting until ready."""
     while True:
-        scan_resp = client.get_read_session(sess_req)
+        try:
+            scan_resp = client.get_read_session(sess_req)
+        except Exception as e:
+            debug_info = _storage_debug_info(client, session_id=sess_req.session_id)
+            logger.error(f"get_read_session failed ({debug_info}): {e!r}")
+            raise
         if scan_resp.session_status == SessionStatus.INIT:
             time.sleep(1)
             continue

tzrec/datasets/odps_dataset_test.py

@@ -14,14 +14,19 @@
 import os
 import time
 import unittest
+from types import SimpleNamespace
+from unittest import mock
 
 import numpy as np
 import pyarrow as pa
+import requests
 from odps import ODPS
+from odps.errors import ODPSError
 from parameterized import parameterized
 from torch import distributed as dist
 from torch.utils.data import DataLoader
 
+from tzrec.datasets import odps_dataset
 from tzrec.datasets.odps_dataset import OdpsDataset, OdpsWriter, _create_odps_account
 from tzrec.features.feature import FgMode, create_features
 from tzrec.protos import data_pb2, feature_pb2, sampler_pb2
@@ -689,5 +694,107 @@ def _writer_worker(rank, port):
             self.assertEqual(reader.count, 1280)
 
 
+class _FailingStorageClient:
+    """Stub storage-api client whose calls raise a connection-reset error."""
+
+    table = SimpleNamespace(full_table_name="test_project.test_table")
+    _quota_name = "test_quota"
+    _rest_endpoint = "http://service.odps.test"
+
+    def _raise(self):
+        raise requests.exceptions.ConnectionError(
+            "Connection aborted.",
+            ConnectionResetError(104, "Connection reset by peer"),
+        )
+
+    def read_rows_arrow(self, read_req):
+        self._raise()
+
+    def get_read_session(self, sess_req):
+        self._raise()
+
+
+class _RetryingStorageClient:
+    """Stub whose read_rows_arrow raises ODPSError `fail_times` times, then succeeds."""
+
+    table = SimpleNamespace(full_table_name="test_project.test_table")
+    _quota_name = "test_quota"
+    _rest_endpoint = "http://service.odps.test"
+
+    def __init__(self, fail_times):
+        self._fail_times = fail_times
+        self.calls = 0
+        self.sentinel = object()
+
+    def read_rows_arrow(self, read_req):
+        self.calls += 1
+        if self.calls <= self._fail_times:
+            raise ODPSError("synthetic ODPS error", request_id="REQ-XYZ")
+        return self.sentinel
+
+
+class OdpsStorageErrorLogTest(unittest.TestCase):
+    def test_read_rows_arrow_logs_session_and_reraises(self):
+        client = _FailingStorageClient()
+        read_req = SimpleNamespace(
+            session_id="sess-read-123",
+            row_index=0,
+            row_count=10,
+            max_batch_rows=100,
+        )
+        with mock.patch.object(odps_dataset, "logger") as m_logger:
+            with self.assertRaises(requests.exceptions.ConnectionError):
+                odps_dataset._read_rows_arrow_with_retry(client, read_req)
+        m_logger.error.assert_called_once()
+        logged = m_logger.error.call_args[0][0]
+        self.assertIn("sess-read-123", logged)
+        self.assertIn("test_project.test_table", logged)
+
+    def test_get_read_session_logs_session_and_reraises(self):
+        client = _FailingStorageClient()
+        sess_req = SimpleNamespace(session_id="sess-scan-456")
+        with mock.patch.object(odps_dataset, "logger") as m_logger:
+            with self.assertRaises(requests.exceptions.ConnectionError):
+                odps_dataset._get_session_record_count(client, sess_req)
+        m_logger.error.assert_called_once()
+        logged = m_logger.error.call_args[0][0]
+        self.assertIn("sess-scan-456", logged)
+
+    def test_read_rows_arrow_retries_odps_error_then_succeeds(self):
+        client = _RetryingStorageClient(fail_times=2)
+        read_req = SimpleNamespace(
+            session_id="sess-retry-1",
+            row_index=0,
+            row_count=10,
+            max_batch_rows=100,
+        )
+        with mock.patch.object(odps_dataset.time, "sleep"):
+            with mock.patch.object(odps_dataset, "logger") as m_logger:
+                reader = odps_dataset._read_rows_arrow_with_retry(client, read_req)
+        self.assertIs(reader, client.sentinel)
+        self.assertEqual(client.calls, 3)  # 2 failures + 1 success
+        self.assertEqual(m_logger.warning.call_count, 2)
+        m_logger.error.assert_not_called()
+
+    def test_read_rows_arrow_odps_error_exhausts_retries_and_raises(self):
+        client = _RetryingStorageClient(fail_times=99)
+        read_req = SimpleNamespace(
+            session_id="sess-retry-2",
+            row_index=5,
+            row_count=10,
+            max_batch_rows=100,
+        )
+        with mock.patch.object(odps_dataset.time, "sleep"):
+            with mock.patch.object(odps_dataset, "logger") as m_logger:
+                with self.assertRaises(ODPSError):
+                    odps_dataset._read_rows_arrow_with_retry(client, read_req)
+        self.assertEqual(client.calls, 4)  # initial + 3 retries
+        self.assertEqual(m_logger.warning.call_count, 3)
+        m_logger.error.assert_called_once()
+        logged = m_logger.error.call_args[0][0]
+        self.assertIn("after 3 retries", logged)
+        self.assertIn("sess-retry-2", logged)
+
+
 if __name__ == "__main__":
     unittest.main()

tzrec/version.py

@@ -9,4 +9,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "1.2.19"
+__version__ = "1.2.20"