mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-08 15:28:26 +01:00
Peter Rowlands (변기호)
71bae7e4af
* fs: use fsspecwrapper in memfs * odb: add memfs staging ODB in ODBManager * odb: add read_only attribute for odb instances * objects: handle memfs paths when staging trees * odb: remove git ODB * output: use staging ODB * RepoDependency: use staging odb for imports * dvcfs: handle case where dir cache is unavailable * fix circular import check * odb: use temp local ODB for staging * objects: use staging ODB in stage() and tree load * objects: add typing for add/check/load/stage * output/dep: use new staging * add unit tests for local odb staging * localfs: account for string paths in move() * include local in odb.by_scheme * tests: count dir cache like any other object in test_pull_git_imports * odb: clear staging on gc() * objects: add tree.get/filter/insert * output: use obj.filter/get for granular dir commit * squash filter logic output * dependency: don't filter imports at the dep level * don't init staging ODB for external cache * use staging ODB when computing tree hash for external outs * remove unneeded force in test granular commit * use unique-per-manager memfs URLs for staging * odb: make staging per-ODBManager * objects: move state dependent staging into ODBManager * objects: support checking multiple odbs - staging dependent checks can be done via ODBManager.check * checkout: do state lookup before explicit stage() in _changed() * output: use odbmanager based staging * RepoDependency: use odbmanager based staging * update obj staging unit tests * checkout: remove check(), handle individual file exceptions on link() * diff: use odbmanager based staging * objects: add ObjectPermissionError for read/write perm errors * odb: move staging out of odbmanager * objects: handle staging ODB in objects.stage() * objects.tree: don't stage anything in tree methods * odb: clear staged objects on gc() * remote: don't filter objs by scheme, objs are already grouped by ODB * diff: update objects.stage usage * output/dep: update stage() usage * update ODB tests * update func tests * remove memfs from FS_MAP * revert unneeded objects.check() changes * remove unused tree.insert() * revert unused odbmanager changes * add separate odb error classes * stage external outs directly in ODB * check for obj existence in main ODB before staging * use .dvc/tmp/staging instead of per-odb staging * gc: only gc loal staging once
91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
import errno
|
|
import itertools
|
|
import logging
|
|
import posixpath
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
from dvc.progress import Tqdm
|
|
from dvc.utils import to_chunks
|
|
|
|
from .base import ObjectDB
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SSHObjectDB(ObjectDB):
|
|
def batch_exists(self, path_infos, callback):
|
|
def _exists(chunk_and_channel):
|
|
chunk, channel = chunk_and_channel
|
|
ret = []
|
|
for path in chunk:
|
|
try:
|
|
channel.stat(path)
|
|
ret.append(True)
|
|
except OSError as exc:
|
|
if exc.errno != errno.ENOENT:
|
|
raise
|
|
ret.append(False)
|
|
callback(path)
|
|
return ret
|
|
|
|
with self.fs.ssh() as ssh:
|
|
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))
|
|
|
|
return results
|
|
|
|
def hashes_exist(self, hashes, jobs=None, name=None):
|
|
"""This is older implementation used in remote/base.py
|
|
We are reusing it in RemoteSSH, because SSH's batch_exists proved to be
|
|
faster than current approach (relying on exists(path_info)) applied in
|
|
remote/base.
|
|
"""
|
|
if not self.fs.CAN_TRAVERSE:
|
|
return list(set(hashes) & set(self.all()))
|
|
|
|
# possibly prompt for credentials before "Querying" progress output
|
|
self.fs.ensure_credentials()
|
|
|
|
with Tqdm(
|
|
desc="Querying "
|
|
+ ("cache in " + name if name else "remote cache"),
|
|
total=len(hashes),
|
|
unit="file",
|
|
) as pbar:
|
|
|
|
def exists_with_progress(chunks):
|
|
return self.batch_exists(chunks, callback=pbar.update_msg)
|
|
|
|
with ThreadPoolExecutor(
|
|
max_workers=jobs or self.fs.jobs
|
|
) as executor:
|
|
path_infos = [self.hash_to_path_info(x) for x in hashes]
|
|
chunks = to_chunks(path_infos, num_chunks=self.fs.jobs)
|
|
results = executor.map(exists_with_progress, chunks)
|
|
in_remote = itertools.chain.from_iterable(results)
|
|
ret = list(itertools.compress(hashes, in_remote))
|
|
return ret
|
|
|
|
def _list_paths(self, prefix=None, progress_callback=None):
|
|
if prefix:
|
|
root = posixpath.join(self.path_info.path, prefix[:2])
|
|
else:
|
|
root = self.path_info.path
|
|
with self.fs.ssh() as ssh:
|
|
if not ssh.exists(root):
|
|
return
|
|
# If we simply return an iterator then with above closes instantly
|
|
if progress_callback:
|
|
for path in ssh.walk_files(root):
|
|
progress_callback()
|
|
yield path
|
|
else:
|
|
yield from ssh.walk_files(root)
|