andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Refactor LargeFile

Browse source
This commit is contained in:
Andras Schmelczer 2022-06-03 10:44:01 +02:00
parent a19a39a3a0
commit a4e16d5739
18 changed files with 583 additions and 396 deletions
Download patch file Download diff file Expand all files Collapse all files

1
great_ai/src/great_ai/large_file/__init__.py Normal file
View file

@ -0,0 +1 @@
from .large_file import LargeFile, LargeFileS3

26
great_ai/src/great_ai/open_s3/__main__.py → great_ai/src/great_ai/large_file/__main__.py
View file

@ -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__":

0
great_ai/src/great_ai/open_s3/helper/__init__.py → great_ai/src/great_ai/large_file/helper/__init__.py
View file

0
great_ai/src/great_ai/open_s3/helper/bytes_to_megabytes.py → great_ai/src/great_ai/large_file/helper/bytes_to_megabytes.py
View file

0
great_ai/src/great_ai/open_s3/helper/human_readable_to_byte.py → great_ai/src/great_ai/large_file/helper/human_readable_to_byte.py
View file

0
great_ai/src/great_ai/open_s3/helper/progress_bar.py → great_ai/src/great_ai/large_file/helper/progress_bar.py
View file

3
great_ai/src/great_ai/large_file/large_file/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .large_file import LargeFile
from .large_file_local import LargeFileLocal
from .large_file_s3 import LargeFileS3

300
great_ai/src/great_ai/large_file/large_file/large_file.py Normal file
View file

@ -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

58
great_ai/src/great_ai/large_file/large_file/large_file_local.py Normal file
View file

@ -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()

169
great_ai/src/great_ai/large_file/large_file/large_file_s3.py Normal file
View file

@ -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)

1
great_ai/src/great_ai/large_file/models/__init__.py Normal file
View file

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

10
great_ai/src/great_ai/large_file/models/data_instance.py Normal file
View file

@ -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"]

11
great_ai/src/great_ai/open_s3/parse_arguments.py → great_ai/src/great_ai/large_file/parse_arguments.py
View file

@ -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(

1
great_ai/src/great_ai/open_s3/__init__.py
View file

@ -1 +0,0 @@
from .large_file import LargeFile

373
great_ai/src/great_ai/open_s3/large_file.py
View file

@ -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)

0
great_ai/tests/open_s3/__init__.py → great_ai/tests/large_file/__init__.py
View file

2
great_ai/tests/open_s3/test_human_readable_to_byte.py → great_ai/tests/large_file/test_human_readable_to_byte.py
View file

@ -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):

24
great_ai/tests/open_s3/test_large_file.py → great_ai/tests/large_file/test_large_file.py
View file

@ -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")