1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
|
#!/usr/bin/env python
"""
Benchmark the uploading of a file using s3transfer. You can also chose how
type of file that is uploaded (i.e. filename, seekable, nonseekable).
Usage
=====
NOTE: Make sure you run ``pip install -r requirements-dev.txt`` before running.
To benchmark with using a temporary file that is generated for you::
./benchmark-upload --file-size 10MB --file-type filename \\
--s3-bucket mybucket
To benchmark with your own local file::
./benchmark-upload --source-file myfile --file-type filename \\
--s3-bucket mybucket
"""
import argparse
import os
import shutil
import subprocess
import tempfile
from botocore.session import get_session
TEMP_FILE = 'temp'
TEMP_KEY = 'temp'
KB = 1024
SIZE_SUFFIX = {
'kb': 1024,
'mb': 1024**2,
'gb': 1024**3,
'tb': 1024**4,
'kib': 1024,
'mib': 1024**2,
'gib': 1024**3,
'tib': 1024**4,
}
def human_readable_to_bytes(value):
"""Converts a human readable size to bytes.
:param value: A string such as "10MB". If a suffix is not included,
then the value is assumed to be an integer representing the size
in bytes.
:returns: The converted value in bytes as an integer
"""
value = value.lower()
if value[-2:] == 'ib':
# Assume IEC suffix.
suffix = value[-3:].lower()
else:
suffix = value[-2:].lower()
has_size_identifier = len(value) >= 2 and suffix in SIZE_SUFFIX
if not has_size_identifier:
try:
return int(value)
except ValueError:
raise ValueError(f"Invalid size value: {value}")
else:
multiplier = SIZE_SUFFIX[suffix]
return int(value[: -len(suffix)]) * multiplier
def create_file(filename, file_size):
with open(filename, 'wb') as f:
for i in range(0, file_size, KB):
f.write(b'a' * i)
def benchmark_upload(args):
source_file = args.source_file
session = get_session()
client = session.create_client('s3')
tempdir = None
try:
# If a source file was not specified, then create a temporary file of
# that size for the user.
if not source_file:
tempdir = tempfile.mkdtemp()
source_file = os.path.join(tempdir, TEMP_FILE)
create_file(source_file, args.file_size)
upload_file_script = (
f'./upload-file --file-name {source_file} --file-type {args.file_type} --s3-bucket {args.s3_bucket} '
f'--s3-key {TEMP_KEY}'
)
benchmark_args = ['./benchmark', upload_file_script]
if args.output_file:
benchmark_args.extend(['--output-file', args.output_file])
subprocess.check_call(benchmark_args)
finally:
if tempdir:
shutil.rmtree(tempdir)
client.delete_object(Bucket=args.s3_bucket, Key=TEMP_KEY)
def main():
parser = argparse.ArgumentParser(usage=__doc__)
source_file_group = parser.add_mutually_exclusive_group(required=True)
source_file_group.add_argument(
'--source-file',
help=(
'The local file to upload. Note this is optional. You can also '
'use --file-size which will create a temporary file for you.'
),
)
source_file_group.add_argument(
'--file-size',
type=human_readable_to_bytes,
help=(
'The size of the temporary file to create. You can also specify '
'your own file with --source-file'
),
)
parser.add_argument(
'--file-type',
choices=['filename', 'seekable', 'nonseekable'],
required=True,
help='The way to represent the file when uploading',
)
parser.add_argument(
'--s3-bucket',
required=True,
help='The S3 bucket to upload the file to',
)
parser.add_argument(
'--output-file',
help=(
'The file to output the data collected to. The default '
'location performance.csv'
),
)
args = parser.parse_args()
benchmark_upload(args)
if __name__ == '__main__':
main()
|
