refactor: revert READER_TYPE_STRING and sequential dcp e2e test

Reverts non-DCP optimized reader changes to make the PR changes clearer:
- Revert fix(tests): resolve e2e test import errors after adding __init__ files
- Revert test: place READER_TYPE_STRING_TO_CLASS in conftest
- Revert a minor test escape sequence fix.

Author: jet-tong

s3torchconnector/tst/__init__.py

@@ -29,7 +29,7 @@
 from typing import Optional
 
 from s3torchconnector.dcp.s3_prefix_strategy import RoundRobinPrefixStrategy
-from ..test_common import _list_folders_in_bucket
+from test_common import _list_folders_in_bucket
 
 
 def generate_random_port():

s3torchconnector/tst/e2e/__init__.py

@@ -10,9 +10,9 @@
 
 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
 
 SIMPLE_MODEL = torch.nn.Sequential(
     nn.Linear(5, 5),
@@ -40,8 +40,15 @@ def __init__(self):
 
 
 @pytest.mark.parametrize("model", [SIMPLE_MODEL, LARGER_MODEL])
+@pytest.mark.parametrize(
+    "reader_class,reader_constructor",
+    [
+        (SequentialS3Reader, S3ReaderConstructor.sequential()),
+        (DCPOptimizedS3Reader, S3ReaderConstructor.dcp_optimized()),
+    ],
+)
 def test_dcp_load_reads_tensors_in_sequential_order(
-    dcp_reader_constructor, checkpoint_directory, model
+    checkpoint_directory, model, reader_class, reader_constructor
 ):
     """
     Test that prepare_local_plan allows dcp.load() to read items in offset order.
@@ -63,11 +70,6 @@ def test_dcp_load_reads_tensors_in_sequential_order(
     storage_writer = S3StorageWriter(region=region, path=s3_uri, overwrite=True)
     dcp.save(state_dict, storage_writer=storage_writer)
 
-    reader_type_string = S3ReaderConstructor.get_reader_type_string(
-        dcp_reader_constructor
-    )
-    reader_class = READER_TYPE_STRING_TO_CLASS[reader_type_string]
-
     read_positions = []
     original_read = reader_class.read
 
@@ -76,10 +78,11 @@ def track_reads(self, size=None):
             read_positions.append(self._position)
         return original_read(self, size)
 
+    # Load with position tracking on read() (called at the start of each torch.load())
     with patch.object(reader_class, "read", track_reads):
         loaded_state_dict = {k: torch.empty_like(v) for k, v in state_dict.items()}
         storage_reader = S3StorageReader(
-            region=region, path=s3_uri, reader_constructor=dcp_reader_constructor
+            region=region, path=s3_uri, reader_constructor=reader_constructor
         )
         dcp.load(loaded_state_dict, storage_reader=storage_reader)
 

s3torchconnector/tst/e2e/dcp/test_e2e_s3_file_system.py

@@ -20,7 +20,7 @@
     from .conftest import BucketPrefixFixture, BucketPrefixData
 
 
-from .test_common import _get_fork_methods, _read_data, _set_start_method
+from test_common import _get_fork_methods, _read_data, _set_start_method
 
 
 start_methods = _get_fork_methods()

s3torchconnector/tst/e2e/dcp/test_e2e_s3_storage_reader.py

@@ -19,8 +19,8 @@
 from s3torchconnector.lightning import S3LightningCheckpoint
 from s3torchconnectorclient import S3Exception, __version__
 
-from .models.net import Net
-from .models.lightning_transformer import LightningTransformer, L
+from models.net import Net
+from models.lightning_transformer import LightningTransformer, L
 
 
 LIGHTNING_ACCELERATOR = "cpu"

s3torchconnector/tst/e2e/test_distributed_training.py

@@ -5,7 +5,7 @@
 import pytest
 
 from s3torchconnector import S3Checkpoint
-from .models.net import Net
+from models.net import Net
 
 
 @pytest.mark.parametrize(

s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py

@@ -5,8 +5,8 @@
 from s3torchconnector._s3client import S3Client
 from s3torchconnectorclient._mountpoint_s3_client import MountpointS3Client
 
-from .test_common import _get_fork_methods
-from .conftest import getenv
+from test_common import _get_fork_methods
+from conftest import getenv
 
 
 NATIVE_S3_CLIENT = None

s3torchconnector/tst/e2e/test_e2e_s3checkpoint.py

@@ -22,7 +22,7 @@
 if TYPE_CHECKING:
     from .conftest import BucketPrefixFixture
 
-from .test_common import _get_fork_methods, _read_data, _set_start_method
+from test_common import _get_fork_methods, _read_data, _set_start_method
 
 
 start_methods = _get_fork_methods()

s3torchconnector/tst/e2e/test_mountpoint_client_parallel_access.py

@@ -28,6 +28,10 @@
 TEST_REGION = "us-east-1"
 S3_PREFIX = f"s3://{TEST_BUCKET}"
 TEST_ENDPOINT = "https://s3.us-east-1.amazonaws.com"
+READER_TYPE_STRING_TO_CLASS = {
+    "sequential": SequentialS3Reader,
+    "range_based": RangedS3Reader,
+}
 
 
 @pytest.mark.parametrize(