Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #2289 from Suor/exists-jobs

Browse Source 
remote: honor --jobs for status collection and other fixes
This commit is contained in:
Ruslan Kuprieiev
2019-07-21 19:24:30 +03:00
committed by GitHub
parent b6bc65fcbf 1f90b426fb
commit a1143c532c
6 changed files with 50 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/remote/base.py
+2 -2
View File
@@ -589,7 +589,7 @@ class RemoteBASE(object):
return self._changed_dir_cache(checksum)
return self.changed_cache_file(checksum)
def cache_exists(self, checksums):
def cache_exists(self, checksums, jobs=None):
"""Check if the given checksums are stored in the remote.
There are two ways of performing this check:
@@ -618,7 +618,7 @@ class RemoteBASE(object):
return self.batch_exists(chunks, callback=progress_callback)
if self.no_traverse and hasattr(self, "batch_exists"):
with ThreadPoolExecutor(max_workers=self.JOBS) as executor:
with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor:
path_infos = [self.checksum_to_path_info(x) for x in checksums]
chunks = to_chunks(path_infos, num_chunks=self.JOBS)
results = executor.map(exists_with_progress, chunks)
dvc/remote/local/__init__.py
+3 -3
View File
@@ -241,7 +241,7 @@ class RemoteLOCAL(RemoteBASE):
move(inp, tmp)
move(tmp, outp)
def cache_exists(self, md5s):
def cache_exists(self, md5s, jobs=None):
return [
checksum
for checksum in progress(md5s)
@@ -306,7 +306,7 @@ class RemoteLOCAL(RemoteBASE):
md5s = list(ret)
logger.info("Collecting information from local cache...")
local_exists = self.cache_exists(md5s)
local_exists = self.cache_exists(md5s, jobs=jobs)
# This is a performance optimization. We can safely assume that,
# if the resources that we want to fetch are already cached,
@@ -316,7 +316,7 @@ class RemoteLOCAL(RemoteBASE):
remote_exists = local_exists
else:
logger.info("Collecting information from remote cache...")
remote_exists = list(remote.cache_exists(md5s))
remote_exists = list(remote.cache_exists(md5s, jobs=jobs))
self._fill_statuses(ret, local_exists, remote_exists)
dvc/remote/pool.py
+4 -4
View File
@@ -4,8 +4,8 @@ from funcy import memoize
@contextmanager
def get_connection(*args, **kwargs):
pool = get_pool(*args, **kwargs)
def get_connection(conn_func, *args, **kwargs):
pool = get_pool(conn_func, *args, **kwargs)
conn = pool.get_connection()
try:
yield conn
@@ -17,8 +17,8 @@ def get_connection(*args, **kwargs):
@memoize
def get_pool(*args, **kwargs):
return Pool(*args, **kwargs)
def get_pool(conn_func, *args, **kwargs):
return Pool(conn_func, *args, **kwargs)
def close_pools():
dvc/remote/ssh/__init__.py
+8 -8
View File
@@ -153,15 +153,15 @@ class RemoteSSH(RemoteBASE):
return ret
with self.ssh(path_infos[0]) as ssh:
with ssh.open_max_sftp_channels() as channels:
max_workers = len(channels)
channels = ssh.open_max_sftp_channels()
max_workers = len(channels)
with ThreadPoolExecutor(max_workers=max_workers) as executor:
paths = [path_info.path for path_info in path_infos]
chunks = to_chunks(paths, num_chunks=max_workers)
chunks_and_channels = zip(chunks, channels)
outcome = executor.map(_exists, chunks_and_channels)
results = list(itertools.chain.from_iterable(outcome))
with ThreadPoolExecutor(max_workers=max_workers) as executor:
paths = [path_info.path for path_info in path_infos]
chunks = to_chunks(paths, num_chunks=max_workers)
chunks_and_channels = zip(chunks, channels)
outcome = executor.map(_exists, chunks_and_channels)
results = list(itertools.chain.from_iterable(outcome))
return results
dvc/remote/ssh/connection.py
+25 -46
View File
@@ -3,7 +3,6 @@ import posixpath
import logging
import errno
import stat
from contextlib import contextmanager
from funcy import cached_property
try:
@@ -61,26 +60,22 @@ class SSHConnection:
self._ssh.connect(host, *args, **kwargs)
self._ssh.get_transport().set_keepalive(10)
self._sftp = None
self._sftp_alive = False
self._sftp_channels = []
def _sftp_connect(self):
if not self._sftp or not self._sftp_alive:
self._sftp = self._ssh.open_sftp()
self._sftp_alive = True
@property
def sftp(self):
if not self._sftp_channels:
self._sftp_channels = [self._ssh.open_sftp()]
return self._sftp_channels[0]
def close(self):
if self._sftp:
self._sftp.close()
self._sftp_alive = False
for sftp in self._sftp_channels:
sftp.close()
self._ssh.close()
def st_mode(self, path):
self._sftp_connect()
with ignore_file_not_found():
return self._sftp.stat(path).st_mode
return self.sftp.stat(path).st_mode
return 0
@@ -97,8 +92,6 @@ class SSHConnection:
return stat.S_ISLNK(self.st_mode(path))
def makedirs(self, path):
self._sftp_connect()
# Single stat call will say whether this is a dir, a file or a link
st_mode = self.st_mode(path)
@@ -117,7 +110,7 @@ class SSHConnection:
if tail:
try:
self._sftp.mkdir(path)
self.sftp.mkdir(path)
except IOError as e:
# Since paramiko errors are very vague we need to recheck
# whether it's because path already exists or something else
@@ -129,11 +122,8 @@ class SSHConnection:
# used as a template.
#
# [1] https://github.com/python/cpython/blob/master/Lib/os.py
self._sftp_connect()
try:
dir_entries = self._sftp.listdir_attr(directory)
dir_entries = self.sftp.listdir_attr(directory)
except IOError as exc:
raise DvcException(
"couldn't get the '{}' remote directory files list".format(
@@ -169,7 +159,7 @@ class SSHConnection:
def _remove_file(self, path):
with ignore_file_not_found():
self._sftp.remove(path)
self.sftp.remove(path)
def _remove_dir(self, path):
for root, dirs, files in self.walk(path, topdown=False):
@@ -181,52 +171,45 @@ class SSHConnection:
for dname in dirs:
path = posixpath.join(root, dname)
with ignore_file_not_found():
self._sftp.rmdir(dname)
self.sftp.rmdir(dname)
with ignore_file_not_found():
self._sftp.rmdir(path)
self.sftp.rmdir(path)
def remove(self, path):
self._sftp_connect()
if self.isdir(path):
self._remove_dir(path)
else:
self._remove_file(path)
def download(self, src, dest, no_progress_bar=False, progress_title=None):
self._sftp_connect()
if no_progress_bar:
self._sftp.get(src, dest)
self.sftp.get(src, dest)
else:
if not progress_title:
progress_title = os.path.basename(src)
self._sftp.get(src, dest, callback=create_cb(progress_title))
self.sftp.get(src, dest, callback=create_cb(progress_title))
progress.finish_target(progress_title)
def move(self, src, dst):
self.makedirs(posixpath.dirname(dst))
self._sftp_connect()
self._sftp.rename(src, dst)
self.sftp.rename(src, dst)
def upload(self, src, dest, no_progress_bar=False, progress_title=None):
self._sftp_connect()
self.makedirs(posixpath.dirname(dest))
tmp_file = tmp_fname(dest)
if no_progress_bar:
self._sftp.put(src, tmp_file)
self.sftp.put(src, tmp_file)
else:
if not progress_title:
progress_title = posixpath.basename(dest)
self._sftp.put(src, tmp_file, callback=create_cb(progress_title))
self.sftp.put(src, tmp_file, callback=create_cb(progress_title))
progress.finish_target(progress_title)
self._sftp.rename(tmp_file, dest)
self.sftp.rename(tmp_file, dest)
def execute(self, cmd):
stdin, stdout, stderr = self._ssh.exec_command(cmd)
@@ -307,18 +290,14 @@ class SSHConnection:
self.makedirs(posixpath.dirname(dest))
self.execute("cp {} {}".format(src, dest))
@contextmanager
def open_max_sftp_channels(self):
try:
channels = []
# If there are more than 1 it means we've already opened max amount
if len(self._sftp_channels) <= 1:
while True:
try:
channels.append(self._ssh.open_sftp())
self._sftp_channels.append(self._ssh.open_sftp())
except paramiko.ssh_exception.ChannelException:
if not channels:
if not self._sftp_channels:
raise
break
yield channels
finally:
for channel in channels:
channel.close()
return self._sftp_channels
tests/conftest.py
+8
View File
@@ -169,3 +169,11 @@ def erepo(repo_dir):
yield repo
finally:
repo.tearDown()
@pytest.fixture(scope="session", autouse=True)
def _close_pools():
from dvc.remote.pool import close_pools
yield
close_pools()