Add additional logging for construction and HeadObject requests (#128)

* Add additional logging to verify which constructor is used for Dataset creation

* Add logging whenever we run an explicit HeadObject request

* Add additional logging for explicit calls to S3Client

* Add extra tests which validate logs are performed when S3Client methods are called

* Add explicit logging for s3_uri in from_prefix constructors


---------

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

Author: muddyfish

s3torchconnector/src/s3torchconnector/_s3client/_s3client.py

@@ -1,6 +1,7 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
 
+import logging
 import os
 from functools import partial
 from typing import Optional, Any
@@ -23,6 +24,9 @@
 """
 
 
+log = logging.getLogger(__name__)
+
+
 def _identity(obj: Any) -> Any:
     return obj
 
@@ -51,6 +55,7 @@ def _client_builder(self) -> MountpointS3Client:
         )
 
     def get_object(self, bucket: str, key: str) -> S3Reader:
+        log.debug(f"GetObject s3://{bucket}/{key}")
         return S3Reader(
             bucket,
             key,
@@ -64,21 +69,25 @@ def _get_object_stream(self, bucket: str, key: str) -> GetObjectStream:
     def put_object(
         self, bucket: str, key: str, storage_class: Optional[str] = None
     ) -> S3Writer:
+        log.debug(f"PutObject s3://{bucket}/{key}")
         return S3Writer(self._client.put_object(bucket, key, storage_class))
 
     # TODO: Probably need a ListObjectResult on dataset side
     def list_objects(
         self, bucket: str, prefix: str = "", delimiter: str = "", max_keys: int = 1000
     ) -> ListObjectStream:
+        log.debug(f"ListObjects s3://{bucket}/{prefix}")
         return self._client.list_objects(bucket, prefix, delimiter, max_keys)
 
     # TODO: We need ObjectInfo on dataset side
     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,

s3torchconnector/src/s3torchconnector/s3iterable_dataset.py

@@ -2,6 +2,7 @@
 #  // SPDX-License-Identifier: BSD
 from functools import partial
 from typing import Iterator, Any, Union, Iterable, Callable
+import logging
 
 import torch.utils.data
 
@@ -14,6 +15,8 @@
     get_objects_from_prefix,
 )
 
+log = logging.getLogger(__name__)
+
 
 class S3IterableDataset(torch.utils.data.IterableDataset):
     """An IterableStyle dataset created from S3 objects.
@@ -58,6 +61,7 @@ def from_objects(
         Raises:
             S3Exception: An error occurred accessing S3.
         """
+        log.info(f"Building {cls.__name__} from_objects")
         return cls(
             region, partial(get_objects_from_uris, object_uris), transform=transform
         )
@@ -83,6 +87,7 @@ def from_prefix(
         Raises:
             S3Exception: An error occurred accessing S3.
         """
+        log.info(f"Building {cls.__name__} from_prefix {s3_uri=}")
         return cls(
             region, partial(get_objects_from_prefix, s3_uri), transform=transform
         )

s3torchconnector/src/s3torchconnector/s3map_dataset.py

@@ -2,6 +2,7 @@
 #  // SPDX-License-Identifier: BSD
 from functools import partial
 from typing import List, Any, Callable, Iterable, Union
+import logging
 
 import torch.utils.data
 from s3torchconnector._s3bucket_key import S3BucketKey
@@ -15,6 +16,8 @@
     identity,
 )
 
+log = logging.getLogger(__name__)
+
 
 class S3MapDataset(torch.utils.data.Dataset):
     """A Map-Style dataset created from S3 objects.
@@ -66,7 +69,7 @@ def from_objects(
         Raises:
             S3Exception: An error occurred accessing S3.
         """
-
+        log.info(f"Building {cls.__name__} from_objects")
         return cls(
             region, partial(get_objects_from_uris, object_uris), transform=transform
         )
@@ -92,6 +95,7 @@ def from_prefix(
         Raises:
             S3Exception: An error occurred accessing S3.
         """
+        log.info(f"Building {cls.__name__} from_prefix {s3_uri=}")
         return cls(
             region, partial(get_objects_from_prefix, s3_uri), transform=transform
         )

s3torchconnector/tst/unit/test_checkpointing.py

@@ -1,7 +1,7 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
-
-from io import BytesIO
+import logging
+from io import BytesIO, SEEK_END
 from operator import eq
 from typing import Any, Callable
 from unittest.mock import patch
@@ -147,6 +147,23 @@ def test_general_checkpointing_untyped_storage_loads_no_modern_pytorch_format(
     )
 
 
+def test_checkpoint_seek_logging(caplog):
+    checkpoint = S3Checkpoint(TEST_REGION)
+
+    # Use MockClient instead of actual client.
+    client = MockS3Client(TEST_REGION, TEST_BUCKET)
+    checkpoint._client = client
+
+    s3_uri = f"s3://{TEST_BUCKET}/{TEST_KEY}"
+    with checkpoint.writer(s3_uri) as writer:
+        writer.write(b"test")
+
+    with checkpoint.reader(s3_uri) as reader:
+        with caplog.at_level(logging.DEBUG):
+            reader.seek(0, SEEK_END)
+    assert f"HeadObject {s3_uri}" in caplog.messages
+
+
 def _test_save(
     data,
     byteorder: str,

s3torchconnector/tst/unit/test_s3_client.py

@@ -0,0 +1,53 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+import logging
+
+import pytest
+
+from s3torchconnector._s3client import S3Client, MockS3Client
+from s3torchconnectorclient._mountpoint_s3_client import ObjectInfo
+
+TEST_BUCKET = "test-bucket"
+TEST_KEY = "test-key"
+TEST_REGION = "us-east-1"
+S3_URI = f"s3://{TEST_BUCKET}/{TEST_KEY}"
+
+
+@pytest.fixture
+def s3_client() -> S3Client:
+    client = MockS3Client(TEST_REGION, TEST_BUCKET)
+    client.add_object(TEST_KEY, b"data")
+    return client
+
+
+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
+
+
+def test_get_object_info_log(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
+
+
+def test_head_object_log(s3_client: S3Client, caplog):
+    with caplog.at_level(logging.DEBUG):
+        s3_client.head_object(TEST_BUCKET, TEST_KEY)
+    assert f"HeadObject {S3_URI}" in caplog.messages
+
+
+def test_put_object_log(s3_client: S3Client, caplog):
+    with caplog.at_level(logging.DEBUG):
+        s3_client.put_object(TEST_BUCKET, TEST_KEY)
+    assert f"PutObject {S3_URI}" in caplog.messages
+
+
+def test_list_objects_log(s3_client: S3Client, caplog):
+    with caplog.at_level(logging.DEBUG):
+        s3_client.list_objects(TEST_BUCKET, TEST_KEY)
+    assert f"ListObjects {S3_URI}" in caplog.messages

s3torchconnector/tst/unit/test_s3iterable_dataset.py

@@ -1,6 +1,6 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
-
+import logging
 from typing import Iterable, Callable, Sequence, Any
 
 import pytest
@@ -15,16 +15,20 @@
 )
 
 
-def test_dataset_creation_from_prefix_with_region():
-    dataset = S3IterableDataset.from_prefix(S3_PREFIX, region=TEST_REGION)
+def test_dataset_creation_from_prefix_with_region(caplog):
+    with caplog.at_level(logging.INFO):
+        dataset = S3IterableDataset.from_prefix(S3_PREFIX, region=TEST_REGION)
     assert isinstance(dataset, S3IterableDataset)
     assert dataset.region == TEST_REGION
+    assert "Building S3IterableDataset from_prefix" in caplog.text
 
 
-def test_dataset_creation_from_objects_with_region():
-    dataset = S3IterableDataset.from_objects([], region=TEST_REGION)
+def test_dataset_creation_from_objects_with_region(caplog):
+    with caplog.at_level(logging.INFO):
+        dataset = S3IterableDataset.from_objects([], region=TEST_REGION)
     assert isinstance(dataset, S3IterableDataset)
     assert dataset.region == TEST_REGION
+    assert "Building S3IterableDataset from_objects" in caplog.text
 
 
 @pytest.mark.parametrize(

s3torchconnector/tst/unit/test_s3mapdataset.py

@@ -1,6 +1,6 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
-
+import logging
 from typing import Sequence, Callable, Any
 
 import pytest
@@ -15,16 +15,20 @@
 )
 
 
-def test_dataset_creation_from_prefix_with_region():
-    dataset = S3MapDataset.from_prefix(S3_PREFIX, region=TEST_REGION)
+def test_dataset_creation_from_prefix_with_region(caplog):
+    with caplog.at_level(logging.INFO):
+        dataset = S3MapDataset.from_prefix(S3_PREFIX, region=TEST_REGION)
     assert isinstance(dataset, S3MapDataset)
     assert dataset.region == TEST_REGION
+    assert "Building S3MapDataset from_prefix" in caplog.text
 
 
-def test_dataset_creation_from_objects_with_region():
-    dataset = S3MapDataset.from_objects([], region=TEST_REGION)
+def test_dataset_creation_from_objects_with_region(caplog):
+    with caplog.at_level(logging.INFO):
+        dataset = S3MapDataset.from_objects([], region=TEST_REGION)
     assert isinstance(dataset, S3MapDataset)
     assert dataset.region == TEST_REGION
+    assert "Building S3MapDataset from_objects" in caplog.text
 
 
 @pytest.mark.parametrize(