Add benchmarking for boto3-using-crt

awslabs · Oct 18, 2023 · 470d3af · 470d3af
1 parent cd5a06b
commit 470d3af
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 18 deletions.
diff --git a/runners/s3-benchrunner-python/main.py b/runners/s3-benchrunner-python/main.py
@@ -11,14 +11,11 @@
     bytes_to_gigabit,
     ns_to_secs,
 )
-from runner.boto3 import Boto3BenchmarkRunner
-from runner.cli import CliBenchmarkRunner
-from runner.crt import CrtBenchmarkRunner
 
 PARSER = argparse.ArgumentParser(
     description='Python benchmark runner. Pick which S3 library to use.')
 PARSER.add_argument('LIB', choices=(
-    'crt', 'boto3-python', 'cli-python', 'cli-crt'))
+    'crt', 'boto3-python', 'boto3-crt', 'cli-python', 'cli-crt'))
 PARSER.add_argument('BENCHMARK')
 PARSER.add_argument('BUCKET')
 PARSER.add_argument('REGION')
@@ -28,12 +25,20 @@
 
 def create_runner_for_lib(lib: str, config: BenchmarkConfig) -> BenchmarkRunner:
     """Factory function. Create appropriate subclass, given the lib."""
+    use_crt = lib.endswith('crt')
+
     if lib == 'crt':
+        from runner.crt import CrtBenchmarkRunner
         return CrtBenchmarkRunner(config)
-    if lib == 'boto3-python':
-        return Boto3BenchmarkRunner(config)
-    if lib.startswith('cli-'):
-        return CliBenchmarkRunner(config, use_crt=lib.endswith('crt'))
+
+    if lib.startswith('boto3'):
+        from runner.boto3 import Boto3BenchmarkRunner
+        return Boto3BenchmarkRunner(config, use_crt)
+
+    if lib.startswith('cli'):
+        from runner.cli import CliBenchmarkRunner
+        return CliBenchmarkRunner(config, use_crt)
+
     else:
         raise ValueError(f'Unknown lib: {lib}')
 

diff --git a/runners/s3-benchrunner-python/runner/boto3.py b/runners/s3-benchrunner-python/runner/boto3.py
@@ -1,45 +1,58 @@
-import boto3  # type: ignore
 from concurrent.futures import ThreadPoolExecutor
+import os
 
 from runner import BenchmarkConfig, BenchmarkRunner
 
 
 class Boto3BenchmarkRunner(BenchmarkRunner):
     """Benchmark runner using boto3.client('s3')"""
 
-    def __init__(self, config: BenchmarkConfig):
+    def __init__(self, config: BenchmarkConfig, use_crt: bool):
         super().__init__(config)
 
+        # boto3 uses CRT if it's installed, unless env-var BOTO_DISABLE_CRT=true
+        # so set the env-var before importing boto3
+        os.environ['BOTO_DISABLE_CRT'] = str(not use_crt).lower()
+        import boto3  # type: ignore
+        import botocore.compat  # type: ignore
+        assert use_crt == botocore.compat.HAS_CRT
+
+        self.use_crt = use_crt
+        if (use_crt):
+            self._verbose('--- boto3-crt ---')
+        else:
+            self._verbose('--- boto3-python ---')
+
         self._s3_client = boto3.client('s3')
 
+    def _verbose(self, msg):
+        if self.config.verbose:
+            print(msg)
+
     def _make_request(self, task_i: int):
         task = self.config.tasks[task_i]
 
         if task.action == 'upload':
             if self.config.files_on_disk:
-                if self.config.verbose:
-                    print(f'boto3 upload_file("{task.key}")')
+                self._verbose(f'upload_file("{task.key}")')
                 self._s3_client.upload_file(
                     task.key, self.config.bucket, task.key)
 
             else:
-                if self.config.verbose:
-                    print(f'boto3 upload_fileobj("{task.key}")')
+                self._verbose(f'upload_fileobj("{task.key}")')
                 upload_stream = self._new_iostream_to_upload_from_ram(
                     task.size)
                 self._s3_client.upload_fileobj(
                     upload_stream, self.config.bucket, task.key)
 
         elif task.action == 'download':
             if self.config.files_on_disk:
-                if self.config.verbose:
-                    print(f'boto3 download_file("{task.key}")')
+                self._verbose(f'download_file("{task.key}")')
                 self._s3_client.download_file(
                     self.config.bucket, task.key, task.key)
 
             else:
-                if self.config.verbose:
-                    print(f'boto3 download_fileobj("{task.key}")')
+                self._verbose(f'download_fileobj("{task.key}")')
                 download_stream = Boto3DownloadFileObj()
                 self._s3_client.download_fileobj(
                     self.config.bucket, task.key, download_stream)