diff --git a/.tool-versions b/.tool-versions index c5cd8bf..c4f2bfc 100644 --- a/.tool-versions +++ b/.tool-versions @@ -1 +1 @@ -python 3.10.12 \ No newline at end of file +python 3.10.12 diff --git a/README.md b/README.md index abe4e73..8333e4a 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. @@ -38,11 +38,16 @@ Or you can directly invoke the script by running python s3_md5/main.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. -## 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/s3_md5/cmd.py b/s3_md5/cmd.py index f5e273b..5a2b9b1 100644 --- a/s3_md5/cmd.py +++ b/s3_md5/cmd.py @@ -1,27 +1,29 @@ '''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 -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, + ) + logger.info(f"md5 hash {md5_hash}") + logger.info(f"took {perf_counter() - start_time} seconds") if __name__ == "__main__": - run() + asyncio_run(run()) diff --git a/s3_md5/src/cli.py b/s3_md5/src/cli.py index c240a92..8caebad 100644 --- a/s3_md5/src/cli.py +++ b/s3_md5/src/cli.py @@ -2,12 +2,33 @@ 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 + + +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 +36,10 @@ 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() + 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..43242c8 --- /dev/null +++ b/s3_md5/src/consumer.py @@ -0,0 +1,36 @@ +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) 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}" + + " " + + f"left {chunk_count - element_to_consume + 1}") + 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..7ba1839 100644 --- a/s3_md5/src/s3_md5.py +++ b/s3_md5/src/s3_md5.py @@ -1,42 +1,77 @@ '''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, Awaitable from mypy_boto3_s3 import S3Client +from setproctitle import setproctitle +from .consumer import consumer from .logger import logger from .s3_file import S3FileHelper +setproctitle('s3-md5') -def parse_file_md5(s3_client: S3Client, - bucket: str, - file_name: str, - chunk_size: int, - workers: int) -> str: + +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): '''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 {file_size} bytes") 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 {chunk_size} bytes") + + chunk_count = file_size // chunk_size + logger.debug(f"chunk count {chunk_count}") + + md5_store = Manager().Value(str, '') + byte_store = Manager().dict() + + 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)) + + async def wrapper(part_number: int): + ranged_bytes_string = s3_file.calculate_range_bytes_from_part_number( + part_number, chunk_size, chunk_count) + logger.debug(f"downloading {ranged_bytes_string}") + ranged_bytes = await s3_file.get_range_bytes(ranged_bytes_string) + logger.debug(f"downloaded {ranged_bytes_string}") + byte_store[part_number] = ranged_bytes + + tasks = [asyncio.create_task(wrapper(part_number)) + for part_number in range(chunk_count)] + try: + await asyncio.gather(*tasks) + # pylint: disable=broad-exception-caught + except Exception as exception: + logger.error(f"parse_file_md5 {exception}") + consumer_process.terminate() + sys.exit(1) + + consumer_process.join() + return md5_store.value diff --git a/setup.py b/setup.py index 504ab79..acbfd1d 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": [ diff --git a/test/test_parse_file_md5.py b/test/test_parse_file_md5.py index 4c09799..72d800c 100644 --- a/test/test_parse_file_md5.py +++ b/test/test_parse_file_md5.py @@ -5,8 +5,8 @@ from s3_md5.src.s3_md5 import parse_file_md5 -def test_get_md5_hash(s3_setup: Tuple[S3Client, str, str, str]): +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 = parse_file_md5(s3_client, test_bucket, test_file_name, 1, 2) assert md5_hash == md5(bytes(test_body, 'utf-8')).hexdigest()