[Bug] Fix download bugs during download benchmark

python/aibrix/aibrix/downloader/base.py

@@ -89,11 +89,11 @@ def download_directory(self, local_path: Path):
 
         if not self._support_range_download():
             # download using multi threads
-            st = time.perf_counter()
             num_threads = envs.DOWNLOADER_NUM_THREADS
             logger.info(
-                f"Downloader {self.__class__.__name__} does not support "
-                f"range download, use {num_threads} threads to download."
+                f"Downloader {self.__class__.__name__} download "
+                f"{len(filtered_files)} files from {self.model_uri} "
+                f"use {num_threads} threads to download."
             )
 
             executor = ThreadPoolExecutor(num_threads)
@@ -108,26 +108,16 @@ def download_directory(self, local_path: Path):
                 for file in filtered_files
             ]
             wait(futures)
-            duration = time.perf_counter() - st
+
+        else:
             logger.info(
                 f"Downloader {self.__class__.__name__} download "
                 f"{len(filtered_files)} files from {self.model_uri} "
-                f"using {num_threads} threads, "
-                f"duration: {duration:.2f} seconds."
+                f"using range support methods."
             )
-
-        else:
-            st = time.perf_counter()
             for file in filtered_files:
                 # use range download to speedup download
                 self.download(local_path, file, self.bucket_name, True)
-            duration = time.perf_counter() - st
-            logger.info(
-                f"Downloader {self.__class__.__name__} download "
-                f"{len(filtered_files)} files from {self.model_uri} "
-                f"using range support methods, "
-                f"duration: {duration:.2f} seconds."
-            )
 
     def download_model(self, local_path: Optional[str] = None):
         if local_path is None:
@@ -139,7 +129,7 @@ def download_model(self, local_path: Optional[str] = None):
         model_path.mkdir(parents=True, exist_ok=True)
 
         # TODO check local file exists
-
+        st = time.perf_counter()
         if self._is_directory():
             self.download_directory(local_path=model_path)
         else:
@@ -149,6 +139,13 @@ def download_model(self, local_path: Optional[str] = None):
                 bucket_name=self.bucket_name,
                 enable_range=self._support_range_download(),
             )
+        duration = time.perf_counter() - st
+        logger.info(
+            f"Downloader {self.__class__.__name__} download "
+            f"from {self.model_uri} "
+            f"duration: {duration:.2f} seconds."
+        )
+
         return model_path
 
     def __hash__(self):

python/aibrix/aibrix/downloader/s3.py

@@ -18,6 +18,7 @@
 
 import boto3
 from boto3.s3.transfer import TransferConfig
+from botocore.config import Config, MAX_POOL_CONNECTIONS
 from tqdm import tqdm
 
 from aibrix import envs
@@ -40,12 +41,24 @@ def __init__(self, model_uri):
         region = envs.DOWNLOADER_AWS_REGION
         bucket_name, bucket_path = _parse_bucket_info_from_uri(model_uri)
 
+        # Avoid warning log "Connection pool is full"
+        # Refs: https://github.com/boto/botocore/issues/619#issuecomment-583511406
+        max_pool_connections = (
+            envs.DOWNLOADER_NUM_THREADS
+            if envs.DOWNLOADER_NUM_THREADS > MAX_POOL_CONNECTIONS
+            else MAX_POOL_CONNECTIONS
+        )
+        client_config = Config(
+            max_pool_connections=max_pool_connections,
+        )
+
         self.client = boto3.client(
             service_name="s3",
             region_name=region,
             endpoint_url=endpoint,
             aws_access_key_id=ak,
             aws_secret_access_key=sk,
+            config=client_config,
         )
 
         super().__init__(

python/aibrix/aibrix/downloader/tos.py

@@ -42,9 +42,12 @@ def __init__(self, model_uri):
         sk = envs.DOWNLOADER_TOS_SECRET_KEY
         endpoint = envs.DOWNLOADER_TOS_ENDPOINT or ""
         region = envs.DOWNLOADER_TOS_REGION or ""
+        enable_crc = envs.DOWNLOADER_TOS_ENABLE_CRC
         bucket_name, bucket_path = _parse_bucket_info_from_uri(model_uri)
 
-        self.client = tos.TosClientV2(ak=ak, sk=sk, endpoint=endpoint, region=region)
+        self.client = tos.TosClientV2(
+            ak=ak, sk=sk, endpoint=endpoint, region=region, enable_crc=enable_crc
+        )
 
         super().__init__(
             model_uri=model_uri,
@@ -109,8 +112,8 @@ def download(
         task_num = envs.DOWNLOADER_NUM_THREADS if enable_range else 1
 
         download_kwargs = {}
-        if envs.DOWNLOADER_PART_THRESHOLD is not None:
-            download_kwargs["multipart_threshold"] = envs.DOWNLOADER_PART_THRESHOLD
+        if envs.DOWNLOADER_PART_CHUNKSIZE is not None:
+            download_kwargs["part_size"] = envs.DOWNLOADER_PART_CHUNKSIZE
 
         # download file
         total_length = meta_data.content_length

python/aibrix/aibrix/envs.py

@@ -64,6 +64,7 @@ def _parse_int_or_none(value: Optional[str]) -> Optional[int]:
 DOWNLOADER_TOS_SECRET_KEY = os.getenv("TOS_SECRET_KEY")
 DOWNLOADER_TOS_ENDPOINT = os.getenv("TOS_ENDPOINT")
 DOWNLOADER_TOS_REGION = os.getenv("TOS_REGION")
+DOWNLOADER_TOS_ENABLE_CRC = _is_true(os.getenv("TOS_ENABLE_CRC"))
 
 # Downloader AWS S3 Envs
 DOWNLOADER_AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY_ID")

runtime.Dockerfile

@@ -5,8 +5,15 @@ FROM ${BASE_IMAGE} AS base
 
 WORKDIR /app
 
-# Install Poetry
 ARG POETRY_VERSION=1.8.3
+
+# Install dependencies
+RUN apt-get update \
+    && apt-get install -y python3-dev build-essential \
+    && apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install Poetry
 RUN python3 -m pip install poetry==${POETRY_VERSION}
 
 # Copy the runtime source