style

python/aibrix/aibrix/downloader/s3.py

@@ -54,7 +54,7 @@ def __init__(
         if model_name is None:
             model_name = infer_model_name(model_uri)
             logger.info(f"model_name is not set, using `{model_name}` as model_name")
-        
+
         auth_config = self._get_auth_config()
         bucket_name, bucket_path = _parse_bucket_info_from_uri(model_uri, scheme=scheme)
 
@@ -71,9 +71,7 @@ def __init__(
         )
 
         self.client = boto3.client(
-            service_name="s3",
-            config=client_config,
-            **auth_config
+            service_name="s3", config=client_config, **auth_config
         )
 
         super().__init__(
@@ -85,26 +83,26 @@ def __init__(
         )  # type: ignore
 
     @abstractmethod
-    def _get_auth_config(self) -> Dict[str, str]:
+    def _get_auth_config(self) -> Dict[str, Optional[str]]:
         """Get auth config for S3 client.
-        
+
         Returns:
             Dict[str, str]: auth config for S3 client, containing following keys:
             - region_name: region name of S3 bucket
             - endpoint_url: endpoint url of S3 bucket
             - aws_access_key_id: access key id of S3 bucket
             - aws_secret_access_key: secret access key of S3 bucket
-            
+
         Example return value:
             {
                 region_name: "region-name",
                 endpoint_url: "URL_ADDRESS3.region-name.com",
                 aws_access_key_id: "AK****",
-                aws_secret_access_key: "SK****",,  
+                aws_secret_access_key: "SK****",,
             }
         """
         pass
-    
+
     @lru_cache()
     def _is_directory(self) -> bool:
         """Check if model_uri is a directory."""
@@ -214,9 +212,12 @@ def _valid_config(self):
             self.client.head_bucket(Bucket=self.bucket_name)
         except Exception as e:
             assert False, f"S3 bucket {self.bucket_name} not exist for {e}."
-
-    def _get_auth_config(self) -> Dict[str, str]:
-        ak, sk = envs.DOWNLOADER_AWS_ACCESS_KEY_ID, envs.DOWNLOADER_AWS_SECRET_ACCESS_KEY
+
+    def _get_auth_config(self) -> Dict[str, Optional[str]]:
+        ak, sk = (
+            envs.DOWNLOADER_AWS_ACCESS_KEY_ID,
+            envs.DOWNLOADER_AWS_SECRET_ACCESS_KEY,
+        )
         assert ak is not None and ak != "", "`AWS_ACCESS_KEY_ID` is not set."
         assert sk is not None and sk != "", "`AWS_SECRET_ACCESS_KEY` is not set."
 
@@ -228,7 +229,6 @@ def _get_auth_config(self) -> Dict[str, str]:
         }
 
 
-
 class TOSDownloader(S3BaseDownloader):
     def __init__(
         self,
@@ -258,11 +258,10 @@ def _valid_config(self):
         except Exception as e:
             assert False, f"TOS bucket {self.bucket_name} not exist for {e}."
 
-    def _get_auth_config(self) -> Dict[str, str]:
+    def _get_auth_config(self) -> Dict[str, Optional[str]]:
         return {
             "region_name": envs.DOWNLOADER_TOS_REGION or "",
             "endpoint_url": envs.DOWNLOADER_TOS_ENDPOINT or "",
             "aws_access_key_id": envs.DOWNLOADER_TOS_ACCESS_KEY or "",
             "aws_secret_access_key": envs.DOWNLOADER_TOS_SECRET_KEY or "",
         }
-

python/aibrix/pyproject.toml

@@ -37,7 +37,6 @@ aibrix_download = 'aibrix.downloader.__main__:main'
 [tool.poetry.dependencies]
 python = ">=3.10,<3.12"
 huggingface-hub = "^0.24.6"
-tos = "2.8.0"
 boto3 = "^1.35.5"
 fastapi = "^0.112.2"
 gunicorn = "^23.0.0"