Add support for canned ACL environment variable (#2729)

mosaicml · Nov 17, 2023 · 2cf73b4 · 2cf73b4
1 parent 1a4e67f
commit 2cf73b4
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 2 deletions.
diff --git a/composer/utils/object_store/s3_object_store.py b/composer/utils/object_store/s3_object_store.py
@@ -135,13 +135,18 @@ def upload_object(self,
         file_size = os.path.getsize(filename)
         cb_wrapper = None if callback is None else lambda bytes_transferred: callback(bytes_transferred, file_size)
 
-        # Validate kwargs
-        if len(kwargs) != 0:
+        # Validate kwargs. Use env var for Canned ACL if present and one has not been passed in.
+        if len(kwargs) == 0:
+            if 'S3_CANNED_ACL' in os.environ:
+                kwargs['ExtraArgs'] = {'ACL': os.environ['S3_CANNED_ACL']}
+        else:
             if len(kwargs) > 1 or 'ExtraArgs' not in kwargs or not isinstance(kwargs['ExtraArgs'], dict):
                 raise ValueError('S3ObjectStore.upload_object only supports an additional ExtraArgs dictionary.')
             for key in kwargs['ExtraArgs']:
                 if key not in S3Transfer.ALLOWED_UPLOAD_ARGS:
                     raise ValueError(f'{key} is not an allowed upload argument.')
+            if 'S3_CANNED_ACL' in os.environ and 'ACL' not in kwargs['ExtraArgs']:
+                kwargs['ExtraArgs']['ACL'] = os.environ['S3_CANNED_ACL']
 
         self.client.upload_file(Bucket=self.bucket,
                                 Key=self.get_key(object_name),

diff --git a/tests/utils/object_store/test_s3_object_store.py b/tests/utils/object_store/test_s3_object_store.py
@@ -4,6 +4,8 @@
 import os
 import pathlib
 import threading
+from unittest import mock
+from unittest.mock import ANY, MagicMock
 
 import pytest
 
@@ -31,3 +33,63 @@ def test_s3_object_store_multi_threads(tmp_path: pathlib.Path, s3_bucket: str):
         threads.append(t)
     for t in threads:
         t.join()
+
+
+def test_s3_upload_object_arguments(tmp_path: pathlib.Path, s3_bucket: str):
+    filename = tmp_path / 'localfile.txt'
+    filename.touch()
+    remote_obj_name = 'remote.txt'
+
+    object_store = S3ObjectStore(bucket=s3_bucket)
+    object_store.client.upload_file = MagicMock()
+
+    with mock.patch.dict('os.environ'):
+        os.environ.pop('S3_CANNED_ACL', None)
+
+        object_store.upload_object(object_name=remote_obj_name, filename=filename)
+        object_store.client.upload_file.assert_called_with(Bucket='my-bucket',
+                                                           Key=remote_obj_name,
+                                                           Filename=filename,
+                                                           Callback=None,
+                                                           Config=ANY)
+
+        object_store.upload_object(object_name=remote_obj_name,
+                                   filename=filename,
+                                   ExtraArgs={'ACL': 'authenticated-read'})
+        object_store.client.upload_file.assert_called_with(Bucket='my-bucket',
+                                                           Key=remote_obj_name,
+                                                           Filename=filename,
+                                                           Callback=None,
+                                                           Config=ANY,
+                                                           ExtraArgs={'ACL': 'authenticated-read'})
+
+        os.environ['S3_CANNED_ACL'] = 'bucket-owner-full-control'
+
+        object_store.upload_object(object_name=remote_obj_name, filename=filename)
+        object_store.client.upload_file.assert_called_with(Bucket='my-bucket',
+                                                           Key=remote_obj_name,
+                                                           Filename=filename,
+                                                           Callback=None,
+                                                           Config=ANY,
+                                                           ExtraArgs={'ACL': 'bucket-owner-full-control'})
+
+        object_store.upload_object(object_name=remote_obj_name, filename=filename, ExtraArgs={'Metadata': {}})
+        object_store.client.upload_file.assert_called_with(Bucket='my-bucket',
+                                                           Key=remote_obj_name,
+                                                           Filename=filename,
+                                                           Callback=None,
+                                                           Config=ANY,
+                                                           ExtraArgs={
+                                                               'ACL': 'bucket-owner-full-control',
+                                                               'Metadata': {}
+                                                           })
+
+        object_store.upload_object(object_name=remote_obj_name,
+                                   filename=filename,
+                                   ExtraArgs={'ACL': 'authenticated-read'})
+        object_store.client.upload_file.assert_called_with(Bucket='my-bucket',
+                                                           Key=remote_obj_name,
+                                                           Filename=filename,
+                                                           Callback=None,
+                                                           Config=ANY,
+                                                           ExtraArgs={'ACL': 'authenticated-read'})