未验证 提交 af819995 编写于 作者: Z zhanghan 提交者: GitHub

fix multi process download model (#662)

上级 2ed67ee5
...@@ -21,6 +21,8 @@ import logging ...@@ -21,6 +21,8 @@ import logging
from tqdm import tqdm from tqdm import tqdm
from pathlib import Path from pathlib import Path
import six import six
import paddle as P
import time
if six.PY2: if six.PY2:
from pathlib2 import Path from pathlib2 import Path
else: else:
...@@ -33,6 +35,8 @@ def _fetch_from_remote(url, ...@@ -33,6 +35,8 @@ def _fetch_from_remote(url,
force_download=False, force_download=False,
cached_dir='~/.paddle-ernie-cache'): cached_dir='~/.paddle-ernie-cache'):
import hashlib, tempfile, requests, tarfile import hashlib, tempfile, requests, tarfile
env = P.distributed.ParallelEnv()
sig = hashlib.md5(url.encode('utf8')).hexdigest() sig = hashlib.md5(url.encode('utf8')).hexdigest()
cached_dir = Path(cached_dir).expanduser() cached_dir = Path(cached_dir).expanduser()
try: try:
...@@ -40,25 +44,36 @@ def _fetch_from_remote(url, ...@@ -40,25 +44,36 @@ def _fetch_from_remote(url,
except OSError: except OSError:
pass pass
cached_dir_model = cached_dir / sig cached_dir_model = cached_dir / sig
if force_download or not cached_dir_model.exists(): done_file = cached_dir_model / 'fetch_done'
cached_dir_model.mkdir() if force_download or not done_file.exists():
tmpfile = cached_dir_model / 'tmp' if env.dev_id == 0:
with tmpfile.open('wb') as f: cached_dir_model.mkdir()
#url = 'https://ernie.bj.bcebos.com/ERNIE_stable.tgz' tmpfile = cached_dir_model / 'tmp'
r = requests.get(url, stream=True) with tmpfile.open('wb') as f:
total_len = int(r.headers.get('content-length')) #url = 'https://ernie.bj.bcebos.com/ERNIE_stable.tgz'
for chunk in tqdm( r = requests.get(url, stream=True)
r.iter_content(chunk_size=1024), total_len = int(r.headers.get('content-length'))
total=total_len // 1024, for chunk in tqdm(
desc='downloading %s' % url, r.iter_content(chunk_size=1024),
unit='KB'): total=total_len // 1024,
if chunk: desc='downloading %s' % url,
f.write(chunk) unit='KB'):
f.flush() if chunk:
log.debug('extacting... to %s' % tmpfile) f.write(chunk)
with tarfile.open(tmpfile.as_posix()) as tf: f.flush()
tf.extractall(path=cached_dir_model.as_posix()) log.debug('extacting... to %s' % tmpfile)
os.remove(tmpfile.as_posix()) with tarfile.open(tmpfile.as_posix()) as tf:
tf.extractall(path=cached_dir_model.as_posix())
os.remove(tmpfile.as_posix())
f = done_file.open('wb')
f.close()
else:
while True:
if done_file.exists():
break
else:
time.sleep(1)
log.debug('%s cached in %s' % (url, cached_dir)) log.debug('%s cached in %s' % (url, cached_dir))
return cached_dir_model return cached_dir_model
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册