perf(s3reader): optimize ListOfRangesS3Reader for sequential DCP workloads

jet-tong · jet-tong · commit 39853e486a20 · 2025-10-14T20:00:33.000+01:00
- Remove dependency on SequentialS3Reader for self-managed streams
- Implement direct stream management with per-group buffering
- Optimize read() method with no BytesIO buffer assuming sequential reading
- We now enforce non-seekable behaviour to force sequential reading patterns

This implementation is now significantly faster for distributed checkpoint
loading patterns while maintaining correctness for sequential access. This relies on
load ordering optimisation which enforces sequential reading with read() operations,
but will not work with readinto() operations since those still have backward seek
patterns.
diff --git a/s3torchconnector/src/s3torchconnector/s3reader/list_of_ranges.py b/s3torchconnector/src/s3torchconnector/s3reader/list_of_ranges.py
@@ -3,22 +3,23 @@
 
 import logging
 from dataclasses import dataclass
-from typing import List, Optional, Callable, Union, Dict
-from io import SEEK_SET
+from typing import List, Optional, Callable, Union, Dict, Iterator
+from io import SEEK_SET, SEEK_CUR
 
 from s3torchconnectorclient._mountpoint_s3_client import (
     ObjectInfo,
     GetObjectStream,
     HeadObjectResult,
 )
 from .s3reader import S3Reader
-from .sequential import SequentialS3Reader
 
 log = logging.getLogger(__name__)
 
 
 @dataclass
 class RangeRequest:
+    """Singular range request; Inclusive start, exclusive end"""
+
     start: int
     end: int
     request_id: Optional[str] = None
@@ -31,6 +32,8 @@ class RangeGroup:
     requests: List[RangeRequest]
 
 
+# TODO: Update name, since it now requires sequential reading and is optimised for DCP
+# TODO: Update docstring to emphasise this requires Load Ordering in prepare_local_plan
 class ListOfRangesS3Reader(S3Reader):
     """Optimized reader with pre-calculated request mapping and batch prefetch."""
 
@@ -42,30 +45,20 @@ def __init__(
         get_object_info: Callable[[], Union[ObjectInfo, HeadObjectResult]],
         get_stream: Callable[[Optional[int], Optional[int]], GetObjectStream],
         max_gap_size: int = 200 * 1024 * 1024,
-        **kwargs,
     ):
         self._bucket = bucket
         self._key = key
         self._get_object_info = get_object_info
         self._get_stream = get_stream
 
         # Calculate range groups using coalescing logic
-        self._range_groups = self._calculate_range_groups(ranges, max_gap_size)
-
-        # Pre-create all readers
-        self._group_readers: Dict[int, SequentialS3Reader] = {}
-        for i, group in enumerate(self._range_groups):
-            reader = SequentialS3Reader(
-                bucket=bucket,
-                key=key,
-                get_object_info=get_object_info,
-                get_stream=get_stream,
-                start_offset=group.start,
-                end_offset=group.end,
-            )
-            # TODO - judge if this is beneficial or not.
-            reader.prefetch()  # Batch prefetch all ranges
-            self._group_readers[i] = reader
+        self._range_groups = self._coalesce_ranges(ranges, max_gap_size)
+        self._current_group_idx: int = 0
+
+        # Per-group stream cache
+        self._streams: Dict[int, Iterator[bytes]] = {}
+        self._stream_positions: Dict[int, int] = {}
+        self._stream_buffers: Dict[int, bytes] = {}
 
         self._position: int = 0
 
@@ -77,76 +70,130 @@ def bucket(self) -> str:
     def key(self) -> str:
         return self._key
 
-    def _calculate_range_groups(
+    def seekable(self) -> bool:
+        """Not seekable — torch/distributed/checkpoint/filesystem.py will use read() instead of readinto()."""
+        return False
+
+    def _coalesce_ranges(
         self, ranges: List[RangeRequest], max_gap_size: int
     ) -> List[RangeGroup]:
-        """Coalescing logic - group ranges within max_gap_size."""
-        # TODO: optimise this logic
+        """Coalescing nearby byte ranges within max_gap_size."""
         if not ranges:
             return []
 
-        # TODO: could be pre-sorted in prepare_local_plan for dcp.load
-        sorted_ranges = sorted(ranges, key=lambda r: r.start)
-        groups = []
-        current_group = [sorted_ranges[0]]
-
-        for i in range(1, len(sorted_ranges)):
-            prev_end = current_group[-1].end
-            curr_start = sorted_ranges[i].start
+        # TODO: could be pre-sorted in prepare_local_plan (small optimisation)
+        ranges = sorted(ranges, key=lambda r: r.start)
+        groups: List[RangeGroup] = []
+        current = [ranges[0]]
 
-            if curr_start - prev_end <= max_gap_size:
-                current_group.append(sorted_ranges[i])
+        for r in ranges[1:]:
+            if r.start - current[-1].end <= max_gap_size:
+                current.append(r)
             else:
-                groups.append(self._create_range_group(current_group))
-                current_group = [sorted_ranges[i]]
+                groups.append(RangeGroup(current[0].start, current[-1].end, current))
+                current = [r]
 
-        groups.append(self._create_range_group(current_group))
+        groups.append(RangeGroup(current[0].start, current[-1].end, current))
         return groups
 
-    def _create_range_group(self, ranges: List[RangeRequest]) -> RangeGroup:
-        """Create range group - always succeeds since we only use gap size."""
-        # TODO remove min/max code by tracking incrementally in _calculate_range_groups
-        # * (was kept since it's easier to understand and test)
-        group_start = min(r.start for r in ranges)
-        group_end = max(r.end for r in ranges)
-        return RangeGroup(start=group_start, end=group_end, requests=ranges)
-
-    def _find_reader_for_offset(self, offset: int) -> Optional[SequentialS3Reader]:
-        """Find reader that contains the given offset."""
-        for i, group in enumerate(self._range_groups):
-            if group.start <= offset < group.end:
-                self._current_reader_index = i
-                return self._group_readers[i]
-            if group.start > offset:  # TODO handle this case properly by raising errors
-                break
-        return None
-
-    def seek(self, offset: int, whence: int = SEEK_SET, /) -> int:
-        self._position = offset
-        reader = self._find_reader_for_offset(offset)
-        if not reader:
-            return self._position
-        return reader.seek(offset, whence)
+    def _get_stream_for_group(self, idx: int) -> Iterator[bytes]:
+        """
+        Returns a cached iterator for the given range group,
+        or creates a new one if not present.
+        """
+        if idx not in self._streams:
+            group = self._range_groups[idx]
+            stream = self._get_stream(group.start, group.end)
+            self._streams[idx] = stream
+            self._stream_positions[idx] = group.start
+            self._stream_buffers[idx] = b""
+        return self._streams[idx]
 
     def read(self, size: Optional[int] = None) -> bytes:
-        reader = self._find_reader_for_offset(self._position)
-        if not reader:
+        """Reads up to `size` bytes sequentially across grouped ranges."""
+        if not size or size <= 0:
             return b""
-        data = reader.read(size)
-        self._position += len(data)
-        return data
+
+        pos = self._position
+
+        # Find group (with cache)
+        if (
+            self._current_group_idx < len(self._range_groups)
+            and self._range_groups[self._current_group_idx].start
+            <= pos
+            < self._range_groups[self._current_group_idx].end
+        ):
+            group_idx = self._current_group_idx
+        else:
+            group_idx = next(
+                (i for i, g in enumerate(self._range_groups) if g.start <= pos < g.end),
+                None,
+            )
+            if group_idx is None:
+                return b""
+            self._current_group_idx = group_idx
+
+        stream = self._get_stream_for_group(group_idx)
+
+        current_pos = self._stream_positions[group_idx]
+        buffer = self._stream_buffers[group_idx]
+        remaining = size
+        chunks: List[bytes] = []
+
+        # 1. Serve from buffered leftover bytes
+        if buffer and current_pos <= pos < current_pos + len(buffer):
+            offset = pos - current_pos
+            end = offset + min(remaining, len(buffer) - offset)
+            chunks.append(buffer[offset:end])
+            remaining -= end - offset
+            current_pos = pos + (end - offset)
+            self._stream_buffers[group_idx] = buffer[end:] if end < len(buffer) else b""
+
+        # 2. Read more data from S3 stream
+        while remaining > 0:
+            try:
+                chunk = next(stream)
+            except StopIteration:
+                break
+
+            # Skip ahead if behind target
+            if current_pos < pos:
+                skip = min(pos - current_pos, len(chunk))
+                chunk = chunk[skip:]
+                current_pos += skip
+
+            # Take needed part of chunk
+            take = min(len(chunk), remaining)
+            chunks.append(chunk[:take])
+            remaining -= take
+            current_pos += take
+
+            # Save leftover bytes
+            if take < len(chunk):
+                self._stream_buffers[group_idx] = chunk[take:]
+                break
+
+        self._stream_positions[group_idx] = current_pos
+        self._position = pos + (size - remaining)
+        return b"".join(chunks)
+
+    def seek(self, offset: int, whence: int = SEEK_SET, /) -> int:
+        if whence == SEEK_SET:
+            self._position = offset
+        elif whence == SEEK_CUR:
+            self._position += offset
+        return self._position
 
     def readinto(self, buf) -> int:
-        reader = self._find_reader_for_offset(self._position)
-        if not reader:
-            return 0
-        bytes_read = reader.readinto(buf)
-        self._position += bytes_read
-        return bytes_read
+        data = self.read(len(buf))
+        n = len(data)
+        buf[:n] = data
+        return n
 
     def tell(self) -> int:
         return self._position
 
     def close(self) -> None:
-        for reader in self._group_readers.values():
-            reader.close()
+        self._streams.clear()
+        self._stream_positions.clear()
+        self._stream_buffers.clear()
diff --git a/s3torchconnector/src/s3torchconnector/s3reader/sequential.py b/s3torchconnector/src/s3torchconnector/s3reader/sequential.py
@@ -1,8 +1,6 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
 
-import os
-import logging
 import io
 from functools import cached_property
 from io import SEEK_CUR, SEEK_END, SEEK_SET
@@ -15,35 +13,6 @@
 )
 from .s3reader import S3Reader
 
-log = logging.getLogger(__name__)
-
-
-class _OffsetReaderView:
-    """Wrapper that translates absolute positions to buffer-relative positions"""
-
-    def __init__(self, buffer: io.BytesIO, start_offset: int = 0):
-        self._buffer = buffer
-        self._start_offset = start_offset
-
-    def seek(self, offset: int, whence: int = SEEK_SET) -> int:
-        if whence == SEEK_SET:
-            buffer_pos = offset - self._start_offset
-            return self._buffer.seek(buffer_pos) + self._start_offset
-        else:
-            return self._buffer.seek(offset, whence) + self._start_offset
-
-    def tell(self) -> int:
-        return self._buffer.tell() + self._start_offset
-
-    def read(self, size=-1):
-        return self._buffer.read(size)
-
-    def readinto(self, buf):
-        return self._buffer.readinto(buf)
-
-    def write(self, data):
-        return self._buffer.write(data)
-
 
 class SequentialS3Reader(S3Reader):
     """Sequential S3 reader implementation
@@ -57,32 +26,19 @@ def __init__(
         bucket: str,
         key: str,
         get_object_info: Callable[[], Union[ObjectInfo, HeadObjectResult]],
-        get_stream: Callable[[Optional[int], Optional[int]], GetObjectStream],
-        start_offset: Optional[int] = None,
-        end_offset: Optional[int] = None,
+        get_stream: Callable[[], GetObjectStream],
     ):
         if not bucket:
             raise ValueError("Bucket should be specified")
         self._bucket = bucket
         self._key = key
-        self._filename = os.path.basename(self._key)
         self._get_object_info = get_object_info
         self._get_stream = get_stream
         self._stream: Optional[Iterator[bytes]] = None
-        self._raw_buffer = io.BytesIO()
-        self._buffer = _OffsetReaderView(self._raw_buffer, start_offset or 0)
+        self._buffer = io.BytesIO()
         self._size: Optional[int] = None
         # Invariant: _position == _buffer._tell() unless _position_at_end()
-        self._position = start_offset or 0
-        self._pid: int = os.getpid()
-
-        self._start_offset = start_offset
-        self._end_offset = end_offset
-
-        # Log start of reading for tracking
-        log.debug(
-            f"file={self._filename}, pid={self._pid}, type=read_initialized, start_offset={start_offset}, end_offset={end_offset}"
-        )
+        self._position = 0
 
     @property
     def bucket(self) -> str:
@@ -104,10 +60,7 @@ def prefetch(self) -> None:
         """
 
         if self._stream is None:
-            if self._start_offset is not None or self._end_offset is not None:
-                self._stream = self._get_stream(self._start_offset, self._end_offset)
-            else:
-                self._stream = self._get_stream(None, None)
+            self._stream = self._get_stream()
 
     def readinto(self, buf) -> int:
         """Read up to len(buf) bytes into a pre-allocated, writable bytes-like object buf.
@@ -120,9 +73,6 @@ def readinto(self, buf) -> int:
             int : numer of bytes read or zero, if no bytes available
         """
         buf_size = len(buf)
-        log.debug(
-            f"file={self._filename}, pid={self._pid}, type=readinto, position={self._position}, size={buf_size}"
-        )
         if self._position_at_end() or buf_size == 0:
             # If no bytes are available or no place to write data, zero should be returned
             return 0
@@ -158,11 +108,6 @@ def read(self, size: Optional[int] = None) -> bytes:
 
         if size is not None and not isinstance(size, int):
             raise TypeError(f"argument should be integer or None, not {type(size)!r}")
-
-        log.debug(
-            f"file={self._filename}, pid={self._pid}, type=read, position={self._position}, size={size}"
-        )
-
         if self._position_at_end():
             # Invariant: if we're at EOF, it doesn't matter what `size` is, we'll always return no data and have no
             # side effect.
@@ -255,15 +200,14 @@ def _position_at_end(self) -> bool:
             # We can never be special cased to EOF if we never saw how long it is.
             # If we _are_ at EOF, we'll just not take the early exits.
             return False
-        end_pos = min(self._end_offset, self._size) if self._end_offset else self._size
-        return self._position >= end_pos
+        return self._position == self._size
 
     def _buffer_size(self) -> int:
-        cur_pos = self._raw_buffer.tell()
-        self._raw_buffer.seek(0, SEEK_END)
-        buffer_size = self._raw_buffer.tell()
-        self._raw_buffer.seek(cur_pos)
-        return buffer_size + (self._start_offset or 0)
+        cur_pos = self._buffer.tell()
+        self._buffer.seek(0, SEEK_END)
+        buffer_size = self._buffer.tell()
+        self._buffer.seek(cur_pos)
+        return buffer_size
 
     def tell(self) -> int:
         """
