sdk/python: add workaround for Streaming-Cold-GET limitation in ObjectFileReader

rkoo19 · rkoo19 · commit 1016bdacbd14 · 2025-05-16T19:11:39.000-07:00
Signed-off-by: Ryan Koo &lt;rbk65@cornell.edu&gt;
diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md
@@ -6,11 +6,13 @@ We structure this changelog in accordance with [Keep a Changelog](https://keepac
 
 ---
 
-## Unreleased
+## [1.13.8] - 2025-05-16
 
 ### Changed
 
 - Improve `ObjectFileReader` logging to include the full exception details and traceback when retrying and resuming.
+- Update `ObjectFileReader` resume logic to accommodate for limitation w/ the `Streaming-Cold-GET` and read range as to not cause a timeout. 
+  - If the remote object is not cached, resuming via stream is not possible and reading must restart from the beginning.
 
 ## [1.13.8] - 2025-05-15
 
diff --git a/python/aistore/sdk/obj/obj_file/object_file.py b/python/aistore/sdk/obj/obj_file/object_file.py
@@ -11,6 +11,7 @@
 from aistore.sdk.utils import get_logger
 from aistore.sdk.obj.obj_file.utils import (
     handle_broken_stream,
+    compute_loop_size,
 )
 
 logger = get_logger(__name__)
@@ -39,12 +40,13 @@ def __init__(self, content_iterator: ContentIterator, max_resume: int):
         self._max_resume = max_resume  # Maximum number of resume attempts allowed
         self._reset()
 
-    def _reset(self):
+    def _reset(self, retain_resumes: bool = False) -> None:
         self._iterable = self._content_iterator.iter()
         self._remainder = None
         self._resume_position = 0
-        self._resume_total = 0
         self._closed = False
+        if not retain_resumes:
+            self._resume_total = 0
 
     @property
     def content_iterator(self) -> ContentIterator:
@@ -82,9 +84,13 @@ def read(self, size: Optional[int] = -1) -> bytes:
             raise ValueError("I/O operation on closed file.")
         if size == 0:
             return b""
+        if size is None:
+            size = -1
 
-        # If size is -1, set it to infinity to read until the end of the stream
-        size = float("inf") if size == -1 else size
+        # Cache original requested size in case of reset
+        original_size = size
+
+        size = compute_loop_size(size)
         result = []
 
         try:
@@ -119,7 +125,16 @@ def read(self, size: Optional[int] = -1) -> bytes:
                         self._max_resume,
                         err,
                     )
+
+                    # If the object is remote and not cached, reset the iterator and restart the read operation
+                    # to avoid timeouts when streaming non-cached remote objects with byte ranges (must wait for
+                    # entire object to be cached in-cluster).
+                    if not self._iterable:
+                        self._reset(retain_resumes=True)
+                        size = compute_loop_size(original_size)
+
                     continue
+
                 except StopIteration:
                     # End of stream, exit loop
                     break
diff --git a/python/aistore/sdk/obj/obj_file/utils.py b/python/aistore/sdk/obj/obj_file/utils.py
@@ -2,7 +2,8 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 #
 
-from typing import Iterator, Tuple
+from sys import maxsize
+from typing import Iterator, Tuple, Optional
 
 from aistore.sdk.obj.content_iterator import ContentIterator
 from aistore.sdk.obj.obj_file.errors import ObjectFileReaderMaxResumeError
@@ -11,6 +12,44 @@
 logger = get_logger(__name__)
 
 
+def compute_loop_size(size: int) -> int:
+    """
+    Compute the size of the loop for reading data and return it.
+    If requested read size is -1, return `sys.maxsize` to loop until `StopIteration`.
+
+    Args:
+        size (int): The requested size to be read.
+
+    Returns:
+        int: The size for the loop.
+
+    """
+    return maxsize if size == -1 else size
+
+
+def get_iterator(
+    content_iterator: ContentIterator, resume_position: int
+) -> Optional[Iterator[bytes]]:
+    """
+    Create a new iterator from the content iterator starting at the specified byte position.
+    Returns None if the object is not cached.
+
+    Args:
+        content_iterator (ContentIterator): The content iterator used to read the data.
+        resume_position (int): The byte position from which to resume reading.
+
+    Returns:
+        Optional[Iterator[bytes]]: A new iterator starting from the specified byte position.
+            None if the object is not cached in the bucket.
+    """
+    # If remote object is not cached, start over
+    if not content_iterator.client.head().present:
+        return None
+    # Otherwise, resume from last known position
+    else:
+        return content_iterator.iter(offset=resume_position)
+
+
 def increment_resume(resume_total: int, max_resume: int, err: Exception) -> int:
     """
     Increment the number of resume attempts and raise an error if the maximum allowed is exceeded.
@@ -38,7 +77,7 @@ def handle_broken_stream(
     resume_total: int,
     max_resume: int,
     err: Exception,
-) -> Tuple[Iterator[bytes], int]:
+) -> Tuple[Optional[Iterator[bytes]], int]:
     """
     Handle the broken stream/iterator by incrementing the resume count, logging a warning,
     and returning a newly instanatiated iterator from the last known position.
@@ -51,7 +90,8 @@ def handle_broken_stream(
         err (Exception): The error that caused the resume attempt.
 
     Returns:
-        Tuple[Iterator[bytes], int]: The new iterator and the updated resume total.
+        Optional[Iterator[bytes]]: The new iterator. None if the object is not cached.
+        int: The updated number of resume attempts.
 
     Raises:
         ObjectFileReaderMaxResumeError: If the maximum number of resume attempts is exceeded.
@@ -66,6 +106,8 @@ def handle_broken_stream(
         exc_info=err,
     )
 
-    # Create a new iterator from the last read position
-    new_iter = content_iterator.iter(offset=resume_position)
+    new_iter = get_iterator(
+        content_iterator=content_iterator, resume_position=resume_position
+    )
+
     return new_iter, resume_total
diff --git a/python/aistore/sdk/obj/object_attributes.py b/python/aistore/sdk/obj/object_attributes.py
@@ -11,6 +11,7 @@
     AIS_ACCESS_TIME,
     AIS_VERSION,
     AIS_CUSTOM_MD,
+    AIS_PRESENT,
 )
 
 
@@ -71,6 +72,13 @@ def custom_metadata(self) -> Dict[str, str]:
             return self._parse_custom(custom_md_header)
         return {}
 
+    @property
+    def present(self) -> bool:
+        """
+        Whether the object is present/cached.
+        """
+        return self._response_headers.get(AIS_PRESENT, "") == "true"
+
     @staticmethod
     def _parse_custom(custom_md_header) -> Dict[str, str]:
         """
diff --git a/python/aistore/version.py b/python/aistore/version.py
@@ -1 +1 @@
-__version__ = "1.13.8"
+__version__ = "1.13.9"
diff --git a/python/tests/unit/sdk/obj/test_object_file_reader.py b/python/tests/unit/sdk/obj/test_object_file_reader.py

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-__version__ = "1.13.8"`
	`1`	`+__version__ = "1.13.9"`