Merge pull request #361 from hssyoo/crt-mode

Add TransferConfig support for CRTTransferManager

Author: hssyoo

.changes/next-release/enhancement-awscrt-49742.json

@@ -0,0 +1,5 @@
+{
+  "type": "feature",
+  "category": "``awscrt``",
+  "description": "``CRTTransferManager`` now supports the following ``TransferConfig`` options - ``multipart_threshold``, ``multipart_chunksize``, ``max_request_concurrency``"
+}

s3transfer/crt.py

@@ -13,6 +13,7 @@
 import logging
 import re
 import threading
+from collections import namedtuple
 from io import BytesIO
 
 import awscrt.http
@@ -184,6 +185,16 @@ def _get_crt_throughput_target_gbps(provided_throughput_target_bytes=None):
     return target_gbps
 
 
+def _has_minimum_crt_version(minimum_version):
+    crt_version_str = awscrt.__version__
+    try:
+        crt_version_ints = map(int, crt_version_str.split("."))
+        crt_version_tuple = tuple(crt_version_ints)
+    except (TypeError, ValueError):
+        return False
+    return crt_version_tuple >= minimum_version
+
+
 class CRTTransferManager:
     ALLOWED_DOWNLOAD_ARGS = TransferManager.ALLOWED_DOWNLOAD_ARGS
     ALLOWED_UPLOAD_ARGS = TransferManager.ALLOWED_UPLOAD_ARGS
@@ -193,7 +204,9 @@ class CRTTransferManager:
 
     _UNSUPPORTED_BUCKET_PATTERNS = TransferManager._UNSUPPORTED_BUCKET_PATTERNS
 
-    def __init__(self, crt_s3_client, crt_request_serializer, osutil=None):
+    def __init__(
+        self, crt_s3_client, crt_request_serializer, osutil=None, config=None
+    ):
         """A transfer manager interface for Amazon S3 on CRT s3 client.
 
         :type crt_s3_client: awscrt.s3.S3Client
@@ -207,12 +220,18 @@ def __init__(self, crt_s3_client, crt_request_serializer, osutil=None):
         :type osutil: s3transfer.utils.OSUtils
         :param osutil: OSUtils object to use for os-related behavior when
             using with transfer manager.
+
+        :type config: s3transfer.manager.TransferConfig
+        :param config: The transfer configuration to be used when
+            making CRT S3 client requests.
         """
         if osutil is None:
             self._osutil = OSUtils()
         self._crt_s3_client = crt_s3_client
         self._s3_args_creator = S3ClientArgsCreator(
-            crt_request_serializer, self._osutil
+            crt_request_serializer,
+            self._osutil,
+            config,
         )
         self._crt_exception_translator = (
             crt_request_serializer.translate_crt_exception
@@ -731,10 +750,72 @@ def set_s3_request(self, s3_request):
         self._crt_future = self._s3_request.finished_future
 
 
+CRTConfigParameter = namedtuple('CRTConfigParameter', ['name', 'min_version'])
+
+
 class S3ClientArgsCreator:
-    def __init__(self, crt_request_serializer, os_utils):
+    _CRT_ARG_TO_CONFIG_PARAM = {
+        'max_active_connections_override': CRTConfigParameter(
+            'max_request_concurrency', (0, 29, 0)
+        ),
+    }
+
+    def __init__(self, crt_request_serializer, os_utils, config=None):
         self._request_serializer = crt_request_serializer
         self._os_utils = os_utils
+        self._config = config
+
+    def _get_crt_transfer_config_options(self, request_type):
+        crt_config = {
+            'part_size': self._config.multipart_chunksize,
+            'max_active_connections_override': self._config.max_request_concurrency,
+        }
+
+        if (
+            self._config.get_deep_attr('multipart_chunksize')
+            is self._config.UNSET_DEFAULT
+        ):
+            # Let CRT dynamically calculate part size.
+            crt_config['part_size'] = None
+        if (
+            self._config.get_deep_attr('max_request_concurrency')
+            is self._config.UNSET_DEFAULT
+        ):
+            crt_config['max_active_connections_override'] = None
+
+        if hasattr(self, f'_get_crt_options_{request_type}'):
+            crt_config.update(
+                getattr(self, f'_get_crt_options_{request_type}')()
+            )
+        self._remove_param_if_not_min_crt_version(crt_config)
+        return crt_config
+
+    def _get_crt_options_put_object(self):
+        return {'multipart_upload_threshold': self._config.multipart_threshold}
+
+    def _remove_param_if_not_min_crt_version(self, crt_config):
+        to_remove = []
+        for request_arg in crt_config:
+            if request_arg not in self._CRT_ARG_TO_CONFIG_PARAM:
+                continue
+            param = self._CRT_ARG_TO_CONFIG_PARAM[request_arg]
+            if _has_minimum_crt_version(param.min_version):
+                continue
+            # Only log the warning if user attempted to explicitly
+            # use the transfer config parameter.
+            if (
+                self._config.get_deep_attr(param.name)
+                is not self._config.UNSET_DEFAULT
+            ):
+                min_ver_str = '.'.join(str(i) for i in param.min_version)
+                logger.warning(
+                    f'Transfer config parameter {param.name} '
+                    f'requires minimum CRT version: {min_ver_str}. '
+                    f'{param.name} will not be used in the request.'
+                )
+            to_remove.append(request_arg)
+        for request_arg in to_remove:
+            del crt_config[request_arg]
 
     def get_make_request_args(
         self, request_type, call_args, coordinator, future, on_done_after_calls
@@ -823,6 +904,10 @@ def _get_make_request_args_put_object(
         )
         make_request_args['send_filepath'] = send_filepath
         make_request_args['checksum_config'] = checksum_config
+        if self._config is not None:
+            make_request_args.update(
+                self._get_crt_transfer_config_options(request_type)
+            )
         return make_request_args
 
     def _get_make_request_args_get_object(
@@ -859,6 +944,10 @@ def _get_make_request_args_get_object(
         make_request_args['recv_filepath'] = recv_filepath
         make_request_args['on_body'] = on_body
         make_request_args['checksum_config'] = checksum_config
+        if self._config is not None:
+            make_request_args.update(
+                self._get_crt_transfer_config_options(request_type)
+            )
         return make_request_args
 
     def _default_get_make_request_args(

s3transfer/manager.py

@@ -50,6 +50,8 @@
 
 
 class TransferConfig:
+    UNSET_DEFAULT = object()
+
     def __init__(
         self,
         multipart_threshold=8 * MB,
@@ -152,12 +154,19 @@ def __init__(
 
     def _validate_attrs_are_nonzero(self):
         for attr, attr_val in self.__dict__.items():
-            if attr_val is not None and attr_val <= 0:
+            if (
+                attr_val is not None
+                and attr_val is not self.UNSET_DEFAULT
+                and attr_val <= 0
+            ):
                 raise ValueError(
                     f'Provided parameter {attr} of value {attr_val} must '
                     'be greater than 0.'
                 )
 
+    def get_deep_attr(self, item):
+        return object.__getattribute__(self, item)
+
 
 class TransferManager:
     ALLOWED_DOWNLOAD_ARGS = ALLOWED_DOWNLOAD_ARGS

tests/functional/test_crt.py

@@ -18,6 +18,8 @@
 
 from botocore.session import Session
 
+from s3transfer.constants import MB
+from s3transfer.manager import TransferConfig
 from s3transfer.subscribers import BaseSubscriber
 from tests import (
     HAS_CRT,
@@ -769,3 +771,98 @@ def test_crt_s3_client_error_handling(self):
         )
         with self.assertRaises(awscrt.exceptions.AwsCrtError):
             future.result()
+
+    def test_transfer_config_used_in_upload_request(self):
+        config = TransferConfig(
+            multipart_threshold=4 * MB,
+            multipart_chunksize=2 * MB,
+            max_request_concurrency=100,
+        )
+        transfer_manager = s3transfer.crt.CRTTransferManager(
+            crt_s3_client=self.s3_crt_client,
+            crt_request_serializer=self.request_serializer,
+            config=config,
+        )
+        future = transfer_manager.upload(
+            self.filename, self.bucket, self.key, {}, []
+        )
+        future.result()
+
+        callargs_kwargs = self.s3_crt_client.make_request.call_args[1]
+        assert callargs_kwargs['multipart_upload_threshold'] == 4 * MB
+        assert callargs_kwargs['part_size'] == 2 * MB
+        assert callargs_kwargs['max_active_connections_override'] == 100
+
+    def test_transfer_config_used_in_download_request(self):
+        config = TransferConfig(
+            multipart_threshold=4 * MB,
+            multipart_chunksize=2 * MB,
+            max_request_concurrency=100,
+        )
+        transfer_manager = s3transfer.crt.CRTTransferManager(
+            crt_s3_client=self.s3_crt_client,
+            crt_request_serializer=self.request_serializer,
+            config=config,
+        )
+        future = transfer_manager.download(
+            self.bucket, self.key, self.filename, {}, []
+        )
+        future.result()
+
+        callargs_kwargs = self.s3_crt_client.make_request.call_args[1]
+        assert callargs_kwargs['part_size'] == 2 * MB
+        assert callargs_kwargs['max_active_connections_override'] == 100
+        # Config option only used for PUT requests.
+        assert 'multipart_upload_threshold' not in callargs_kwargs
+
+    def test_unset_part_size_defaults_to_none_in_upload_request(self):
+        config = TransferConfig(
+            multipart_chunksize=TransferConfig.UNSET_DEFAULT,
+        )
+        transfer_manager = s3transfer.crt.CRTTransferManager(
+            crt_s3_client=self.s3_crt_client,
+            crt_request_serializer=self.request_serializer,
+            config=config,
+        )
+        future = transfer_manager.upload(
+            self.filename, self.bucket, self.key, {}, []
+        )
+        future.result()
+
+        callargs_kwargs = self.s3_crt_client.make_request.call_args[1]
+        assert callargs_kwargs['part_size'] is None
+
+    def test_unset_max_concurrency_defaults_to_none(self):
+        config = TransferConfig(
+            max_request_concurrency=TransferConfig.UNSET_DEFAULT,
+        )
+        transfer_manager = s3transfer.crt.CRTTransferManager(
+            crt_s3_client=self.s3_crt_client,
+            crt_request_serializer=self.request_serializer,
+            config=config,
+        )
+        future = transfer_manager.upload(
+            self.filename, self.bucket, self.key, {}, []
+        )
+        future.result()
+
+        callargs_kwargs = self.s3_crt_client.make_request.call_args[1]
+        assert callargs_kwargs['max_active_connections_override'] is None
+
+    @mock.patch('awscrt.__version__', '0.28.0')
+    def test_args_removed_if_not_min_awscrt_version(self):
+        config = TransferConfig(
+            max_request_concurrency=TransferConfig.UNSET_DEFAULT,
+        )
+        transfer_manager = s3transfer.crt.CRTTransferManager(
+            crt_s3_client=self.s3_crt_client,
+            crt_request_serializer=self.request_serializer,
+            config=config,
+        )
+        future = transfer_manager.upload(
+            self.filename, self.bucket, self.key, {}, []
+        )
+        future.result()
+
+        callargs_kwargs = self.s3_crt_client.make_request.call_args[1]
+        assert 'max_active_connections_override' not in callargs_kwargs