diff --git a/runners/s3-benchrunner-python/main.py b/runners/s3-benchrunner-python/main.py index fc345358..1f20109d 100755 --- 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 index 33e43787..4ecc0a39 100644 --- a/runners/s3-benchrunner-python/runner/boto3.py +++ b/runners/s3-benchrunner-python/runner/boto3.py @@ -1,5 +1,5 @@ -import boto3 # type: ignore from concurrent.futures import ThreadPoolExecutor +import os from runner import BenchmarkConfig, BenchmarkRunner @@ -7,24 +7,39 @@ 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( @@ -32,14 +47,12 @@ def _make_request(self, task_i: int): 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)