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..1e95559 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. +- `-b` or block size that determines how many concurrent requests the application should send to s3. This is to prevent s3 from throttling the application due to rate limiting. By default it is set to **10**. Please increase this as this is dependent on your aws account s3 api request limit. -## 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..086046e 100644 --- a/s3_md5/cmd.py +++ b/s3_md5/cmd.py @@ -1,27 +1,32 @@ '''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() + asyncio_run(run()) 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/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..190be7c 100644 --- a/s3_md5/src/s3_md5.py +++ b/s3_md5/src/s3_md5.py @@ -1,42 +1,74 @@ '''module uses threads to download file from s3 and generates md5 hash''' -from concurrent.futures import ThreadPoolExecutor +import asyncio +import sys from hashlib import md5 from mypy_boto3_s3 import S3Client +from setproctitle import setproctitle +from tqdm import tqdm 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: + +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}') + chunk_size = file_size + logger.info(f"chunk size {bytes_to_mega_bytes(chunk_size)} megabyte(s)") + + chunk_count = file_size // chunk_size + logger.debug(f"chunk count {chunk_count}") + + logger.info(f"block size {block_size}") + + md5_store = md5() + byte_store = {} + next_ingest_part_number = 0 + progress_bar = tqdm(total=chunk_count) + semaphore = asyncio.Semaphore(block_size) - with ThreadPoolExecutor(max_workers=workers) as thread_executor: - def wrapper(part_number: int): + async def wrapper(part_number: int): + async with semaphore: 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() + 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 + + nonlocal next_ingest_part_number + while next_ingest_part_number < chunk_count: + logger.debug(f"checking {next_ingest_part_number}") + potential_bytes = byte_store.get(next_ingest_part_number) + if potential_bytes is None: + logger.debug( + f"expected part number {next_ingest_part_number} not available") + break + md5_store.update(potential_bytes) + progress_bar.update(1) + logger.debug(f"ingested {next_ingest_part_number}") + del byte_store[next_ingest_part_number] + next_ingest_part_number += 1 + + 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}") + sys.exit(1) + + return md5_store.hexdigest() 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..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()