From 805eb908d0340d6a7eaa27c7020e645b7fb60fec Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Mon, 16 Jul 2018 20:48:56 +0300 Subject: [PATCH] import: support external urls Fixes #891 Signed-off-by: Ruslan Kuprieiev --- dvc/command/imp.py | 2 +- dvc/dependency/base.py | 4 ++++ dvc/dependency/local.py | 4 ++-- dvc/project.py | 6 +----- dvc/remote/gs.py | 28 ++++++++++++++++++---------- dvc/remote/hdfs.py | 6 +++++- dvc/remote/s3.py | 26 ++++++++++++++++++-------- dvc/remote/ssh.py | 17 ++++++++++++----- dvc/stage.py | 12 ++++++------ tests/test_repro.py | 6 +++++- 10 files changed, 72 insertions(+), 39 deletions(-) diff --git a/dvc/command/imp.py b/dvc/command/imp.py index f5119e7be..abe497838 100644 --- a/dvc/command/imp.py +++ b/dvc/command/imp.py @@ -7,6 +7,6 @@ class CmdImport(CmdBase): try: self.project.imp(self.args.url, self.args.out) except DvcException as ex: - self.project.logger.error('Failed to import {}', ex) + self.project.logger.error('Failed to import {}'.format(self.args.url), ex) return 1 return 0 diff --git a/dvc/dependency/base.py b/dvc/dependency/base.py index 9b887e765..1fef1a63a 100644 --- a/dvc/dependency/base.py +++ b/dvc/dependency/base.py @@ -49,6 +49,10 @@ class DependencyBase(object): def sep(self): return '/' + @property + def exists(self): + return self.remote.exists([self.path_info]) + def changed(self): raise NotImplementedError diff --git a/dvc/dependency/local.py b/dvc/dependency/local.py index 47a903929..20705f75d 100644 --- a/dvc/dependency/local.py +++ b/dvc/dependency/local.py @@ -47,7 +47,7 @@ class DependencyLOCAL(DependencyBase): return os.path.relpath(self.path) def changed(self): - if not os.path.exists(self.path): + if not self.exists: return True info = self.remote.save_info(self.path_info) @@ -55,7 +55,7 @@ class DependencyLOCAL(DependencyBase): return self.info != info def save(self): - if not os.path.exists(self.path): + if not self.exists: raise self.DoesNotExistError(self.rel_path) if not os.path.isfile(self.path) and not os.path.isdir(self.path): diff --git a/dvc/project.py b/dvc/project.py index 66c0d8366..ec4945ece 100644 --- a/dvc/project.py +++ b/dvc/project.py @@ -211,14 +211,10 @@ class Project(object): return stage def imp(self, url, out): - stage_fname = out + Stage.STAGE_FILE_SUFFIX - cwd = os.path.dirname(os.path.abspath(out)) stage = Stage.loads(project=self, cmd=None, deps=[url], - outs=[out], - fname=stage_fname, - cwd=cwd) + outs=[out]) stage.run() stage.dump() diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index 79cb5d1a0..aa2b43500 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -55,18 +55,24 @@ class RemoteGS(RemoteBase): return {self.PARAM_ETAG: self.get_etag(path_info['bucket'], path_info['key'])} + def _copy(from_info, to_info, gs=None): + gs = gs if gs else self.gs + + blob = gs.bucket(from_info['bucket']).get_blob(from_info['key']) + if not blob: + raise DvcException('{} doesn\'t exist in the cloud'.format(from_info['key'])) + + self.gs.bucket(to_info['bucket']).copy_blob(blob, self.gs.bucket(to_info['bucket']), new_name=to_info['key']) + def save(self, path_info): if path_info['scheme'] != 'gs': raise NotImplementedError etag = self.get_etag(path_info['bucket'], path_info['key']) - dest_key = posixpath.join(self.prefix, etag[0:2], etag[2:]) + key = posixpath.join(self.prefix, etag[0:2], etag[2:]) + to_info = {'scheme': 'gs', 'bucket': self.bucket, 'key': key} - blob = self.gs.bucket(path_info['bucket']).get_blob(path_info['key']) - if not blob: - raise DvcException('{} doesn\'t exist in the cloud'.format(path_info['key'])) - - self.gs.bucket(self.bucket).copy_blob(blob, self.gs.bucket(path_info['bucket']), new_name=dest_key) + self._copy(path_info, to_info) return {self.PARAM_ETAG: etag} @@ -79,11 +85,9 @@ class RemoteGS(RemoteBase): return key = posixpath.join(self.prefix, etag[0:2], etag[2:]) - blob = self.gs.bucket(self.bucket).get_blob(key) - if not blob: - raise DvcException('{} doesn\'t exist in the cloud'.format(key)) + from_info = {'scheme': 'gs', 'bucket': self.bucket, 'key': key} - self.gs.bucket(path_info['bucket']).copy_blob(blob, self.gs.bucket(self.bucket), new_name=path_info['key']) + self._copy(from_info, path_info) def remove(self, path_info): if path_info['scheme'] != 'gs': @@ -156,6 +160,10 @@ class RemoteGS(RemoteBase): if from_info['scheme'] != 'gs': raise NotImplementedError + if from_info['scheme'] == 'gs': + self._copy(from_info, to_info, gs=gs) + continue + if to_info['scheme'] != 'local': raise NotImplementedError diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index 91f53cd7f..9c73c0b0e 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -57,7 +57,7 @@ class RemoteHDFS(RemoteBase): def cp(self, from_info, to_info): self.hadoop_fs('mkdir -p {}'.format(posixpath.dirname(to_info['url'])), user=to_info['user']) - self.hadoop_fs('cp {} {}'.format(from_info['url'], to_info['url']), user=to_info['user']) + self.hadoop_fs('cp -f {} {}'.format(from_info['url'], to_info['url']), user=to_info['user']) def rm(self, path_info): self.hadoop_fs('rm {}'.format(path_info['url']), user=path_info['user']) @@ -156,6 +156,10 @@ class RemoteHDFS(RemoteBase): if from_info['scheme'] != 'hdfs': raise NotImplementedError + if to_info['scheme'] == 'hdfs': + self.cp(from_info, to_info) + continue + if to_info['scheme'] != 'local': raise NotImplementedError diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 097bcb328..522d45b26 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -18,6 +18,7 @@ from dvc.progress import progress from dvc.config import Config from dvc.remote.base import RemoteBase from dvc.remote.local import RemoteLOCAL +from dvc.exceptions import DvcException class Callback(object): @@ -64,7 +65,7 @@ class RemoteS3(RemoteBase): try: obj = self.s3.Object(bucket, key).get() except Exception: - return None + raise DvcException('s3://{}/{} does not exist'.format(bucket, key)) return obj['ETag'].strip('"') @@ -74,16 +75,22 @@ class RemoteS3(RemoteBase): return {self.PARAM_ETAG: self.get_etag(path_info['bucket'], path_info['key'])} + def _copy(self, from_info, to_info, s3=None): + s3 = s3 if s3 else self.s3 + + source = {'Bucket': from_info['bucket'], + 'Key': from_info['key']} + self.s3.Bucket(to_info['bucket']).copy(source, to_info['key']) + def save(self, path_info): if path_info['scheme'] != 's3': raise NotImplementedError etag = self.get_etag(path_info['bucket'], path_info['key']) - dest_key = posixpath.join(self.prefix, etag[0:2], etag[2:]) + key = posixpath.join(self.prefix, etag[0:2], etag[2:]) + to_info = {'scheme': 's3', 'bucket': self.bucket, 'key': key} - source = {'Bucket': path_info['bucket'], - 'Key': path_info['key']} - self.s3.Bucket(self.bucket).copy(source, dest_key) + self._copy(path_info, to_info) return {self.PARAM_ETAG: etag} @@ -96,10 +103,9 @@ class RemoteS3(RemoteBase): return key = posixpath.join(self.prefix, etag[0:2], etag[2:]) - source = {'Bucket': self.bucket, - 'Key': key} + from_info = {'scheme': 's3', 'bucket': self.bucket, 'key': key} - self.s3.Bucket(path_info['bucket']).copy(source, path_info['key']) + self._copy(from_info, path_info) def remove(self, path_info): if path_info['scheme'] != 's3': @@ -197,6 +203,10 @@ class RemoteS3(RemoteBase): if from_info['scheme'] != 's3': raise NotImplementedError + if to_info['scheme'] == 's3': + self._copy(from_info, to_info, s3=s3) + continue + if to_info['scheme'] != 'local': raise NotImplementedError diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py index bf2cf1fef..c64fd57e3 100644 --- a/dvc/remote/ssh.py +++ b/dvc/remote/ssh.py @@ -122,20 +122,21 @@ class RemoteSSH(RemoteBase): return md5 - def cp(self, from_info, to_info): + def cp(self, from_info, to_info, ssh=None): if from_info['scheme'] != 'ssh' or to_info['scheme'] != 'ssh': raise NotImplementedError assert from_info['host'] == to_info['host'] assert from_info['user'] == to_info['user'] - ssh = self.ssh(host=from_info['host'], user=from_info['user']) + s = ssh if ssh else self.ssh(host=from_info['host'], user=from_info['user']) dname = posixpath.dirname(to_info['path']) - self._exec(ssh, 'mkdir -p {}'.format(dname)) - self._exec(ssh, 'cp {} {}'.format(from_info['path'], to_info['path'])) + self._exec(s, 'mkdir -p {}'.format(dname)) + self._exec(s, 'cp {} {}'.format(from_info['path'], to_info['path'])) - ssh.close() + if not ssh: + s.close() def save_info(self, path_info): if path_info['scheme'] != 'ssh': @@ -189,6 +190,12 @@ class RemoteSSH(RemoteBase): if from_info['scheme'] != 'ssh': raise NotImplementedError + if to_info['scheme'] == 'ssh': + assert from_info['host'] == to_info['host'] + assert from_info['user'] == to_info['user'] + self.cp(from_info, to_info, ssh=ssh) + continue + if to_info['scheme'] != 'local': raise NotImplementedError diff --git a/dvc/stage.py b/dvc/stage.py index 6e869cbb1..8c247273d 100644 --- a/dvc/stage.py +++ b/dvc/stage.py @@ -103,8 +103,7 @@ class Stage(object): def is_import(self): return not self.cmd and \ len(self.deps) == 1 and \ - len(self.outs) == 1 and \ - self.outs[0].path_info['scheme'] == 'local' + len(self.outs) == 1 def changed(self): ret = False @@ -280,7 +279,7 @@ class Stage(object): self.check_missing_outputs() elif self.is_import: msg = u'Importing \'{}\' -> \'{}\'' - self.project.logger.info(msg.format(self.deps[0].path, self.outs[0].rel_path)) + self.project.logger.info(msg.format(self.deps[0].path, self.outs[0].path)) self.deps[0].download(self.outs[0].path_info) elif self.is_data_source: @@ -301,9 +300,10 @@ class Stage(object): self.save() def check_missing_outputs(self): - missing_outs = [out.rel_path for out in self.outs if not os.path.exists(out.rel_path)] - if missing_outs: - raise MissingDataSource(missing_outs) + outs = [out for out in self.outs if not out.exists] + paths = [out.path if out.path_info['scheme'] != 'local' else out.rel_path for out in outs] + if paths: + raise MissingDataSource(paths) def checkout(self): for out in self.outs: diff --git a/tests/test_repro.py b/tests/test_repro.py index e88819960..5fa913723 100644 --- a/tests/test_repro.py +++ b/tests/test_repro.py @@ -337,6 +337,10 @@ class TestReproExternalBase(TestDvc): self.assertTrue(os.path.exists('import')) self.assertTrue(filecmp.cmp('import', self.FOO)) + import_remote_stage = self.dvc.imp(out_foo_path, out_foo_path + '_imported') + + add_stage = self.dvc.add(out_foo_path) + cmd_stage = self.dvc.run(outs=[out_bar_path], deps=[out_foo_path], cmd=self.cmd(foo_path, bar_path)) @@ -346,7 +350,7 @@ class TestReproExternalBase(TestDvc): sleep() stages = self.dvc.reproduce(import_stage.path) - self.assertEqual(len(stages), 1) + self.assertEqual(len(stages), 2) self.assertTrue(os.path.exists('import')) self.assertTrue(filecmp.cmp('import', self.BAR))