feat(dcp): use constructor pattern for ListOfRanges optimization

Add DCPListOfRangesConstructor and dcp_list_of_ranges() factory method to enable
DCP range optimization through reader_constructor parameter. Includes better range
injection logic and support for both direct ListOfRanges usage and DCP optimization.

Users can now opt-in via: reader_constructor=S3ReaderConstructor.dcp_list_of_ranges()

Author: jet-tong

s3torchconnector/src/s3torchconnector/dcp/s3_file_system.py

@@ -28,7 +28,11 @@
 
 from s3torchconnector._s3client import S3Client
 from s3torchconnector._s3dataset_common import parse_s3_uri
-from ..s3reader import S3ReaderConstructor, S3ReaderConstructorProtocol
+from ..s3reader import (
+    S3ReaderConstructor,
+    DCPListOfRangesConstructor,
+    S3ReaderConstructorProtocol,
+)
 from .. import S3ClientConfig
 from .s3_prefix_strategy import S3PrefixStrategyBase, DefaultPrefixStrategy
 from .._user_agent import UserAgent
@@ -79,7 +83,6 @@ def create_stream(
         self,
         path: Union[str, os.PathLike],
         mode: str,
-        reader_constructor: Optional[S3ReaderConstructorProtocol] = None,
     ) -> Generator[io.IOBase, None, None]:
         """
         Create a stream for reading or writing to S3.
@@ -102,18 +105,8 @@ def create_stream(
             with self._client.put_object(bucket, key) as stream:
                 yield stream
         elif mode == "rb":  # read mode
-            logger.debug("create_stream readable for %s", path_str)
-            relative_path = os.path.relpath(path, self._path)
-            if self.file_ranges and relative_path in self.file_ranges:
-                ranges = self.file_ranges[relative_path]
-                # ! Force use list_of_ranges reader for now
-                # TODO: (Important) improve this by passing in ranges parameter properly
-                reader_constructor = S3ReaderConstructor.list_of_ranges(ranges)
-
-            # Use provided reader_constructor or fall back to default
-            constructor = reader_constructor or self._reader_constructor
             with self._client.get_object(
-                bucket, key, reader_constructor=constructor
+                bucket, key, reader_constructor=self._reader_constructor
             ) as stream:
                 yield stream
         else:
@@ -380,20 +373,21 @@ def prepare_local_plan(self, plan: LoadPlan) -> LoadPlan:
             LoadPlan: The same plan with items sorted by storage offset.
         """
 
-        # Calculate ranges per file
-        per_file_ranges = {}
-        for read_item in plan.items:
-            item_md = self.storage_data[read_item.storage_index]
-            path = item_md.relative_path
-            if path not in per_file_ranges:
-                per_file_ranges[path] = []
-            per_file_ranges[path].append(
-                RangeRequest(start=item_md.offset, end=item_md.offset + item_md.length)
-            )
-
-        # Store ranges in filesystem
-        # TODO find a better place to handle this information
-        self.fs.file_ranges = per_file_ranges
+        # Inject ranges if using DCP list-of-ranges reader constructor
+        if isinstance(self.fs._reader_constructor, DCPListOfRangesConstructor):
+            # Calculate ranges per file
+            per_file_ranges = {}
+            for read_item in plan.items:
+                item_md = self.storage_data[read_item.storage_index]
+                path = item_md.relative_path
+                if path not in per_file_ranges:
+                    per_file_ranges[path] = []
+                per_file_ranges[path].append(
+                    RangeRequest(
+                        start=item_md.offset, end=item_md.offset + item_md.length
+                    )
+                )
+            self.fs._reader_constructor.set_file_ranges(per_file_ranges)
 
         # Sort items in plan based on their offset in checkpoints shards
         plan.items.sort(key=lambda item: self.storage_data[item.storage_index].offset)

s3torchconnector/src/s3torchconnector/s3reader/__init__.py

@@ -2,14 +2,16 @@
 #  // SPDX-License-Identifier: BSD
 
 from .s3reader import S3Reader
-from .constructor import S3ReaderConstructor
+from .constructor import S3ReaderConstructor, DCPListOfRangesConstructor
 from .sequential import SequentialS3Reader
 from .ranged import RangedS3Reader
+from .list_of_ranges import ListOfRangesS3Reader
 from .protocol import GetStreamCallable, S3ReaderConstructorProtocol
 
 __all__ = [
     "S3Reader",
     "S3ReaderConstructor",
     "SequentialS3Reader",
     "RangedS3Reader",
+    "ListOfRangesS3Reader",
 ]

s3torchconnector/src/s3torchconnector/s3reader/constructor.py

@@ -2,14 +2,35 @@
 #  // SPDX-License-Identifier: BSD
 
 from functools import partial
-from typing import Optional, List
+from typing import Optional, List, Dict
 
+from .s3reader import S3Reader
 from .protocol import S3ReaderConstructorProtocol
 from .sequential import SequentialS3Reader
 from .ranged import RangedS3Reader
 from .list_of_ranges import ListOfRangesS3Reader, RangeRequest
 
 
+class DCPListOfRangesConstructor:
+    def __init__(self) -> None:
+        self._file_ranges: Dict[str, List[RangeRequest]] = {}
+
+    def set_file_ranges(self, file_ranges: Dict[str, List[RangeRequest]]) -> None:
+        self._file_ranges = file_ranges
+
+    def __call__(
+        self, bucket: str, key: str, get_object_info, get_stream, **kwargs
+    ) -> S3Reader:
+        if self._file_ranges:
+            # TODO: Check if using filename only works with prefix strategies
+            filename = key.split("/")[-1]
+            return ListOfRangesS3Reader(
+                bucket, key, self._file_ranges[filename], get_object_info, get_stream
+            )
+        # Fallback to SequentialS3Reader if no file_ranges yet (e.g. when reading .metadata)
+        return SequentialS3Reader(bucket, key, get_object_info, get_stream, **kwargs)
+
+
 class S3ReaderConstructor:
     """Constructor for creating ``partial(S3Reader)`` instances.
 
@@ -93,6 +114,12 @@ def list_of_ranges(ranges: List[RangeRequest]) -> S3ReaderConstructorProtocol:
         # TODO update docstring
         return partial(ListOfRangesS3Reader, ranges=ranges)
 
+    @staticmethod
+    def dcp_list_of_ranges() -> S3ReaderConstructorProtocol:
+        """Creates a DCP-optimized constructor that uses ListOfRanges when ranges are available"""
+        # TODO update docstring
+        return DCPListOfRangesConstructor()
+
     @staticmethod
     def default() -> S3ReaderConstructorProtocol:
         """Creates default reader constructor (sequential)
@@ -112,10 +139,11 @@ def get_reader_type_string(
                 S3ReaderConstructor.default()
             )
 
-        if not isinstance(constructor, partial):
+        if isinstance(constructor, DCPListOfRangesConstructor):
+            return "dcp_list_of_ranges"
+        elif not isinstance(constructor, partial):
             return "unknown"
-
-        if constructor.func == RangedS3Reader:
+        elif constructor.func == RangedS3Reader:
             return "range_based"
         elif constructor.func == SequentialS3Reader:
             return "sequential"