Move folders

src/__main__.py

@@ -14,7 +14,7 @@ from uvicorn.supervisors.basereload import BaseReload
 from watchdog.events import FileSystemEvent, PatternMatchingEventHandler
 from watchdog.observers import Observer
 
-from utilities import get_logger
+from great_ai.utilities import get_logger
 
 from .great_ai.constants import SERVER_NAME
 from .great_ai.context import _is_in_production_mode

src/great_ai/constants.py

@@ -1,4 +1,4 @@
-from large_file import LargeFileMongo, LargeFileS3
+from great_ai.large_file import LargeFileMongo, LargeFileS3
 
 from .persistence.mongodb_driver import MongodbDriver
 

src/great_ai/context.py

@@ -6,8 +6,8 @@ from typing import Any, Dict, Optional, Type, cast
 
 from pydantic import BaseModel
 
-from large_file import LargeFileBase, LargeFileLocal
+from great_ai.large_file import LargeFileBase, LargeFileLocal
-from utilities import get_logger
+from great_ai.utilities import get_logger
 
 from .constants import (
     DEFAULT_LARGE_FILE_CONFIG_PATHS,

src/great_ai/deploy/great_ai.py

@@ -16,7 +16,7 @@ from typing import (
 from fastapi import APIRouter, FastAPI, status
 from pydantic import BaseModel, create_model
 
-from utilities import parallel_map
+from great_ai.utilities import parallel_map
 
 from ..constants import DASHBOARD_PATH
 from ..context import get_context

src/great_ai/deploy/routes/dashboard/create_dash_app.py

@@ -8,7 +8,7 @@ from dash import Dash, dcc, html
 from dash.dependencies import Input, Output
 from flask import Flask
 
-from utilities import unique
+from great_ai.utilities import unique
 
 from ....constants import DASHBOARD_PATH, ONLINE_TAG_NAME
 from ....context import get_context

src/great_ai/large_file/__init__.py

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

src/great_ai/large_file/__main__.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+
+from argparse import Namespace
+from pathlib import Path
+from typing import Mapping, Type
+
+from great_ai.large_file import (
+    LargeFileBase,
+    LargeFileLocal,
+    LargeFileMongo,
+    LargeFileS3,
+)
+from great_ai.utilities import get_logger
+
+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)

src/great_ai/large_file/helper/__init__.py

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

src/great_ai/large_file/helper/bytes_to_megabytes.py

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

src/great_ai/large_file/helper/human_readable_to_byte.py

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

src/great_ai/large_file/helper/progress_bar.py

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

src/great_ai/large_file/large_file/__init__.py

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

src/great_ai/large_file/large_file/large_file_base.py

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

src/great_ai/large_file/large_file/large_file_local.py

@@ -0,0 +1,59 @@
+from pathlib import Path
+from typing import Any, List, Optional
+
+from great_ai.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()

src/great_ai/large_file/large_file/large_file_mongo.py

@@ -0,0 +1,122 @@
+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 great_ai.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])

src/great_ai/large_file/large_file/large_file_s3.py

@@ -0,0 +1,152 @@
+from functools import cached_property
+from pathlib import Path
+from typing import Any, List, Mapping, Optional
+
+import boto3
+
+from great_ai.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)

src/great_ai/large_file/models/__init__.py

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

src/great_ai/large_file/models/data_instance.py

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

src/great_ai/large_file/parse_arguments.py

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

src/great_ai/persistence/tracing_database_driver.py

@@ -3,7 +3,7 @@ from datetime import datetime
 from pathlib import Path
 from typing import List, Optional, Sequence, Tuple, Union
 
-from utilities import ConfigFile
+from great_ai.utilities import ConfigFile
 
 from ..views import Filter, SortBy, Trace