Ensure HeadObject is not explicitly called when seeking to the end of a file when dataset is created with from_prefix (#145)

Co-authored-by: Simon Beal <[email protected]>

Author: muddyfish

s3torchconnector/src/s3torchconnector/_s3_bucket_iterable.py

@@ -11,8 +11,8 @@
     ListObjectStream,
 )
 
+from ._s3bucket_key_data import S3BucketKeyData
 from ._s3client import S3Client
-from . import S3Reader
 
 
 class S3BucketIterable:
@@ -21,7 +21,7 @@ def __init__(self, client: S3Client, bucket: str, prefix: str):
         self._bucket = bucket
         self._prefix = prefix
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[S3BucketKeyData]:
         # This allows us to iterate multiple times by re-creating the `_list_stream`
         return iter(S3BucketIterator(self._client, self._bucket, self._prefix))
 
@@ -32,10 +32,9 @@ def __init__(self, client: S3Client, bucket: str, prefix: str):
         self._bucket = bucket
         self._list_stream = _PickleableListObjectStream(client, bucket, prefix)
 
-    def __iter__(self) -> Iterator[S3Reader]:
-        return map(
-            partial(self._client.from_bucket_and_object_info, self._bucket),
-            chain.from_iterable(map(_extract_object_info, self._list_stream)),
+    def __iter__(self) -> Iterator[S3BucketKeyData]:
+        return chain.from_iterable(
+            map(partial(_extract_list_results, self._bucket), self._list_stream)
         )
 
 
@@ -66,5 +65,11 @@ def __setstate__(self, state):
         self._list_stream = ListObjectStream._from_state(**state)
 
 
-def _extract_object_info(list_result: ListObjectResult) -> List[ObjectInfo]:
-    return list_result.object_info
+def _extract_list_results(
+    bucket: str, list_result: ListObjectResult
+) -> Iterator[S3BucketKeyData]:
+    return map(partial(_extract_object_info, bucket), list_result.object_info)
+
+
+def _extract_object_info(bucket: str, object_info: ObjectInfo) -> S3BucketKeyData:
+    return S3BucketKeyData(bucket=bucket, key=object_info.key, object_info=object_info)

…or/src/s3torchconnector/_s3bucket_key.py …c/s3torchconnector/_s3bucket_key_data.pys3torchconnector/src/s3torchconnector/_s3bucket_key.py renamed to s3torchconnector/src/s3torchconnector/_s3bucket_key_data.py s3torchconnector/src/s3torchconnector/_s3bucket_key.py renamed to s3torchconnector/src/s3torchconnector/_s3bucket_key_data.py

@@ -1,10 +1,13 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
-from typing import NamedTuple
+from typing import NamedTuple, Optional
 
+from s3torchconnectorclient._mountpoint_s3_client import ObjectInfo
 
-class S3BucketKey(NamedTuple):
+
+class S3BucketKeyData(NamedTuple):
     """Read-only information about object stored in S3."""
 
     bucket: str
     key: str
+    object_info: Optional[ObjectInfo] = None

s3torchconnector/src/s3torchconnector/_s3client/_s3client.py

@@ -57,12 +57,19 @@ def _client_builder(self) -> MountpointS3Client:
             user_agent_prefix=user_agent_prefix,
         )
 
-    def get_object(self, bucket: str, key: str) -> S3Reader:
-        log.debug(f"GetObject s3://{bucket}/{key}")
+    def get_object(
+        self, bucket: str, key: str, *, object_info: Optional[ObjectInfo] = None
+    ) -> S3Reader:
+        log.debug(f"GetObject s3://{bucket}/{key}, {object_info is None=}")
+        if object_info is None:
+            get_object_info = partial(self.head_object, bucket, key)
+        else:
+            get_object_info = partial(_identity, object_info)
+
         return S3Reader(
             bucket,
             key,
-            get_object_info=partial(self.head_object, bucket, key),
+            get_object_info=get_object_info,
             get_stream=partial(self._get_object_stream, bucket, key),
         )
 
@@ -86,14 +93,3 @@ def list_objects(
     def head_object(self, bucket: str, key: str) -> ObjectInfo:
         log.debug(f"HeadObject s3://{bucket}/{key}")
         return self._client.head_object(bucket, key)
-
-    def from_bucket_and_object_info(
-        self, bucket: str, object_info: ObjectInfo
-    ) -> S3Reader:
-        log.debug(f"GetObjectWithInfo s3://{bucket}/{object_info.key}")
-        return S3Reader(
-            bucket,
-            object_info.key,
-            get_object_info=partial(_identity, object_info),
-            get_stream=partial(self._get_object_stream, bucket, object_info.key),
-        )

s3torchconnector/src/s3torchconnector/_s3dataset_common.py

@@ -6,7 +6,7 @@
 from ._s3_bucket_iterable import S3BucketIterable
 from ._s3client import S3Client
 from . import S3Reader
-from ._s3bucket_key import S3BucketKey
+from ._s3bucket_key_data import S3BucketKeyData
 
 """
 _s3dataset_common.py
@@ -38,18 +38,15 @@ def parse_s3_uri(uri: str) -> Tuple[str, str]:
 
 def get_objects_from_uris(
     object_uris: Union[str, Iterable[str]], client: S3Client
-) -> Iterable[S3BucketKey]:
+) -> Iterable[S3BucketKeyData]:
     if isinstance(object_uris, str):
         object_uris = [object_uris]
     # TODO: We should be consistent with URIs parsing. Revise if we want to do this upfront or lazily.
     bucket_key_pairs = [parse_s3_uri(uri) for uri in object_uris]
 
-    return (S3BucketKey(bucket, key) for bucket, key in bucket_key_pairs)
+    return (S3BucketKeyData(bucket, key) for bucket, key in bucket_key_pairs)
 
 
-def get_objects_from_prefix(s3_uri: str, client: S3Client) -> Iterable[S3BucketKey]:
+def get_objects_from_prefix(s3_uri: str, client: S3Client) -> Iterable[S3BucketKeyData]:
     bucket, prefix = parse_s3_uri(s3_uri)
-    s3objects = S3BucketIterable(client, bucket, prefix)
-    bucket_key_pairs = (S3BucketKey(obj.bucket, obj.key) for obj in s3objects)
-
-    return bucket_key_pairs
+    return iter(S3BucketIterable(client, bucket, prefix))

s3torchconnector/src/s3torchconnector/s3iterable_dataset.py

@@ -7,7 +7,7 @@
 import torch.utils.data
 
 from . import S3Reader
-from ._s3bucket_key import S3BucketKey
+from ._s3bucket_key_data import S3BucketKeyData
 from ._s3client import S3Client
 from ._s3dataset_common import (
     identity,
@@ -28,7 +28,7 @@ class S3IterableDataset(torch.utils.data.IterableDataset):
     def __init__(
         self,
         region: str,
-        get_dataset_objects: Callable[[S3Client], Iterable[S3BucketKey]],
+        get_dataset_objects: Callable[[S3Client], Iterable[S3BucketKeyData]],
         endpoint: str = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
@@ -113,9 +113,11 @@ def _get_client(self):
             self._client = S3Client(self.region, self.endpoint)
         return self._client
 
-    def _get_transformed_object(self, bucket_key: S3BucketKey) -> Any:
+    def _get_transformed_object(self, bucket_key: S3BucketKeyData) -> Any:
         return self._transform(
-            self._get_client().get_object(bucket_key.bucket, bucket_key.key)
+            self._get_client().get_object(
+                bucket_key.bucket, bucket_key.key, object_info=bucket_key.object_info
+            )
         )
 
     def __iter__(self) -> Iterator[Any]:

s3torchconnector/src/s3torchconnector/s3map_dataset.py

@@ -5,7 +5,7 @@
 import logging
 
 import torch.utils.data
-from s3torchconnector._s3bucket_key import S3BucketKey
+from s3torchconnector._s3bucket_key_data import S3BucketKeyData
 
 from ._s3client import S3Client
 from . import S3Reader
@@ -29,7 +29,7 @@ class S3MapDataset(torch.utils.data.Dataset):
     def __init__(
         self,
         region: str,
-        get_dataset_objects: Callable[[S3Client], Iterable[S3BucketKey]],
+        get_dataset_objects: Callable[[S3Client], Iterable[S3BucketKeyData]],
         endpoint: str = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
@@ -49,7 +49,7 @@ def endpoint(self):
         return self._endpoint
 
     @property
-    def _dataset_bucket_key_pairs(self) -> List[S3BucketKey]:
+    def _dataset_bucket_key_pairs(self) -> List[S3BucketKeyData]:
         if self._bucket_key_pairs is None:
             self._bucket_key_pairs = list(self._get_dataset_objects(self._get_client()))
         return self._bucket_key_pairs
@@ -123,7 +123,9 @@ def _get_client(self):
 
     def _get_object(self, i: int) -> S3Reader:
         bucket_key = self._dataset_bucket_key_pairs[i]
-        return self._get_client().get_object(bucket_key.bucket, bucket_key.key)
+        return self._get_client().get_object(
+            bucket_key.bucket, bucket_key.key, object_info=bucket_key.object_info
+        )
 
     def __getitem__(self, i: int) -> Any:
         return self._transform(self._get_object(i))

s3torchconnector/tst/unit/test_s3_client.py

@@ -1,11 +1,11 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
 import logging
+from unittest.mock import MagicMock
 
 import pytest
 
 from s3torchconnector._s3client import S3Client, MockS3Client
-from s3torchconnectorclient._mountpoint_s3_client import ObjectInfo
 
 TEST_BUCKET = "test-bucket"
 TEST_KEY = "test-key"
@@ -23,16 +23,13 @@ def s3_client() -> S3Client:
 def test_get_object_log(s3_client: S3Client, caplog):
     with caplog.at_level(logging.DEBUG):
         s3_client.get_object(TEST_BUCKET, TEST_KEY)
-    assert f"GetObject {S3_URI}" in caplog.messages
+    assert f"GetObject {S3_URI}, object_info is None=True" in caplog.messages
 
 
-def test_get_object_info_log(s3_client: S3Client, caplog):
+def test_get_object_log_with_info(s3_client: S3Client, caplog):
     with caplog.at_level(logging.DEBUG):
-        s3_client.from_bucket_and_object_info(
-            TEST_BUCKET,
-            ObjectInfo(TEST_KEY, "", 0, 0, None, None),
-        )
-    assert f"GetObjectWithInfo {S3_URI}" in caplog.messages
+        s3_client.get_object(TEST_BUCKET, TEST_KEY, object_info=MagicMock())
+    assert f"GetObject {S3_URI}, object_info is None=False" in caplog.messages
 
 
 def test_head_object_log(s3_client: S3Client, caplog):

s3torchconnector/tst/unit/test_s3iterable_dataset.py

@@ -1,11 +1,14 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
 import logging
+from io import SEEK_END
 from typing import Iterable, Callable, Sequence, Any
+from unittest.mock import patch
 
 import pytest
 
 from s3torchconnector import S3IterableDataset, S3Reader
+from s3torchconnector._s3client import MockS3Client
 
 from test_s3dataset_common import (
     TEST_BUCKET,
@@ -200,6 +203,21 @@ def test_dataset_creation_from_prefix_with_region_and_endpoint():
     assert dataset.endpoint == TEST_ENDPOINT
 
 
+def test_from_prefix_seek_no_head():
+    dataset = S3IterableDataset.from_prefix(S3_PREFIX, region=TEST_REGION)
+
+    # use mock client for unit testing
+    client = _create_mock_client_with_dummy_objects(TEST_BUCKET, ["foo"])
+    dataset._client = client
+
+    with patch.object(
+        MockS3Client, "head_object", wraps=client.head_object
+    ) as head_object:
+        s3_object = next(iter(dataset))
+        s3_object.seek(0, SEEK_END)
+    head_object.assert_not_called()
+
+
 def _verify_dataset(
     dataset: S3IterableDataset,
     expected_keys: Sequence[str],

s3torchconnector/tst/unit/test_s3mapdataset.py

@@ -1,11 +1,14 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
 import logging
+from io import SEEK_END
 from typing import Sequence, Callable, Any
+from unittest.mock import patch
 
 import pytest
 
 from s3torchconnector import S3MapDataset, S3Reader
+from s3torchconnector._s3client import MockS3Client
 
 from test_s3dataset_common import (
     TEST_BUCKET,
@@ -159,6 +162,21 @@ def test_transform_from_objects(
     assert list(dataset) == [expected]
 
 
+def test_from_prefix_seek_no_head():
+    dataset = S3MapDataset.from_prefix(S3_PREFIX, region=TEST_REGION)
+
+    # use mock client for unit testing
+    client = _create_mock_client_with_dummy_objects(TEST_BUCKET, ["foo"])
+    dataset._client = client
+
+    with patch.object(
+        MockS3Client, "head_object", wraps=client.head_object
+    ) as head_object:
+        s3_object = next(iter(dataset))
+        s3_object.seek(0, SEEK_END)
+    head_object.assert_not_called()
+
+
 @pytest.mark.parametrize(
     "keys, length",
     [