fix: move to chunk based writing

This commit is contained in:
Félix Voituret
2019-11-19 12:11:47 +01:00
parent c2f412b1d2
commit ab099f2256

View File

@@ -18,7 +18,6 @@ import tarfile
from os import environ from os import environ
from tempfile import TemporaryFile from tempfile import TemporaryFile
from shutil import copyfileobj
import requests import requests
@@ -60,14 +59,18 @@ class GithubModelProvider(ModelProvider):
self._release, self._release,
name) name)
get_logger().info('Downloading model archive %s', url) get_logger().info('Downloading model archive %s', url)
response = requests.get(url, stream=True) with requests.get(url, stream=True) as response:
if response.status_code != 200: # Note: check for error logging here or upstream ?
raise IOError(f'Resource {url} not found') response.raise_for_status()
with TemporaryFile() as stream: with TemporaryFile() as stream:
copyfileobj(response.raw, stream) # Note: check for chunk size parameters ?
get_logger().info('Extracting downloaded %s archive', name) for chunk in response.iter_content(chunk_size=8192):
stream.seek(0) if chunk:
tar = tarfile.open(fileobj=stream) stream.write(chunk)
tar.extractall(path=path) get_logger().info('Extracting downloaded %s archive', name)
tar.close() stream.seek(0)
tar = tarfile.open(fileobj=stream)
tar.extractall(path=path)
tar.close()
# TODO: perform checksum control
get_logger().info('%s model file(s) extracted', name) get_logger().info('%s model file(s) extracted', name)