diff --git a/.github/workflows/package.yaml b/.github/workflows/package.yaml index e76cc5a..7e36e54 100644 --- a/.github/workflows/package.yaml +++ b/.github/workflows/package.yaml @@ -22,7 +22,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v4 with: - python-version: "3.10" + python-version: "3.12" - name: Create Tag id: tag_generator diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index a78391a..fff1d3a 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -11,7 +11,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v4 with: - python-version: "3.10" + python-version: "3.12" - name: Install Dependencies run: pip install ".[develop]" diff --git a/.tool-versions b/.tool-versions index c5cd8bf..69e5cf7 100644 --- a/.tool-versions +++ b/.tool-versions @@ -1 +1 @@ -python 3.10.12 \ No newline at end of file +python 3.12.2 diff --git a/README.md b/README.md index abe4e73..4f123da 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,12 @@ # s3-md5 -Get fast md5 hashes for an s3 file. +Get fast md5 hashes for an s3 file. This works by utilizing a process to fetch chunks into by doing them in threads while having another process consuming the shared dictionary. Given MD5 hashes need to be processed sequentially, it keeps looking for the expected chunk to ensure the order is right. -## installation +## Requirements - python 3.10 -## how to use +## Usage You can use the tool as a command line argument. You can download the latest release from [here](https://github.com/sakibstark11/s3-md5-python/releases). You can also build the wheel file yourself by running the following command. @@ -35,14 +35,19 @@ s3-md5 Or you can directly invoke the script by running ```sh -python s3_md5/main.py +python s3_md5/cmd.py ``` +### Arguments + There are two _optional_ arguments that you may want to provide -- `-w` or workers sets the number of python threads to use for downloading purposes, by default its set to the following equation `number of cpu cores * 2 - 1` -- `-c` or chunk size in **bytes** sets the individual download size on each get request sent to s3, by default its set to `1000000` +- `-c` or chunk size in **bytes** sets the individual download size on each get request sent to s3, by default it will use [speedtest-cli](https://pypi.org/project/speedtest-cli/) to determine the network speed. +- `-b` or block size to determine the number of maximum concurrent requests sent to s3 to protect against rate limiting. By default it is set to **10**. Please change this as this is related to your aws account s3 api rate limits. -## caveats +### Example -- File size can not be smaller than the default chunk size of `1000000`, if yes, then the chunk size must be manually provided or it will raise an assertion error. +for a file size of `1048576000` bytes +on a 250 mpbs bandwidth +on a macbook m1 8 core cpu +a chunk size of `4000000` works the best as it completes it within ~100 seconds diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..4088045 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +asyncio_mode=auto diff --git a/s3_md5/cmd.py b/s3_md5/cmd.py index f5e273b..0f13384 100644 --- a/s3_md5/cmd.py +++ b/s3_md5/cmd.py @@ -1,27 +1,36 @@ '''driver''' +from asyncio import run as asyncio_run from time import perf_counter -from boto3 import client +from aioboto3 import Session from s3_md5.src.cli import parse_args from s3_md5.src.logger import logger from s3_md5.src.s3_md5 import parse_file_md5 +from s3_md5.src.utils import seconds_to_minutes -def run(): +async def run(): + '''runs the script''' start_time = perf_counter() args = parse_args() - main_s3_client = client('s3') - md5_hash = parse_file_md5( - main_s3_client, - args.bucket, - args.file_name, - args.chunk_size, - args.workers - ) - logger.info(f'md5 hash {md5_hash}') - logger.info(f'took {perf_counter() - start_time} seconds') + main_s3_session = Session() + async with main_s3_session.client('s3') as s3_client: + md5_hash = await parse_file_md5( + s3_client, + args.bucket, + args.file_name, + args.chunk_size, + args.block_size + ) + logger.info(f"md5 hash {md5_hash}") + logger.info( + f"took {seconds_to_minutes(perf_counter() - start_time)} minute(s)") if __name__ == "__main__": - run() + try: + asyncio_run(run()) + # pylint: disable=broad-exception-caught + except Exception as exception: + logger.error(f"cmd {exception}") diff --git a/s3_md5/src/cli.py b/s3_md5/src/cli.py index c240a92..7205b63 100644 --- a/s3_md5/src/cli.py +++ b/s3_md5/src/cli.py @@ -2,12 +2,34 @@ from argparse import ArgumentParser from multiprocessing import cpu_count +from speedtest import Speedtest + +from .logger import logger + +DEFAULT_WORKERS = cpu_count() * 2 - 1 +BIT_IN_BYTE = 0.125 +DEFAULT_CHUNK_SIZE = 1000000 +DEFAULT_BLOCK_SIZE = 10 + + +def get_download_speed(): + '''uses speed test to get download speed''' + logger.info("picking chunk size") + try: + speed_test = Speedtest() + download_speed = speed_test.download(threads=1) + chunk_size = int(download_speed * BIT_IN_BYTE) + return chunk_size + # pylint: disable=broad-exception-caught + except Exception as exception: + logger.debug(f"get_download_speed {exception}") + logger.warning( + "will use default chunk size as automatic chunk size calculation failed") + return DEFAULT_CHUNK_SIZE + def parse_args(): '''parses command line arguments''' - DEFAULT_WORKERS = cpu_count() * 2 - 1 - DEFAULT_CHUNK_SIZE = 1000000 - parser = ArgumentParser(description='parse md5 of an s3 object') parser.add_argument('bucket', type=str, @@ -15,10 +37,13 @@ def parse_args(): parser.add_argument('file_name', help='file name', type=str) - parser.add_argument('-w', '--workers', type=int, - default=DEFAULT_WORKERS, - help='number of cpu threads to use for downloading') parser.add_argument('-c', '--chunk_size', type=int, - default=DEFAULT_CHUNK_SIZE, + default=None, help='chunk size to download on each request') - return parser.parse_args() + parser.add_argument('-b', '--block_size', type=int, + default=DEFAULT_BLOCK_SIZE, + help='maximum concurrent request') + parsed_args = parser.parse_args() + if parsed_args.chunk_size is None: + parsed_args.chunk_size = get_download_speed() + return parsed_args diff --git a/s3_md5/src/consumer.py b/s3_md5/src/consumer.py new file mode 100644 index 0000000..f81405b --- /dev/null +++ b/s3_md5/src/consumer.py @@ -0,0 +1,37 @@ +import sys +from hashlib import md5 +from multiprocessing.managers import ValueProxy +from typing import Dict + +from tqdm import tqdm + +from .logger import logger + + +def consumer(store: Dict[int, bytes], variable: ValueProxy[str], chunk_count: int): + '''a process that subscribes to the queue and processes md5''' + hasher = md5() + logger.debug("consumer started") + element_to_consume = 0 + with tqdm(total=chunk_count, position=1, desc="consumed") as progress_bar: + while element_to_consume < chunk_count: + try: + potential_item = store.get(element_to_consume) + if potential_item is not None: + hasher.update(potential_item) + logger.debug( + f"consumed chunk {element_to_consume + 1}" + + " " + + f"left {chunk_count - (element_to_consume + 1)}") + del store[element_to_consume] + element_to_consume += 1 + progress_bar.update(1) + # pylint: disable=broad-exception-caught + except Exception as exception: + logger.error(f"consumer {exception}") + sys.exit(1) + + logger.debug("calculating md5 hash") + md5_hash = hasher.hexdigest() + variable.value = md5_hash + sys.exit() diff --git a/s3_md5/src/logger.py b/s3_md5/src/logger.py index 0d5741c..12062e8 100644 --- a/s3_md5/src/logger.py +++ b/s3_md5/src/logger.py @@ -1,9 +1,18 @@ '''creates a logger''' import logging +import os import sys +LOG_LEVELS = { + 'CRITICAL': logging.CRITICAL, + 'WARNING': logging.WARNING, + 'ERROR': logging.ERROR, + 'DEBUG': logging.DEBUG, + 'INFO': logging.INFO +} logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) +log_level = LOG_LEVELS[os.getenv('LOG_LEVEL', 'INFO')] +logger.setLevel(log_level) stream_handler = logging.StreamHandler(sys.stdout) formatter = logging.Formatter( diff --git a/s3_md5/src/s3_file.py b/s3_md5/src/s3_file.py index d38fae8..7c75c2c 100644 --- a/s3_md5/src/s3_file.py +++ b/s3_md5/src/s3_file.py @@ -1,4 +1,6 @@ '''S3 file helper module''' +from typing import Awaitable + from mypy_boto3_s3 import S3Client @@ -11,10 +13,10 @@ def __init__(self, s3_client: S3Client, bucket: str, file_name: str) -> None: self.bucket = bucket self.file_name = file_name - def get_file_size(self) -> int: + async def get_file_size(self) -> Awaitable[int]: '''makes a head object request to get file size in bytes''' - s3_object = self.s3_client.head_object(Bucket=self.bucket, - Key=self.file_name) + s3_object = await self.s3_client.head_object(Bucket=self.bucket, + Key=self.file_name) self.__file_size = s3_object['ContentLength'] return self.__file_size @@ -33,10 +35,12 @@ def calculate_range_bytes_from_part_number(self, part_number: int, end_bytes: int = self.__file_size if part_number + \ 1 == file_chunk_count else (((part_number * chunk_size) + chunk_size) - 1) - return f'bytes={start_bytes}-{end_bytes}' + return f"bytes={start_bytes}-{end_bytes}" - def get_range_bytes(self, range_string: str) -> bytes: + async def get_range_bytes(self, range_string: str) -> bytes: '''fetches the range bytes requested from s3''' - return self.s3_client.get_object(Bucket=self.bucket, - Key=self.file_name, - Range=range_string)['Body'].read() + s3_object = await self.s3_client.get_object(Bucket=self.bucket, + Key=self.file_name, + Range=range_string) + async with s3_object['Body'] as stream: + return await stream.read() diff --git a/s3_md5/src/s3_md5.py b/s3_md5/src/s3_md5.py index eeb276a..afe3b9a 100644 --- a/s3_md5/src/s3_md5.py +++ b/s3_md5/src/s3_md5.py @@ -1,42 +1,84 @@ -'''module uses threads to download file from s3 and generates md5 hash''' -from concurrent.futures import ThreadPoolExecutor -from hashlib import md5 +import asyncio +import sys +from multiprocessing import Manager, Process +from signal import SIGCHLD, signal +from typing import Any from mypy_boto3_s3 import S3Client +from setproctitle import setproctitle +from tqdm import tqdm +from .consumer import consumer from .logger import logger from .s3_file import S3FileHelper +from .utils import bytes_to_mega_bytes +setproctitle('s3-md5') -def parse_file_md5(s3_client: S3Client, - bucket: str, - file_name: str, - chunk_size: int, - workers: int) -> str: - '''main function to orchestrate the md5 generation of s3 object''' + +def consumer_death_strategy(signal_number: int, + stack: Any, + process: Process): + '''Handler to call when consumer process dies''' + if process.exitcode != 0: + logger.error( + f"Consumer died with signal number {signal_number} exit code {process.exitcode}") + logger.error(f"Consumer stack {stack}") + logger.warning("Will exit") + process.terminate() + sys.exit(1) + logger.debug("Consumer process finished") + + +async def parse_file_md5(s3_client: S3Client, + bucket: str, + file_name: str, + chunk_size: int, + block_size: int): + '''Main function to orchestrate the MD5 generation of S3 object''' s3_file = S3FileHelper(s3_client, bucket, file_name) - file_size = s3_file.get_file_size() + file_size = await s3_file.get_file_size() + logger.info(f"File size {bytes_to_mega_bytes(file_size)} megabyte(s)") if file_size < chunk_size: - raise AssertionError('file size cannot be smaller than chunk size') - logger.info(f'file size {file_size} bytes') - file_chunk_count = file_size // chunk_size - logger.info(f'file chunk count {file_chunk_count}') - - with ThreadPoolExecutor(max_workers=workers) as thread_executor: - def wrapper(part_number: int): - ranged_bytes_string = s3_file.calculate_range_bytes_from_part_number( - part_number, chunk_size, file_chunk_count) - logger.info(f"downloading {ranged_bytes_string}") - ranged_bytes = s3_file.get_range_bytes(ranged_bytes_string) - logger.info(f"downloaded {ranged_bytes_string}") - return ranged_bytes - - logger.info('downloading file') - results = thread_executor.map(wrapper, - range(file_chunk_count)) - - hash_object = md5() - for result in results: - hash_object.update(result) - return hash_object.hexdigest() + chunk_size = file_size + logger.info(f"Chunk size {bytes_to_mega_bytes(chunk_size)} megabyte(s)") + + chunk_count = file_size // chunk_size + logger.info(f"Chunk count {chunk_count}") + + logger.info(f"Block size {block_size}") + + md5_store = Manager().Value(str, '') + byte_store = Manager().dict() + semaphore = asyncio.Semaphore(block_size) + + consumer_process = Process(target=consumer, args=( + byte_store, md5_store, chunk_count)) + consumer_process.start() + + signal(SIGCHLD, lambda signal_number, stack: consumer_death_strategy( + signal_number, stack, consumer_process)) + + with tqdm(total=chunk_count, position=0, desc="downloaded") as progress_bar: + async with asyncio.TaskGroup() as task_group: + async def wrapper(part_number: int): + ranged_bytes_string = s3_file.calculate_range_bytes_from_part_number( + part_number, chunk_size, chunk_count) + async with semaphore: + logger.debug( + f"Downloading {part_number + 1} {ranged_bytes_string}") + ranged_bytes = await s3_file.get_range_bytes(ranged_bytes_string) + logger.debug( + f"Downloaded {part_number + 1} {ranged_bytes_string}") + progress_bar.update(1) + byte_store[part_number] = ranged_bytes + + # Process tasks in blocks + for i in range(0, chunk_count, block_size): + block_end = min(i + block_size, chunk_count) + for part_number in range(i, block_end): + task_group.create_task(wrapper(part_number)) + + consumer_process.join() + return md5_store.value diff --git a/s3_md5/src/utils.py b/s3_md5/src/utils.py new file mode 100644 index 0000000..a2ae509 --- /dev/null +++ b/s3_md5/src/utils.py @@ -0,0 +1,11 @@ +'''basic converter utilities''' + + +def bytes_to_mega_bytes(value: int) -> float: + '''convert bytes to megabytes''' + return value / (1000 * 1000) + + +def seconds_to_minutes(value: float) -> float: + '''convert seconds to minutes''' + return value / 60 diff --git a/setup.py b/setup.py index 504ab79..d325893 100644 --- a/setup.py +++ b/setup.py @@ -12,8 +12,11 @@ author_email="16sakib@gmail.com", license="MIT", install_requires=[ - "boto3==1.26.41", "boto3-stubs[s3]", + "setproctitle==1.3.3", + "tqdm==4.66.2", + "speedtest-cli==2.1.3", + "aioboto3==12.3.0" ], extras_require={ "develop": [ @@ -21,6 +24,7 @@ "moto==4.0.12", "pytest==7.2.0", "pylint==3.1.0", + "pytest-asyncio==0.23.8" ], "release": ["wheel==0.43.0"] }, diff --git a/test/conftest.py b/test/conftest.py index 2cb6019..6d7c38e 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,20 +1,24 @@ -'''common resource for testing''' -from boto3 import client +import aioboto3 from moto import mock_s3 from pytest import fixture @fixture -def s3_setup(): +async def s3_setup(): with mock_s3(): - s3_client = client('s3') - test_bucket = 'bucket' - test_file_name = 'key' - test_body = '0123456789' - s3_client.create_bucket(Bucket=test_bucket, CreateBucketConfiguration={ - 'LocationConstraint': 'ap-east-1', - },) - s3_client.put_object(Bucket=test_bucket, - Key=test_file_name, - Body=test_body) - yield s3_client, test_bucket, test_file_name, test_body + async with aioboto3.Session().client('s3', region_name='ap-east-1') as s3_client: + test_bucket = 'bucket' + test_file_name = 'key' + test_body = '0123456789' + + await s3_client.create_bucket( + Bucket=test_bucket, + CreateBucketConfiguration={'LocationConstraint': 'ap-east-1'} + ) + await s3_client.put_object( + Bucket=test_bucket, + Key=test_file_name, + Body=test_body + ) + + yield s3_client, test_bucket, test_file_name, test_body diff --git a/test/test_calculate_range_bytes_from_part_number.py b/test/test_calculate_range_bytes_from_part_number.py index 16f4c4a..1c198c9 100644 --- a/test/test_calculate_range_bytes_from_part_number.py +++ b/test/test_calculate_range_bytes_from_part_number.py @@ -1,14 +1,16 @@ '''tests range byte calculator function''' from typing import Tuple + from mypy_boto3_s3 import S3Client + from s3_md5.src.s3_file import S3FileHelper -def test_calculate_range_bytes_from_part_number(s3_setup: Tuple[S3Client, str, str, str]): +async def test_calculate_range_bytes_from_part_number(s3_setup: Tuple[S3Client, str, str, str]): '''test function''' s3_client, test_bucket, test_file_name, _ = s3_setup s3_file = S3FileHelper(s3_client, test_bucket, test_file_name) - s3_file.get_file_size() + await s3_file.get_file_size() part_number = 1 chunk_size = 1000000 file_chunk_count = 10 diff --git a/test/test_get_file_size.py b/test/test_get_file_size.py index 22d5e77..975f691 100644 --- a/test/test_get_file_size.py +++ b/test/test_get_file_size.py @@ -1,12 +1,14 @@ '''tests file size method''' from typing import Tuple + from mypy_boto3_s3 import S3Client + from s3_md5.src.s3_file import S3FileHelper -def test_get_file_size(s3_setup: Tuple[S3Client, str, str, str]): +async def test_get_file_size(s3_setup: Tuple[S3Client, str, str, str]): '''test function''' s3_client, test_bucket, test_file_name, test_body = s3_setup s3_file = S3FileHelper(s3_client, test_bucket, test_file_name) - file_size = s3_file.get_file_size() + file_size = await s3_file.get_file_size() assert file_size == len(test_body) diff --git a/test/test_get_range_bytes.py b/test/test_get_range_bytes.py index 6afb26c..cb884ce 100644 --- a/test/test_get_range_bytes.py +++ b/test/test_get_range_bytes.py @@ -1,14 +1,16 @@ '''tests range fetching function''' from typing import Tuple + from mypy_boto3_s3 import S3Client + from s3_md5.src.s3_file import S3FileHelper -def test_get_range_bytes(s3_setup: Tuple[S3Client, str, str, str]): +async def test_get_range_bytes(s3_setup: Tuple[S3Client, str, str, str]): '''test function''' s3_client, test_bucket, test_file_name, test_body = s3_setup s3_file = S3FileHelper(s3_client, test_bucket, test_file_name) - s3_file.get_file_size() + await s3_file.get_file_size() range_string = 'bytes=0-0' - data = s3_file.get_range_bytes(range_string) + data = await s3_file.get_range_bytes(range_string) assert data == bytes(test_body[0], 'utf-8') diff --git a/test/test_parse_file_md5.py b/test/test_parse_file_md5.py index 4c09799..cba25a8 100644 --- a/test/test_parse_file_md5.py +++ b/test/test_parse_file_md5.py @@ -1,12 +1,15 @@ '''tests the driver function''' from hashlib import md5 from typing import Tuple + from mypy_boto3_s3 import S3Client + from s3_md5.src.s3_md5 import parse_file_md5 -def test_get_md5_hash(s3_setup: Tuple[S3Client, str, str, str]): +async def test_parse_file_md5(s3_setup: Tuple[S3Client, str, str, str]): '''test function''' s3_client, test_bucket, test_file_name, test_body = s3_setup - md5_hash = parse_file_md5(s3_client, test_bucket, test_file_name, 1, 1) + md5_hash = await parse_file_md5( + s3_client, test_bucket, test_file_name, 1, 2) assert md5_hash == md5(bytes(test_body, 'utf-8')).hexdigest()