test(dcp): add dcp optimized reader e2e test for coalescing behaviour

Add e2e integration test for DCPOptimizedS3Reader range coalescing
behaviour with full and partial loading patterns and different max_gap_sizes.

Author: jet-tong

s3torchconnector/tst/e2e/dcp/test_e2e_s3_storage_reader.py

@@ -10,7 +10,7 @@
 
 from s3torchconnector import S3ReaderConstructor
 from s3torchconnector.dcp import S3StorageWriter, S3StorageReader
-from s3torchconnector.s3reader import SequentialS3Reader, DCPOptimizedS3Reader
+from s3torchconnector._s3client import S3Client
 
 from ...conftest import READER_TYPE_STRING_TO_CLASS
 
@@ -93,3 +93,89 @@ def track_reads(self, size=None):
     assert loaded_state_dict.keys() == state_dict.keys()
     for key in state_dict:
         assert torch.equal(loaded_state_dict[key], state_dict[key])
+
+
+@pytest.mark.parametrize("model", [SIMPLE_MODEL, LARGER_MODEL])
+@pytest.mark.parametrize(
+    "max_gap_size,load_filter,filter_name,expected_streams",
+    [
+        # Full load - all tensors are consecutive, so always 1 stream
+        (0, lambda k: True, "Full", 1),
+        (float("inf"), lambda k: True, "Full", 1),
+        # Weights only - scattered by biases, so stream count depends on max_gap_size
+        (0, lambda k: k.endswith(".weight"), "Weights", 3),
+        (float("inf"), lambda k: k.endswith(".weight"), "Weights", 1),
+        # Layer 2 only - their bias+weight tensors are consecutive, so always 1 stream
+        (0, lambda k: "2." in k, "Layer 2", 1),
+        (float("inf"), lambda k: "2." in k, "Layer 2", 1),
+    ],
+)
+def test_dcp_optimized_loading_patterns(
+    checkpoint_directory,
+    model,
+    max_gap_size,
+    load_filter,
+    filter_name,
+    expected_streams,
+):
+    """Test DCPOptimized reader with full and partial loading patterns and different max_gap_size.
+
+    Validates that full loads use 1 stream, and partial load stream usage depends
+    on max_gap_size and whether tensors are consecutive / neighbours.  
+    
+    SIMPLE_MODEL tensors: ['0.bias', '0.weight', '1.bias', '1.weight', '2.bias', '2.weight']
+    LARGER_MODEL tensors: ['linear_relu_stack.0.bias', 'linear_relu_stack.0.weight', 'linear_relu_stack.2.bias',
+                           'linear_relu_stack.2.weight', 'linear_relu_stack.4.bias', 'linear_relu_stack.4.weight']
+    """
+    region = checkpoint_directory.region
+    s3_uri = checkpoint_directory.s3_uri
+
+    state_dict = model.state_dict()
+    dcp.save(state_dict, storage_writer=S3StorageWriter(region, s3_uri, overwrite=True))
+
+    # Print model structure (once per model)
+    all_keys = list(state_dict.keys())
+    if max_gap_size == 0 and filter_name == "Full":
+        print(f"\nTensors: {sorted(all_keys)}")
+
+    # Apply filter for partial load
+    filtered_keys = [k for k in all_keys if load_filter(k)]
+    excluded_keys = [k for k in all_keys if not load_filter(k)]
+    assert filtered_keys, f"No keys match {filter_name} filter for this model"
+    filtered_dict = {k: torch.empty_like(state_dict[k]) for k in filtered_keys}
+
+    # Load full / partial checkpoint with stream call tracker
+    stream_calls = []
+    original_get_object_stream = S3Client._get_object_stream
+    def track_get_object_stream(self, bucket, key, start=None, end=None):
+        if not key.endswith(".metadata"):
+            stream_calls.append((start, end))
+        return original_get_object_stream(self, bucket, key, start=start, end=end)
+
+    with patch.object(S3Client, "_get_object_stream", track_get_object_stream):
+        reader_constructor = S3ReaderConstructor.dcp_optimized(max_gap_size)
+        reader = S3StorageReader(region, s3_uri, reader_constructor=reader_constructor)
+        dcp.load(filtered_dict, storage_reader=reader)
+
+    # Verify correctness
+    assert len(filtered_dict) == len(filtered_keys)
+    for k, v in filtered_dict.items():
+        assert torch.equal(v, state_dict[k])
+        assert load_filter(k) 
+
+    # Verify excluded keys are not loaded
+    for k in excluded_keys:
+        assert k not in filtered_dict, f"Key {k} should not be in {filter_name} load"
+
+    # Verify expected stream count
+    assert len(stream_calls) == expected_streams
+    if len(stream_calls) > 1:
+        for i in range(1, len(stream_calls)):
+            assert stream_calls[i][0] >= stream_calls[i - 1][1]
+            assert stream_calls[i][0] - stream_calls[i - 1][1] >= max_gap_size
+
+    # Print number of stream calls
+    coalesce = "no coalesce" if max_gap_size == 0 else "full coalesce"
+    print(
+        f"{filter_name} load, {coalesce}: {len(stream_calls)} streams, {len(filtered_keys)} tensors"
+    )