Make Gluon download function to be atomic

python/mxnet/gluon/utils.py

@@ -22,7 +22,9 @@
            'check_sha1', 'download']
 
 import os
+import sys
 import hashlib
+import uuid
 import warnings
 import collections
 import weakref
@@ -195,6 +197,39 @@ def check_sha1(filename, sha1_hash):
     return sha1.hexdigest() == sha1_hash
 
 
+if sys.platform != 'win32':
+    # refer to https://github.com/untitaker/python-atomicwrites
+    def _replace_atomic(src, dst):
+        """Implement atomic os.replace with linux and OSX. Internal Usa only"""
+        os.rename(src, dst)
+else:
+    from ctypes import windll, WinError
+
+    _MOVEFILE_REPLACE_EXISTING = 0x1
+    _MOVEFILE_WRITE_THROUGH = 0x8
+    _windows_default_flags = _MOVEFILE_WRITE_THROUGH
+
+    text_type = unicode if sys.version_info[0] == 2 else str  # noqa
+
+    def _path_to_unicode(x):
+        """Handle text decoding. Internal use only"""
+        if not isinstance(x, text_type):
+            return x.decode(sys.getfilesystemencoding())
+        return x
+
+    def _handle_errors(rv):
+        """Handle WinError. Internal use only"""
+        if not rv:
+            raise WinError()
+
+    def _replace_atomic(src, dst):
+        """Implement atomic os.replace with windows. Internal use only"""
+        _handle_errors(windll.kernel32.MoveFileExW(
+            _path_to_unicode(src), _path_to_unicode(dst),
+            _windows_default_flags | _MOVEFILE_REPLACE_EXISTING
+        ))
+
+
 def download(url, path=None, overwrite=False, sha1_hash=None, retries=5, verify_ssl=True):
     """Download an given URL
 
@@ -231,7 +266,7 @@ def download(url, path=None, overwrite=False, sha1_hash=None, retries=5, verify_
             fname = os.path.join(path, url.split('/')[-1])
         else:
             fname = path
-    assert retries >= 0, "Number of retries should be at least 0"
+    assert retries >= 0, 'Number of retries should be at least 0'
 
     if not verify_ssl:
         warnings.warn(
@@ -242,23 +277,34 @@ def download(url, path=None, overwrite=False, sha1_hash=None, retries=5, verify_
         dirname = os.path.dirname(os.path.abspath(os.path.expanduser(fname)))
         if not os.path.exists(dirname):
             os.makedirs(dirname)
-        while retries+1 > 0:
+        while retries + 1 > 0:
             # Disable pyling too broad Exception
             # pylint: disable=W0703
             try:
-                print('Downloading %s from %s...'%(fname, url))
+                print('Downloading {} from {}...'.format(fname, url))
                 r = requests.get(url, stream=True, verify=verify_ssl)
                 if r.status_code != 200:
-                    raise RuntimeError("Failed downloading url %s"%url)
-                with open(fname, 'wb') as f:
+                    raise RuntimeError('Failed downloading url {}'.format(url))
+                # create uuid for temporary files
+                random_uuid = str(uuid.uuid4())
+                with open('{}.{}'.format(fname, random_uuid), 'wb') as f:
+                    # create uuid for temporary files
                     for chunk in r.iter_content(chunk_size=1024):
                         if chunk: # filter out keep-alive new chunks
                             f.write(chunk)
+                # if the target file exists(created by other processes),
+                # delete the temporary file
+                if os.path.exists(fname):
+                    os.remove('{}.{}'.format(fname, random_uuid))
+                else:
+                    # atmoic operation in the same file system
+                    _replace_atomic('{}.{}'.format(fname, random_uuid), fname)
                 if sha1_hash and not check_sha1(fname, sha1_hash):
-                    raise UserWarning('File {} is downloaded but the content hash does not match.'\
-                                      ' The repo may be outdated or download may be incomplete. '\
-                                      'If the "repo_url" is overridden, consider switching to '\
-                                      'the default repo.'.format(fname))
+                    raise UserWarning(
+                        'File {} is downloaded but the content hash does not match.'
+                        ' The repo may be outdated or download may be incomplete. '
+                        'If the "repo_url" is overridden, consider switching to '
+                        'the default repo.'.format(fname))
                 break
             except Exception as e:
                 retries -= 1

tests/python/unittest/test_gluon_utils.py

@@ -19,6 +19,9 @@
 import os
 import tempfile
 import warnings
+import glob
+import shutil
+import multiprocessing as mp
 
 try:
     from unittest import mock
@@ -46,22 +49,53 @@ def test_download_retries():
 
 @mock.patch(
     'requests.get',
-    mock.Mock(side_effect=
-              lambda *args, **kwargs: MockResponse(200, 'MOCK CONTENT' * 100)))
+    mock.Mock(side_effect=lambda *args, **kwargs: MockResponse(200, 'MOCK CONTENT' * 100)))
+def _download_successful(tmp):
+    """ internal use for testing download successfully """
+    mx.gluon.utils.download(
+        "https://raw.githubusercontent.com/apache/incubator-mxnet/master/README.md",
+        path=tmp)
+
+
 def test_download_successful():
+    """ test download with one process """
     tmp = tempfile.mkdtemp()
     tmpfile = os.path.join(tmp, 'README.md')
-    mx.gluon.utils.download(
-        "https://raw.githubusercontent.com/apache/incubator-mxnet/master/README.md",
-        path=tmpfile)
-    assert os.path.getsize(tmpfile) > 100
+    _download_successful(tmpfile)
+    assert os.path.getsize(tmpfile) > 100, os.path.getsize(tmpfile)
+    pattern = os.path.join(tmp, 'README.md*')
+    # check only one file we want left
+    assert len(glob.glob(pattern)) == 1, glob.glob(pattern)
+    # delete temp dir
+    shutil.rmtree(tmp)
+
+
+def test_multiprocessing_download_successful():
+    """ test download with multiprocessing """
+    tmp = tempfile.mkdtemp()
+    tmpfile = os.path.join(tmp, 'README.md')
+    process_list = []
+    # test it with 10 processes
+    for i in range(10):
+        process_list.append(mp.Process(
+            target=_download_successful, args=(tmpfile,)))
+        process_list[i].start()
+    for i in range(10):
+        process_list[i].join()
+    assert os.path.getsize(tmpfile) > 100, os.path.getsize(tmpfile)
+    # check only one file we want left
+    pattern = os.path.join(tmp, 'README.md*')
+    assert len(glob.glob(pattern)) == 1, glob.glob(pattern)
+    # delete temp dir
+    shutil.rmtree(tmp)
 
 
 @mock.patch(
     'requests.get',
     mock.Mock(
         side_effect=lambda *args, **kwargs: MockResponse(200, 'MOCK CONTENT')))
 def test_download_ssl_verify():
+    """ test download verify_ssl parameter """
     with warnings.catch_warnings(record=True) as warnings_:
         mx.gluon.utils.download(
             "https://mxnet.incubator.apache.org/index.html", verify_ssl=False)