forked from mumrah/s3-multipart
/
s3-mp-download.py
executable file
·133 lines (111 loc) · 4.41 KB
/
s3-mp-download.py
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
#!/usr/bin/env python
import argparse
import logging
from math import ceil
from multiprocessing import Pool
import os
import time
import urlparse
import boto
parser = argparse.ArgumentParser(description="Download a file from S3 in parallel",
prog="s3-mp-download")
parser.add_argument("src", help="The S3 key to download")
parser.add_argument("dest", help="The destination file")
parser.add_argument("-np", "--num-processes", help="Number of processors to use",
type=int, default=2)
parser.add_argument("-f", "--force", help="Overwrite an existing file",
action="store_true")
log = logging.getLogger("s3-mp-download")
def do_part_download(args):
"""
Download a part of an S3 object using Range header
We utilize the existing S3 GET request implemented by Boto and tack on the
Range header. We then read in 1Mb chunks of the file and write out to the
correct position in the target file
:type args: tuple of (string, string, int, int)
:param args: The actual arguments of this method. Due to lameness of
multiprocessing, we have to extract these outside of the
function definition.
The arguments are: S3 Bucket name, S3 key, local file name,
chunk size, and part number
"""
bucket_name, key_name, fname, min_byte, max_byte = args
conn = boto.connect_s3()
# Make the S3 request
resp = conn.make_request("GET", bucket=bucket_name,
key=key_name, headers={'Range':"bytes=%d-%d" % (min_byte, max_byte)})
# Open the target file, seek to byte offset
fd = os.open(fname, os.O_WRONLY)
log.info("Opening file descriptor %d, seeking to %d" % (fd, min_byte))
os.lseek(fd, min_byte, os.SEEK_SET)
chunk_size = min((max_byte-min_byte), 32*1024*1024)
log.info("Reading HTTP stream in %dM chunks" % (chunk_size/1024./1024))
t1 = time.time()
s = 0
while True:
data = resp.read(chunk_size)
if data == "":
break
os.write(fd, data)
s += len(data)
t2 = time.time() - t1
os.close(fd)
s = s / 1024 / 1024.
log.info("Downloaded %0.2fM in %0.2fs at %0.2fMbps" % (s, t2, s/t2))
def gen_byte_ranges(size, num_parts):
part_size = int(ceil(1. * size / num_parts))
for i in range(num_parts):
yield (part_size*i, min(part_size*(i+1)-1, size-1))
def main():
logging.basicConfig(level=logging.INFO)
args = parser.parse_args()
log.debug("Got args: %s" % args)
# Check that src is a valid S3 url
split_rs = urlparse.urlsplit(args.src)
if split_rs.scheme != "s3":
raise ValueError("'%s' is not an S3 url" % args.src)
# Check that dest does not exist
if os.path.exists(args.dest):
if args.force:
os.remove(args.dest)
else:
raise ValueError("Destination file '%s' exists, specify -f to"
" overwrite" % args.dest)
# Split out the bucket and the key
s3 = boto.connect_s3()
bucket = s3.lookup(split_rs.netloc)
key = bucket.get_key(split_rs.path)
# Determine the total size and calculate byte ranges
conn = boto.connect_s3()
resp = conn.make_request("HEAD", bucket=bucket, key=key)
size = int(resp.getheader("content-length"))
logging.info("Got headers: %s" % resp.getheaders())
# Skipping multipart if file is less than 1mb
if size < 1024 * 1024:
t1 = time.time()
key.get_contents_to_filename(args.dest)
t2 = time.time() - t1
log.info("Finished single-part download of %0.2fM in %0.2fs (%0.2fMbps)" %
(size, t2, size/t2))
else:
# Touch the file
fd = os.open(args.dest, os.O_CREAT)
os.close(fd)
num_parts = args.num_processes
def arg_iterator(num_parts):
for min_byte, max_byte in gen_byte_ranges(size, num_parts):
yield (bucket.name, key.name, args.dest, min_byte, max_byte)
s = size / 1024 / 1024.
try:
t1 = time.time()
pool = Pool(processes=args.num_processes)
pool.map_async(do_part_download, arg_iterator(num_parts)).get(9999999)
t2 = time.time() - t1
log.info("Finished downloading %0.2fM in %0.2fs (%0.2fMbps)" %
(s, t2, s/t2))
except KeyboardInterrupt:
log.info("User terminated")
except Exception, err:
log.error(err)
if __name__ == "__main__":
main()