uses asyncio

.tool-versions

@@ -1 +1 @@
-python 3.10.12
+python 3.10.12

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 <bucket_name> <file_name>
 ```
 
+### 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

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

s3_md5/src/cli.py

@@ -2,23 +2,44 @@
 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,
                         help='bucket name')
     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

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

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(

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

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

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": [

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