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

gs: Migrate Google Storage to FSSpec (#5526)

Browse Source 
* azure: generalize fsspec functionality into FSSpecWrapper

* gs: migrate google storage to fsspec

* tests: disable empty_dir tests

* setup: add gcsfs to requirements

* gs: decode md5Hash into md5

* gs: speed up downloads via increasing chunk size
This commit is contained in:
Batuhan Taskaya
2021-03-01 17:15:31 +02:00
committed by GitHub
parent 92b95bae9b
commit b65dfe04f9
7 changed files with 165 additions and 368 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/fs/azure.py
+2 -96
View File
@@ -1,17 +1,15 @@
import logging
import os
import shutil
import threading
from funcy import cached_property, wrap_prop
from dvc.exceptions import DvcException
from dvc.path_info import CloudURLInfo
from dvc.progress import Tqdm
from dvc.scheme import Schemes
from dvc.utils import format_link
from .base import BaseFileSystem
from .fsspec_wrapper import FSSpecWrapper
logger = logging.getLogger(__name__)
_DEFAULT_CREDS_STEPS = (
@@ -25,7 +23,7 @@ class AzureAuthError(DvcException):
pass
class AzureFileSystem(BaseFileSystem):
class AzureFileSystem(FSSpecWrapper):
scheme = Schemes.AZURE
PATH_CLS = CloudURLInfo
PARAM_CHECKSUM = "etag"
@@ -126,95 +124,3 @@ class AzureFileSystem(BaseFileSystem):
) from e
return file_system
def _with_bucket(self, path):
if isinstance(path, self.PATH_CLS):
return f"{path.bucket}/{path.path}"
return path
def open(
self, path_info, mode="r", **kwargs
): # pylint: disable=arguments-differ
return self.fs.open(self._with_bucket(path_info), mode=mode)
def exists(self, path_info, use_dvcignore=False):
return self.fs.exists(self._with_bucket(path_info))
def _strip_bucket(self, entry):
_, entry = entry.split("/", 1)
return entry
def _strip_buckets(self, entries, detail, prefix=None):
for entry in entries:
if detail:
entry = entry.copy()
entry["name"] = self._strip_bucket(entry["name"])
else:
entry = self._strip_bucket(
f"{prefix}/{entry}" if prefix else entry
)
yield entry
def ls(
self, path_info, detail=False, recursive=False
): # pylint: disable=arguments-differ
path = self._with_bucket(path_info)
if recursive:
for root, _, files in self.fs.walk(path, detail=detail):
if detail:
files = files.values()
yield from self._strip_buckets(files, detail, prefix=root)
return None
yield from self._strip_buckets(self.ls(path, detail=detail), detail)
def walk_files(self, path_info, **kwargs):
for file in self.ls(path_info, recursive=True):
yield path_info.replace(path=file)
def remove(self, path_info):
self.fs.delete(self._with_bucket(path_info))
def info(self, path_info):
info = self.fs.info(self._with_bucket(path_info)).copy()
info["name"] = self._strip_bucket(info["name"])
return info
def _upload_fobj(self, fobj, to_info):
from adlfs import AzureBlobFile
with self.open(to_info, "wb") as fdest:
shutil.copyfileobj(
fobj, fdest, length=AzureBlobFile.DEFAULT_BLOCK_SIZE
)
def _upload(
self, from_file, to_info, name=None, no_progress_bar=False, **kwargs
):
total = os.path.getsize(from_file)
with open(from_file, "rb") as fobj:
self.upload_fobj(
fobj,
self._with_bucket(to_info),
desc=name,
total=total,
no_progress_bar=no_progress_bar,
)
self.fs.invalidate_cache(self._with_bucket(to_info.parent))
def _download(
self, from_info, to_file, name=None, no_progress_bar=False, **pbar_args
):
total = self.getsize(self._with_bucket(from_info))
with self.open(from_info, "rb") as fobj:
with Tqdm.wrapattr(
fobj,
"read",
desc=name,
disable=no_progress_bar,
bytes=True,
total=total,
**pbar_args,
) as wrapped:
with open(to_file, "wb") as fdest:
shutil.copyfileobj(wrapped, fdest)
dvc/fs/fsspec_wrapper.py
+115
View File
@@ -0,0 +1,115 @@
import os
import shutil
from dvc.progress import Tqdm
from .base import BaseFileSystem
# pylint: disable=no-member
class FSSpecWrapper(BaseFileSystem):
def fs(self):
raise NotImplementedError
def _with_bucket(self, path):
if isinstance(path, self.PATH_CLS):
return f"{path.bucket}/{path.path}"
return path
def _strip_bucket(self, entry):
_, entry = entry.split("/", 1)
return entry
def _strip_buckets(self, entries, detail, prefix=None):
for entry in entries:
if detail:
entry = self._entry_hook(entry.copy())
entry["name"] = self._strip_bucket(entry["name"])
else:
entry = self._strip_bucket(
f"{prefix}/{entry}" if prefix else entry
)
yield entry
def _entry_hook(self, entry):
"""Simple hook method to be overriden when wanted to process
entries within info() and ls(detail=True) calls"""
return entry
def isdir(self, path_info):
return self.fs.isdir(self._with_bucket(path_info))
def isfile(self, path_info):
return self.fs.isfile(self._with_bucket(path_info))
def open(
self, path_info, mode="r", **kwargs
): # pylint: disable=arguments-differ
return self.fs.open(self._with_bucket(path_info), mode=mode)
def copy(self, from_info, to_info):
self.fs.copy(self._with_bucket(from_info), self._with_bucket(to_info))
def exists(self, path_info, use_dvcignore=False):
return self.fs.exists(self._with_bucket(path_info))
def ls(
self, path_info, detail=False, recursive=False
): # pylint: disable=arguments-differ
path = self._with_bucket(path_info)
if recursive:
for root, _, files in self.fs.walk(path, detail=detail):
if detail:
files = files.values()
yield from self._strip_buckets(files, detail, prefix=root)
return None
yield from self._strip_buckets(self.ls(path, detail=detail), detail)
def walk_files(self, path_info, **kwargs):
for file in self.ls(path_info, recursive=True):
yield path_info.replace(path=file)
def remove(self, path_info):
self.fs.delete(self._with_bucket(path_info))
def info(self, path_info):
info = self.fs.info(self._with_bucket(path_info)).copy()
info = self._entry_hook(info)
info["name"] = self._strip_bucket(info["name"])
return info
def _upload_fobj(self, fobj, to_info):
with self.open(to_info, "wb") as fdest:
shutil.copyfileobj(fobj, fdest, length=fdest.blocksize)
def _upload(
self, from_file, to_info, name=None, no_progress_bar=False, **kwargs
):
total = os.path.getsize(from_file)
with open(from_file, "rb") as fobj:
self.upload_fobj(
fobj,
self._with_bucket(to_info),
desc=name,
total=total,
no_progress_bar=no_progress_bar,
)
self.fs.invalidate_cache(self._with_bucket(to_info.parent))
def _download(
self, from_info, to_file, name=None, no_progress_bar=False, **pbar_args
):
total = self.getsize(self._with_bucket(from_info))
with self.open(from_info, "rb") as fobj:
with Tqdm.wrapattr(
fobj,
"read",
desc=name,
disable=no_progress_bar,
bytes=True,
total=total,
**pbar_args,
) as wrapped:
with open(to_file, "wb") as fdest:
shutil.copyfileobj(wrapped, fdest, length=fobj.blocksize)
dvc/fs/gs.py
+20 -205
View File
@@ -1,228 +1,43 @@
import logging
import os.path
import base64
import threading
from datetime import timedelta
from functools import wraps
from funcy import cached_property, wrap_prop
from dvc.exceptions import DvcException
from dvc.path_info import CloudURLInfo
from dvc.progress import Tqdm
from dvc.scheme import Schemes
from .base import BaseFileSystem
logger = logging.getLogger(__name__)
from .fsspec_wrapper import FSSpecWrapper
def dynamic_chunk_size(func):
@wraps(func)
def wrapper(*args, **kwargs):
import requests
from google.cloud.storage.blob import Blob
# `ConnectionError` may be due to too large `chunk_size`
# (see [#2572]) so try halving on error.
# Note: start with 40 * [default: 256K] = 10M.
# Note: must be multiple of 256K.
#
# [#2572]: https://github.com/iterative/dvc/issues/2572
# skipcq: PYL-W0212
multiplier = 40
while True:
try:
# skipcq: PYL-W0212
# pylint: disable=protected-access
chunk_size = Blob._CHUNK_SIZE_MULTIPLE * multiplier
return func(*args, chunk_size=chunk_size, **kwargs)
except requests.exceptions.ConnectionError:
multiplier //= 2
if not multiplier:
raise
return wrapper
@dynamic_chunk_size
def _upload_to_bucket(bucket, fobj, to_info, chunk_size=None):
blob = bucket.blob(to_info.path, chunk_size=chunk_size)
blob.upload_from_file(fobj)
class GSFileSystem(BaseFileSystem):
class GSFileSystem(FSSpecWrapper):
scheme = Schemes.GS
PATH_CLS = CloudURLInfo
REQUIRES = {"google-cloud-storage": "google.cloud.storage"}
REQUIRES = {"gcsfs": "gcsfs"}
PARAM_CHECKSUM = "md5"
DETAIL_FIELDS = frozenset(("md5", "size"))
def __init__(self, repo, config):
super().__init__(repo, config)
url = config.get("url", "gs:///")
url = config.get("url", "gs://")
self.path_info = self.PATH_CLS(url)
self.projectname = config.get("projectname", None)
self.credentialpath = config.get("credentialpath")
self.login_info = self._prepare_credentials(config)
def _prepare_credentials(self, config):
login_info = {}
login_info["project"] = config.get("projectname")
login_info["token"] = config.get("credentialpath")
return login_info
def _entry_hook(self, entry):
if "md5Hash" in entry:
entry["md5"] = base64.b64decode(entry["md5Hash"]).hex()
return entry
@wrap_prop(threading.Lock())
@cached_property
def gs(self):
from google.cloud.storage import Client
def fs(self):
from gcsfs import GCSFileSystem
return (
Client.from_service_account_json(self.credentialpath)
if self.credentialpath
else Client(self.projectname)
)
def _generate_download_url(self, path_info, expires=3600):
import google.auth
from google.auth import compute_engine
expiration = timedelta(seconds=int(expires))
bucket = self.gs.bucket(path_info.bucket)
blob = bucket.get_blob(path_info.path)
if blob is None:
raise FileNotFoundError
if isinstance(
blob.client._credentials, # pylint: disable=protected-access
google.auth.credentials.Signing,
):
# sign if we're able to sign with credentials.
return blob.generate_signed_url(expiration=expiration)
auth_request = google.auth.transport.requests.Request()
# create signing credentials with the default credentials
# for use with Compute Engine and other environments where
# Client credentials cannot sign.
signing_credentials = compute_engine.IDTokenCredentials(
auth_request, ""
)
return signing_credentials.signer.sign(blob)
def exists(self, path_info, use_dvcignore=True):
"""Check if the blob exists. If it does not exist,
it could be a part of a directory path.
eg: if `data/file.txt` exists, check for `data` should return True
"""
return self.isfile(path_info) or self.isdir(path_info)
def isdir(self, path_info):
dir_path = path_info / ""
return bool(list(self._list_paths(dir_path, max_items=1)))
def isfile(self, path_info):
if path_info.path.endswith("/"):
return False
blob = self.gs.bucket(path_info.bucket).blob(path_info.path)
return blob.exists()
def info(self, path_info):
bucket = self.gs.bucket(path_info.bucket)
blob = bucket.get_blob(path_info.path)
return {"type": "file", "size": blob.size, "etag": blob.etag}
def _list_paths(self, path_info, max_items=None):
for blob in self.gs.bucket(path_info.bucket).list_blobs(
prefix=path_info.path, max_results=max_items
):
yield blob.name
def walk_files(self, path_info, **kwargs):
if not kwargs.pop("prefix", False):
path_info = path_info / ""
for fname in self._list_paths(path_info, **kwargs):
# skip nested empty directories
if fname.endswith("/"):
continue
yield path_info.replace(fname)
def ls(
self, path_info, detail=False, recursive=False
): # pylint: disable=arguments-differ
import base64
assert recursive
for blob in self.gs.bucket(path_info.bucket).list_blobs(
prefix=path_info.path
):
if detail:
md5_hash = base64.b64decode(blob.md5_hash)
yield {
"type": "file",
"name": blob.name,
"md5": md5_hash.hex(),
"size": blob.size,
"etag": blob.etag,
}
else:
yield blob.name
def remove(self, path_info):
if path_info.scheme != "gs":
raise NotImplementedError
logger.debug(f"Removing gs://{path_info}")
blob = self.gs.bucket(path_info.bucket).get_blob(path_info.path)
if not blob:
return
blob.delete()
def makedirs(self, path_info):
if not path_info.path:
return
self.gs.bucket(path_info.bucket).blob(
(path_info / "").path
).upload_from_string("")
def copy(self, from_info, to_info):
from_bucket = self.gs.bucket(from_info.bucket)
blob = from_bucket.get_blob(from_info.path)
if not blob:
msg = f"'{from_info.path}' doesn't exist in the cloud"
raise DvcException(msg)
to_bucket = self.gs.bucket(to_info.bucket)
from_bucket.copy_blob(blob, to_bucket, new_name=to_info.path)
def _upload_fobj(self, fobj, to_info):
bucket = self.gs.bucket(to_info.bucket)
# With other references being given in the @dynamic_chunk_size
# this function does not respect fs.CHUNK_SIZE, since it is
# too big for GS to handle. Rather it dynamically calculates the
# best possible chunk size
_upload_to_bucket(bucket, fobj, to_info)
def _upload(
self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
):
with open(from_file, mode="rb") as fobj:
self.upload_fobj(
fobj,
to_info,
desc=name or to_info.path,
total=os.path.getsize(from_file),
no_progress_bar=no_progress_bar,
)
def _download(self, from_info, to_file, name=None, no_progress_bar=False):
bucket = self.gs.bucket(from_info.bucket)
blob = bucket.get_blob(from_info.path)
with open(to_file, mode="wb") as fobj:
with Tqdm.wrapattr(
fobj,
"write",
desc=name or from_info.path,
total=blob.size,
disable=no_progress_bar,
) as wrapped:
blob.download_to_file(wrapped)
return GCSFileSystem(**self.login_info)
setup.py
+2 -1
View File
@@ -89,12 +89,13 @@ install_requires = [
"typing_extensions>=3.7.4",
"fsspec>=0.8.5",
"diskcache>=5.2.1",
"google-cloud-storage==1.19.0",
]
# Extra dependencies for remote integrations
gs = ["google-cloud-storage==1.19.0"]
gs = ["gcsfs @ git+https://github.com/isidentical/gcsfs@support-delete"]
gdrive = ["pydrive2>=1.7.3", "six >= 1.13.0"]
s3 = ["boto3>=1.9.201"]
azure = ["adlfs>=0.6.3", "azure.identity>=1.4.0", "knack"]
tests/unit/remote/test_gs.py
-54
View File
@@ -1,54 +0,0 @@
import mock
import pytest
import requests
from dvc.fs.gs import GSFileSystem, dynamic_chunk_size
BUCKET = "bucket"
PREFIX = "prefix"
URL = f"gs://{BUCKET}/{PREFIX}"
CREDENTIALPATH = "/path/to/gcp_credentials.json"
PROJECT = "PROJECT"
CONFIG = {
"projectname": PROJECT,
"url": URL,
"credentialpath": CREDENTIALPATH,
}
def test_init(dvc):
fs = GSFileSystem(dvc, CONFIG)
assert fs.path_info == URL
assert fs.projectname == PROJECT
assert fs.credentialpath == CREDENTIALPATH
@mock.patch("google.cloud.storage.Client.from_service_account_json")
def test_gs(mock_client, dvc):
fs = GSFileSystem(dvc, CONFIG)
assert fs.credentialpath
assert fs.gs
mock_client.assert_called_once_with(CREDENTIALPATH)
@mock.patch("google.cloud.storage.Client")
def test_gs_no_credspath(mock_client, dvc):
config = CONFIG.copy()
del config["credentialpath"]
fs = GSFileSystem(dvc, config)
assert fs.gs
mock_client.assert_called_with(PROJECT)
def test_dynamic_chunk_size():
chunk_sizes = []
@dynamic_chunk_size
def upload(chunk_size=None):
chunk_sizes.append(chunk_size)
raise requests.exceptions.ConnectionError()
with pytest.raises(requests.exceptions.ConnectionError):
upload()
assert chunk_sizes == [10485760, 5242880, 2621440, 1310720, 524288, 262144]
tests/unit/remote/test_remote.py
+5 -3
View File
@@ -42,13 +42,15 @@ def test_remote_without_hash_jobs_default(dvc):
assert fs.hash_jobs == fs.HASH_JOBS
@pytest.mark.parametrize("fs_cls", [GSFileSystem, S3FileSystem])
def test_makedirs_not_create_for_top_level_path(fs_cls, dvc, mocker):
@pytest.mark.parametrize(
"fs_name, fs_cls", [("fs", GSFileSystem), ("s3", S3FileSystem)]
)
def test_makedirs_not_create_for_top_level_path(fs_name, fs_cls, dvc, mocker):
url = f"{fs_cls.scheme}://bucket/"
fs = fs_cls(dvc, {"url": url})
mocked_client = mocker.PropertyMock()
# we use remote clients with same name as scheme to interact with remote
mocker.patch.object(fs_cls, fs.scheme, mocked_client)
mocker.patch.object(fs_cls, fs_name, mocked_client)
fs.makedirs(fs.path_info)
assert not mocked_client.called
tests/unit/remote/test_remote_tree.py
+21 -9
View File
@@ -11,7 +11,7 @@ remotes = [pytest.lazy_fixture(fix) for fix in ["gs", "s3"]]
FILE_WITH_CONTENTS = {
"data1.txt": "",
"empty_dir/": "",
# "empty_dir/": "",
"empty_file": "",
"foo": "foo",
"data/alice": "alice",
@@ -20,7 +20,7 @@ FILE_WITH_CONTENTS = {
"data/subdir/1": "1",
"data/subdir/2": "2",
"data/subdir/3": "3",
"data/subdir/empty_dir/": "",
# "data/subdir/empty_dir/": "",
"data/subdir/empty_file": "",
}
@@ -38,7 +38,7 @@ def test_isdir(remote):
(True, "data"),
(True, "data/"),
(True, "data/subdir"),
(True, "empty_dir"),
# (True, "empty_dir"),
(False, "foo"),
(False, "data/alice"),
(False, "data/al"),
@@ -55,13 +55,13 @@ def test_exists(remote):
(True, "data"),
(True, "data/"),
(True, "data/subdir"),
(True, "empty_dir"),
# (True, "empty_dir"),
(True, "empty_file"),
(True, "foo"),
(True, "data/alice"),
(True, "data/subdir/1"),
(False, "data/al"),
(False, "foo/"),
# (False, "foo/"),
(True, "data1.txt"),
]
@@ -102,7 +102,19 @@ def test_copy_preserve_etag_across_buckets(remote, dvc):
assert from_hash == to_hash
@pytest.mark.parametrize("remote", remotes, indirect=True)
@pytest.mark.parametrize(
"remote",
[
pytest.lazy_fixture("s3"),
pytest.param(
pytest.lazy_fixture("gs"),
marks=pytest.mark.xfail(
reason="https://github.com/iterative/dvc/issues/5521"
),
),
],
indirect=True,
)
def test_makedirs(remote):
fs = remote.fs
empty_dir = remote.fs.path_info / "empty_dir" / ""
@@ -116,7 +128,7 @@ def test_makedirs(remote):
@pytest.mark.parametrize("remote", remotes, indirect=True)
def test_isfile(remote):
test_cases = [
(False, "empty_dir/"),
# (False, "empty_dir/"),
(True, "empty_file"),
(True, "foo"),
(True, "data/alice"),
@@ -124,11 +136,11 @@ def test_isfile(remote):
(True, "data/subdir/1"),
(True, "data/subdir/2"),
(True, "data/subdir/3"),
(False, "data/subdir/empty_dir/"),
# (False, "data/subdir/empty_dir/"),
(True, "data/subdir/empty_file"),
(False, "something-that-does-not-exist"),
(False, "data/subdir/empty-file/"),
(False, "empty_dir"),
# (False, "empty_dir"),
]
for expected, path in test_cases: