diff --git a/dvc/command/imp.py b/dvc/command/imp.py index 1e125ecd0..a0292bb05 100644 --- a/dvc/command/imp.py +++ b/dvc/command/imp.py @@ -15,7 +15,7 @@ class CmdImport(CmdBase): out = self.args.out or default_out - self.project.imp(self.args.url, out) + self.project.imp(self.args.url, out, self.args.resume) except DvcException: logger.error( "failed to import {}. You could also try downloading " @@ -48,5 +48,11 @@ def add_parser(subparsers, parent_parser): "'remote://myremote/path/to/file'(see " "`dvc remote` commands). ", ) + import_parser.add_argument( + "--resume", + action="store_true", + default=False, + help="Resume previously started download.", + ) import_parser.add_argument("out", nargs="?", help="Output.") import_parser.set_defaults(func=CmdImport) diff --git a/dvc/output/base.py b/dvc/output/base.py index 6303507e4..71c9e7e7a 100644 --- a/dvc/output/base.py +++ b/dvc/output/base.py @@ -164,8 +164,8 @@ class OutputBase(object): return ret - def download(self, to_info): - self.remote.download([self.path_info], [to_info]) + def download(self, to_info, resume=False): + self.remote.download([self.path_info], [to_info], resume=resume) def checkout(self, force=False): if not self.use_cache: diff --git a/dvc/project.py b/dvc/project.py index 5a7c3b4aa..2ea02a527 100644 --- a/dvc/project.py +++ b/dvc/project.py @@ -413,7 +413,7 @@ class Project(object): return stage - def imp(self, url, out): + def imp(self, url, out, resume=False): from dvc.stage import Stage stage = Stage.create(project=self, cmd=None, deps=[url], outs=[out]) @@ -425,7 +425,7 @@ class Project(object): self.files_to_git_add = [] with self.state: - stage.run() + stage.run(resume=resume) stage.dump() diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index 16e69a7f8..077f0f742 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -145,7 +145,12 @@ class RemoteAzure(RemoteBase): progress.finish_target(name) def download( - self, from_infos, to_infos, no_progress_bar=False, names=None + self, + from_infos, + to_infos, + no_progress_bar=False, + names=None, + resume=False, ): names = self._verify_path_args(from_infos, to_infos, names) diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 8af5cc3fd..d515defec 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -167,7 +167,14 @@ class RemoteBase(object): def save(self, path_info): raise NotImplementedError - def download(self, from_infos, to_infos, no_progress_bar=False, name=None): + def download( + self, + from_infos, + to_infos, + no_progress_bar=False, + name=None, + resume=False, + ): raise NotImplementedError def upload(self, from_infos, to_infos, path_info, names=None): diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index 756829848..0752d782c 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -162,7 +162,12 @@ class RemoteGS(RemoteBase): progress.finish_target(name) def download( - self, from_infos, to_infos, no_progress_bar=False, names=None + self, + from_infos, + to_infos, + no_progress_bar=False, + names=None, + resume=False, ): names = self._verify_path_args(from_infos, to_infos, names) diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index f62eede8a..38bda2e72 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -145,7 +145,12 @@ class RemoteHDFS(RemoteBase): self.hadoop_fs(cmd, user=to_info["user"]) def download( - self, from_infos, to_infos, no_progress_bar=False, names=None + self, + from_infos, + to_infos, + no_progress_bar=False, + names=None, + resume=False, ): names = self._verify_path_args(from_infos, to_infos, names) diff --git a/dvc/remote/http.py b/dvc/remote/http.py index 7444ca87e..2b27d4f39 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -28,7 +28,7 @@ class RemoteHTTP(RemoteBase): scheme = "http" REGEX = r"^https?://.*$" REQUEST_TIMEOUT = 10 - CHUNK_SIZE = 1000000 # Megabyte + CHUNK_SIZE = 2 ** 16 PARAM_CHECKSUM = "etag" def __init__(self, project, config): @@ -42,7 +42,12 @@ class RemoteHTTP(RemoteBase): return self.cache_dir def download( - self, from_infos, to_infos, no_progress_bar=False, names=None + self, + from_infos, + to_infos, + no_progress_bar=False, + names=None, + resume=False, ): names = self._verify_path_args(to_infos, from_infos, names) @@ -58,7 +63,6 @@ class RemoteHTTP(RemoteBase): ) logger.debug(msg) - tmp_file = self.tmp_file(to_info["path"]) if not name: name = os.path.basename(to_info["path"]) @@ -72,14 +76,18 @@ class RemoteHTTP(RemoteBase): cb = ProgressBarCallback(name, total) try: - self._download_to(from_info["path"], tmp_file, callback=cb) + self._download_to( + from_info["path"], + to_info["path"], + callback=cb, + resume=resume, + ) + except Exception: msg = "failed to download '{}'".format(from_info["path"]) logger.error(msg) continue - os.rename(tmp_file, to_info["path"]) - if not no_progress_bar: progress.finish_target(name) @@ -125,18 +133,54 @@ class RemoteHTTP(RemoteBase): return etag - def _download_to(self, url, file, callback=None): - r = self._request("GET", url, stream=True) + def _download_to(self, url, target_file, callback=None, resume=False): + request = self._request("GET", url, stream=True) + partial_file = target_file + ".part" - with open(file, "wb") as fd: - bytes_transfered = 0 + mode, transferred_bytes = self._determine_mode_get_transferred_bytes( + partial_file, resume + ) - for chunk in r.iter_content(chunk_size=self.CHUNK_SIZE): - fd.write(chunk) + self._validate_existing_file_size(transferred_bytes, partial_file) + + self._write_request_content( + mode, partial_file, request, transferred_bytes, callback + ) + + os.rename(partial_file, target_file) + + def _write_request_content( + self, mode, partial_file, request, transferred_bytes, callback=None + ): + with open(partial_file, mode) as fd: + + for index, chunk in enumerate( + request.iter_content(chunk_size=self.CHUNK_SIZE) + ): + chunk_number = index + 1 + if chunk_number * self.CHUNK_SIZE > transferred_bytes: + fd.write(chunk) + fd.flush() + transferred_bytes += len(chunk) if callback: - bytes_transfered += len(chunk) - callback(bytes_transfered) + callback(transferred_bytes) + + def _validate_existing_file_size(self, bytes_transferred, partial_file): + if bytes_transferred % self.CHUNK_SIZE != 0: + raise DvcException( + "File {}, might be corrupted, please remove " + "it and retry importing".format(partial_file) + ) + + def _determine_mode_get_transferred_bytes(self, partial_file, resume): + if os.path.exists(partial_file) and resume: + mode = "ab" + bytes_transfered = os.path.getsize(partial_file) + else: + mode = "wb" + bytes_transfered = 0 + return mode, bytes_transfered def _request(self, method, url, **kwargs): kwargs.setdefault("allow_redirects", True) diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 8ee326fcd..21127ae21 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -487,7 +487,12 @@ class RemoteLOCAL(RemoteBase): ) def download( - self, from_infos, to_infos, no_progress_bar=False, names=None + self, + from_infos, + to_infos, + no_progress_bar=False, + names=None, + resume=False, ): names = self._verify_path_args(from_infos, to_infos, names) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index fc0f88325..0c05c047a 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -207,7 +207,12 @@ class RemoteS3(RemoteBase): progress.finish_target(name) def download( - self, from_infos, to_infos, no_progress_bar=False, names=None + self, + from_infos, + to_infos, + no_progress_bar=False, + names=None, + resume=False, ): names = self._verify_path_args(from_infos, to_infos, names) diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py index f49199536..2fff01ddc 100644 --- a/dvc/remote/ssh.py +++ b/dvc/remote/ssh.py @@ -279,7 +279,12 @@ class RemoteSSH(RemoteBase): ssh.close() def download( - self, from_infos, to_infos, no_progress_bar=False, names=None + self, + from_infos, + to_infos, + no_progress_bar=False, + names=None, + resume=False, ): names = self._verify_path_args(from_infos, to_infos, names) diff --git a/dvc/stage.py b/dvc/stage.py index 5f7c9cdf8..2400b5b19 100644 --- a/dvc/stage.py +++ b/dvc/stage.py @@ -564,7 +564,7 @@ class Stage(object): if p.returncode != 0: raise StageCmdFailedError(self) - def run(self, dry=False): + def run(self, dry=False, resume=False): if self.locked: logger.info( "Verifying outputs in locked stage '{stage}'".format( @@ -584,7 +584,9 @@ class Stage(object): if self._already_cached(): self.outs[0].checkout() else: - self.deps[0].download(self.outs[0].path_info) + self.deps[0].download( + self.outs[0].path_info, resume=resume + ) elif self.is_data_source: msg = "Verifying data sources in '{}'".format(self.relpath) diff --git a/tests/test_import.py b/tests/test_import.py index 289631402..5fc3c2799 100644 --- a/tests/test_import.py +++ b/tests/test_import.py @@ -4,10 +4,12 @@ import os from uuid import uuid4 from dvc import logger +from dvc.utils.compat import urljoin from dvc.exceptions import DvcException from dvc.main import main -from mock import patch +from mock import patch, mock_open, call from tests.basic_env import TestDvc +from tests.utils.httpd import StaticFileServer from dvc.utils.compat import StringIO @@ -66,3 +68,72 @@ class TestFailedImportMessage(TestDvc): "command.", logger.logger.handlers[1].stream.getvalue(), ) + + +class TestInterruptedDownload(TestDvc): + @property + def remote(self): + return "http://localhost:8000/" + + def _prepare_interrupted_download(self): + import_url = urljoin(self.remote, self.FOO) + import_output = "imported_file" + tmp_file_name = import_output + ".part" + tmp_file_path = os.path.join(self._root_dir, tmp_file_name) + self._import_with_interrupt(import_output, import_url) + self.assertTrue(os.path.exists(tmp_file_name)) + self.assertFalse(os.path.exists(import_output)) + return import_output, import_url, tmp_file_path + + def _import_with_interrupt(self, import_output, import_url): + def interrupting_generator(): + yield self.FOO[0].encode("utf8") + raise KeyboardInterrupt + + with patch( + "requests.models.Response.iter_content", + return_value=interrupting_generator(), + ): + with patch( + "dvc.remote.http.RemoteHTTP._content_length", return_value=3 + ): + result = main(["import", import_url, import_output]) + self.assertEqual(result, 252) + + +class TestShouldResumeDownload(TestInterruptedDownload): + @patch("dvc.remote.http.RemoteHTTP.CHUNK_SIZE", 1) + def test(self): + with StaticFileServer(): + import_output, import_url, tmp_file_path = ( + self._prepare_interrupted_download() + ) + + m = mock_open() + with patch("dvc.remote.http.open", m): + result = main( + ["import", "--resume", import_url, import_output] + ) + self.assertEqual(result, 0) + m.assert_called_once_with(tmp_file_path, "ab") + m_handle = m() + expected_calls = [call(b"o"), call(b"o")] + m_handle.write.assert_has_calls(expected_calls, any_order=False) + + +class TestShouldNotResumeDownload(TestInterruptedDownload): + @patch("dvc.remote.http.RemoteHTTP.CHUNK_SIZE", 1) + def test(self): + with StaticFileServer(): + import_output, import_url, tmp_file_path = ( + self._prepare_interrupted_download() + ) + + m = mock_open() + with patch("dvc.remote.http.open", m): + result = main(["import", import_url, import_output]) + self.assertEqual(result, 0) + m.assert_called_once_with(tmp_file_path, "wb") + m_handle = m() + expected_calls = [call(b"f"), call(b"o"), call(b"o")] + m_handle.write.assert_has_calls(expected_calls, any_order=False) diff --git a/tests/utils/httpd.py b/tests/utils/httpd.py index 4d11d6913..b92bc5768 100644 --- a/tests/utils/httpd.py +++ b/tests/utils/httpd.py @@ -1,5 +1,6 @@ import hashlib import os +import socket import threading from dvc.utils.compat import HTTPServer, SimpleHTTPRequestHandler @@ -30,13 +31,26 @@ class ContentMD5Handler(SimpleHTTPRequestHandler): class StaticFileServer: - __server_lock = threading.Lock() + _server_lock = threading.Lock() def __init__(self, handler="etag"): + self._server_lock.acquire() handler_class = ETagHandler if handler == "etag" else ContentMD5Handler + self.bind_port_by_any_means(handler_class) - self.__server_lock.acquire() - self.httpd = HTTPServer(("localhost", 8000), handler_class) + def bind_port_by_any_means(self, handler_class): + import time + + # shutdowning/closing socket/server does not unbind port in time, + # locking the server also does not bring results, hence + # this method + for i in range(10000): + try: + self.httpd = HTTPServer(("localhost", 8000), handler_class) + except Exception: + time.sleep(0.01) + continue + break def __enter__(self): self.server_thread = threading.Thread(target=self.httpd.serve_forever) @@ -44,6 +58,7 @@ class StaticFileServer: self.server_thread.start() def __exit__(self, *args): + self.httpd.socket.close() self.httpd.shutdown() self.httpd.server_close() - self.__server_lock.release() + self._server_lock.release()