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

Format

Browse source 
Signed-off-by: András Schmelczer <andras@schmelczer.dev>
This commit is contained in:
Andras Schmelczer 2022-04-02 13:57:16 +02:00
parent 889e79174b
commit 60cd55c0cd
No known key found for this signature in database
GPG key ID: 39260B5B0614A13E
13 changed files with 168 additions and 116 deletions
Download patch file Download diff file Expand all files Collapse all files

2
good_ai/src/open_s3/helper/__init__.py
View file

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

2
good_ai/src/open_s3/helper/human_readable_to_byte.py
View file

@ -27,5 +27,5 @@ def human_readable_to_byte(size: str) -> int:
scalar = float(results["scalar"])
idx = possible_units.index(results["unit"].upper())
factor = 1024 ** idx
factor = 1024**idx
return round(scalar * factor)

19
good_ai/src/open_s3/helper/progress_bar.py
View file

@ -1,40 +1,35 @@
import os
import threading
from logging import Logger
from typing import IO, Any, Optional
from tqdm.auto import tqdm
from pathlib import Path
import os
import sys
import threading
class ProgressBar:
def __init__(self, file_size: int, logger: Logger, prefix: str):
self._file_size = file_size
self._logger = logger
self._prefix=prefix
self._prefix = prefix
self._seen_so_far = 0
self._lock = threading.Lock()
def __call__(self, bytes_amount: int):
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
size_length = len(str(self._file_size))
progress = str(self._seen_so_far).rjust(size_length)
self._logger.info(f"{self._prefix} {progress}/{self._file_size} bytes ({percentage:.1f}%)")
self._logger.info(
f"{self._prefix} {progress}/{self._file_size} bytes ({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}')
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}')
super().__init__(file_size=size, logger=logger, prefix=f"Uploading {path.name}")

66
good_ai/src/open_s3/large_file.py
View file

@ -6,9 +6,8 @@ import tempfile
from pathlib import Path
from types import TracebackType
from typing import IO, Any, Dict, List, Optional, Type, Union
import boto3
from botocore.exceptions import ClientError
import boto3
from helper import DownloadProgressBar, UploadProgressBar, human_readable_to_byte
logger = logging.getLogger("open_s3")
@ -62,7 +61,7 @@ class LargeFile:
newline: Optional[str] = None,
version: Optional[int] = None,
keep_last_n: Optional[int] = None,
offline_mode: bool = False
offline_mode: bool = False,
):
self._name: str = name
self._version = version
@ -79,7 +78,7 @@ class LargeFile:
self._find_versions()
self._check_mode_and_set_version()
@classmethod
def configure_credentials(
cls,
@ -154,7 +153,7 @@ class LargeFile:
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]
@ -175,8 +174,17 @@ class LargeFile:
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(size=size, name=key, logger=logger))
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(size=size, name=key, 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
@ -193,7 +201,7 @@ class LargeFile:
with tempfile.TemporaryDirectory() as tmp:
if path.is_file():
logger.info(f"Copying file for {self._local_name}")
copy = shutil.copy
copy: Any = shutil.copy
else:
logger.info(f"Copying directory for {self._local_name}")
copy = shutil.copytree
@ -216,7 +224,14 @@ class LargeFile:
logger.info(f"Uploading {self._local_name} to S3 from {path}")
file_to_be_uploaded = Path(tmp) / 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(file_to_be_uploaded, logger=logger))
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(file_to_be_uploaded, logger=logger),
)
self.clean_up()
@ -239,15 +254,21 @@ class LargeFile:
"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)
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:
self._fetch_versions_from_cache()
else:
self._create_client()
self._fetch_versions_from_s3()
if self._versions:
logger.info(f"Found versions: {self.version_ids}")
else:
@ -257,8 +278,7 @@ class LargeFile:
logger.info(f"Fetching offline versions of {self._name}")
self._versions = [
path
for path in self.cache_path.glob(f'{self._local_name}-*')
path for path in self.cache_path.glob(f"{self._local_name}-*")
]
def _fetch_versions_from_s3(self) -> None:
@ -267,10 +287,7 @@ class LargeFile:
Bucket=self.bucket_name, Prefix=self._name
)
self._versions = (
sorted(
o["Key"]
for o in found_objects["Contents"]
)
sorted(o["Key"] for o in found_objects["Contents"])
if "Contents" in found_objects
else []
)
@ -309,22 +326,25 @@ class LargeFile:
@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):
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)
@ -343,4 +363,4 @@ class LargeFile:
logger.info(
f"Deleting file from cache to meet quota (max_cache_size={self.max_cache_size}): {file}"
)
os.unlink(file)
os.unlink(file)

8
good_ai/tests/open_s3/test_human_readable_to_byte.py
View file

@ -4,22 +4,22 @@ from src.open_s3.helper import human_readable_to_byte
class TestHumanReadableToByte(unittest.TestCase):
def test_simple_cases(self):
def test_simple_cases(self) -> None:
self.assertEqual(human_readable_to_byte("1KB"), 1024)
self.assertEqual(human_readable_to_byte("2KB"), 2048)
def test_fractions(self):
def test_fractions(self) -> None:
self.assertEqual(human_readable_to_byte("0.5KB"), 512)
self.assertEqual(human_readable_to_byte("20.5KB"), 1024 * 20 + 512)
def test_formating(self):
def test_formating(self) -> None:
self.assertEqual(human_readable_to_byte(" 1MB"), 1024 * 1024)
self.assertEqual(human_readable_to_byte(" 2 MB"), 1024 * 1024 * 2)
self.assertEqual(human_readable_to_byte(" 4 MB "), 1024 * 1024 * 4)
self.assertEqual(human_readable_to_byte("8MB "), 1024 * 1024 * 8)
self.assertEqual(human_readable_to_byte(" 1.5 MB "), 1024 * 1024 * 1.5)
def test_casing(self):
def test_casing(self) -> None:
self.assertEqual(human_readable_to_byte("0.5GB"), 0.5 * 1024 * 1024 * 1024)
self.assertEqual(human_readable_to_byte("0.5gB"), 0.5 * 1024 * 1024 * 1024)
self.assertEqual(human_readable_to_byte("0.5Gb"), 0.5 * 1024 * 1024 * 1024)

12
good_ai/tests/open_s3/test_large_file.py
View file

@ -1,5 +1,6 @@
from pathlib import Path
import unittest
from pathlib import Path
from typing import Any
from unittest.mock import Mock, create_autospec, patch
import botocore.session
@ -14,23 +15,22 @@ credentials = {
"aws_access_key_id": "YOUR_ACCESS_KEY_ID",
"aws_secret_access_key": "YOUR_VERY_SECRET_ACCESS_KEY",
"large_files_bucket_name": "create_a_bucket_and_put_its_name_here",
"other_key": 23,
"endpoint_url": "this is optional, for backblaze, use this: https://s3.us-west-002.backblazeb2.com",
}
class TestLargeFile(unittest.TestCase):
def test_uninitialized(self):
def test_uninitialized(self) -> None:
self.assertRaises(ValueError, LargeFile, "test-file")
def test_bad_file_modes(self):
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")
@patch("botocore.session")
def test_initialized_with_dict(self, session):
def test_initialized_with_dict(self, session) -> None:
session_mock = Mock()
session.get_session = create_autospec(
botocore.session.get_session, return_value=session_mock
@ -79,7 +79,7 @@ class TestLargeFile(unittest.TestCase):
self.assertEqual(lf._s3_name, "test-file/2")
@patch("botocore.session")
def test_initialized_with_file(self, session):
def test_initialized_with_file(self, session: Any) -> None:
session_mock = Mock()
session.get_session = create_autospec(
botocore.session.get_session, return_value=session_mock