intellif-aihub 0.1.14__py3-none-any.whl → 0.1.15__py3-none-any.whl

aihub/__init__.py

@@ -1 +1 @@
-__version__ = "0.1.14"
+__version__ = "0.1.15"

aihub/models/artifact.py

@@ -14,6 +14,7 @@ from pydantic import BaseModel, Field
 
 class ArtifactType(str, Enum):
     """制品类型枚举："dataset"-数据集类型；"model"-模型类型；"metrics"-指标类型；"log"-日志类型；"checkpoint"-检查点类型；"image"-图像类型；"prediction"-预测结果类型；"other"-其他类型"""
+
     dataset = "dataset"  # 数据集类型
     model = "model"  # 模型类型
     metrics = "metrics"  # 指标类型
@@ -26,24 +27,29 @@ class ArtifactType(str, Enum):
 
 class CreateArtifactsReq(BaseModel):
     """创建制品请求"""
+
     entity_id: str = Field(alias="entity_id", description="实体ID，通常是运行ID，用于关联制品与特定运行")
-    entity_type: ArtifactType = Field(default=ArtifactType.other, alias="entity_type",
-                                      description="制品类型，指定制品的类型，默认为other")
+    entity_type: ArtifactType = Field(
+        default=ArtifactType.other, alias="entity_type", description="制品类型，指定制品的类型，默认为other"
+    )
     src_path: str = Field(alias="src_path", description="源路径，制品在系统中的路径标识")
-    is_dir: bool = Field(default=False, alias="is_dir",
-                         description="是否为目录，True表示制品是一个目录，False表示是单个文件")
+    is_dir: bool = Field(
+        default=False, alias="is_dir", description="是否为目录，True表示制品是一个目录，False表示是单个文件"
+    )
 
     model_config = {"use_enum_values": True}
 
 
 class CreateArtifactsResponseData(BaseModel):
     """创建制品响应数据"""
+
     id: int = Field(description="制品ID")
     s3_path: str = Field(alias="s3_path", description="S3存储路径")
 
 
 class CreateArtifactsResponseModel(BaseModel):
     """创建制品响应模型"""
+
     code: int = Field(description="响应码，0表示成功")
     msg: str = Field(default="", description="响应消息")
     data: Optional[CreateArtifactsResponseData] = Field(default=None, description="响应数据")
@@ -51,6 +57,7 @@ class CreateArtifactsResponseModel(BaseModel):
 
 class CreateEvalReq(BaseModel):
     """创建评估请求"""
+
     dataset_id: int = Field(alias="dataset_id", description="数据集ID")
     dataset_version_id: int = Field(alias="dataset_version_id", description="数据集版本ID")
     prediction_artifact_path: str = Field(alias="prediction_artifact_path", description="预测结果制品路径")
@@ -62,6 +69,7 @@ class CreateEvalReq(BaseModel):
 
 class ArtifactResp(BaseModel):
     """制品响应模型，表示一个制品的详细信息"""
+
     id: int = Field(description="制品ID")
     entity_type: str = Field(alias="entity_type", description="实体类型")
     entity_id: str = Field(alias="entity_id", description="实体ID")
@@ -72,6 +80,7 @@ class ArtifactResp(BaseModel):
 
 class ArtifactRespData(BaseModel):
     """制品分页数据"""
+
     total: int = Field(description="总记录数")
     page_size: int = Field(alias="page_size", description="每页大小")
     page_num: int = Field(alias="page_num", description="页码")
@@ -80,6 +89,7 @@ class ArtifactRespData(BaseModel):
 
 class ArtifactRespModel(BaseModel):
     """获取制品响应模型"""
+
     code: int = Field(description="响应码，0表示成功")
     msg: str = Field(default="", description="响应消息")
     data: ArtifactRespData = Field(description="响应数据")
@@ -91,8 +101,10 @@ InfinityPageSize = 10000 * 100
 
 class StsResp(BaseModel):
     """STS 临时凭证"""
+
     access_key_id: Optional[str] = Field(default=None, alias="access_key_id", description="访问密钥ID")
     secret_access_key: Optional[str] = Field(default=None, alias="secret_access_key", description="秘密访问密钥")
     session_token: Optional[str] = Field(default=None, alias="session_token", description="会话令牌")
     expiration: Optional[int] = Field(default=None, alias="expiration", description="过期时间")
     endpoint: Optional[str] = Field(default=None, alias="endpoint", description="端点URL")
+    bucket: Optional[str] = Field(default=None, alias="bucket", description="存储桶名称")

aihub/models/dataset_management.py

@@ -8,6 +8,7 @@ from pydantic import BaseModel, Field
 
 class DatasetVersionStatus(IntEnum):
     """数据集版本状态：1-等待中；2-运行中；3-成功；4-失败；5-加载meta；6-构建index"""
+
     Waiting = 1  # 等待中
     Running = 2  # 运行中
     Success = 3  # 成功
@@ -18,6 +19,7 @@ class DatasetVersionStatus(IntEnum):
 
 class UploadType(IntEnum):
     """上传类型：1-本地上传；3-服务器路径上传；4-Labelfree；5-数据接入"""
+
     LOCAL = 1  # 本地上传
     SERVER_PATH = 3  # 服务器路径上传
     LABELFREE = 4  # Labelfree
@@ -26,6 +28,7 @@ class UploadType(IntEnum):
 
 class CreateDatasetRequest(BaseModel):
     """创建数据集请求"""
+
     name: str = Field(description="数据集名称")
     description: str = Field(description="数据集描述")
     tags: List[int] = Field(description="标签ID列表，通过标签管理系统查询")
@@ -37,11 +40,13 @@ class CreateDatasetRequest(BaseModel):
 
 class CreateDatasetResponse(BaseModel):
     """创建数据集返回"""
+
     id: int = Field(alias="id", description="数据集ID")
 
 
 class DatasetVersionBase(BaseModel):
     """数据集版本概要"""
+
     id: int = Field(description="版本ID")
     version: int = Field(description="版本号")
     status: DatasetVersionStatus = Field(description="版本状态")
@@ -53,6 +58,7 @@ class DatasetVersionBase(BaseModel):
 
 class DatasetDetail(BaseModel):
     """数据集详情"""
+
     id: int = Field(description="数据集 ID")
     name: str = Field(description="名称")
     description: str = Field(description="描述")
@@ -69,6 +75,7 @@ class DatasetDetail(BaseModel):
 
 class ExtInfo(BaseModel):
     """扩展信息"""
+
     rec_file_path: Optional[str] = Field(None, alias="rec_file_path", description="rec文件路径")
     idx_file_path: Optional[str] = Field(None, alias="idx_file_path", description="idx文件路径")
     json_file_path: Optional[str] = Field(None, alias="json_file_path", description="json文件路径")
@@ -77,6 +84,7 @@ class ExtInfo(BaseModel):
 
 class CreateDatasetVersionRequest(BaseModel):
     """创建版本请求"""
+
     upload_path: str = Field(alias="upload_path", description="上传路径")
     description: Optional[str] = Field(None, description="版本描述")
     dataset_id: int = Field(alias="dataset_id", description="数据集ID")
@@ -91,11 +99,13 @@ class CreateDatasetVersionRequest(BaseModel):
 
 class CreateDatasetVersionResponse(BaseModel):
     """创建版本返回"""
+
     id: int = Field(alias="id", description="版本ID")
 
 
 class UploadDatasetVersionRequest(BaseModel):
     """上传数据集版本请求"""
+
     upload_path: str = Field(alias="upload_path", description="上传目录")
     upload_type: UploadType = Field(alias="upload_type", description="上传类型")
     dataset_id: int = Field(alias="dataset_id", description="数据集ID")
@@ -107,11 +117,13 @@ class UploadDatasetVersionRequest(BaseModel):
 
 class UploadDatasetVersionResponse(BaseModel):
     """上传数据集版本返回"""
+
     id: int = Field(alias="id", description="版本ID")
 
 
 class DatasetVersionDetail(BaseModel):
     """数据集版本详情"""
+
     id: int = Field(description="版本ID")
     version: int = Field(description="版本号")
     dataset_id: int = Field(alias="dataset_id", description="数据集ID")
@@ -133,6 +145,7 @@ class DatasetVersionDetail(BaseModel):
 
 class FileUploadData(BaseModel):
     """文件上传数据"""
+
     path: str = Field(description="路径")
     url: str = Field(description="URL")
 
@@ -203,3 +216,13 @@ class ListDatasetVersionResp(BaseModel):
     page_size: int = Field(alias="page_size", description="每页大小")
     page_num: int = Field(alias="page_num", description="当前页码")
     data: List[ListDatasetVersionItem] = Field(description="数据集版本列表")
+
+
+class CreateDatasetVersionByDataIngestReqV2(BaseModel):
+    """通过数据集成创建数据集版本请求"""
+
+    description: Optional[str] = Field(None, description="描述")
+    dataset_id: int = Field(..., description="数据集ID")
+    s3_object_sheet: str = Field(..., description="S3对象表")
+    object_cnt: Optional[int] = Field(None, description="对象数量")
+    data_size: Optional[int] = Field(None, description="数据大小")

aihub/services/artifact.py

@@ -98,9 +98,7 @@ class ArtifactService:
         """
         return self._artifact.get_sts()
 
-    def get_by_run_id(
-            self, run_id: str, artifact_path: Optional[str] = None
-    ) -> List[ArtifactResp]:
+    def get_by_run_id(self, run_id: str, artifact_path: Optional[str] = None) -> List[ArtifactResp]:
         """根据运行ID获取制品列表
 
         Args:
@@ -116,11 +114,11 @@ class ArtifactService:
         return self._artifact.get_by_run_id(run_id, artifact_path)
 
     def create_artifact(
-            self,
-            local_path: str,
-            artifact_path: Optional[str] = None,
-            run_id: Optional[str] = None,
-            artifact_type: ArtifactType = ArtifactType.other,
+        self,
+        local_path: str,
+        artifact_path: Optional[str] = None,
+        run_id: Optional[str] = None,
+        artifact_type: ArtifactType = ArtifactType.other,
     ) -> None:
         """创建单个文件制品并上传
 
@@ -171,11 +169,11 @@ class ArtifactService:
         return
 
     def create_artifacts(
-            self,
-            local_dir: str,
-            artifact_path: Optional[str] = None,
-            run_id: Optional[str] = None,
-            artifact_type: ArtifactType = ArtifactType.other,
+        self,
+        local_dir: str,
+        artifact_path: Optional[str] = None,
+        run_id: Optional[str] = None,
+        artifact_type: ArtifactType = ArtifactType.other,
     ) -> None:
         """创建目录制品并上传
 
@@ -223,9 +221,7 @@ class ArtifactService:
         logger.info(f"log artifact done: {artifact_path}")
         return
 
-    def download_artifacts(
-            self, run_id: str, artifact_path: Optional[str], local_dir: str
-    ) -> None:
+    def download_artifacts(self, run_id: str, artifact_path: Optional[str], local_dir: str) -> None:
         """下载制品
 
         Args:
@@ -252,9 +248,7 @@ class ArtifactService:
             if artifact_item.is_dir:
                 download_dir_from_s3(self.s3_client, bucket, object_name, local_dir)
             else:
-                self.s3_client.fget_object(
-                    bucket, object_name, str(Path(local_dir) / artifact_item.src_path)
-                )
+                self.s3_client.fget_object(bucket, object_name, str(Path(local_dir) / artifact_item.src_path))
 
         logger.info(f"download artifact done: {artifact_path}")
         return
@@ -311,9 +305,7 @@ class _Artifact:
             raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
         return
 
-    def get_by_run_id(
-            self, run_id: str, artifact_path: Optional[str]
-    ) -> List[ArtifactResp]:
+    def get_by_run_id(self, run_id: str, artifact_path: Optional[str]) -> List[ArtifactResp]:
         """根据运行ID获取制品列表
 
         Args:
@@ -326,18 +318,12 @@ class _Artifact:
         Raises:
             APIError: 当API调用失败时抛出
         """
-        resp = self._http.get(
-            f"{_Base}/artifacts?entity_id={run_id}&page_num=1&page_size={InfinityPageSize}"
-        )
+        resp = self._http.get(f"{_Base}/artifacts?entity_id={run_id}&page_num=1&page_size={InfinityPageSize}")
         wrapper = APIWrapper[ArtifactRespData].model_validate(resp.json())
         if wrapper.code != 0:
             raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
         if artifact_path:
-            return [
-                artifact
-                for artifact in wrapper.data.data
-                if artifact.src_path == artifact_path
-            ]
+            return [artifact for artifact in wrapper.data.data if artifact.src_path == artifact_path]
         else:
             return wrapper.data.data
 

aihub/services/dataset_management.py

@@ -17,7 +17,6 @@ from __future__ import annotations
 import mimetypes
 import os
 import pathlib
-import tempfile
 import time
 import uuid
 
@@ -25,6 +24,7 @@ import httpx
 from loguru import logger
 
 from ..exceptions import APIError
+from ..models.artifact import StsResp
 from ..models.common import APIWrapper
 from ..models.dataset_management import (
     CreateDatasetRequest,
@@ -40,9 +40,12 @@ from ..models.dataset_management import (
     ListDatasetResp,
     ListDatasetVersionReq,
     ListDatasetVersionResp,
+    CreateDatasetVersionByDataIngestReqV2,
+    UploadType,
 )
 from ..models.dataset_management import DatasetVersionStatus
-from ..utils.download import dataset_download, zip_dir
+from ..utils.di import SimpleS3Client, DataUploader
+from ..utils.download import dataset_download
 
 _BASE = "/dataset-mng/api/v2"
 
@@ -138,20 +141,29 @@ class DatasetManagementService:
     def dataset(self) -> _Dataset:
         return self._dataset
 
+    def _get_sts(self) -> StsResp:
+        return self.dataset_version.get_sts()
+
     @property
     def dataset_version(self) -> _DatasetVersion:
         return self._dataset_version
 
+    def upload_by_data_ingest(
+        self,
+        req: CreateDatasetVersionByDataIngestReqV2,
+    ) -> CreateDatasetVersionResponse:
+        return self.dataset_version.upload_by_data_ingest(req)
+
     def create_dataset_and_version(
-            self,
-            *,
-            dataset_name: str,
-            dataset_description: str = "",
-            is_local_upload: bool,
-            local_file_path: str | None = None,
-            server_file_path: str | None = None,
-            version_description: str = "",
-            timeout: int = 1_800,
+        self,
+        *,
+        dataset_name: str,
+        dataset_description: str = "",
+        is_local_upload: bool,
+        local_file_path: str | None = None,
+        server_file_path: str | None = None,
+        version_description: str = "",
+        timeout: int = 1_800,
     ) -> tuple[int, int, str]:
         """创建数据集及其版本，并等待版本状态变为 *Success*。
 
@@ -169,17 +181,51 @@ class DatasetManagementService:
 
         Returns:
            tuple[int, int, str]: 一个三元组，包含：[数据集 ID,数据集版本 ID, 数据集版本标签（格式为 <dataset_name>/V<version_number>)]
+
+        Raises:
+            ValueError: 当参数不满足要求时
+            APIError: 当后端返回错误时
+            TimeoutError: 当等待超时时
         """
+        # 参数校验
+        self._validate_create_params(is_local_upload, local_file_path, server_file_path)
+
+        # 创建数据集
+        dataset_id = self._create_dataset(dataset_name, dataset_description)
+        logger.info(f"创建数据集成功，名称为 {dataset_name} ,开始准备创建版本、上传数据")
+
+        # 创建数据集版本
+        version_id = self._create_dataset_version(
+            dataset_id=dataset_id,
+            is_local_upload=is_local_upload,
+            local_file_path=local_file_path,
+            server_file_path=server_file_path,
+            version_description=version_description,
+        )
+
+        # 获取版本标签
+        version_tag = self._get_version_tag(dataset_id, version_id)
+        logger.info(f"数据集版本创建成功，名称为 {version_tag}，开始轮询状态…")
+
+        # 轮询等待版本状态变为成功
+        self._wait_for_version_success(version_id, timeout)
+
+        return dataset_id, version_id, version_tag
+
+    def _validate_create_params(
+        self, is_local_upload: bool, local_file_path: str | None, server_file_path: str | None
+    ) -> None:
+        """验证创建数据集和版本所需的参数"""
         if is_local_upload:
             if not local_file_path:
                 raise ValueError("is_local_upload=True 时必须提供 local_file_path")
-            upload_type = 1
         else:
             if not server_file_path:
                 raise ValueError("is_local_upload=False 时必须提供 server_file_path")
-            upload_type = 3
 
-        dataset_id = self._dataset.create(
+    def _create_dataset(self, dataset_name: str, dataset_description: str) -> int:
+        """创建数据集"""
+        return self._dataset.create(
             CreateDatasetRequest(
                 name=dataset_name,
                 description=dataset_description,
@@ -190,39 +236,96 @@ class DatasetManagementService:
                 access_user_ids=None,
             )
         )
-        logger.info(
-            f"创建数据集成功，名称为 {dataset_name} ,开始准备创建版本、上传数据"
-        )
 
+    def _create_dataset_version(
+        self,
+        dataset_id: int,
+        is_local_upload: bool,
+        local_file_path: str | None,
+        server_file_path: str | None,
+        version_description: str,
+    ) -> int:
+        """根据上传类型创建数据集版本"""
         if is_local_upload:
-            # 上传文件，检查是文件夹还是zip
-            local_file_path = pathlib.Path(local_file_path)
-            if local_file_path.is_dir():
-                # 把文件夹打包为一个 zip
-                temp_zip_path = (
-                        pathlib.Path(tempfile.mkdtemp()) / f" {uuid.uuid4().hex}.zip"
-                )
-                zip_dir(local_file_path, temp_zip_path)
-                upload_data = self._upload.upload_file(temp_zip_path)
-                os.remove(temp_zip_path)
-            else:
-                upload_data = self._upload.upload_file(local_file_path)
-
-            upload_path = upload_data.path
+            return self._create_local_dataset_version(dataset_id, local_file_path, version_description)
         else:
-            upload_path = server_file_path
-        logger.info(f"文件上传成功：{local_file_path}")
+            return self._create_server_dataset_version(dataset_id, server_file_path, version_description)
+
+    def _create_local_dataset_version(
+        self, dataset_id: int, local_file_path: str | None, version_description: str
+    ) -> int:
+        """创建本地文件数据集版本"""
+        if pathlib.Path(local_file_path).is_dir():
+            return self._create_local_dir_dataset_version(dataset_id, local_file_path)
+        elif pathlib.Path(local_file_path).is_file():
+            return self._create_local_file_dataset_version(dataset_id, local_file_path, version_description)
+        else:
+            raise ValueError(f"本地路径既不是文件也不是目录: {local_file_path}")
+
+    def _create_local_dir_dataset_version(self, dataset_id: int, local_file_path: str) -> int:
+        """处理本地目录上传"""
+        sts = self._get_sts()
+        s3_client = SimpleS3Client(
+            sts.endpoint, sts.access_key_id, sts.secret_access_key, session_token=sts.session_token
+        )
+        uid = uuid.uuid4().hex
+        s3_target = f"s3://{sts.bucket}/dataset_workspace/{dataset_id}/{uid}"
+        s3_csv_path = f"s3://{sts.bucket}/dataset_workspace/{dataset_id}/{uid}.csv"
+        s3_status_path = f"s3://{sts.bucket}/dataset_workspace/{dataset_id}/{uid}.json"
+
+        # 创建上传器并执行
+        uploader = DataUploader(
+            task_id=dataset_id,
+            local_path=str(local_file_path),
+            s3_target=s3_target,
+            csv_path=s3_csv_path,
+            status_path=s3_status_path,
+            num_workers=40,
+        )
 
-        version_id = self._dataset_version.upload(
+        upload_stats = uploader.run(s3_client)
+        req = CreateDatasetVersionByDataIngestReqV2(
+            description=f"sdk 上传",
+            dataset_id=dataset_id,
+            s3_object_sheet=s3_csv_path,
+            object_cnt=upload_stats.uploaded_count,
+            data_size=upload_stats.uploaded_size,
+        )
+        return self.upload_by_data_ingest(req).id
+
+    def _create_local_file_dataset_version(
+        self, dataset_id: int, local_file_path: str, version_description: str
+    ) -> int:
+        """处理本地文件上传"""
+        upload_data = self._upload.upload_file(local_file_path)
+        upload_path = upload_data.path
+        logger.info(f"文件上传成功：{local_file_path}")
+        return self._dataset_version.upload(
             UploadDatasetVersionRequest(
                 upload_path=upload_path,
-                upload_type=upload_type,
+                upload_type=UploadType.LOCAL,  # 本地上传类型
+                dataset_id=dataset_id,
+                description=version_description,
+                parent_version_id=0,
+            )
+        )
+
+    def _create_server_dataset_version(
+        self, dataset_id: int, server_file_path: str | None, version_description: str
+    ) -> int:
+        """创建服务器文件数据集版本"""
+        return self._dataset_version.upload(
+            UploadDatasetVersionRequest(
+                upload_path=server_file_path,
+                upload_type=UploadType.SERVER_PATH,  # 服务器文件上传类型
                 dataset_id=dataset_id,
                 description=version_description,
                 parent_version_id=0,
             )
         )
 
+    def _get_version_tag(self, dataset_id: int, version_id: int) -> str:
+        """获取版本标签"""
         detail = self._dataset.get(dataset_id)
         ver_num = next(
             (v.version for v in detail.versions if v.id == version_id),
@@ -231,9 +334,10 @@ class DatasetManagementService:
         if ver_num is None:
             ver_num = 1
 
-        version_tag = f"{detail.name}/V{ver_num}"
-        logger.info(f"数据集版本创建成功，名称为 {version_tag}，开始轮询状态…")
+        return f"{detail.name}/V{ver_num}"
 
+    def _wait_for_version_success(self, version_id: int, timeout: int) -> None:
+        """轮询等待版本状态变为成功"""
         start_ts = time.time()
         poll_interval = 10
 
@@ -255,8 +359,6 @@ class DatasetManagementService:
             logger.debug(f"已等待 {elapsed:.0f}s，继续轮询…")
             time.sleep(poll_interval)
 
-        return dataset_id, version_id, version_tag
-
     def run_download(self, dataset_version_name: str, local_dir: str, worker: int = 4) -> None:
         """根据数据集版本名称下载对应的数据集文件。
 
@@ -400,9 +502,7 @@ class _DatasetVersion:
         return wrapper.data
 
     def get_by_name(self, version_name: str) -> DatasetVersionDetail:
-        resp = self._http.get(
-            f"{_BASE}/dataset-versions-detail", params={"name": version_name}
-        )
+        resp = self._http.get(f"{_BASE}/dataset-versions-detail", params={"name": version_name})
         wrapper = APIWrapper[DatasetVersionDetail].model_validate(resp.json())
         if wrapper.code != 0:
             raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
@@ -417,6 +517,40 @@ class _DatasetVersion:
             raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
         return wrapper.data
 
+    def get_sts(self) -> StsResp:
+        """获取STS临时凭证
+
+        获取用于访问S3存储的临时凭证。
+
+        Returns:
+            StsResp: STS临时凭证信息
+
+        Raises:
+            APIError: 当API调用失败时抛出
+        """
+        resp = self._http.get(f"{_BASE}/dataset-versions/get-sts")
+        logger.info(f"get sts: {resp.text}")
+        wrapper = APIWrapper[StsResp].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def upload_by_data_ingest(self, req: CreateDatasetVersionByDataIngestReqV2) -> CreateDatasetVersionResponse:
+        """上传数据集版本（数据集导入）
+        Args:
+            req
+
+        """
+        resp = self._http.post(
+            f"{_BASE}/dataset-versions/data-ingest",
+            json=req.model_dump(),
+        )
+        logger.debug(f"upload_by_data_ingest: {resp.text}")
+        wrapper = APIWrapper[CreateDatasetVersionResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
 
 class _Upload:
     def __init__(self, http: httpx.Client):

aihub/utils/di.py

@@ -0,0 +1,337 @@
+import argparse
+import csv
+import hashlib
+import json
+import os
+import queue
+import sys
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Dict, Any, Tuple
+
+import minio
+from loguru import logger
+
+
+class UploadStatus:
+    """上传状态类"""
+
+    def __init__(self):
+        self.uploaded_count = 0
+        self.uploaded_size = 0
+        self.updated_at = int(time.time() * 1000)
+
+    def update(self, count: int, size: int):
+        self.uploaded_count += count
+        self.uploaded_size += size
+        self.updated_at = int(time.time() * 1000)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "uploaded_count": self.uploaded_count,
+            "uploaded_size": self.uploaded_size,
+            "updated_at": self.updated_at,
+        }
+
+
+class SimpleS3Client:
+    """简化的S3客户端"""
+
+    def __init__(self, endpoint: str, access_key: str, secret_key: str, session_token: str):
+        self.client = minio.Minio(
+            endpoint, access_key=access_key, secret_key=secret_key, secure=False, session_token=session_token
+        )
+
+    def upload_file(self, local_path: str, bucket: str, object_name: str) -> Tuple[str, int]:
+        """上传文件并返回哈希和大小"""
+        file_size = os.path.getsize(local_path)
+
+        # 计算文件哈希
+        sha256_hash = hashlib.sha256()
+        with open(local_path, "rb") as f:
+            for chunk in iter(lambda: f.read(8192), b""):
+                sha256_hash.update(chunk)
+
+        file_hash = sha256_hash.hexdigest()
+
+        # 上传文件
+        with open(local_path, "rb") as f:
+            self.client.put_object(bucket, object_name, f, file_size)
+
+        return file_hash, file_size
+
+    def upload_json(self, data: Dict[str, Any], bucket: str, object_name: str):
+        """上传JSON数据"""
+        json_str = json.dumps(data)
+        json_bytes = json_str.encode("utf-8")
+
+        from io import BytesIO
+
+        self.client.put_object(
+            bucket, object_name, BytesIO(json_bytes), len(json_bytes), content_type="application/json"
+        )
+
+
+class DataUploader:
+    """数据上传器"""
+
+    def __init__(
+        self, task_id: int, local_path: str, s3_target: str, csv_path: str, status_path: str, num_workers: int = 10
+    ):
+        self.task_id = task_id
+        self.local_path = local_path
+        self.num_workers = num_workers
+
+        # 解析S3路径
+        self.target_bucket, self.target_prefix = self._parse_s3_path(s3_target)
+        self.csv_bucket, self.csv_key = self._parse_s3_path(csv_path)
+        self.status_bucket, self.status_key = self._parse_s3_path(status_path)
+
+        # 创建工作目录
+        self.work_dir = Path.home() / ".di_workspace" / str(task_id)
+        self.work_dir.mkdir(parents=True, exist_ok=True)
+        self.csv_file = self.work_dir / "upload_records.csv"
+
+        # CSV记录队列
+        self.csv_queue = queue.Queue()
+        self.processed_files = set()
+        self.total_files = 0
+
+    def _parse_s3_path(self, s3_path: str) -> Tuple[str, str]:
+        """解析S3路径"""
+        if s3_path.startswith("s3://"):
+            parts = s3_path[5:].split("/", 1)
+            bucket = parts[0]
+            key = parts[1] if len(parts) > 1 else ""
+            return bucket, key
+        return "", ""
+
+    def _collect_files(self) -> list:
+        """收集需要上传的文件"""
+        files = []
+
+        if os.path.isfile(self.local_path):
+            files.append(self.local_path)
+            self.total_files += 1
+        else:
+            for root, _, filenames in os.walk(self.local_path):
+                for filename in filenames:
+                    file_path = os.path.join(root, filename)
+                    if not os.path.islink(file_path):  # 跳过符号链接
+                        files.append(file_path)
+                        self.total_files += 1
+
+        # 过滤已处理的文件
+        base_path = os.path.dirname(self.local_path) if os.path.isfile(self.local_path) else self.local_path
+        unprocessed_files = []
+
+        for file_path in files:
+            rel_path = os.path.relpath(file_path, base_path)
+            if rel_path not in self.processed_files:
+                unprocessed_files.append(file_path)
+
+        return unprocessed_files
+
+    def _csv_writer_worker(self):
+        """CSV写入工作器"""
+        # 初始化CSV文件
+        uploaded_count = 0
+        file_exists = os.path.exists(self.csv_file)
+
+        with open(self.csv_file, "a", newline="", encoding="utf-8") as f:
+            writer = csv.writer(f)
+            if not file_exists:
+                writer.writerow(["local_path", "sha256", "s3path", "file_size"])
+
+            while True:
+                try:
+                    record = self.csv_queue.get(timeout=1)
+                    if record is None:  # 结束信号
+                        break
+
+                    writer.writerow(
+                        [record["local_path"], record["file_hash"], record["s3_path"], str(record["file_size"])]
+                    )
+
+                    f.flush()  # 确保数据写入磁盘
+                    self.csv_queue.task_done()
+                    uploaded_count += 1
+                    # 每上传100个文件，打印进度
+                    if uploaded_count % 1000 == 0:
+                        logger.info(f"已上传 {uploaded_count} 个文件")
+
+                except queue.Empty:
+                    continue
+                except Exception as e:
+                    logger.error(f"Failed to write CSV record: {e}")
+                    self.csv_queue.task_done()
+
+    def _upload_worker(self, s3_client: SimpleS3Client, file_queue: queue.Queue, base_path: str):
+        """上传工作器"""
+        while True:
+            try:
+                file_path = file_queue.get(timeout=1)
+                if file_path is None:  # 结束信号
+                    break
+
+                try:
+                    # 计算相对路径和S3对象名
+                    rel_path = os.path.relpath(file_path, base_path)
+
+                    object_name = os.path.join(self.target_prefix, rel_path).replace("\\", "/")
+
+                    # 上传文件
+                    file_hash, file_size = s3_client.upload_file(file_path, self.target_bucket, object_name)
+
+                    # 将记录放入CSV队列
+                    s3_path = f"s3://{self.target_bucket}/{object_name}"
+                    record = {
+                        "local_path": os.path.join("/", rel_path),
+                        "file_hash": file_hash,
+                        "s3_path": s3_path,
+                        "file_size": file_size,
+                    }
+                    self.csv_queue.put(record)
+
+                    logger.debug(f"Uploaded: {rel_path}")
+
+                except Exception as e:
+                    logger.error(f"Failed to upload {file_path}: {e}")
+                finally:
+                    file_queue.task_done()
+
+            except queue.Empty:
+                break
+
+    def _calculate_final_stats(self) -> UploadStatus:
+        """从CSV文件计算最终统计信息"""
+        stats = UploadStatus()
+        if not os.path.exists(self.csv_file):
+            return stats
+
+        total_count = 0
+        total_size = 0
+
+        try:
+            with open(self.csv_file, "r", encoding="utf-8") as f:
+                reader = csv.DictReader(f)
+                for row in reader:
+                    total_count += 1
+                    total_size += int(row["file_size"])
+        except Exception as e:
+            logger.error(f"Failed to calculate stats: {e}")
+
+        stats.update(total_count, total_size)
+
+        return stats
+
+    def run(self, s3_client: SimpleS3Client) -> UploadStatus:
+        """执行上传任务"""
+        # 收集文件
+        files = self._collect_files()
+        if not files:
+            logger.info("No files to upload")
+            return UploadStatus()
+
+        logger.info(f"Found {len(files)} files to upload")
+
+        # 准备文件队列
+        file_queue = queue.Queue()
+        for file_path in files:
+            file_queue.put(file_path)
+
+        base_path = os.path.dirname(self.local_path) if os.path.isfile(self.local_path) else self.local_path
+
+        # 启动CSV写入线程
+        csv_thread = threading.Thread(target=self._csv_writer_worker)
+        csv_thread.daemon = True
+        csv_thread.start()
+
+        try:
+            # 启动上传工作器
+            with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
+                futures = []
+                for i in range(self.num_workers):
+                    future = executor.submit(self._upload_worker, s3_client, file_queue, base_path)
+                    futures.append(future)
+
+                # 等待所有任务完成
+                for future in as_completed(futures):
+                    future.result()
+
+            # 等待CSV队列处理完成
+            self.csv_queue.join()
+
+            # 发送结束信号给CSV写入线程
+            self.csv_queue.put(None)
+            csv_thread.join()
+
+            # 上传记录文件到S3
+            if os.path.exists(self.csv_file):
+                s3_client.upload_file(str(self.csv_file), self.csv_bucket, self.csv_key)
+                logger.info("Upload records saved to S3")
+
+            # 计算并上传最终统计信息
+            stats = self._calculate_final_stats()
+            s3_client.upload_json(stats.to_dict(), self.status_bucket, self.status_key)
+            logger.info(f"Upload completed: {stats.uploaded_count} files, {stats.uploaded_size} bytes")
+
+        finally:
+            # 清理工作目录
+            try:
+                import shutil
+
+                shutil.rmtree(self.work_dir)
+            except Exception as e:
+                logger.warning(f"Failed to cleanup workspace: {e}")
+        return stats
+
+
+def main():
+    """主函数"""
+    parser = argparse.ArgumentParser(description="简化的数据摄取工具")
+    parser.add_argument("-e", "--endpoint", default="192.168.13.160:9008", help="S3端点")
+    parser.add_argument("-ak", "--access-key", default="admin2024", help="访问密钥")
+    parser.add_argument("-sk", "--secret-key", default="root@23452024", help="秘密密钥")
+    parser.add_argument("-t", "--target", default="s3://testbucket/test_ok11", help="目标S3路径")
+    parser.add_argument("-l", "--local", default="./test_data", help="本地路径")
+    parser.add_argument("-o", "--object-sheet", default="s3://testbucket/records/123.csv", help="记录文件S3路径")
+    parser.add_argument("-s", "--status", default="s3://testbucket/status/123.json", help="状态文件S3路径")
+    parser.add_argument("-i", "--task-id", type=int, default=123, help="任务ID")
+    parser.add_argument("-n", "--num-workers", type=int, default=10, help="工作线程数")
+
+    args = parser.parse_args()
+
+    # 检查本地路径
+    if not os.path.exists(args.local):
+        logger.error(f"Local path does not exist: {args.local}")
+        sys.exit(1)
+
+    logger.info(f"Starting upload: {args.local} -> {args.target}")
+
+    try:
+        # 创建S3客户端
+        s3_client = SimpleS3Client(args.endpoint, args.access_key, args.secret_key)
+
+        # 创建上传器并执行
+        uploader = DataUploader(
+            task_id=args.task_id,
+            local_path=args.local,
+            s3_target=args.target,
+            csv_path=args.object_sheet,
+            status_path=args.status,
+            num_workers=args.num_workers,
+        )
+
+        uploader.run(s3_client)
+        logger.info("Upload completed successfully")
+
+    except Exception as e:
+        logger.error(f"Upload failed: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

aihub/utils/download.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-import concurrent.futures
 import os
 import tempfile
 import zipfile
 from typing import List, TypedDict
 
 import pyarrow.parquet as pq
-from tqdm import tqdm
+from tqdm.contrib.concurrent import thread_map
 
-from .http import http_download_file
+from .http import http_download_file, http_download_file_wrapper
 from .s3 import s3_to_url
 
 
@@ -59,18 +58,7 @@ def dataset_download(index_url: str, local_dir: str, worker: int = 4) -> None:
     if worker < 1:
         worker = 1
 
-    with (
-        tqdm(total=len(files), desc="Downloading dataset") as bar,
-        concurrent.futures.ThreadPoolExecutor(max_workers=worker) as pool,
-    ):
-
-        def _one(flocal: str, furl: str):
-            http_download_file(furl, flocal)
-            bar.update()
-
-        futures = [pool.submit(_one, p, u) for p, u in files]
-        for fut in concurrent.futures.as_completed(futures):
-            fut.result()
+    thread_map(http_download_file_wrapper, files, max_workers=worker)
 
 
 def zip_dir(dir_path: str, zip_path: str):

aihub/utils/http.py

@@ -5,6 +5,12 @@ import os
 import requests
 
 
+def http_download_file_wrapper(item):
+    """Wrapper function"""
+    dst_path, url = item
+    return http_download_file(url, dst_path)
+
+
 def http_download_file(url: str, dst_path: str, chunk: int = 1 << 16) -> None:
     os.makedirs(os.path.dirname(dst_path), exist_ok=True)
     with requests.get(url, timeout=None, stream=True) as r:

{intellif_aihub-0.1.14.dist-info → intellif_aihub-0.1.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: intellif-aihub
-Version: 0.1.14
+Version: 0.1.15
 Summary: Intellif AI-hub SDK.
 Author-email: Platform Team <aihub@example.com>
 License-Expression: Apache-2.0

{intellif_aihub-0.1.14.dist-info → intellif_aihub-0.1.15.dist-info}/RECORD

@@ -1,11 +1,11 @@
-aihub/__init__.py,sha256=PIBqEOI-nqKFL9oJAWQQwlHuujG9Cd7EmdxDrThNQto,23
+aihub/__init__.py,sha256=qb0TalpSt1CbprnFyeLUKqgrqNtmnk9IoQQ7umAoXVY,23
 aihub/client.py,sha256=nVELjkyVOG6DKJjurYn59fCoT5JsSayUweiH7bvKcAo,5547
 aihub/exceptions.py,sha256=l2cMAvipTqQOio3o11fXsCCSCevbuK4PTsxofkobFjk,500
 aihub/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aihub/models/artifact.py,sha256=4xIWV3xfZuZWcCdmZGEZ5k_rvV4oc5C_3gapw5O-2vQ,4586
+aihub/models/artifact.py,sha256=F-r7DJY9A09yIQJqWol6gLRu6y7NGjRa6-BxkMEluxU,4655
 aihub/models/common.py,sha256=qmabc2LkAdQJXIcpT1P35zxd0Lc8yDYdD4ame1iF4Bs,241
 aihub/models/data_warehouse.py,sha256=zXvWwg7ySoFJMdqQ_1UMTNEKDMhu1hDHlWdBAXdizBk,3905
-aihub/models/dataset_management.py,sha256=etvMOwKO3Oc8h-pJ-adoGCOLWIe2YYzx2smZZ7eCdWw,10327
+aihub/models/dataset_management.py,sha256=4DuQ0zM7jv73SJiqvieHLtn2Y-T6FIFV9r7bgzyCtDo,10790
 aihub/models/document_center.py,sha256=od9bzx6krAS6ktIA-ChxeqGcch0v2wsS1flY2vuHXBc,1340
 aihub/models/eval.py,sha256=4Gon4Sg4dOkyCx3KH2mO5ip3AhrBwrPC0UZA447HeoQ,910
 aihub/models/labelfree.py,sha256=nljprYO6ECuctTVbHqriQ73N5EEyYURhBrnU28Ngfvc,1589
@@ -17,9 +17,9 @@ aihub/models/task_center.py,sha256=HE21Q4Uj-vt9LHezHnqBYgdinhrh4iJPOq8VXbSMllU,5
 aihub/models/user_system.py,sha256=0L_pBkWL9v1tv_mclOyRgCyWibuuj_XU3mPoe2v48vQ,12216
 aihub/models/workflow_center.py,sha256=4xtI1WZ38ceXJ8gwDBj-QNjOiRlLO_8kGiQybdudJPY,20121
 aihub/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aihub/services/artifact.py,sha256=PtCGhYFpFK_hppQr7A0bdXvceXzmYSwEpRj-PE2rIcQ,11473
+aihub/services/artifact.py,sha256=lfOrgOT2AlH1w-75NLcQGOhVWdhmJcWD1gESPpUzqUw,11257
 aihub/services/data_warehouse.py,sha256=awvlJdggo8ph6sXweXXVp4GLRuUSD46LoD0QQksXRts,2964
-aihub/services/dataset_management.py,sha256=i9Q_Mt8XZeFQR3_5JqIeXu3Y01ay-fPXbykkrfs0OmQ,15961
+aihub/services/dataset_management.py,sha256=R7mFsJ1dNOI_p5yNj_rQdLolRC0UKEN4WejE7uOjVlE,21379
 aihub/services/document_center.py,sha256=dG67Ji-DOnzL2t-4x4gVfMt9fbSj_IjVHCLw5R-VTkQ,1813
 aihub/services/eval.py,sha256=V1nBISIyYWg9JJO24xzy4-kit9NsaCYp1EWIX_fgJkQ,2128
 aihub/services/labelfree.py,sha256=xua62UWhVXTxJjHRyy86waaormnJjmpQwepcARBy_h0,1450
@@ -32,11 +32,12 @@ aihub/services/task_center.py,sha256=rVQG7q2_GN0501w5KHsOOlSVFX9ovpRMGX5hskCqggw
 aihub/services/user_system.py,sha256=IqWL4bnsKyyzuGT5l6adnw0qNXlH9PSo1-C_pFyOSzA,18868
 aihub/services/workflow_center.py,sha256=caKxOlba0J1s1RUK6RUm1ndJSwAcZXEakRanu3sGKPU,17468
 aihub/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aihub/utils/download.py,sha256=yC3SoL5uE68pMB-IsNz233wj-gFrHB7D7ALzQA5JkFM,2155
-aihub/utils/http.py,sha256=SvEWB4BxvwaHYqMVE4B0Go3OWGAD4xyQnUXDZ16yOSo,410
+aihub/utils/di.py,sha256=vFUzno5WbRKu6-pj8Hnz9IqT7xb9UDZQ4qpOFH1YAtM,11812
+aihub/utils/download.py,sha256=ZZVbcC-PnN3PumV7ZiJ_-srkt4HPPovu2F6Faa2RrPE,1830
+aihub/utils/http.py,sha256=AmfHHNjptuuSFx2T1twWCnerR_hLN_gd0lUs8z36ERA,547
 aihub/utils/s3.py,sha256=ISIBP-XdBPkURpXnN56ZnIWokOOg2SRUh_qvxJk-G1Q,2187
-intellif_aihub-0.1.14.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-intellif_aihub-0.1.14.dist-info/METADATA,sha256=f1WbrM_51scvCQQZsNdhgYpweM3tVtdMftIif8zq7yU,2949
-intellif_aihub-0.1.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-intellif_aihub-0.1.14.dist-info/top_level.txt,sha256=vIvTtSIN73xv46BpYM-ctVGnyOiUQ9EWP_6ngvdIlvw,6
-intellif_aihub-0.1.14.dist-info/RECORD,,
+intellif_aihub-0.1.15.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+intellif_aihub-0.1.15.dist-info/METADATA,sha256=Hz8Z3sB06pNTJF8lygzDU37da2bCgXCrzJ1-CRAlN7Y,2949
+intellif_aihub-0.1.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+intellif_aihub-0.1.15.dist-info/top_level.txt,sha256=vIvTtSIN73xv46BpYM-ctVGnyOiUQ9EWP_6ngvdIlvw,6
+intellif_aihub-0.1.15.dist-info/RECORD,,