fix: Use sagemaker session's s3 resource for downloading folders

Author: Darren Stack

src/sagemaker/utils.py

@@ -384,8 +384,7 @@ def download_folder(bucket_name, prefix, target, sagemaker_session):
         sagemaker_session (sagemaker.session.Session): a sagemaker session to
             interact with S3.
     """
-    boto_session = sagemaker_session.boto_session
-    s3 = boto_session.resource("s3", region_name=boto_session.region_name)
+    s3 = sagemaker_session.s3_resource
 
     prefix = prefix.lstrip("/")
 
@@ -657,9 +656,8 @@ def download_file(bucket_name, path, target, sagemaker_session):
             interact with S3.
     """
     path = path.lstrip("/")
-    boto_session = sagemaker_session.boto_session
 
-    s3 = boto_session.resource("s3", region_name=sagemaker_session.boto_region_name)
+    s3 = sagemaker_session.s3_resource
     bucket = s3.Bucket(bucket_name)
     bucket.download_file(path, target)
 

tests/unit/test_utils.py

@@ -1032,10 +1032,10 @@ def __init__(self, tmp):
         self.current_bucket = None
         self.object_mock = MagicMock()
 
-        self.sagemaker_session.boto_session.resource().Bucket().download_file.side_effect = (
+        self.sagemaker_session.s3_resource.Bucket().download_file.side_effect = (
             self.download_file
         )
-        self.sagemaker_session.boto_session.resource().Bucket.side_effect = self.bucket
+        self.sagemaker_session.s3_resource.Bucket.side_effect = self.bucket
         self.fake_upload_path = self.mock_s3_upload()
 
     def bucket(self, name):
@@ -1070,6 +1070,7 @@ def upload_file(self, target, **kwargs):
                 object_mock.upload_file(target, **kwargs)
 
         self.sagemaker_session.boto_session.resource().Object = MockS3Object
+        self.sagemaker_session.s3_resource.Object = MockS3Object
         return dst