Commit 24725ef5 authored by Dorota Czaplejewicz's avatar Dorota Czaplejewicz

Devkit: Resume image download on errors

parent ff65b589
......@@ -12,6 +12,7 @@ import shutil
import subprocess
import sys
import tempfile
import time
import tqdm
import yaml
......@@ -51,12 +52,17 @@ FB: Done
'''
BLOCK_SIZE = 8192
UNCOMPRESSED_SIZE = 3600000000
MAX_ATTEMPTS = 10
class VerifyImageException(Exception):
pass
class PrematureEndException(Exception):
pass
def verify_image(image, meta):
m = hashlib.sha256()
size = int(meta['image']['size'])
......@@ -80,6 +86,42 @@ def verify_image(image, meta):
"does not match {}".format(m.hexdigest(), hexdigest))
def resuming_stream(url, expected_size):
position = 0
for i in range(MAX_ATTEMPTS):
try:
resp = requests.get(url,
stream=True,
headers={'Range': 'bytes={}-'.format(position)})
resp.raise_for_status()
if resp.status_code != requests.codes.partial_content:
position = 0
logging.info('Proceeding from {} bytes'.format(position))
for data in resp.iter_content(BLOCK_SIZE):
position += len(data)
yield data
if position < expected_size:
raise PrematureEndException()
return
except (requests.exceptions.ConnectionError, PrematureEndException):
if i < MAX_ATTEMPTS - 1:
logging.info("Connection error, retrying")
time.sleep(5)
else:
logging.error("Max connection errors reached, aborting")
raise
def stream_file(url):
resp = requests.head(url, stream=True)
resp.raise_for_status()
ts = int(resp.headers.get('content-length', 0))
return resuming_stream(url, ts), ts
def download_image(url, target):
decomp = lzma.LZMADecompressor()
logging.info("Downloading image from {}".format(url))
......@@ -98,9 +140,7 @@ def download_image(url, target):
meta = None
uncompressed_size = UNCOMPRESSED_SIZE
resp = requests.get(url, stream=True)
resp.raise_for_status()
ts = int(resp.headers.get('content-length', 0))
stream, ts = stream_file(url)
download_bar = tqdm.tqdm(total=ts,
desc='Download',
leave=False)
......@@ -108,7 +148,7 @@ def download_image(url, target):
desc='Decompr.',
leave=False)
with open(target, 'wb+') as f:
for data in resp.iter_content(BLOCK_SIZE):
for data in stream:
if data:
out = decomp.decompress(data)
decompress_bar.update(len(out))
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment