refact: use bucket path

Author: brosoul

python/aibrix/aibrix/downloader/base.py

@@ -20,9 +20,10 @@
 from pathlib import Path
 from typing import List, Optional
 
-from aibrix import envs
 from aibrix.logger import init_logger
 
+from aibrix import envs
+
 logger = init_logger(__name__)
 
 
@@ -32,6 +33,8 @@ class BaseDownloader(ABC):
 
     model_uri: str
     model_name: str
+    bucket_path: str
+    bucket_name: Optional[str]
     allow_file_suffix: Optional[List[str]] = field(
         default_factory=lambda: envs.DOWNLOADER_ALLOW_FILE_SUFFIX
     )
@@ -59,7 +62,13 @@ def _support_range_download(self) -> bool:
         pass
 
     @abstractmethod
-    def download(self, filename: str, local_path: Path, enable_range: bool = True):
+    def download(
+        self,
+        local_path: Path,
+        bucket_path: str,
+        bucket_name: str = None,
+        enable_range: bool = True,
+    ):
         pass
 
     def download_directory(self, local_path: Path):
@@ -68,9 +77,9 @@ def download_directory(self, local_path: Path):
         directory method for ``Downloader``. Otherwise, the following logic will be
         used to download the directory.
         """
-        directory_list = self._directory_list(self.model_uri)
+        directory_list = self._directory_list(self.bucket_path)
         if self.allow_file_suffix is None:
-            logger.info(f"All files from {self.model_uri} will be downloaded.")
+            logger.info(f"All files from {self.bucket_path} will be downloaded.")
             filtered_files = directory_list
         else:
             filtered_files = [
@@ -92,8 +101,9 @@ def download_directory(self, local_path: Path):
             futures = [
                 executor.submit(
                     self.download,
-                    filename=file,
                     local_path=local_path,
+                    bucket_path=file,
+                    bucket_name=self.bucket_name,
                     enable_range=False,
                 )
                 for file in filtered_files
@@ -111,7 +121,7 @@ def download_directory(self, local_path: Path):
             st = time.perf_counter()
             for file in filtered_files:
                 # use range download to speedup download
-                self.download(file, local_path, True)
+                self.download(local_path, file, self.bucket_name, True)
             duration = time.perf_counter() - st
             logger.info(
                 f"Downloader {self.__class__.__name__} download "
@@ -132,11 +142,12 @@ def download_model(self, local_path: Optional[str] = None):
         # TODO check local file exists
 
         if self._is_directory():
-            self.download_directory(model_path)
+            self.download_directory(local_path=model_path)
         else:
             self.download(
-                filename=self.model_uri,
                 local_path=model_path,
+                bucket_path=self.bucket_path,
+                bucket_name=self.bucket_name,
                 enable_range=self._support_range_download(),
             )
 

python/aibrix/aibrix/downloader/huggingface.py

@@ -15,11 +15,11 @@
 from pathlib import Path
 from typing import List, Optional
 
+from aibrix.downloader.base import BaseDownloader
+from aibrix.logger import init_logger
 from huggingface_hub import HfApi, hf_hub_download, snapshot_download
 
 from aibrix import envs
-from aibrix.downloader.base import BaseDownloader
-from aibrix.logger import init_logger
 
 logger = init_logger(__name__)
 
@@ -36,12 +36,19 @@ def __init__(self, model_uri: str, model_name: Optional[str] = None):
             else:
                 model_name = _parse_model_name_from_uri(model_uri)
 
-        super().__init__(model_uri=model_uri, model_name=model_name)  # type: ignore
-
         self.hf_token = envs.DOWNLOADER_HF_TOKEN
         self.hf_endpoint = envs.DOWNLOADER_HF_ENDPOINT
         self.hf_revision = envs.DOWNLOADER_HF_REVISION
 
+        super().__init__(
+            model_uri=model_uri,
+            model_name=model_name,
+            bucket_path=model_uri,
+            bucket_name=None,
+        )  # type: ignore
+
+        # Dependent on the attributes generated in the base class,
+        # so place it after the super().__init__() call.
         self.allow_patterns = (
             None
             if self.allow_file_suffix is None
@@ -58,7 +65,7 @@ def _valid_config(self):
         assert (
             len(self.model_uri.split("/")) == 2
         ), "Model uri must be in `repo/name` format."
-
+        assert self.bucket_name is None, "Bucket name is empty in HuggingFace."
         assert self.model_name is not None, "Model name is not set."
 
     def _is_directory(self) -> bool:
@@ -74,10 +81,16 @@ def _directory_list(self, path: str) -> List[str]:
     def _support_range_download(self) -> bool:
         return False
 
-    def download(self, filename: str, local_path: Path, enable_range: bool = True):
+    def download(
+        self,
+        local_path: Path,
+        bucket_path: str,
+        bucket_name: str = None,
+        enable_range: bool = True,
+    ):
         hf_hub_download(
             repo_id=self.model_uri,
-            filename=filename,
+            filename=bucket_path,
             local_dir=local_path,
             revision=self.hf_revision,
             token=self.hf_token,

python/aibrix/aibrix/downloader/s3.py

@@ -17,11 +17,11 @@
 from urllib.parse import urlparse
 
 import boto3
+from aibrix.downloader.base import BaseDownloader
 from boto3.s3.transfer import TransferConfig
 from tqdm import tqdm
 
 from aibrix import envs
-from aibrix.downloader.base import BaseDownloader
 
 
 def _parse_bucket_info_from_uri(uri: str) -> Tuple[str, str]:
@@ -39,8 +39,6 @@ def __init__(self, model_uri):
         endpoint = envs.DOWNLOADER_AWS_ENDPOINT
         region = envs.DOWNLOADER_AWS_REGION
         bucket_name, bucket_path = _parse_bucket_info_from_uri(model_uri)
-        self.bucket_name = bucket_name
-        self.bucket_path = bucket_path
 
         self.client = boto3.client(
             service_name="s3",
@@ -50,7 +48,12 @@ def __init__(self, model_uri):
             aws_secret_access_key=sk,
         )
 
-        super().__init__(model_uri=model_uri, model_name=model_name)  # type: ignore
+        super().__init__(
+            model_uri=model_uri,
+            model_name=model_name,
+            bucket_path=bucket_path,
+            bucket_name=bucket_name,
+        )  # type: ignore
 
     def _valid_config(self):
         assert (
@@ -79,26 +82,28 @@ def _is_directory(self) -> bool:
 
     def _directory_list(self, path: str) -> List[str]:
         objects_out = self.client.list_objects_v2(
-            Bucket=self.bucket_name, Delimiter="/", Prefix=self.bucket_path
+            Bucket=self.bucket_name, Delimiter="/", Prefix=path
         )
         contents = objects_out.get("Contents", [])
         return [content.get("Key") for content in contents]
 
     def _support_range_download(self) -> bool:
         return True
 
-    def download(self, filename: str, local_path: Path, enable_range: bool = True):
-        # filename should extract from model_uri when it is not a directory
-        if not self._is_directory():
-            filename = self.bucket_path
-
+    def download(
+        self,
+        local_path: Path,
+        bucket_path: str,
+        bucket_name: str = None,
+        enable_range: bool = True,
+    ):
         # check if file exist
         try:
-            meta_data = self.client.head_object(Bucket=self.bucket_name, Key=filename)
+            meta_data = self.client.head_object(Bucket=bucket_name, Key=bucket_path)
         except Exception as e:
-            raise ValueError(f"TOS file {filename} not exist for {e}.")
+            raise ValueError(f"S3 bucket path {bucket_path} not exist for {e}.")
 
-        _file_name = filename.split("/")[-1]
+        _file_name = bucket_path.split("/")[-1]
         # S3 client does not support Path, convert it to str
         local_file = str(local_path.joinpath(_file_name).absolute())
 
@@ -122,8 +127,8 @@ def download_progress(bytes_transferred):
                 pbar.update(bytes_transferred)
 
             self.client.download_file(
-                Bucket=self.bucket_name,
-                Key=filename,
+                Bucket=bucket_name,
+                Key=bucket_path,
                 Filename=local_file,
                 Config=config,
                 Callback=download_progress,

python/aibrix/aibrix/downloader/tos.py

@@ -18,11 +18,11 @@
 from urllib.parse import urlparse
 
 import tos
+from aibrix.downloader.base import BaseDownloader
 from tos import DataTransferType
 from tqdm import tqdm
 
 from aibrix import envs
-from aibrix.downloader.base import BaseDownloader
 
 tos_logger = logging.getLogger("tos")
 tos_logger.setLevel(logging.WARNING)
@@ -40,18 +40,18 @@ def __init__(self, model_uri):
         model_name = envs.DOWNLOADER_MODEL_NAME
         ak = envs.DOWNLOADER_TOS_ACCESS_KEY
         sk = envs.DOWNLOADER_TOS_SECRET_KEY
-        endpoint = envs.DOWNLOADER_TOS_ENDPOINT or ''
-        region = envs.DOWNLOADER_TOS_REGION or ''
+        endpoint = envs.DOWNLOADER_TOS_ENDPOINT or ""
+        region = envs.DOWNLOADER_TOS_REGION or ""
         bucket_name, bucket_path = _parse_bucket_info_from_uri(model_uri)
-        self.bucket_name = bucket_name
-        self.bucket_path = bucket_path
-        
-        self.client = tos.TosClientV2(ak=ak,
-                                      sk=sk,
-                                      endpoint=endpoint,
-                                      region=region)
-        
-        super().__init__(model_uri=model_uri, model_name=model_name)  # type: ignore
+
+        self.client = tos.TosClientV2(ak=ak, sk=sk, endpoint=endpoint, region=region)
+
+        super().__init__(
+            model_uri=model_uri,
+            model_name=model_name,
+            bucket_path=bucket_path,
+            bucket_name=bucket_name,
+        )  # type: ignore
 
     def _valid_config(self):
         assert (
@@ -83,26 +83,27 @@ def _is_directory(self) -> bool:
     def _directory_list(self, path: str) -> List[str]:
         # TODO cache list_objects_type2 result to avoid too many requests
         objects_out = self.client.list_objects_type2(
-            self.bucket_name, prefix=self.bucket_path, delimiter="/"
+            self.bucket_name, prefix=path, delimiter="/"
         )
-
         return [obj.key for obj in objects_out.contents]
 
     def _support_range_download(self) -> bool:
         return True
 
-    def download(self, filename: str, local_path: Path, enable_range: bool = True):
-        # filename should extract from model_uri when it is not a directory
-        if not self._is_directory():
-            filename = self.bucket_path
-
+    def download(
+        self,
+        local_path: Path,
+        bucket_path: str,
+        bucket_name: str = None,
+        enable_range: bool = True,
+    ):
         # check if file exist
         try:
-            meta_data = self.client.head_object(bucket=self.bucket_name, key=filename)
+            meta_data = self.client.head_object(bucket=bucket_name, key=bucket_path)
         except Exception as e:
-            raise ValueError(f"TOS file {filename} not exist for {e}.")
+            raise ValueError(f"TOS bucket path {bucket_path} not exist for {e}.")
 
-        _file_name = filename.split("/")[-1]
+        _file_name = bucket_path.split("/")[-1]
         # TOS client does not support Path, convert it to str
         local_file = str(local_path.joinpath(_file_name).absolute())
         task_num = envs.DOWNLOADER_NUM_THREADS if enable_range else 1
@@ -121,8 +122,8 @@ def download_progress(
                 pbar.update(rw_once_bytes)
 
             self.client.download_file(
-                bucket=self.bucket_name,
-                key=filename,
+                bucket=bucket_name,
+                key=bucket_path,
                 file_path=local_file,
                 task_num=task_num,
                 data_transfer_listener=download_progress,