Move files

This commit is contained in:
Andras Schmelczer 2022-07-04 19:31:15 +02:00
parent 3cf28379e8
commit 00cc8225c5
159 changed files with 31 additions and 49 deletions

View file

@ -0,0 +1 @@
from .large_file import LargeFileBase, LargeFileLocal, LargeFileMongo, LargeFileS3

View file

@ -0,0 +1,71 @@
#!/usr/bin/env python3
from argparse import Namespace
from pathlib import Path
from typing import Mapping, Type
from ..utilities import get_logger
from .large_file import LargeFileBase, LargeFileLocal, LargeFileMongo, LargeFileS3
from .parse_arguments import parse_arguments
logger = get_logger("large_file")
def main() -> None:
parser, args = parse_arguments()
large_file = get_class(args)
if not args.cache and not args.push and not args.delete:
logger.warning("No action required.")
parser.print_help()
if args.cache:
for c in args.cache:
split = c.split(":")
file_name = split[0]
version = None if len(split) == 1 else int(split[1])
large_file(file_name, "r", version=version).get()
if args.push:
for p in args.push:
path = Path(p)
large_file(path.name, "w").push(path)
if args.delete:
for f in args.delete:
large_file(f).delete()
def get_class(args: Namespace) -> Type[LargeFileBase]:
factory: Mapping[str, Type[LargeFileBase]] = {
"s3": LargeFileS3,
"local": LargeFileLocal,
"mongodb": LargeFileMongo,
}
if args.backend not in factory:
raise ValueError(
f"Backend {args.backend} does not exits, available options: {' ,'.join(factory.keys())}"
)
large_file = factory[args.backend]
if args.backend != "local":
if args.secrets is None:
raise ValueError(
"Providing a credentials file is required when the backend mode is not `local`."
)
large_file.configure_credentials_from_file(args.secrets)
return large_file
if __name__ == "__main__":
try:
main()
except KeyboardInterrupt:
logger.warning("Exiting")
exit()
except Exception as e:
logger.exception(e)

View file

@ -0,0 +1,2 @@
from .human_readable_to_byte import human_readable_to_byte
from .progress_bar import DownloadProgressBar, UploadProgressBar

View file

@ -0,0 +1,2 @@
def bytes_to_megabytes(bytes: int) -> str:
return f"{round(bytes / 1000 / 1000, 2):.2f}"

View file

@ -0,0 +1,31 @@
import re
def human_readable_to_byte(size: str) -> int:
"""Case is ignored, kb, kB, Kb, and KB are all treated as kilobyte."""
if size.strip() == "0":
return 0
possible_units = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
units_re = "|".join(possible_units)
regex = re.compile(
rf"""
\s* # trim
(?P<scalar>\d+(.\d+)?) # get scalar, it might be a float
\s* # ignore optional whitespace
(?P<unit>{units_re}) # capture the unit
""",
flags=re.VERBOSE | re.IGNORECASE,
)
match = regex.match(size)
if not match:
raise ValueError(f'Could not find values in "{size}"')
results = match.groupdict()
scalar = float(results["scalar"])
idx = possible_units.index(results["unit"].upper())
factor = 1024**idx
return round(scalar * factor)

View file

@ -0,0 +1,51 @@
import os
import threading
from logging import Logger
from pathlib import Path
from .bytes_to_megabytes import bytes_to_megabytes
class ProgressBar:
min_progress_percentage_change = 10
def __init__(self, file_size: int, logger: Logger, prefix: str):
self._file_size = file_size
self._logger = logger
self._prefix = prefix
self._last_percentage: float = 0
self._seen_so_far = 0
self._lock = threading.Lock()
def __call__(self, bytes_amount: int) -> None:
with self._lock:
self._seen_so_far += bytes_amount
percentage = (self._seen_so_far / float(self._file_size)) * 100
if (
percentage != 100
and percentage - self._last_percentage
< self.min_progress_percentage_change
):
return
self._last_percentage += self.min_progress_percentage_change
file_size_mb = bytes_to_megabytes(self._file_size)
seen_so_far_mb = bytes_to_megabytes(self._seen_so_far)
progress = seen_so_far_mb.rjust(len(file_size_mb))
self._logger.info(
f"{self._prefix} {progress}/{file_size_mb} MB ({percentage:.1f}%)"
)
class DownloadProgressBar(ProgressBar):
def __init__(self, name: str, size: int, logger: Logger):
super().__init__(file_size=size, logger=logger, prefix=f"Downloading {name}")
class UploadProgressBar(ProgressBar):
def __init__(self, path: Path, logger: Logger):
size = os.path.getsize(path)
super().__init__(file_size=size, logger=logger, prefix=f"Uploading {path.name}")

View file

@ -0,0 +1,4 @@
from .large_file_base import LargeFileBase
from .large_file_local import LargeFileLocal
from .large_file_mongo import LargeFileMongo
from .large_file_s3 import LargeFileS3

View file

@ -0,0 +1,316 @@
import os
import shutil
import tempfile
from abc import ABC, abstractmethod
from functools import lru_cache
from pathlib import Path
from types import TracebackType
from typing import IO, Any, List, Optional, Type, Union, cast
from great_ai.utilities import ConfigFile, get_logger
from ..helper import human_readable_to_byte
from ..models import DataInstance
logger = get_logger("large_file")
CACHE_NAME_VERSION_SEPARATOR = "-"
COMPRESSION_ALGORITHM = "gztar"
ARCHIVE_EXTENSION = ".tar.gz"
class LargeFileBase(ABC):
"""
Store large files remotely. Use local cache for speed up.
Examples:
```
with LargeFile("test.txt", "w", keep_last_n=3) as f:
for i in range(1000000):
f.write('test\n')
with LargeFile("test.txt", "r") as f:
print(f.readlines()[0])
path_to_cached_text_file = LargeFile("test.txt", version=0).get()
```
By default, files are stored in the ".cache" folder and the
least recently use is deleted after the overall size reaches 30 GBs.
Change it with the following properties.
```
LargeFile.cache_path = Path(".cache")
LargeFile.max_cache_size = "30GB"
```
"""
initialized = False
cache_path = Path(".cache")
max_cache_size: Optional[str] = "30GB"
def __init__(
self,
name: str,
mode: str = "r",
*,
buffering: int = -1,
encoding: Optional[str] = None,
errors: Optional[str] = None,
newline: Optional[str] = None,
version: Optional[int] = None,
keep_last_n: Optional[int] = None,
cache_only_mode: bool = False,
):
self._name = name
self._version = version
self._mode = mode
self._keep_last_n = keep_last_n
self._cache_only_mode = cache_only_mode
self._buffering = buffering
self._encoding = encoding
self._errors = errors
self._newline = newline
LargeFileBase.cache_path.mkdir(parents=True, exist_ok=True)
self._find_instances()
self._check_mode_and_set_version()
@classmethod
def configure_credentials_from_file(
cls,
secrets_path: Union[Path, str],
) -> None:
cls.configure_credentials(**ConfigFile(secrets_path))
@classmethod
def configure_credentials(
cls,
) -> None:
cls.initialized = True
def __enter__(self) -> IO:
self._file: IO[Any] = (
tempfile.NamedTemporaryFile(
mode=self._mode,
buffering=self._buffering,
encoding=self._encoding,
newline=self._newline,
errors=self._errors,
delete=False,
prefix="large_file-",
)
if "w" in self._mode
else open(
self.get(),
mode=self._mode,
buffering=self._buffering,
encoding=self._encoding,
newline=self._newline,
errors=self._errors,
)
)
return self._file
def __exit__(
self,
type: Optional[Type[BaseException]],
exc: Optional[BaseException],
traceback: Optional[TracebackType],
) -> bool:
self._file.close()
if type is None:
if "w" in self._mode:
self.push(Path(self._file.name))
os.unlink(self._file.name)
else:
logger.exception("Could not finish operation.")
return True
@property
def _local_name(self) -> str:
return f"{self._name}{CACHE_NAME_VERSION_SEPARATOR}{self.version}"
@property
def version(self) -> int:
return cast(int, self._version)
@lru_cache(1)
def get(self, hide_progress: bool = False) -> Path:
remote_path = next(
i.remote_path for i in self._instances if i.version == self._version
)
destination = self.cache_path / self._local_name
if not destination.exists():
logger.info(f"File {self._local_name} does not exist locally")
with tempfile.TemporaryDirectory() as tmp:
local_root_path = Path(tmp)
tmp_file_archive = (
local_root_path / f"{self._local_name}{ARCHIVE_EXTENSION}"
)
self._download(
remote_path, tmp_file_archive, hide_progress=hide_progress
)
logger.info(f"Decompressing {self._local_name}")
shutil.unpack_archive(str(tmp_file_archive), tmp, COMPRESSION_ALGORITHM)
shutil.move(str(local_root_path / self._local_name), str(destination))
else:
logger.info(f"File {self._local_name} found in cache")
return destination
def push(self, path: Union[Path, str], hide_progress: bool = False) -> None:
if isinstance(path, str):
path = Path(path)
with tempfile.TemporaryDirectory() as tmp:
if path.is_file():
logger.info(f"Copying file for {self._local_name}")
copy: Any = shutil.copy
else:
logger.info(f"Copying directory for {self._local_name}")
copy = shutil.copytree
try:
# Make local copy in the cache
shutil.rmtree(self.cache_path / self._local_name, ignore_errors=True)
copy(str(path), str(self.cache_path / self._local_name))
except shutil.SameFileError:
pass # No worries
copy(str(path), str(Path(tmp) / self._local_name))
with tempfile.TemporaryDirectory() as tmp2:
# A directory has to be zipped and it cannot contain the output of the zipping
logger.info(f"Compressing {self._local_name}")
shutil.make_archive(
str(Path(tmp2) / self._local_name),
COMPRESSION_ALGORITHM,
tmp,
)
file_to_be_uploaded = (
Path(tmp2) / f"{self._local_name}{ARCHIVE_EXTENSION}"
)
self._upload(file_to_be_uploaded, hide_progress=hide_progress)
self.clean_up()
def delete(self) -> None:
self._keep_last_n = 0
self._delete_old_remote_versions()
def _find_instances(self) -> None:
if self._cache_only_mode:
self._instances = self._find_instances_from_cache()
else:
self._instances = self._find_remote_instances()
self._instances = sorted(self._instances, key=lambda i: i.version)
def _find_instances_from_cache(self) -> List[DataInstance]:
logger.info(f"Fetching cached versions of {self._name}")
candidates = [
DataInstance(
name=CACHE_NAME_VERSION_SEPARATOR.join(
f.name.split(CACHE_NAME_VERSION_SEPARATOR)[:-1]
),
version=int(f.name.split(CACHE_NAME_VERSION_SEPARATOR)[-1]),
remote_path=f,
)
for f in self.cache_path.glob(
f"{self._name}{CACHE_NAME_VERSION_SEPARATOR}*"
)
]
return [c for c in candidates if c.name == self._name]
def _check_mode_and_set_version(self) -> None:
if "+" in self._mode:
raise ValueError(
f"File mode `{self._mode}` is not allowed3, remove the `+`."
)
if "w" in self._mode:
if self._version is not None:
raise ValueError("Providing a version is not allowed in write mode.")
self._version = self._instances[-1].version + 1 if self._instances else 0
elif "r" in self._mode:
if not self._instances:
raise FileNotFoundError(
f"File {self._name} not found. No versions are available."
)
if self._version is None:
self._version = self._instances[-1].version
logger.info(
f"Latest version of {self._name} is {self._version} "
+ f"(from versions: {self.versions_pretty})"
)
elif self._version not in [i.version for i in self._instances]:
raise FileNotFoundError(
f"File {self._name} not found with version {self._version}. "
+ f"(from versions: {self.versions_pretty})"
)
else:
raise ValueError("Unsupported file mode.")
@property
def versions_pretty(self) -> str:
return ", ".join((str(i.version) for i in self._instances))
def clean_up(self) -> None:
self._delete_old_remote_versions()
self._prune_cache()
def _prune_cache(self) -> None:
self.cache_path.mkdir(parents=True, exist_ok=True)
if self.max_cache_size is None:
return
allowed_size = human_readable_to_byte(self.max_cache_size)
assert allowed_size >= 0
least_recently_read = sorted(
[f for f in self.cache_path.glob("*")], key=lambda f: f.stat().st_atime
)
while sum(os.path.getsize(f) for f in least_recently_read) > allowed_size:
file = least_recently_read.pop(0)
logger.info(
f"Deleting file from cache to meet quota (max_cache_size={self.max_cache_size}): {file}"
)
os.unlink(file)
@abstractmethod
def _find_remote_instances(self) -> List[DataInstance]:
pass
@abstractmethod
def _download(
self, remote_path: Any, local_path: Path, hide_progress: bool
) -> None:
pass
@abstractmethod
def _upload(self, local_path: Path, hide_progress: bool) -> None:
pass
@abstractmethod
def _delete_old_remote_versions(self) -> None:
pass

View file

@ -0,0 +1,58 @@
from pathlib import Path
from typing import Any, List, Optional
from ...utilities import get_logger
from ..models import DataInstance
from .large_file_base import LargeFileBase
logger = get_logger("large_file")
class LargeFileLocal(LargeFileBase):
def __init__(
self,
name: str,
mode: str = "r",
*,
buffering: int = -1,
encoding: Optional[str] = None,
errors: Optional[str] = None,
newline: Optional[str] = None,
version: Optional[int] = None,
keep_last_n: Optional[int] = None,
):
super().__init__(
name,
mode,
buffering=buffering,
encoding=encoding,
errors=errors,
newline=newline,
version=version,
keep_last_n=keep_last_n,
cache_only_mode=True,
)
super().configure_credentials()
def _find_remote_instances(self) -> List[DataInstance]:
return []
def _download(
self, remote_path: Any, local_path: Path, hide_progress: bool
) -> None:
raise NotImplementedError()
def _upload(self, local_path: Path, hide_progress: bool) -> None:
pass # the "upload" is already done py the parent's caching mechanism
def _delete_old_remote_versions(self) -> None:
if self._keep_last_n is not None:
for i in (
self._instances[: -self._keep_last_n]
if self._keep_last_n > 0
else self._instances
):
logger.info(
f"Removing old version (keep_last_n={self._keep_last_n}): {i.remote_path}"
)
i.remote_path.unlink()

View file

@ -0,0 +1,121 @@
import re
from functools import cached_property
from pathlib import Path
from typing import Any, List, Mapping
from gridfs import DEFAULT_CHUNK_SIZE, Database, GridFSBucket
from pymongo import MongoClient
from ...utilities import get_logger
from ..helper import DownloadProgressBar, UploadProgressBar
from ..models import DataInstance
from .large_file_base import LargeFileBase
logger = get_logger("large_file")
MONGO_NAME_VERSION_SEPARATOR = "_"
class LargeFileMongo(LargeFileBase):
mongo_connection_string = None
mongo_database = None
@classmethod
def configure_credentials( # type: ignore
cls,
*,
mongo_connection_string: str,
mongo_database: str,
**_: Mapping[str, Any],
) -> None:
cls.mongo_connection_string = mongo_connection_string
cls.mongo_database = mongo_database
super().configure_credentials()
@cached_property
def _client(self) -> GridFSBucket:
if self.mongo_connection_string is None or self.mongo_database is None:
raise ValueError(
"Please configure the MongoDB access options by calling LargeFileMongo.configure_credentials or set offline_mode=True in the constructor."
)
db: Database = MongoClient(self.mongo_connection_string)[self.mongo_database]
return GridFSBucket(db)
def _find_remote_instances(self) -> List[DataInstance]:
logger.debug(f"Fetching Mongo (GridFS) versions of {self._name}")
return [
DataInstance(
name=MONGO_NAME_VERSION_SEPARATOR.join(
f.name.split(MONGO_NAME_VERSION_SEPARATOR)[:-1]
),
version=int(f.name.split(MONGO_NAME_VERSION_SEPARATOR)[-1]),
remote_path=(f._id, f.length),
origin="mongodb",
)
for f in self._client.find(
{
"filename": re.compile(
re.escape(self._name + MONGO_NAME_VERSION_SEPARATOR) + ".*"
)
}
)
]
def _download(
self, remote_path: Any, local_path: Path, hide_progress: bool
) -> None:
logger.info(f"Downloading {remote_path[0]} from Mongo (GridFS)")
progress = (
DownloadProgressBar(
name=str(remote_path[0]), size=remote_path[1], logger=logger
)
if not hide_progress
else None
)
with self._client.open_download_stream(remote_path[0]) as stream:
with open(local_path, "wb") as f:
while True:
content = stream.read(DEFAULT_CHUNK_SIZE)
f.write(content)
if progress:
progress(len(content))
if len(content) < DEFAULT_CHUNK_SIZE:
break
def _upload(self, local_path: Path, hide_progress: bool) -> None:
logger.info(f"Uploading {local_path} to Mongo (GridFS)")
progress = (
UploadProgressBar(path=local_path, logger=logger)
if not hide_progress
else None
)
with self._client.open_upload_stream(
f"{self._name}{MONGO_NAME_VERSION_SEPARATOR}{self.version}"
) as stream:
with open(local_path, "rb") as f:
while True:
content = f.read(DEFAULT_CHUNK_SIZE)
stream.write(content)
if progress:
progress(len(content))
if len(content) < DEFAULT_CHUNK_SIZE:
break
def _delete_old_remote_versions(self) -> None:
if self._keep_last_n is not None:
for i in (
self._instances[: -self._keep_last_n]
if self._keep_last_n > 0
else self._instances
):
logger.info(
f"Removing old version from MongoDB (GridFS) (keep_last_n={self._keep_last_n}): {i.name}{MONGO_NAME_VERSION_SEPARATOR}{i.version}"
)
self._client.delete(i.remote_path[0])

View file

@ -0,0 +1,151 @@
from functools import cached_property
from pathlib import Path
from typing import Any, List, Mapping, Optional
import boto3
from ...utilities import get_logger
from ..helper import DownloadProgressBar, UploadProgressBar
from ..models import DataInstance
from .large_file_base import LargeFileBase
logger = get_logger("large_file")
S3_NAME_VERSION_SEPARATOR = "/"
class LargeFileS3(LargeFileBase):
"""
Store large files in S3. Use local cache for speed up.
Examples:
```
with LargeFile("test.txt", "w", keep_last_n=3) as f:
for i in range(1000000):
f.write('test\n')
with LargeFile("test.txt", "r") as f:
print(f.readlines()[0])
path_to_cached_text_file = LargeFile("test.txt", version=0).get()
```
By default, files are stored in the ".cache" folder and the
least recently use is deleted after the overall size reaches 30 GBs.
Change it with the following properties.
```
LargeFile.cache_path = Path(".cache")
LargeFile.max_cache_size = "30GB"
```
"""
region_name = None
access_key_id = None
secret_access_key = None
bucket_name = None
endpoint_url = None
@classmethod
def configure_credentials( # type: ignore
cls,
*,
aws_region_name: str,
aws_access_key_id: str,
aws_secret_access_key: str,
large_files_bucket_name: str,
aws_endpoint_url: Optional[str] = None,
**_: Mapping[str, Any],
) -> None:
cls.region_name = aws_region_name
cls.access_key_id = aws_access_key_id
cls.secret_access_key = aws_secret_access_key
cls.bucket_name = large_files_bucket_name
cls.endpoint_url = aws_endpoint_url
super().configure_credentials()
@cached_property
def _client(self) -> boto3.client:
if (
self.region_name is None
or self.access_key_id is None
or self.secret_access_key is None
or self.bucket_name is None
):
raise ValueError(
"Please configure the S3 access options by calling LargeFileS3.configure_credentials or set offline_mode=True in the constructor."
)
return boto3.client(
"s3",
aws_access_key_id=self.access_key_id,
aws_secret_access_key=self.secret_access_key,
region_name=self.region_name,
endpoint_url=self.endpoint_url,
)
def _find_remote_instances(self) -> List[DataInstance]:
logger.debug(f"Fetching S3 versions of {self._name}")
found_objects = self._client.list_objects_v2(
Bucket=self.bucket_name, Prefix=self._name
)
return (
[
DataInstance(
name=o["Key"].split(S3_NAME_VERSION_SEPARATOR)[0],
version=int(o["Key"].split(S3_NAME_VERSION_SEPARATOR)[-1]),
remote_path=o["Key"],
)
for o in found_objects["Contents"]
if o["Key"].split(S3_NAME_VERSION_SEPARATOR)[0] == self._name
]
if "Contents" in found_objects
else []
)
def _download(
self, remote_path: Any, local_path: Path, hide_progress: bool
) -> None:
logger.info(f"Downloading {remote_path} from S3")
size = self._client.head_object(Bucket=self.bucket_name, Key=remote_path)[
"ContentLength"
]
self._client.download_file(
Bucket=self.bucket_name,
Key=remote_path,
Filename=str(local_path),
Callback=None
if hide_progress
else DownloadProgressBar(name=str(remote_path), size=size, logger=logger),
)
def _upload(self, local_path: Path, hide_progress: bool) -> None:
key = f"{self._name}/{self.version}"
logger.info(f"Uploading {self._local_name} to S3 as {key}")
self._client.upload_file(
Filename=str(local_path),
Bucket=self.bucket_name,
Key=key,
Callback=None
if hide_progress
else UploadProgressBar(path=local_path, logger=logger),
)
def _delete_old_remote_versions(self) -> None:
if self._keep_last_n is not None:
for i in (
self._instances[: -self._keep_last_n]
if self._keep_last_n > 0
else self._instances
):
logger.info(
f"Removing old version from S3 (keep_last_n={self._keep_last_n}): {i.remote_path}"
)
self._client.delete_object(Bucket=self.bucket_name, Key=i.remote_path)

View file

@ -0,0 +1 @@
from .data_instance import DataInstance

View file

@ -0,0 +1,9 @@
from typing import Any
from pydantic import BaseModel
class DataInstance(BaseModel):
name: str
version: int
remote_path: Any

View file

@ -0,0 +1,56 @@
from argparse import ArgumentParser, Namespace
from typing import Tuple
def parse_arguments() -> Tuple[ArgumentParser, Namespace]:
parser = ArgumentParser(
description="Store and version large files in S3; open them like regular files. Caching included.",
)
parser.add_argument(
"-b",
"--backend",
type=str,
help="choose which backend to use, available options: `local`, `s3`, `mongodb`",
required=True,
)
parser.add_argument(
"-s",
"--secrets",
type=str,
help="path to an .ini configuration file with your S3 credentials",
required=False,
)
parser.add_argument(
"-c",
"--cache",
nargs="+",
type=str,
help="download file into local cache, example: file_name.txt:version",
required=False,
)
parser.add_argument(
"-p",
"--push",
nargs="+",
type=str,
help="push a local file into S3 and set it as the most recent version",
required=False,
)
parser.add_argument(
"-d",
"--delete",
nargs="+",
type=str,
help="delete every version of file from S3",
required=False,
)
parser.print_usage = parser.print_help # type: ignore
args = parser.parse_args()
return parser, args