From a4e16d5739705c88d2b537dd0bd17d2d26174cff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= Date: Fri, 3 Jun 2022 10:44:01 +0200 Subject: [PATCH] Refactor LargeFile --- great_ai/src/great_ai/large_file/__init__.py | 1 + .../{open_s3 => large_file}/__main__.py | 26 +- .../helper/__init__.py | 0 .../helper/bytes_to_megabytes.py | 0 .../helper/human_readable_to_byte.py | 0 .../helper/progress_bar.py | 0 .../large_file/large_file/__init__.py | 3 + .../large_file/large_file/large_file.py | 300 ++++++++++++++ .../large_file/large_file/large_file_local.py | 58 +++ .../large_file/large_file/large_file_s3.py | 169 ++++++++ .../great_ai/large_file/models/__init__.py | 1 + .../large_file/models/data_instance.py | 10 + .../parse_arguments.py | 11 +- great_ai/src/great_ai/open_s3/__init__.py | 1 - great_ai/src/great_ai/open_s3/large_file.py | 373 ------------------ .../tests/{open_s3 => large_file}/__init__.py | 0 .../test_human_readable_to_byte.py | 2 +- .../test_large_file.py | 24 +- 18 files changed, 583 insertions(+), 396 deletions(-) create mode 100644 great_ai/src/great_ai/large_file/__init__.py rename great_ai/src/great_ai/{open_s3 => large_file}/__main__.py (51%) rename great_ai/src/great_ai/{open_s3 => large_file}/helper/__init__.py (100%) rename great_ai/src/great_ai/{open_s3 => large_file}/helper/bytes_to_megabytes.py (100%) rename great_ai/src/great_ai/{open_s3 => large_file}/helper/human_readable_to_byte.py (100%) rename great_ai/src/great_ai/{open_s3 => large_file}/helper/progress_bar.py (100%) create mode 100644 great_ai/src/great_ai/large_file/large_file/__init__.py create mode 100644 great_ai/src/great_ai/large_file/large_file/large_file.py create mode 100644 great_ai/src/great_ai/large_file/large_file/large_file_local.py create mode 100644 great_ai/src/great_ai/large_file/large_file/large_file_s3.py create mode 100644 great_ai/src/great_ai/large_file/models/__init__.py create mode 100644 great_ai/src/great_ai/large_file/models/data_instance.py rename great_ai/src/great_ai/{open_s3 => large_file}/parse_arguments.py (80%) delete mode 100644 great_ai/src/great_ai/open_s3/__init__.py delete mode 100644 great_ai/src/great_ai/open_s3/large_file.py rename great_ai/tests/{open_s3 => large_file}/__init__.py (100%) rename great_ai/tests/{open_s3 => large_file}/test_human_readable_to_byte.py (94%) rename great_ai/tests/{open_s3 => large_file}/test_large_file.py (82%) diff --git a/great_ai/src/great_ai/large_file/__init__.py b/great_ai/src/great_ai/large_file/__init__.py new file mode 100644 index 0000000..437e068 --- /dev/null +++ b/great_ai/src/great_ai/large_file/__init__.py @@ -0,0 +1 @@ +from .large_file import LargeFile, LargeFileS3 diff --git a/great_ai/src/great_ai/open_s3/__main__.py b/great_ai/src/great_ai/large_file/__main__.py similarity index 51% rename from great_ai/src/great_ai/open_s3/__main__.py rename to great_ai/src/great_ai/large_file/__main__.py index 113e3a2..9881a22 100644 --- a/great_ai/src/great_ai/open_s3/__main__.py +++ b/great_ai/src/great_ai/large_file/__main__.py @@ -4,16 +4,30 @@ from pathlib import Path from great_ai.utilities.logger import get_logger -from .large_file import LargeFile +from .large_file import LargeFileLocal, LargeFileS3 from .parse_arguments import parse_arguments -logger = get_logger("open_s3") +logger = get_logger("large_file") def main() -> None: parser, args = parse_arguments() - LargeFile.configure_credentials_from_file(args.secrets) + factory = {"s3": LargeFileS3, "local": LargeFileLocal} + + if args.driver not in factory: + raise ValueError( + f"Driver {args.driver} does not exits, available options: {' ,'.join(factory.keys())}" + ) + + large_file = factory[args.driver] + + if args.driver != "local": + if args.secrets is None: + raise ValueError( + "Providing a credentials file is required when the driver mode is not `local`." + ) + large_file.configure_credentials_from_file(args.secrets) # type: ignore if not args.cache and not args.push and not args.delete: logger.warning("No action required.") @@ -24,16 +38,16 @@ def main() -> None: split = c.split(":") file_name = split[0] version = None if len(split) == 1 else int(split[1]) - LargeFile(file_name, "r", version=version).get() + large_file(file_name, "r", version=version).get() if args.push: for p in args.push: path = Path(p) - LargeFile(path.name, "w").push(path) + large_file(path.name, "w").push(path) if args.delete: for f in args.delete: - LargeFile(f).delete() + large_file(f).delete() if __name__ == "__main__": diff --git a/great_ai/src/great_ai/open_s3/helper/__init__.py b/great_ai/src/great_ai/large_file/helper/__init__.py similarity index 100% rename from great_ai/src/great_ai/open_s3/helper/__init__.py rename to great_ai/src/great_ai/large_file/helper/__init__.py diff --git a/great_ai/src/great_ai/open_s3/helper/bytes_to_megabytes.py b/great_ai/src/great_ai/large_file/helper/bytes_to_megabytes.py similarity index 100% rename from great_ai/src/great_ai/open_s3/helper/bytes_to_megabytes.py rename to great_ai/src/great_ai/large_file/helper/bytes_to_megabytes.py diff --git a/great_ai/src/great_ai/open_s3/helper/human_readable_to_byte.py b/great_ai/src/great_ai/large_file/helper/human_readable_to_byte.py similarity index 100% rename from great_ai/src/great_ai/open_s3/helper/human_readable_to_byte.py rename to great_ai/src/great_ai/large_file/helper/human_readable_to_byte.py diff --git a/great_ai/src/great_ai/open_s3/helper/progress_bar.py b/great_ai/src/great_ai/large_file/helper/progress_bar.py similarity index 100% rename from great_ai/src/great_ai/open_s3/helper/progress_bar.py rename to great_ai/src/great_ai/large_file/helper/progress_bar.py diff --git a/great_ai/src/great_ai/large_file/large_file/__init__.py b/great_ai/src/great_ai/large_file/large_file/__init__.py new file mode 100644 index 0000000..cd28af0 --- /dev/null +++ b/great_ai/src/great_ai/large_file/large_file/__init__.py @@ -0,0 +1,3 @@ +from .large_file import LargeFile +from .large_file_local import LargeFileLocal +from .large_file_s3 import LargeFileS3 diff --git a/great_ai/src/great_ai/large_file/large_file/large_file.py b/great_ai/src/great_ai/large_file/large_file/large_file.py new file mode 100644 index 0000000..ca3a434 --- /dev/null +++ b/great_ai/src/great_ai/large_file/large_file/large_file.py @@ -0,0 +1,300 @@ +import os +import shutil +import tempfile +from abc import ABC, abstractmethod +from pathlib import Path +from types import TracebackType +from typing import IO, Any, List, Optional, Type, Union, cast + +from great_ai.utilities.logger import 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 LargeFile(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" + ``` + """ + + 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 + + LargeFile.cache_path.mkdir(parents=True, exist_ok=True) + + self._find_instances() + self._check_mode_and_set_version() + + 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) + + 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 + 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, + origin="filesystem", + ) + 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 diff --git a/great_ai/src/great_ai/large_file/large_file/large_file_local.py b/great_ai/src/great_ai/large_file/large_file/large_file_local.py new file mode 100644 index 0000000..0b76749 --- /dev/null +++ b/great_ai/src/great_ai/large_file/large_file/large_file_local.py @@ -0,0 +1,58 @@ +from pathlib import Path +from typing import Any, List, Optional + +from great_ai.utilities.logger import get_logger + +from ..models import DataInstance +from .large_file import LargeFile + +logger = get_logger("large_file") + + +class LargeFileLocal(LargeFile): + 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, + ) + + 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() diff --git a/great_ai/src/great_ai/large_file/large_file/large_file_s3.py b/great_ai/src/great_ai/large_file/large_file/large_file_s3.py new file mode 100644 index 0000000..5c7a8e0 --- /dev/null +++ b/great_ai/src/great_ai/large_file/large_file/large_file_s3.py @@ -0,0 +1,169 @@ +import configparser +from functools import cached_property +from pathlib import Path +from typing import Any, List, Mapping, Optional, Union + +import boto3 + +from great_ai.utilities.logger import get_logger + +from ..helper import DownloadProgressBar, UploadProgressBar +from ..models import DataInstance +from .large_file import LargeFile + +logger = get_logger("large_file") + + +S3_NAME_VERSION_SEPARATOR = "/" + + +class LargeFileS3(LargeFile): + """ + 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_from_file( + cls, + secrets_path: Union[Path, str], + ) -> None: + if isinstance(secrets_path, str): + secrets_path = Path(secrets_path) + + if not secrets_path.exists(): + raise FileNotFoundError(secrets_path.resolve()) + + credentials = configparser.ConfigParser() + credentials.read(secrets_path) + credentials.default_section + cls.configure_credentials(**credentials[credentials.default_section]) + + @classmethod + def configure_credentials( + 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 + + @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 LargeFile.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"], + origin="s3", + ) + 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=hide_progress, + 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) diff --git a/great_ai/src/great_ai/large_file/models/__init__.py b/great_ai/src/great_ai/large_file/models/__init__.py new file mode 100644 index 0000000..a5cb501 --- /dev/null +++ b/great_ai/src/great_ai/large_file/models/__init__.py @@ -0,0 +1 @@ +from .data_instance import DataInstance diff --git a/great_ai/src/great_ai/large_file/models/data_instance.py b/great_ai/src/great_ai/large_file/models/data_instance.py new file mode 100644 index 0000000..7de533c --- /dev/null +++ b/great_ai/src/great_ai/large_file/models/data_instance.py @@ -0,0 +1,10 @@ +from typing import Any, Literal + +from pydantic import BaseModel + + +class DataInstance(BaseModel): + name: str + version: int + remote_path: Any + origin: Literal["filesystem", "mongodb", "s3"] diff --git a/great_ai/src/great_ai/open_s3/parse_arguments.py b/great_ai/src/great_ai/large_file/parse_arguments.py similarity index 80% rename from great_ai/src/great_ai/open_s3/parse_arguments.py rename to great_ai/src/great_ai/large_file/parse_arguments.py index acb6966..18c395a 100644 --- a/great_ai/src/great_ai/open_s3/parse_arguments.py +++ b/great_ai/src/great_ai/large_file/parse_arguments.py @@ -7,12 +7,19 @@ def parse_arguments() -> Tuple[ArgumentParser, Namespace]: description="Store and version large files in S3; open them like regular files. Caching included.", ) + parser.add_argument( + "--driver", + 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 configration file with your S3 credentials", - required=True, + help="path to an .ini configuration file with your S3 credentials", + required=False, ) parser.add_argument( diff --git a/great_ai/src/great_ai/open_s3/__init__.py b/great_ai/src/great_ai/open_s3/__init__.py deleted file mode 100644 index 3e78b20..0000000 --- a/great_ai/src/great_ai/open_s3/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .large_file import LargeFile diff --git a/great_ai/src/great_ai/open_s3/large_file.py b/great_ai/src/great_ai/open_s3/large_file.py deleted file mode 100644 index 271d994..0000000 --- a/great_ai/src/great_ai/open_s3/large_file.py +++ /dev/null @@ -1,373 +0,0 @@ -import configparser -import os -import shutil -import tempfile -from pathlib import Path -from types import TracebackType -from typing import IO, Any, List, Mapping, Optional, Type, Union, cast - -import boto3 - -from great_ai.utilities.logger import get_logger - -from .helper import DownloadProgressBar, UploadProgressBar, human_readable_to_byte - -logger = get_logger("open_s3") - - -class LargeFile: - """ - 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 - - 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, - offline_mode: bool = False, - ): - self._name: str = name - self._version: int = cast(int, version) - self._mode: str = mode - self._keep_last_n = keep_last_n - self._offline_mode = offline_mode - - self._buffering = buffering - self._encoding = encoding - self._errors = errors - self._newline = newline - - LargeFile.cache_path.mkdir(parents=True, exist_ok=True) - - self._find_versions() - self._check_mode_and_set_version() - - @classmethod - def configure_credentials_from_file( - cls, - secrets_path: Union[Path, str], - ) -> None: - if isinstance(secrets_path, str): - secrets_path = Path(secrets_path) - - if not secrets_path.exists(): - raise FileNotFoundError(secrets_path.resolve()) - - credentials = configparser.ConfigParser() - credentials.read(secrets_path) - credentials.default_section - cls.configure_credentials(**credentials[credentials.default_section]) - - @classmethod - def configure_credentials( - 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 - - 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 version_ids(self) -> List[int]: - return [self._get_version_from_key(key) for key in self._versions] - - @property - def version(self) -> int: - return self._version - - def get(self, hide_progress: bool = False) -> Path: - key = next( - key - for key in self._versions - if self._get_version_from_key(key) == self._version - ) - - destination = self.cache_path / self._local_name - if not destination.exists(): - logger.info( - f"File {self._local_name} does not exist locally, starting download from S3" - ) - - with tempfile.TemporaryDirectory() as tmp: - tmp_file_archive = Path(tmp) / f"{self._local_name}.tar.gz" - - size = self._client.head_object(Bucket=self.bucket_name, Key=key)[ - "ContentLength" - ] - self._client.download_file( - Bucket=self.bucket_name, - Key=key, - Filename=str(tmp_file_archive), - Callback=None - if hide_progress - else DownloadProgressBar(name=str(key), size=size, logger=logger), - ) - logger.info(f"Decompressing {self._local_name}") - shutil.unpack_archive(str(tmp_file_archive), tmp, "gztar") - tmp_file = Path(tmp) / self._local_name - shutil.move(str(tmp_file), 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 - 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), - "gztar", - tmp, - ) - - logger.info(f"Uploading {self._local_name} to S3 from {path}") - - file_to_be_uploaded = Path(tmp2) / f"{self._local_name}.tar.gz" - self._client.upload_file( - Filename=str(file_to_be_uploaded), - Bucket=self.bucket_name, - Key=self._s3_name, - Callback=None - if hide_progress - else UploadProgressBar(path=file_to_be_uploaded, logger=logger), - ) - - self.clean_up() - - def delete(self) -> None: - self._keep_last_n = 0 - self._delete_old_versions_from_s3() - - def clean_up(self) -> None: - self._delete_old_versions_from_s3() - self._delete_old_versions_from_disk() - - def _create_client(self) -> None: - 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 LargeFile.configure_credentials or set offline_mode=True in the constructor." - ) - - self._client = 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_versions(self) -> None: - if self._offline_mode: - versions: Union[List[str], List[Path]] = self._fetch_versions_from_cache() - else: - self._create_client() - versions = self._fetch_versions_from_s3() - - self._versions = sorted(versions, key=self._get_version_from_key) - - def _fetch_versions_from_cache(self) -> List[Path]: - logger.debug(f"Fetching offline versions of {self._name}") - - return list(self.cache_path.glob(f"{self._local_name}-*")) - - def _fetch_versions_from_s3(self) -> List[str]: - logger.debug(f"Fetching online versions of {self._name}") - - found_objects = self._client.list_objects_v2( - Bucket=self.bucket_name, Prefix=self._name - ) - return ( - [o["Key"] for o in found_objects["Contents"]] - if "Contents" in found_objects - else [] - ) - - def _check_mode_and_set_version(self) -> None: - if "+" in self._mode: - raise ValueError("Read-write mode is not allowed.") - - 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.version_ids[-1] + 1 if self.version_ids else 0 - - elif "r" in self._mode: - if not self.version_ids: - raise FileNotFoundError( - f"File {self._name} not found. No versions are available." - ) - - if self._version is None: - self._version = self.version_ids[-1] - logger.info( - f"Latest version of {self._name} is {self._version} " - + f"(from versions: {', '.join((str(v) for v in self.version_ids))})" - ) - - elif self._version not in self.version_ids: - raise FileNotFoundError( - f"File {self._name} not found with version {self._version}. " - + f"(from versions: {', '.join((str(v) for v in self.version_ids))})" - ) - else: - raise ValueError("Unsupported file mode.") - - @property - def _local_name(self) -> str: - return f"{self._name}-{self._version}" - - @property - def _s3_name(self) -> str: - return f"{self._name}/{self._version}" - - @staticmethod - def _get_version_from_key(key: Union[str, Path]) -> int: - if isinstance(key, Path): - return int(key.name.split("-")[-1]) - return int(key.split("/")[-1]) - - def _delete_old_versions_from_s3(self) -> None: - if self._keep_last_n is not None: - for key in ( - self._versions[: -self._keep_last_n] - if self._keep_last_n > 0 - else self._versions - ): - logger.info( - f"Removing old version (keep_last_n={self._keep_last_n}): {key}" - ) - self._client.delete_object(Bucket=self.bucket_name, Key=key) - - def _delete_old_versions_from_disk(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) diff --git a/great_ai/tests/open_s3/__init__.py b/great_ai/tests/large_file/__init__.py similarity index 100% rename from great_ai/tests/open_s3/__init__.py rename to great_ai/tests/large_file/__init__.py diff --git a/great_ai/tests/open_s3/test_human_readable_to_byte.py b/great_ai/tests/large_file/test_human_readable_to_byte.py similarity index 94% rename from great_ai/tests/open_s3/test_human_readable_to_byte.py rename to great_ai/tests/large_file/test_human_readable_to_byte.py index 64576bf..a7290b5 100644 --- a/great_ai/tests/open_s3/test_human_readable_to_byte.py +++ b/great_ai/tests/large_file/test_human_readable_to_byte.py @@ -1,6 +1,6 @@ import unittest -from src.great_ai.open_s3.helper import human_readable_to_byte +from src.great_ai.large_file.helper import human_readable_to_byte class TestHumanReadableToByte(unittest.TestCase): diff --git a/great_ai/tests/open_s3/test_large_file.py b/great_ai/tests/large_file/test_large_file.py similarity index 82% rename from great_ai/tests/open_s3/test_large_file.py rename to great_ai/tests/large_file/test_large_file.py index 99ff39c..5b163ed 100644 --- a/great_ai/tests/open_s3/test_large_file.py +++ b/great_ai/tests/large_file/test_large_file.py @@ -8,7 +8,7 @@ import boto3 PATH = Path(__file__).parent.resolve() -from src.great_ai import LargeFile +from src.great_ai import LargeFileS3 credentials = { "aws_region_name": "your_region_like_eu-west-2", @@ -19,15 +19,15 @@ credentials = { } -class TestLargeFile(unittest.TestCase): +class TestLargeFileS3(unittest.TestCase): def test_uninitialized(self) -> None: - self.assertRaises(ValueError, LargeFile, "test-file") + self.assertRaises(ValueError, LargeFileS3, "test-file") def test_bad_file_modes(self) -> None: - self.assertRaises(ValueError, LargeFile, "test-file", "w", version=3) - self.assertRaises(ValueError, LargeFile, "test-file", "wb", version=3) - self.assertRaises(ValueError, LargeFile, "test-file", "w+r") - self.assertRaises(ValueError, LargeFile, "test-file", "test") + self.assertRaises(ValueError, LargeFileS3, "test-file", "w", version=3) + self.assertRaises(ValueError, LargeFileS3, "test-file", "wb", version=3) + self.assertRaises(ValueError, LargeFileS3, "test-file", "w+r") + self.assertRaises(ValueError, LargeFileS3, "test-file", "test") @patch.object(boto3, "client") def test_initialized_with_dict(self, client: Any) -> None: @@ -52,14 +52,14 @@ class TestLargeFile(unittest.TestCase): ) boto3.client = Mock(return_value=s3) - LargeFile.configure_credentials( + LargeFileS3.configure_credentials( aws_region_name=credentials["aws_region_name"], aws_access_key_id=credentials["aws_access_key_id"], aws_secret_access_key=credentials["aws_secret_access_key"], large_files_bucket_name=credentials["large_files_bucket_name"], aws_endpoint_url=credentials["aws_endpoint_url"], ) - lf = LargeFile("test-file") + lf = LargeFileS3("test-file") boto3.client.assert_called_once_with( "s3", @@ -75,7 +75,6 @@ class TestLargeFile(unittest.TestCase): self.assertEqual(lf._version, 2) self.assertEqual(lf._local_name, "test-file-2") - self.assertEqual(lf._s3_name, "test-file/2") @patch.object(boto3, "client") def test_initialized_with_file(self, client: Any) -> None: @@ -101,8 +100,8 @@ class TestLargeFile(unittest.TestCase): boto3.client = Mock(return_value=s3) - LargeFile.configure_credentials_from_file(PATH / "../../example_secrets.ini") - lf = LargeFile("test-file") + LargeFileS3.configure_credentials_from_file(PATH / "../../example_secrets.ini") + lf = LargeFileS3("test-file") boto3.client.assert_called_once_with( "s3", @@ -119,4 +118,3 @@ class TestLargeFile(unittest.TestCase): self.assertEqual(lf._version, 2) self.assertEqual(lf._local_name, "test-file-2") - self.assertEqual(lf._s3_name, "test-file/2")