Move files
This commit is contained in:
parent
3cf28379e8
commit
00cc8225c5
159 changed files with 31 additions and 49 deletions
1
great_ai/large_file/__init__.py
Normal file
1
great_ai/large_file/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .large_file import LargeFileBase, LargeFileLocal, LargeFileMongo, LargeFileS3
|
||||
71
great_ai/large_file/__main__.py
Normal file
71
great_ai/large_file/__main__.py
Normal 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)
|
||||
2
great_ai/large_file/helper/__init__.py
Normal file
2
great_ai/large_file/helper/__init__.py
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
from .human_readable_to_byte import human_readable_to_byte
|
||||
from .progress_bar import DownloadProgressBar, UploadProgressBar
|
||||
2
great_ai/large_file/helper/bytes_to_megabytes.py
Normal file
2
great_ai/large_file/helper/bytes_to_megabytes.py
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
def bytes_to_megabytes(bytes: int) -> str:
|
||||
return f"{round(bytes / 1000 / 1000, 2):.2f}"
|
||||
31
great_ai/large_file/helper/human_readable_to_byte.py
Normal file
31
great_ai/large_file/helper/human_readable_to_byte.py
Normal 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)
|
||||
51
great_ai/large_file/helper/progress_bar.py
Normal file
51
great_ai/large_file/helper/progress_bar.py
Normal 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}")
|
||||
4
great_ai/large_file/large_file/__init__.py
Normal file
4
great_ai/large_file/large_file/__init__.py
Normal 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
|
||||
316
great_ai/large_file/large_file/large_file_base.py
Normal file
316
great_ai/large_file/large_file/large_file_base.py
Normal 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
|
||||
58
great_ai/large_file/large_file/large_file_local.py
Normal file
58
great_ai/large_file/large_file/large_file_local.py
Normal 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()
|
||||
121
great_ai/large_file/large_file/large_file_mongo.py
Normal file
121
great_ai/large_file/large_file/large_file_mongo.py
Normal 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])
|
||||
151
great_ai/large_file/large_file/large_file_s3.py
Normal file
151
great_ai/large_file/large_file/large_file_s3.py
Normal 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)
|
||||
1
great_ai/large_file/models/__init__.py
Normal file
1
great_ai/large_file/models/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .data_instance import DataInstance
|
||||
9
great_ai/large_file/models/data_instance.py
Normal file
9
great_ai/large_file/models/data_instance.py
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class DataInstance(BaseModel):
|
||||
name: str
|
||||
version: int
|
||||
remote_path: Any
|
||||
56
great_ai/large_file/parse_arguments.py
Normal file
56
great_ai/large_file/parse_arguments.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue