dvc: support ssh deps/outs/cache

Signed-off-by: Ruslan Kuprieiev <[email protected]>

dvc/config.py

@@ -42,6 +42,7 @@ class Config(object):
         schema.Optional(SECTION_CACHE_S3): str,
         schema.Optional(SECTION_CACHE_GS): str,
         schema.Optional(SECTION_CACHE_HDFS): str,
+        schema.Optional(SECTION_CACHE_SSH): str,
 
         # backward compatibility
         schema.Optional(SECTION_CACHE_DIR, default='cache'): str,

dvc/dependency/__init__.py

@@ -13,6 +13,7 @@ from dvc.dependency.s3 import DependencyS3
 from dvc.dependency.gs import DependencyGS
 from dvc.dependency.local import DependencyLOCAL
 from dvc.dependency.hdfs import DependencyHDFS
+from dvc.dependency.ssh import DependencySSH
 
 from dvc.remote import Remote
 from dvc.remote.local import RemoteLOCAL
@@ -21,9 +22,10 @@ from dvc.remote.gs import RemoteGS
 from dvc.remote.ssh import RemoteSSH
 from dvc.remote.hdfs import RemoteHDFS
 
-DEPS = [DependencyHDFS, DependencyS3, DependencyGS, DependencyLOCAL]
+DEPS = [DependencyHDFS, DependencyS3, DependencyGS, DependencySSH, DependencyLOCAL]
 
 DEP_MAP = {'': DependencyLOCAL,
+           'ssh': DependencySSH,
            's3': DependencyS3,
            'gs': DependencyGS,
            'hdfs': DependencyHDFS,}

dvc/dependency/ssh.py

@@ -0,0 +1,47 @@
+import posixpath
+
+try:
+    from urlparse import urlparse
+except ImportError:
+    from urllib.parse import urlparse
+
+from dvc.dependency.base import DependencyBase
+from dvc.remote.ssh import RemoteSSH
+
+
+class DependencySSH(DependencyBase):
+    REGEX = RemoteSSH.REGEX
+
+    def __init__(self, stage, path, info=None, remote=None):
+        super(DependencySSH, self).__init__(stage, path)
+        self.info = info
+        self.remote = remote if remote else RemoteSSH(stage.project, {})
+
+        host = remote.host if remote else self.group('host')
+        user = remote.user if remote else self.group('user')
+        if remote:
+            path = posixpath.join(remote.prefix, urlparse(path).path.lstrip('/'))
+        else:
+            path = self.match(self.path).group('path')
+
+        self.path_info = {'scheme': 'ssh',
+                          'host': host,
+                          'user': user,
+                          'path': path}
+
+    def group(self, name):
+        match = self.match(self.path)
+        if not match:
+            return None
+        return match.group(name)
+
+    def changed(self):
+        return self.info != self.remote.save_info(self.path_info)
+
+    def save(self):
+        self.info = self.remote.save_info(self.path_info)
+
+    def dumpd(self):
+        ret = self.info
+        ret[self.PARAM_PATH] = self.path
+        return ret

dvc/output/__init__.py

@@ -9,15 +9,17 @@ from dvc.output.s3 import OutputS3
 from dvc.output.gs import OutputGS
 from dvc.output.local import OutputLOCAL
 from dvc.output.hdfs import OutputHDFS
+from dvc.output.ssh import OutputSSH
 
 from dvc.remote import Remote
 
 
-OUTS = [OutputHDFS, OutputS3, OutputGS, OutputLOCAL]
+OUTS = [OutputHDFS, OutputS3, OutputGS, OutputSSH, OutputLOCAL]
 
 OUTS_MAP = {'hdfs': OutputHDFS,
             's3': OutputS3,
             'gs': OutputGS,
+            'ssh': OutputSSH,
             '': OutputLOCAL}
 
 SCHEMA[schema.Optional(OutputLOCAL.PARAM_CACHE)] = bool

dvc/output/ssh.py

@@ -0,0 +1,46 @@
+from dvc.dependency.ssh import DependencySSH
+from dvc.exceptions import DvcException
+
+
+class OutputSSH(DependencySSH):
+    PARAM_CACHE = 'cache'
+    PARAM_METRIC = 'metric'
+
+    def __init__(self, stage, path, info=None, remote=None, cache=True, metric=False):
+        super(OutputSSH, self).__init__(stage, path, info, remote=remote)
+        self.use_cache = cache
+        self.metric = metric
+        if cache and self.project.cache.ssh == None:
+            raise DvcException("No cache location setup for \'ssh\' outputs.")
+
+    def dumpd(self):
+        ret = super(OutputSSH, self).dumpd()
+        ret[self.PARAM_CACHE] = self.use_cache
+        ret[self.PARAM_METRIC] = self.metric
+        return ret
+
+    def changed(self):
+        if super(OutputSSH, self).changed():
+            return True
+
+        if self.use_cache and self.info != self.project.cache.ssh.save_info(self.path_info):
+            return True
+
+        return False
+
+    def checkout(self):
+        if not self.use_cache:
+            return
+
+        self.project.cache.ssh.checkout(self.path_info, self.info)
+
+    def save(self):
+        super(OutputSSH, self).save()
+
+        if not self.use_cache:
+            return
+
+        self.info = self.project.cache.ssh.save(self.path_info)
+
+    def remove(self, ignore_remove=False):
+        self.remote.remove(self.path_info)

dvc/remote/ssh.py

@@ -1,5 +1,7 @@
 import os
+import re
 import errno
+import getpass
 import filecmp
 import paramiko
 import posixpath
@@ -36,19 +38,25 @@ def create_cb(name):
 class RemoteSSH(RemoteBase):
     scheme = 'ssh'
 
-    #NOTE: ~/ paths are temporarily forbidden
-    REGEX = r'^(?P<user>.*)@(?P<host>.*):(?P<path>/+.*)$'
+    #NOTE: temporarily only absolute paths are allowed
+    REGEX = r'^ssh://((?P<user>.*)@)?(?P<host>[^/]*):(?P<path>/.*)$'
+
+    PARAM_MD5 = 'md5'
 
     def __init__(self, project, config):
         self.project = project
-        storagepath = config.get(Config.SECTION_AWS_STORAGEPATH, '/')
-        self.url = config.get(Config.SECTION_REMOTE_URL, storagepath)
+        self.url = config.get(Config.SECTION_REMOTE_URL, '/')
         self.host = self.group('host')
         self.user = self.group('user')
+        if not self.user:
+            self.user = config.get(Config.SECTION_REMOTE_USER, getpass.getuser())
         self.prefix = self.group('path')
 
     def group(self, group):
-        return self.match(self.url).group(group)
+        m = self.match(self.url)
+        if not m:
+            return None
+        return m.group(group)
 
     def cache_file_key(self, path):
         relpath = os.path.relpath(os.path.abspath(path), self.project.cache.local.cache_dir)
@@ -116,6 +124,67 @@ class RemoteSSH(RemoteBase):
         self._makedirs_remote(ret)
         return ret
 
+    def md5(self, path_info):
+        if path_info['scheme'] != 'ssh':
+            raise NotImplementedError
+
+        ssh = self.ssh(host=path_info['host'], user=path_info['user'])
+        stdin, stdout, stderr = ssh.exec_command('md5sum {}'.format(path_info['path']))
+        md5 = re.match(r'^(?P<md5>.*)  .*$', stdout.read().decode('utf-8')).group('md5')
+        ssh.close()
+
+        return md5
+
+    def cp(self, from_info, to_info):
+        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'])
+        ssh.exec_command('cp {} {}'.format(from_info['path'], to_info['path']))
+        ssh.close()
+
+    def save_info(self, path_info):
+        if path_info['scheme'] != 'ssh':
+            raise NotImplementedError
+
+        return {self.PARAM_MD5: self.md5(path_info)}
+
+    def save(self, path_info):
+        if path_info['scheme'] != 'ssh':
+            raise NotImplementedError
+
+        md5 = self.md5(path_info)
+        dest = path_info.copy()
+        dest['path'] = posixpath.join(self.prefix, md5[0:2], md5[2:])
+
+        self.cp(path_info, dest)
+
+        return {self.PARAM_MD5: md5}
+
+    def checkout(self, path_info, checksum_info):
+        if path_info['scheme'] != 'ssh':
+            raise NotImplementedError
+
+        md5 = checksum_info.get(self.PARAM_MD5, None)
+        if not md5:
+            return
+
+        src = path_info.copy()
+        src['path'] = posixpath.join(self.prefix, md5[0:2], md5[2:])
+
+        self.cp(src, path_info)
+
+    def remove(self, path_info):
+        if path_info['scheme'] != 'ssh':
+            raise NotImplementedError
+
+        ssh = self.ssh(host=path_info['host'], user=path_info['user'])
+        ssh.open_sftp().remove(path_info['path'])
+        ssh.close()
+
     def download(self, path_info, path, no_progress_bar=False, name=None):
         if path_info['scheme'] != 'ssh':
             raise NotImplementedError

tests/test_data_cloud.py

@@ -5,10 +5,11 @@ import uuid
 import shutil
 import getpass
 import tempfile
+import platform
 
 from dvc.main import main
 from dvc.config import Config, ConfigError
-from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteLOCAL
+from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteLOCAL, RemoteSSH
 from dvc.remote.base import STATUS_UNKNOWN, STATUS_OK, STATUS_MODIFIED, STATUS_NEW, STATUS_DELETED
 
 from tests.basic_env import TestDvc
@@ -54,8 +55,8 @@ def _should_test_ssh():
     if os.getenv("DVC_TEST_SSH") == "true":
         return True
 
-    #FIXME: enable on windows
-    if os.name == 'nt':
+    #FIXME: enable on windows and osx
+    if platform.system() != 'Linux':
         return False
 
     assert os.system('ssh 127.0.0.1 ls &> /dev/null') == 0
@@ -72,7 +73,7 @@ def get_local_url():
 
 
 def get_ssh_url():
-    return '{}@127.0.0.1:{}'.format(getpass.getuser(), get_local_storagepath())
+    return 'ssh://{}@127.0.0.1:{}'.format(getpass.getuser(), get_local_storagepath())
 
 
 def get_aws_storagepath():
@@ -108,7 +109,10 @@ class TestDataCloud(TestDvc):
     def test(self):
         config = TEST_CONFIG
 
-        for scheme, cl in [('s3://', RemoteS3), ('gs://', RemoteGS), (tempfile.mkdtemp(), RemoteLOCAL)]:
+        for scheme, cl in [('s3://', RemoteS3),
+                           ('gs://', RemoteGS),
+                           ('ssh://user@localhost:/', RemoteSSH),
+                           (tempfile.mkdtemp(), RemoteLOCAL)]:
             config[TEST_SECTION][Config.SECTION_REMOTE_URL] = scheme + str(uuid.uuid4())
             self._test_cloud(config, cl)
 
@@ -252,6 +256,20 @@ class TestRemoteLOCAL(TestDataCloudBase):
         self.assertTrue(os.path.isdir(self.dname))
 
 
+class TestRemoteSSH(TestDataCloudBase):
+    def _should_test(self):
+        return _should_test_ssh()
+
+    def _get_url(self):
+        return get_ssh_url()
+
+    def _get_cloud_class(self):
+        return RemoteSSH
+
+    def test(self):
+        super(TestRemoteSSH, self).test()
+
+
 class TestDataCloudCLIBase(TestDvc):
     def main(self, args):
         ret = main(args)

tests/test_repro.py

@@ -3,7 +3,9 @@ import yaml
 import stat
 import shutil
 import filecmp
+import getpass
 import tempfile
+import posixpath
 from subprocess import Popen, PIPE
 
 import boto3
@@ -20,6 +22,7 @@ from dvc.stage import Stage
 from tests.basic_env import TestDvc
 from tests.test_data_cloud import _should_test_aws, TEST_AWS_REPO_BUCKET
 from tests.test_data_cloud import _should_test_gcp, TEST_GCP_REPO_BUCKET
+from tests.test_data_cloud import _should_test_ssh
 
 
 class TestRepro(TestDvc):
@@ -298,7 +301,7 @@ class TestReproExternalBase(TestDvc):
 
         cmd_stage = self.dvc.run(outs=[out_bar_path],
                              deps=[out_foo_path],
-                             cmd='{} {} {}'.format(self.cmd, foo_path, bar_path))
+                             cmd=self.cmd(foo_path, bar_path))
 
         self.write(self.bucket, foo_key, 'bar')
 
@@ -323,9 +326,8 @@ class TestReproExternalS3(TestReproExternalBase):
     def bucket(self):
         return TEST_AWS_REPO_BUCKET
 
-    @property
-    def cmd(self):
-        return 'aws s3 cp'
+    def cmd(self, i, o):
+        return 'aws s3 cp {} {}'.format(i, o)
 
     def write(self, bucket, key, body):
         s3 = boto3.resource('s3')
@@ -347,9 +349,8 @@ class TestReproExternalGS(TestReproExternalBase):
     def bucket(self):
         return TEST_GCP_REPO_BUCKET
 
-    @property
-    def cmd(self):
-        return 'gsutil cp'
+    def cmd(self, i, o):
+        return 'gsutil cp {} {}'.format(i, o)
 
     def write(self, bucket, key, body):
         client = gc.Client()
@@ -373,9 +374,8 @@ class TestReproExternalHDFS(TestReproExternalBase):
     def bucket(self):
         return os.getenv('HADOOP_CONTAINER_IP')
 
-    @property
-    def cmd(self):
-        return 'HADOOP_USER_NAME=root hadoop fs -cp'
+    def cmd(self, i, o):
+        return 'HADOOP_USER_NAME=root hadoop fs -cp {} {}'.format(i, o)
 
     def write(self, bucket, key, body):
         url = self.scheme + '://' + bucket + '/' + key
@@ -396,6 +396,57 @@ class TestReproExternalHDFS(TestReproExternalBase):
         p.communicate()
         self.assertEqual(p.returncode, 0)
 
+
+class TestReproExternalSSH(TestReproExternalBase):
+    _dir = None
+
+    def should_test(self):
+        return _should_test_ssh()
+
+    @property
+    def scheme(self):
+        return 'ssh'
+
+    @property
+    def bucket(self):
+        if not self._dir:
+            self._dir = tempfile.mkdtemp()
+        return '{}@127.0.0.1:{}'.format(getpass.getuser(), self._dir)
+
+    def cmd(self, i, o):
+        i = i.strip('ssh://')
+        o = o.strip('ssh://')
+        return 'scp {} {}'.format(i, o)
+
+    def write(self, bucket, key, body):
+        dest = '{}@127.0.0.1'.format(getpass.getuser())
+        path = posixpath.join(self._dir, key)
+        p = Popen('ssh {} rm {}'.format(dest, path),
+                  shell=True,
+                  executable=os.getenv('SHELL'),
+                  stdin=PIPE,
+                  stdout=PIPE,
+                  stderr=PIPE)
+        p.communicate()
+
+        p = Popen('ssh {} "mkdir -p $(dirname {})"'.format(dest, path),
+                  shell=True,
+                  executable=os.getenv('SHELL'),
+                  stdin=PIPE,
+                  stdout=PIPE,
+                  stderr=PIPE)
+        p.communicate()
+        self.assertEqual(p.returncode, 0)
+
+        p = Popen('echo "{}" | ssh {} "tr -d \'\\n\' > {}"'.format(body, dest, path),
+                  shell=True,
+                  executable=os.getenv('SHELL'),
+                  stdin=PIPE,
+                  stdout=PIPE,
+                  stderr=PIPE)
+        p.communicate()
+        self.assertEqual(p.returncode, 0)
+
  
 class TestReproExternalLOCAL(TestReproExternalBase):
     def setUp(self):
@@ -425,11 +476,10 @@ class TestReproExternalLOCAL(TestReproExternalBase):
     def bucket(self):
         return self.tmpdir
 
-    @property
-    def cmd(self):
+    def cmd(self, i, o):
         if os.name == 'nt':
-            return 'copy'
-        return 'cp'
+            return 'copy {} {}'.format(i, o)
+        return 'cp {} {}'.format(i, o)
 
     def write(self, bucket, key, body):
         path = os.path.join(bucket, key)