sdk/python: objectfile error handling fixes

rkoo19 · rkoo19 · commit 8503952f7520 · 2025-05-15T21:34:24.000-07:00
Signed-off-by: Ryan Koo &lt;rbk65@cornell.edu&gt;
diff --git a/docs/_posts/2024-09-26-resilient-streaming-with-object-file.md b/docs/_posts/2024-09-26-resilient-streaming-with-object-file.md
@@ -39,7 +39,7 @@ During each `read` operation, `ObjectFile` will attempt to resume a disconnected
 
 ## Background
 
-The basic `object.get()` function in the AIS SDK returns an `ObjectReader`, which enables efficient streaming of data by providing an iterator that fetches data in chunks (using the [requests library's stream option](https://requests.readthedocs.io/en/latest/user/advanced/#body-content-workflow)). However, it doesn’t account for network interruptions, such as unexpected stream closures that may lead to a `ChunkedEncodingError`. In the case of a failure, the entire stream may need to be manually restarted, resulting in major inefficiencies and redundant data transfer.
+The basic `object.get()` function in the AIS SDK returns an `ObjectReader`, which enables efficient streaming of data by providing an iterator that fetches data in chunks (using the [requests library's stream option](https://requests.readthedocs.io/en/latest/user/advanced/#body-content-workflow)). However, it doesn’t account for network interruptions, such as unexpected stream closures or timeouts. In the case of a failure, the entire stream may need to be manually restarted, resulting in major inefficiencies and redundant data transfer.
 
 `ObjectFile` addresses this problem by wrapping `ObjectReader` with a file-like interface that incorporates buffer management and logic to handle unexpected interruptions to the object stream. By resuming streams from the exact byte where an interruption occurred, `ObjectFile` eliminates the need to re-fetch previously read data, ensuring uninterrupted data access even in unstable network conditions.
 
@@ -65,7 +65,7 @@ The basic `object.get()` function in the AIS SDK returns an `ObjectReader`, whic
 
 When a `read` operation is initiated, `ObjectFile` first checks if the requested data is already in the buffer. If sufficient data is available, it reads directly from the buffer. If not, it fetches additional chunks from the object stream using the iterator and refills the buffer.
 
-If a network interruption occurs and the iterator raises a `ChunkedEncodingError`, `ObjectFile` automatically attempts to re-establish the connection to the object in AIS. It does this by issuing an [HTTP range request](https://developer.mozilla.org/en-US/docs/Web/HTTP/Range_requests) starting at the next byte that hasn’t yet been loaded into the buffer. This ensures that no previously accessed data is fetched again, and the stream resumes seamlessly from where it left off.
+In case of unexpected stream interruptions (e.g. `ChunkedEncodingError`, `ConnectionError`) or timeouts (e.g. `ReadTimeout`), `ObjectFile` automatically attempts to re-establish the connection to the object in AIS. It does this by issuing an [HTTP range request](https://developer.mozilla.org/en-US/docs/Web/HTTP/Range_requests) starting at the next byte that hasn’t yet been loaded into the buffer. This ensures that no previously accessed data is fetched again, and the stream resumes seamlessly from where it left off.
 
 Once the buffer has sufficient data, `ObjectFile` reads directly from it to complete the `read` operation.
 
diff --git a/docs/_posts/2024-11-21-enhancing-object-file-performance-with-zero-copy-techniques.md b/docs/_posts/2024-11-21-enhancing-object-file-performance-with-zero-copy-techniques.md
@@ -8,7 +8,7 @@ categories: aistore python resilient streaming benchmark optimization enhancemen
 
 # Enhancing ObjectFile Performance with Zero-Copy Techniques
 
-In our [previous blog post](https://aistore.nvidia.com/blog/2024/09/26/resilient-streaming-with-object-file), we introduced `ObjectFile`, a resilient, file-like interface in the AIStore Python SDK designed to handle large, streamed datasets efficiently. `ObjectFile` addresses the challenges of network interruptions and unexpected node failures by automatically handling `ChunkedEncodingError` exceptions. This allows streams to resume seamlessly from the last successful byte, ensuring uninterrupted data access even in unstable network conditions — a critical feature for long-running tasks like training models on massive datasets.
+In our [previous blog post](https://aistore.nvidia.com/blog/2024/09/26/resilient-streaming-with-object-file), we introduced `ObjectFile`, a resilient, file-like interface in the AIStore Python SDK designed to handle large, streamed datasets efficiently. `ObjectFile` addresses the challenges of network interruptions and unexpected node failures by automatically handling unexpected stream interruptions (e.g. `ChunkedEncodingError`,`ConnectionError`) or timeouts (e.g. `ReadTimeout`). This allows streams to resume seamlessly from the last successful byte, ensuring uninterrupted data access even in unstable network conditions — a critical feature for long-running tasks like training models on massive datasets.
 
 While `ObjectFile` provided significant improvements in reliability by handling intermittent failures gracefully, we identified opportunities to optimize its performance. In this post, we'll delve into recent enhancements that leverage zero-copy techniques using `memoryview`, resulting in substantial performance gains when handling large-scale data.
 
diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md
@@ -16,7 +16,9 @@ We structure this changelog in accordance with [Keep a Changelog](https://keepac
 - Add `urllib3.exceptions.TimeoutError` to retry config to handle additional timeout scenarios.
 
 ### Changed
+
 - Enhance object file logging to include the object path for better traceability.
+- Refactored `ObjectFileReader` to handle additional relevant stream and connection level errors and fixed incorrect assumptions about when certain errors occur (during iteration rather than iterator instantiation).
 
 ## [1.13.7] - 2025-05-06
 
diff --git a/python/aistore/sdk/obj/obj_file/object_file.py b/python/aistore/sdk/obj/obj_file/object_file.py
@@ -1,14 +1,17 @@
 #
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 #
 
-import requests
+from requests.exceptions import ChunkedEncodingError, ConnectionError, ReadTimeout
+from urllib3.exceptions import ProtocolError, ReadTimeoutError
 from io import BufferedIOBase, BufferedWriter
 from typing import Optional
 from overrides import override
 from aistore.sdk.obj.content_iterator import ContentIterator
-from aistore.sdk.obj.obj_file.utils import handle_chunked_encoding_error
 from aistore.sdk.utils import get_logger
+from aistore.sdk.obj.obj_file.utils import (
+    handle_broken_stream,
+)
 
 logger = get_logger(__name__)
 
@@ -22,9 +25,9 @@ class ObjectFileReader(BufferedIOBase):
     data is insufficient to satisfy the request, the `read()` method fetches additional chunks from the provided
     `content_iterator` as needed, until the requested size is fulfilled or the end of the stream is reached.
 
-    In case of stream interruptions (e.g., `ChunkedEncodingError`), the `read()` method automatically retries and
-    resumes fetching data from the last successfully retrieved chunk. The `max_resume` parameter controls how many
-    retry attempts are made before an error is raised.
+    In case of unexpected stream interruptions (e.g. `ChunkedEncodingError`, `ConnectionError`) or timeouts (e.g.
+    `ReadTimeout`), the `read()` method automatically retries and resumes fetching data from the last successfully
+    retrieved chunk. The `max_resume` parameter controls how many retry attempts are made before an error is raised.
 
     Args:
         content_iterator (ContentIterator): An iterator that can fetch object data from AIS in chunks.
@@ -33,11 +36,14 @@ class ObjectFileReader(BufferedIOBase):
 
     def __init__(self, content_iterator: ContentIterator, max_resume: int):
         self._content_iterator = content_iterator
-        self._iterable = self._content_iterator.iter()
         self._max_resume = max_resume  # Maximum number of resume attempts allowed
-        self._remainder = None  # Remainder from the last chunk as a memoryview
-        self._resume_position = 0  # Tracks the current position in the stream
-        self._resume_total = 0  # Tracks the number of resume attempts
+        self._reset()
+
+    def _reset(self):
+        self._iterable = self._content_iterator.iter()
+        self._remainder = None
+        self._resume_position = 0
+        self._resume_total = 0
         self._closed = False
 
     @property
@@ -47,11 +53,7 @@ def content_iterator(self) -> ContentIterator:
 
     @override
     def __enter__(self):
-        self._iterable = self._content_iterator.iter()
-        self._remainder = None
-        self._resume_position = 0
-        self._resume_total = 0
-        self._closed = False
+        self._reset()
         return self
 
     @override
@@ -101,19 +103,26 @@ def read(self, size: Optional[int] = -1) -> bytes:
             while size:
                 try:
                     chunk = memoryview(next(self._iterable))
-                except StopIteration:
-                    # End of stream, exit loop
-                    break
-                except requests.exceptions.ChunkedEncodingError as err:
-                    # Handle ChunkedEncodingError and attempt to resume the stream
-                    self._iterable, self._resume_total = handle_chunked_encoding_error(
+                except (
+                    ConnectionError,
+                    ChunkedEncodingError,
+                    ProtocolError,
+                    ReadTimeout,
+                    ReadTimeoutError,
+                ) as err:
+                    # If the stream is broken (e.g. TCP reset, dropped connection, malformed or incomplete chunk)
+                    # or there is a timeout, retry with a new iterator from the last known position
+                    self._iterable, self._resume_total = handle_broken_stream(
                         self._content_iterator,
                         self._resume_position,
                         self._resume_total,
                         self._max_resume,
                         err,
                     )
                     continue
+                except StopIteration:
+                    # End of stream, exit loop
+                    break
 
                 # Track the position of the stream by adding the length of each fetched chunk
                 self._resume_position += len(chunk)
diff --git a/python/aistore/sdk/obj/obj_file/utils.py b/python/aistore/sdk/obj/obj_file/utils.py
@@ -1,44 +1,16 @@
 #
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 #
 
-from typing import Iterator, Optional, Tuple
+from typing import Iterator, Tuple
 
 from aistore.sdk.obj.content_iterator import ContentIterator
+from aistore.sdk.obj.obj_file.errors import ObjectFileReaderMaxResumeError
 from aistore.sdk.utils import get_logger
-from aistore.sdk.obj.obj_file.errors import (
-    ObjectFileReaderStreamError,
-    ObjectFileReaderMaxResumeError,
-)
 
 logger = get_logger(__name__)
 
 
-def reset_iterator(
-    content_iterator: ContentIterator,
-    resume_position: Optional[int] = 0,
-) -> Iterator[bytes]:
-    """
-    Return a new iterator for establishing an object stream and reading chunks of data from
-    byte position `resume_position`.
-
-    Args:
-        content_iterator (ContentIterator): An instance of `ContentIterator` to read data from.
-        resume_position (int, optional): The byte position to resume reading from. Defaults to 0.
-
-    Returns:
-        Iterator[bytes]: An iterator to read chunks of data from the object stream.
-
-    Raises:
-        ObjectFileReaderStreamError if a connection cannot be made.
-    """
-    try:
-        return content_iterator.iter(offset=resume_position)
-    except Exception as err:
-        logger.error("Error establishing object stream: (%s)", err)
-        raise ObjectFileReaderStreamError(err) from err
-
-
 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.
@@ -60,16 +32,16 @@ def increment_resume(resume_total: int, max_resume: int, err: Exception) -> int:
     return resume_total
 
 
-def handle_chunked_encoding_error(
+def handle_broken_stream(
     content_iterator: ContentIterator,
     resume_position: int,
     resume_total: int,
     max_resume: int,
     err: Exception,
 ) -> Tuple[Iterator[bytes], int]:
     """
-    Handle the chunked encoding error by incrementing the resume count, logging a warning,
-    and resetting the iterator from the last known position.
+    Handle the broken stream/iterator by incrementing the resume count, logging a warning,
+    and returning a newly instanatiated iterator from the last known position.
 
     Args:
         content_iterator (ContentIterator): The content iterator used to read the data.
@@ -93,5 +65,6 @@ def handle_chunked_encoding_error(
         resume_total,
         max_resume,
     )
-    chunk_iterator = reset_iterator(content_iterator, resume_position)
-    return chunk_iterator, resume_total
+    # Create a new iterator from the last read position
+    new_iter = content_iterator.iter(offset=resume_position)
+    return new_iter, resume_total
diff --git a/python/aistore/sdk/obj/object_reader.py b/python/aistore/sdk/obj/object_reader.py
@@ -98,7 +98,8 @@ def as_file(
         """
         Create a read-only, non-seekable `ObjectFileReader` instance for streaming object data in chunks.
         This file-like object primarily implements the `read()` method to retrieve data sequentially,
-        with automatic retry/resumption in case of stream interruptions such as `ChunkedEncodingError`.
+        with automatic retry/resumption in case of stream interruptions (e.g. `ChunkedEncodingError`,
+        `ConnectionError`) or timeouts (e.g. `ReadTimeout`, `ConnectTimeout`).
 
         Args:
             buffer_size (int, optional): Currently unused; retained for backward compatibility and future
diff --git a/python/examples/sdk/resilient-streaming-object-file.ipynb b/python/examples/sdk/resilient-streaming-object-file.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "# Resilient Object Streaming in AIStore\n",
     "\n",
-    "The following demo shows how to use `ObjectFileReader` (`aistore.sdk.obj.obj_file.object_file`) to stream large objects amidst potential instances of `ChunkedEncodingError` due to momentary issues with the cluster or its availability mid-read:"
+    "The following demo shows how to use `ObjectFileReader` (`aistore.sdk.obj.obj_file.object_file`) to stream large objects amidst unexpected cases of stream interruptions (e.g. `ChunkedEncodingError`, `ConnectionError`) or timeouts (e.g. `ReadTimeout`) mid-read:"
    ]
   },
   {

Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@`
`6`	`6`	`"source": [`
`7`	`7`	`"# Resilient Object Streaming in AIStore\n",`
`8`	`8`	`"\n",`
`9`		- "The following demo shows how to use `ObjectFileReader` (`aistore.sdk.obj.obj_file.object_file`) to stream large objects amidst potential instances of `ChunkedEncodingError` due to momentary issues with the cluster or its availability mid-read:"
	`9`	+ "The following demo shows how to use `ObjectFileReader` (`aistore.sdk.obj.obj_file.object_file`) to stream large objects amidst unexpected cases of stream interruptions (e.g. `ChunkedEncodingError`, `ConnectionError`) or timeouts (e.g. `ReadTimeout`) mid-read:"
`10`	`10`	`]`
`11`	`11`	`},`
`12`	`12`	`{`