Repository: vllm-project/aibrix
File: python/aibrix/aibrix/runtime/downloaders.py (Python)

Code before:
                logger.info(f"Downloaded TOS file to {destination_file}")

            return local_path

        except (TosClientError, TosServerError) as e:
            logger.error(f"TOS download error: {e}")
            status_code = getattr(e, "status_code", None)
            if status_code == 404 or "404" in str(e):
                raise FileNotFoundError(f"TOS object not found: {source_url}")
            if status_code == 403 or "403" in str(e):
                raise PermissionError(f"Access denied to TOS object: {source_url}")
            raise RuntimeError(f"TOS download error: {e}")

    def _download_tos_directory(
        self, client, bucket_name: str, prefix: str, local_path: str
    ) -> None:
        resp = client.list_objects_type2(bucket_name, prefix=prefix)
        for obj in getattr(resp, "contents", []) or []:
            key = getattr(obj, "key", None)
            if not key or key == prefix:
                continue
            relative_path = key[len(prefix) :]
            if not relative_path:
                continue

            if key.endswith("/") or relative_path.endswith("/"):
                dir_path = os.path.join(local_path, relative_path.rstrip("/"))
                if dir_path:
                    os.makedirs(dir_path, exist_ok=True)
                continue

            local_file_path = os.path.join(local_path, relative_path)
            os.makedirs(os.path.dirname(local_file_path), exist_ok=True)

            tmp_file = local_file_path + ".part"
            client.download_file(bucket=bucket_name, key=key, file_path=tmp_file)
            os.replace(tmp_file, local_file_path)

        logger.info(f"Downloaded TOS directory {prefix} to {local_path}")


class HuggingFaceArtifactDownloader(ArtifactDownloader):
    """Download artifacts from HuggingFace Hub."""

    async def download(
        self, source_url: str, local_path: str, credentials: Optional[Dict] = None
    ) -> str:
        """
        Download from HuggingFace Hub.

        source_url format: huggingface://repo_id (e.g., huggingface://meta-llama/Llama-2-7b or huggingface://my-org/my-model)

Review comment (suggestion):
seems response has `is_truncated`? seems it doesn't handle the pagination. this is ok since the model weights won't be that large but let's leave a comment here

Code after:
                    bucket=bucket_name,
                    key=object_key,
                    file_path=tmp_file,
                )
                os.replace(tmp_file, destination_file)
                logger.info(f"Downloaded TOS file to {destination_file}")

            return local_path

        except (TosClientError, TosServerError) as e:
            logger.error(f"TOS download error: {e}")
            status_code = getattr(e, "status_code", None)
            if status_code == 404 or "404" in str(e):
                raise FileNotFoundError(f"TOS object not found: {source_url}")
            if status_code == 403 or "403" in str(e):
                raise PermissionError(f"Access denied to TOS object: {source_url}")
            raise RuntimeError(f"TOS download error: {e}")

    def _download_tos_directory(
        self, client, bucket_name: str, prefix: str, local_path: str
    ) -> None:
        # NOTE: list_objects_type2 may paginate results (e.g., resp.is_truncated / continuation token).
        # We currently only fetch the first page. This is OK for small model artifacts, but may miss objects
        # if the prefix contains many keys.
        resp = client.list_objects_type2(bucket_name, prefix=prefix)
        for obj in getattr(resp, "contents", []) or []:
            key = getattr(obj, "key", None)
            if not key or key == prefix:
                continue
            relative_path = key[len(prefix) :]
            if not relative_path:
                continue

            if key.endswith("/") or relative_path.endswith("/"):
                dir_path = os.path.join(local_path, relative_path.rstrip("/"))
                if dir_path:
                    os.makedirs(dir_path, exist_ok=True)
                continue

            local_file_path = os.path.join(local_path, relative_path)
            os.makedirs(os.path.dirname(local_file_path), exist_ok=True)

            tmp_file = local_file_path + ".part"
            client.download_file(bucket=bucket_name, key=key, file_path=tmp_file)
            os.replace(tmp_file, local_file_path)

        logger.info(f"Downloaded TOS directory {prefix} to {local_path}")


class HuggingFaceArtifactDownloader(ArtifactDownloader):
    """Download artifacts from HuggingFace Hub."""