Merge pull request #482 from nbcsm/retry

support retry for url get

Author: guschmue

setup.py

@@ -35,7 +35,7 @@ def initialize_options(self):
 
     def finalize_options(self):
         pass
-        
+
     def run(self):
         with open(os.path.join(SRC_DIR, 'version.py'), 'w') as f:
             f.write(dedent('''
@@ -74,11 +74,11 @@ def run(self):
     version=VersionInfo.version,
     description='Tensorflow to ONNX converter',
     setup_requires=['pytest-runner'],
-    tests_require=['graphviz', 'requests', 'parameterized', 'pytest', 'pytest-cov', 'pyyaml'],
+    tests_require=['graphviz', 'parameterized', 'pytest', 'pytest-cov', 'pyyaml'],
     cmdclass=cmdclass,
     packages=find_packages(),
     author='[email protected]',
     author_email='[email protected]',
     url='https://github.com/onnx/tensorflow-onnx',
-    install_requires=['numpy>=1.14.1', 'onnx>=1.4.1', 'six']
+    install_requires=['numpy>=1.14.1', 'onnx>=1.4.1', 'requests', 'six']
 )

tests/run_pretrained_models.py

@@ -17,7 +17,6 @@
 
 import PIL.Image
 import numpy as np
-import requests
 import six
 import tensorflow as tf
 # contrib ops are registered only when the module is imported, the following import statement is needed,
@@ -119,11 +118,7 @@ def download_file(self):
         os.makedirs(dir_name, exist_ok=True)
         fpath = os.path.join(dir_name, fname)
         if not os.path.exists(fpath):
-            response = requests.get(url)
-            if response.status_code not in [200]:
-                response.raise_for_status()
-            with open(fpath, "wb") as f:
-                f.write(response.content)
+            utils.get_url(url, fpath)
         model_path = os.path.join(dir_name, self.local)
         if not os.path.exists(model_path):
             if ftype == 'tgz':

tf2onnx/utils.py

@@ -13,13 +13,18 @@
 import re
 import shutil
 import tempfile
+
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 import six
 import numpy as np
 import tensorflow as tf
 from tensorflow.core.framework import types_pb2, tensor_pb2
 from google.protobuf import text_format
 import onnx
 from onnx import helper, onnx_pb, defs, numpy_helper
+
 from . import constants
 
 #
@@ -468,3 +473,23 @@ def set_debug_mode(enabled):
 
 def get_max_value(np_dtype):
     return np.iinfo(np_dtype).max
+
+
+def get_url(url, path, max_retries=5):
+    """ Download url and save to path. """
+    retries = Retry(total=max_retries, backoff_factor=0.1, status_forcelist=[500, 502, 503, 504])
+    adapter = HTTPAdapter(max_retries=retries)
+    session = requests.Session()
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+
+    response = session.get(url, allow_redirects=True)
+    if response.status_code not in [200]:
+        response.raise_for_status()
+
+    dir_name = os.path.dirname(path)
+    if dir_name:
+        os.makedirs(dir_name, exist_ok=True)
+
+    with open(path, "wb") as f:
+        f.write(response.content)