feat: implement range-based S3 Reader (#339)

- Add RangedS3Reader for efficient byte range requests with adaptive buffering
- Update Rust components to support range parameters in mountpoint client
- Introduce S3ReaderConstructor to configure reader types and parameters,
  maintaining sequential reader as default for backwards compatibility
- Expose reader_constructor in S3Client, datasets, and DCP interfaces
- Update user agent to include dataset and reader types
- Extend test coverage with parametrized tests for both reader implementations
- Add configurable S3Reader to s3torchbenchmarking module
- Update documentation with usage examples and performance considerations

Author: jet-tong

.gitignore

@@ -79,3 +79,6 @@ s3torchconnectorclient/**/*.rs.bk
 
 # JetBrains
 .idea/
+
+# Sphinx documentation
+s3torchconnector/docs/

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## TBD
 
 ### New features
+* Implement range-based S3 Reader for byte range requests, enabling efficient random read patterns (#339)
 
 ### Bug fixes
 * Enable multiple CRT clients per process with different configs (#340)

README.md

@@ -115,13 +115,16 @@ For example, assuming the following directory bucket name `my-test-bucket--usw2-
 usw2-az1, then the URI used will look like: `s3://my-test-bucket--usw2-az1--x-s3/<PREFIX>` (**please note that the 
 prefix for Amazon S3 Express One Zone should end with '/'**), paired with region us-west-2.
 
+
 ## Distributed checkpoints
 
 ### Overview
 
 Amazon S3 Connector for PyTorch provides robust support for PyTorch distributed checkpoints. This feature includes:
 
-- `S3StorageWriter` and `S3StorageReader`: Implementations of PyTorch's StorageWriter and StorageReader interfaces.
+- `S3StorageWriter`: Implementation of PyTorch's StorageWriter interface.
+
+- `S3StorageReader`: Implementation of PyTorch's StorageReader interface. Supports configurable reading strategies via the `reader_constructor` parameter (see [Reader Configurations](#reader-configurations)).
 - `S3FileSystem`: An implementation of PyTorch's FileSystemBase.
 
 These tools enable seamless integration of Amazon S3 with 
@@ -187,7 +190,7 @@ the load across multiple S3 partitions.
 #### 1. RoundRobinPrefixStrategy
 Distributes checkpoints across specified prefixes in a round-robin fashion, ideal for balancing data across multiple storage locations.
 
-```python
+```py
 from s3torchconnector.dcp import RoundRobinPrefixStrategy, S3StorageWriter
 
 model = torchvision.models.resnet18()
@@ -234,7 +237,7 @@ CHECKPOINT_URI
 
 Generates binary (base-2) prefixes for optimal partitioning in distributed environments.
 
-```python
+```py
 from s3torchconnector.dcp import BinaryPrefixStrategy
 
 strategy = BinaryPrefixStrategy(
@@ -261,7 +264,7 @@ s3://my-bucket/checkpoints/
 #### 3. HexPrefixStrategy
 
 Uses hexadecimal (base-16) prefixes for a balance of efficiency and readability.
-```
+```py
 from s3torchconnector.dcp import HexPrefixStrategy
 
 strategy = HexPrefixStrategy(
@@ -288,7 +291,7 @@ s3://my-bucket/checkpoints/
 ### Creating Custom Strategies
 
 You can implement custom prefix strategies by extending the S3PrefixStrategyBase class:
-```
+```py
 from s3torchconnector.dcp import S3PrefixStrategyBase
 
 class CustomPrefixStrategy(S3PrefixStrategyBase):
@@ -312,7 +315,7 @@ The S3IterableDataset can be directly passed to PyTorch's DataLoader for paralle
 By default, all worker processes will share the same list of training objects. However, 
 if you need each worker to have access to a unique portion of the dataset for better parallelization, 
 you can enable dataset sharding using the `enable_sharding` parameter. 
-```
+```py
 dataset = S3IterableDataset.from_prefix(DATASET_URI, region=REGION, enable_sharding=True)
 dataloader = DataLoader(dataset, num_workers=4)
 ```
@@ -324,7 +327,7 @@ Each worker, regardless of its host, will load and process a distinct subset of
 For the S3MapDataset, you need to pass it to DataLoader along with a [DistributedSampler](https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler) wrapped around it. 
 The DistributedSampler ensures that each worker or node receives a unique subset of the dataset, 
 enabling efficient parallel and distributed training.
-```
+```py
 dataset = S3MapDataset.from_prefix(DATASET_URI, region=REGION)
 sampler = DistributedSampler(dataset)
 dataloader = DataLoader(dataset, sampler=sampler, num_workers=4)
@@ -371,6 +374,118 @@ To enable versioning on an S3 bucket, see [Enabling versioning on buckets](https
 
 S3 Versioning and S3 Lifecycle are not supported by S3 Express One Zone.
 
+
+## Direct S3Client Usage
+
+For advanced use cases, you can use the S3Client directly for custom streaming patterns and integration with existing pipelines.
+
+```py
+from s3torchconnector._s3client import S3Client
+
+REGION = "us-east-1"
+BUCKET_NAME = "my-bucket"
+OBJECT_KEY = "large_object.bin"
+
+s3_client = S3Client(region=REGION)
+
+# Writing data to S3
+data = b"content" * 1048576
+s3writer = s3_client.put_object(bucket=BUCKET_NAME, key=OBJECT_KEY)
+s3writer.write(data)
+s3writer.close()
+
+# Reading data from S3
+s3reader = s3_client.get_object(bucket=BUCKET_NAME, key=OBJECT_KEY)
+data = s3reader.read()
+```
+
+## Reader Configurations
+
+Amazon S3 Connector for PyTorch supports two types of readers, configurable through `S3ReaderConstructor`.
+
+### Reader Types
+
+#### 1. Sequential Reader (Default)
+
+- Downloads and buffers the entire S3 object in memory.
+- Prioritizes performance over memory usage by buffering entire objects.
+
+#### 2. Range-based Reader
+
+- Performs byte-range requests to read specific portions of S3 objects without downloading the entire file.
+- Prioritizes memory efficiency, with performance gains only for sparse partial reads.
+- Features adaptive buffering with forward overlap handling:
+  - **Small reads** (< `buffer_size`): Use internal buffer to reduce S3 API calls.
+  - **Large reads** (≥ `buffer_size`): Bypass buffer for direct transfer.
+
+### When to Use Each Reader
+
+- **Sequential Reader**: For processing entire files, and when repeated access to the data is required. Best for most general use cases.
+- **Range-based Reader**: For larger objects (100MB+) that require sparse partial reads, and in memory-constrained environments. 
+
+**Note**: S3Reader instances are not thread-safe and should not be shared across threads. For multiprocessing with DataLoader, each worker process creates its own S3Reader instance automatically.
+
+### Examples
+
+Direct method - `S3Client` usage with range-based reader without buffer:
+```py
+# Direct S3Client usage for zero-copy partial reads into pre-allocated buffers, for memory efficiency and fast data transfer
+from s3torchconnector._s3client import S3Client
+from s3torchconnector import S3ReaderConstructor
+
+s3_client = S3Client(region=REGION)
+reader_constructor = S3ReaderConstructor.range_based(
+    buffer_size=0  # No buffer, for direct transfer
+)
+s3reader = s3_client.get_object(
+    bucket=BUCKET_NAME, 
+    key=OBJECT_NAME, 
+    reader_constructor=reader_constructor
+)
+
+buffer = bytearray(10 * 1024 * 1024)  # 10MB buffer
+s3reader.seek(100 * 1024 * 1024)   # Skip to 100MB offset
+bytes_read = s3reader.readinto(buffer)  # Direct read into buffer
+```
+
+DCP interface - `S3StorageReader` usage with range-based reader with buffer:
+```py
+# Load distributed checkpoint with range-based reader to optimize memory usage for large checkpoint files
+from s3torchconnector.dcp import S3StorageReader
+from s3torchconnector import S3ReaderConstructor
+
+reader_constructor = S3ReaderConstructor.range_based(
+    buffer_size=16*1024*1024  # 16MB buffer
+)
+s3_storage_reader = S3StorageReader(
+    region=REGION, 
+    path=CHECKPOINT_URI,
+    reader_constructor=reader_constructor
+)
+DCP.load(
+    state_dict=model_state_dict,
+    storage_reader=s3_storage_reader,
+)
+```
+
+Dataset interface - `S3MapDataset` usage with sequential reader:
+```py
+# Use sequential reader for optimal performance when reading entire objects
+from s3torchconnector import S3MapDataset, S3ReaderConstructor
+
+dataset = S3MapDataset.from_prefix(
+    DATASET_URI, 
+    region=REGION,
+    reader_constructor=S3ReaderConstructor.sequential()
+)
+
+for item in dataset:
+    content = item.read()
+    ...
+```
+
+For `S3ReaderConstructor` usage details, please refer to the [`S3ReaderConstructor` documentation](https://awslabs.github.io/s3-connector-for-pytorch/autoapi/s3torchconnector/s3reader/constructor/index.html).
+
 ## Contributing
 
 We welcome contributions to Amazon S3 Connector for PyTorch. Please see [CONTRIBUTING](CONTRIBUTING.md) for more

examples/dcp/stateful_example.py

@@ -17,7 +17,7 @@
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
-from s3torchconnector import S3ClientConfig
+from s3torchconnector import S3ClientConfig, S3ReaderConstructor
 from s3torchconnector.dcp import S3StorageWriter, S3StorageReader
 from s3torchconnector.dcp.s3_prefix_strategy import RoundRobinPrefixStrategy
 
@@ -170,9 +170,15 @@ def run(rank, world_size, region, s3_uri, device="cuda"):
         device, model, rank, world_size
     )
     print(f"Load previously saved checkpoint on rank:{rank}")
+    # Configure S3 reader constructor (sequential or range_based)
+    reader_constructor = S3ReaderConstructor.sequential()
+    # reader_constructor = S3ReaderConstructor.range_based(buffer_size=16 * 1024 * 1024)
     # initialize S3StorageReader with region and bucket name, before passing to dcp.load as reader
     storage_reader = S3StorageReader(
-        region=region, path=s3_uri, s3client_config=s3config
+        region=region,
+        path=s3_uri,
+        s3client_config=s3config,
+        reader_constructor=reader_constructor,
     )
     dcp.load(
         state_dict={"model": modified_model, "optimizer": modified_optim},

s3torchbenchmarking/conf/dataloader/s3iterabledataset.yaml

@@ -1,3 +1,8 @@
 kind: s3iterabledataset
 batch_size: 128
-num_workers: 8
+num_workers: 8
+s3reader:
+  # s3reader type: sequential or range_based
+  type: sequential # default
+  # buffer_size (bytes): only used with range_based s3reader type
+  buffer_size: 8*1024*1024 # default

s3torchbenchmarking/conf/dataloader/s3mapdataset.yaml

@@ -1,3 +1,8 @@
 kind: s3mapdataset
 batch_size: 128
-num_workers: 8
+num_workers: 8
+s3reader:
+  # s3reader type: sequential or range_based
+  type: sequential # default
+  # buffer_size (bytes): only used with range_based s3reader type
+  buffer_size: 8*1024*1024 # default

s3torchbenchmarking/conf/dataset.yaml

@@ -31,6 +31,12 @@ hydra:
       # Name of a model (valid options: "entitlement", "vit").
       +model: entitlement
       # Kind of the dataloader (valid options: "fsspec", "s3iterabledataset", "mountpoint", "mountpointcache").
+      # For dataloader kind specific options, see specific conf/dataloader/{dataloader-kind}.yaml
       +dataloader: fsspec, s3iterabledataset, mountpoint, mountpointcache
       # Dataset name (corresponds to the name of a folder in S3); will be used to build an S3 URI
       +dataset: 100k_496x387_images
+      # S3 reader sweeps (only applies to s3iterabledataset/s3mapdataset)
+      # s3reader type: sequential or range_based
+      dataloader.s3reader.type: sequential
+      # buffer_size (bytes): only used with range_based s3reader
+      dataloader.s3reader.buffer_size: 8*1024*1024