intellif-aihub 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

aihub/models/task_center.py

@@ -20,18 +20,18 @@ class TaskCenterPriorityEnum(Enum):
 class LabelProjectTypeEnum(Enum):
     """
     任务类型枚举
+     1 - 目标检测 2 - 语义分割 3 - 图片分类 4 - 实例分割 5 - 视频标注 6 - 人类偏好文本标注 7- 敏感预料文本标注 8 - 文本标注 9 - 关键点标注
     """
 
-    IMAGE_CLASSIFICATION = 1
-    OBJECT_DETECTION = 2
-    SEMANTIC_SEGMENTATION = 3
-    TEXT_CLASSIFICATION = 4
-    TEXT_SEGMENTATION = 5
-    TEXT_RECOGNITION = 6
-    IMAGE_CAPTIONING = 7
-    IMAGE_STYLE_TRANSFER = 8
-    IMAGE_COLORIZATION = 9
-    IMAGE_RESTORATION = 10
+    OBJECT_DETECTION = 1
+    SEGMENTATION = 2
+    IMAGE_CLASSIFICATION = 3
+    INSTANCE_SEGMENTATION = 4
+    VIDEO_LABELING = 5
+    HUMAN_PREFERENCE_TEXT_LABELING = 6
+    SENSITIVE_TEXT_LABELING = 7
+    TEXT_LABELING = 8
+    KEYPOINT_LABELING = 9
 
 
 class CreateTaskOtherInfo(BaseModel):

aihub/models/user_system.py

@@ -0,0 +1,262 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+# ======================================================================
+# COMMON
+# ======================================================================
+
+class Role(BaseModel):
+    id: int
+    name: str
+    role_type: int = Field(alias="role_type")
+    menu_ids: Optional[List[int]] = Field(None, alias="menu_ids")
+
+
+class Menu(BaseModel):
+    id: int
+    name: str
+    parent: int
+    auth: str
+
+
+class TreeMenu(BaseModel):
+    id: int
+    name: str
+    parent: int
+    auth: str
+    children: Optional[List["TreeMenu"]] = None
+    roles: Optional[List[Role]] = None
+
+
+class TagBrief(BaseModel):
+    id: int
+    name: str
+
+
+# ======================================================================
+# ------------------------------- AUTH ---------------------------------
+# ======================================================================
+
+class LoginRequest(BaseModel):
+    username: str = Field(alias="username")
+    password: str = Field(alias="password")
+
+
+class LoginResponse(BaseModel):
+    id: int = Field(alias="id")
+    token: str = Field(alias="token")
+
+
+class SignupRequest(BaseModel):
+    username: str = Field(alias="username")
+    password: str = Field(alias="password")
+    nickname: str = Field(alias="nickname")
+    email: str = Field(alias="email")
+    role_ids: List[int] = Field(alias="role_ids")
+
+
+class SignupResponse(BaseModel):
+    id: int = Field(alias="id")
+
+
+# ======================================================================
+# ------------------------------- MENU ---------------------------------
+# ======================================================================
+class ListMenusRequest(BaseModel):
+    need_roles: Optional[bool] = Field(None, alias="need_roles")
+
+
+class ListMenusResponse(BaseModel):
+    menus: List[TreeMenu] = Field(None, alias="menus")
+
+
+class CreateMenuRequest(BaseModel):
+    name: str = Field(alias="name")
+    parent: int = Field(alias="parent")
+    auth: str = Field(alias="auth")
+    role_ids: Optional[List[int]] = Field(None, alias="role_ids")
+
+
+class CreateMenuResponse(BaseModel):
+    id: int = Field(alias="id")
+
+
+class UpdateMenuRequest(BaseModel):
+    name: Optional[str] = Field(None, alias="name")
+    parent: Optional[int] = Field(None, alias="parent")
+    auth: str = Field(alias="auth")
+    role_ids: Optional[List[int]] = Field(None, alias="role_ids")
+
+
+class GetMenuRolesResponse(BaseModel):
+    role_ids: List[int] = Field(alias="role_ids")
+
+
+class SetMenuRolesRequest(BaseModel):
+    role_ids: List[int] = Field(alias="role_ids")
+
+
+class SearchMenusRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+    name: Optional[str] = None
+    parent_ids: Optional[List[int]] = Field(None, alias="parent_ids")
+    auth: Optional[str] = None
+    menu_ids: Optional[List[int]] = Field(None, alias="menu_ids")
+
+
+class SearchMenusResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[Menu]
+
+
+# ======================================================================
+# ------------------------------- ROLE ---------------------------------
+# ======================================================================
+
+
+class CreateRoleRequest(BaseModel):
+    id: Optional[int] = Field(None, alias="id")
+    name: str
+    role_type: int = Field(alias="role_type")
+    menu_ids: Optional[List[int]] = Field(None, alias="menu_ids")
+
+
+class CreateRoleResponse(BaseModel):
+    id: int
+
+
+class UpdateRoleRequest(BaseModel):
+    name: Optional[str] = None
+    role_type: Optional[int] = Field(None, alias="role_type")
+    menu_ids: Optional[List[int]] = Field(None, alias="menu_ids")
+
+
+class GetRoleMenusResponse(BaseModel):
+    menu_ids: List[int] = Field(alias="menu_ids")
+
+
+class SetRoleMenusRequest(BaseModel):
+    menu_ids: List[int] = Field(alias="menu_ids")
+
+
+class ListRolesRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+    role_type: Optional[int] = Field(None, alias="role_type")
+
+
+class ListRolesResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[Role]
+
+
+class SearchRolesRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+    name: Optional[str] = None
+    role_ids: Optional[List[int]] = Field(None, alias="role_ids")
+    menu_ids: Optional[List[int]] = Field(None, alias="menu_ids")
+
+
+class SearchRolesResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[Role]
+
+
+# ======================================================================
+# ------------------------------- USER ---------------------------------
+# ======================================================================
+
+class User(BaseModel):
+    id: int
+    username: str
+    nickname: str
+    email: str
+    roles: Optional[List[Role]] = Field(None, alias="roles")
+    status: int
+    tags: Optional[List[TagBrief]] = Field(None, alias="tags")
+    created_at: int = Field(alias="created_at")
+
+
+class ListUsersRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+    search_key: Optional[str] = Field(None, alias="search_key")
+
+
+class ListUsersResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[User]
+
+
+class CreateUserRequest(BaseModel):
+    id: int
+    username: str
+    password: str
+    nickname: str
+    email: str
+    role_ids: Optional[List[int]] = Field(None, alias="role_ids")
+    created_at: Optional[int] = Field(None, alias="created_at")
+    updated_at: Optional[int] = Field(None, alias="updated_at")
+    status: Optional[int] = None
+    tag_ids: Optional[List[int]] = Field(None, alias="tag_ids")
+
+
+class CreateUserResponse(BaseModel):
+    id: int
+
+
+class UpdateUserRequest(BaseModel):
+    username: Optional[str] = None
+    nickname: Optional[str] = None
+    email: Optional[str] = None
+    password: Optional[str] = None
+    role_ids: Optional[List[int]] = Field(default_factory=list, alias="role_ids")
+    status: Optional[int] = None
+    tag_ids: Optional[List[int]] = Field(default_factory=list, alias="tag_ids")
+
+
+class SetUserRolesRequest(BaseModel):
+    role_ids: List[int] = Field(alias="role_ids")
+
+
+class GetUserMenusResponse(BaseModel):
+    menus: List[TreeMenu]
+
+
+class SearchUsersRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+    username: Optional[str] = None
+    nickname: Optional[str] = None
+    email: Optional[str] = None
+    user_ids: Optional[List[int]] = Field(None, alias="user_ids")
+    role_ids: Optional[List[int]] = Field(None, alias="role_ids")
+    role_names: Optional[List[str]] = Field(None, alias="role_names")
+    status: Optional[int] = None
+
+
+class SearchUsersResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[User]
+
+
+# 此行放在文件末尾，否则序列化报错
+TreeMenu.update_forward_refs()

aihub/services/artifact.py

@@ -52,14 +52,8 @@ class ArtifactService:
         """
         self._http = http
         self._Artifact = _Artifact(http)
-        self.sts = self._get_sts()
-        self.s3_client = minio.Minio(
-            self.sts.endpoint,
-            access_key=self.sts.access_key_id,
-            secret_key=self.sts.secret_access_key,
-            secure=False,
-            session_token=self.sts.session_token,
-        )
+        self.sts = None
+        self.s3_client = None
 
     @property
     def _artifact(self) -> _Artifact:
@@ -142,6 +136,15 @@ class ArtifactService:
             APIError: 当API调用失败时抛出
         """
         logger.info(f"log artifact: {artifact_path},local path: {local_path} ")
+        if self.s3_client is None:
+            self.sts = self._get_sts()
+            self.s3_client = minio.Minio(
+                self.sts.endpoint,
+                access_key=self.sts.access_key_id,
+                secret_key=self.sts.secret_access_key,
+                session_token=self.sts.session_token,
+                secure=False,
+            )
 
         # 检查文件是否存在
         if not os.path.exists(local_path):
@@ -187,6 +190,15 @@ class ArtifactService:
             ValueError: 当本地目录不存在时抛出
             APIError: 当API调用失败时抛出
         """
+        if self.s3_client is None:
+            self.sts = self._get_sts()
+            self.s3_client = minio.Minio(
+                self.sts.endpoint,
+                access_key=self.sts.access_key_id,
+                secret_key=self.sts.secret_access_key,
+                session_token=self.sts.session_token,
+                secure=False,
+            )
 
         logger.info(f"log artifact: {artifact_path},local path: {local_dir} ")
         if not os.path.exists(local_dir):
@@ -217,6 +229,15 @@ class ArtifactService:
         Raises:
             APIError: 当API调用失败时抛出
         """
+        if self.s3_client is None:
+            self.sts = self._get_sts()
+            self.s3_client = minio.Minio(
+                self.sts.endpoint,
+                access_key=self.sts.access_key_id,
+                secret_key=self.sts.secret_access_key,
+                session_token=self.sts.session_token,
+                secure=False,
+            )
         artifacts = self.get_by_run_id(run_id, artifact_path)
 
         for artifact_item in artifacts:

aihub/services/dataset_management.py

@@ -131,7 +131,7 @@ class DatasetManagementService:
         return dataset_id, version_id, version_tag
 
     def run_download(
-        self, *, dataset_version_name: str, local_dir: str, worker: int = 4
+        self, dataset_version_name: str, local_dir: str, worker: int = 4
     ) -> None:
         """根据数据集版本名称下载对应的数据集文件。
 

aihub/services/eval.py

@@ -0,0 +1,68 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+import httpx
+
+from ..exceptions import APIError
+from ..models.common import APIWrapper
+from ..models.eval import CreatEvalReq, CreatEvalResp
+
+_BASE = "/eval-platform/api/v1"
+
+
+class EvalService:
+    """评测服务"""
+
+    def __init__(self, http: httpx.Client):
+        self._http = http
+        self._eval = _Eval(http)
+
+    def create(
+        self,
+        dataset_version_name: str,
+        prediction_artifact_path: str,
+        evaled_artifact_path: str,
+        report_json: dict,
+        run_id,
+    ) -> int:
+        """创建评测报告
+
+        Args:
+            run_id (str): RUN ID
+            report_json (dict): 报告内容
+            evaled_artifact_path:   评测结果制品路径
+            prediction_artifact_path: 推理结果制品路径
+            dataset_version_name (str): 数据集名称
+
+
+        Returns:
+            id (int): 评测报告id
+
+        """
+        from .dataset_management import DatasetManagementService
+
+        dataset_service = DatasetManagementService(self._http)
+        dataset_version = dataset_service.get_dataset_version_by_name(
+            dataset_version_name
+        )
+        payload = CreatEvalReq(
+            dataset_id=dataset_version.dataset_id,
+            dataset_version_id=dataset_version.id,
+            evaled_artifact_path=evaled_artifact_path,
+            prediction_artifact_path=prediction_artifact_path,
+            report=report_json,
+            run_id=run_id,
+        )
+
+        return self._eval.create(payload)
+
+
+class _Eval:
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def create(self, payload: CreatEvalReq) -> int:
+        resp = self._http.post(f"{_BASE}/RUN/", json=payload.model_dump())
+        wrapper = APIWrapper[CreatEvalResp].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data.id

aihub/services/labelfree.py

@@ -4,7 +4,7 @@ import httpx
 
 from ..exceptions import APIError
 from ..models.common import APIWrapper
-from ..models.labelfree import *
+from ..models.labelfree import GetGlobalStatsResponse
 
 _BASE = "/labelfree/api/v2"
 

aihub/services/model_training_platform.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import httpx
+
+from ..exceptions import APIError
+from ..models.common import APIWrapper
+from ..models.model_training_platform import (
+    CreateTrainingRequest,
+    ListTrainingsRequest,
+    ListTrainingsResponse,
+    ListTrainingPodsRequest,
+    GetTrainingPodLogsNewResponse,
+    GetTrainingPodEventsResponse,
+    Training,
+    CreateTrainingResponse,
+    ListTrainingUsersRequest,
+    ListTrainingPodsResponse,
+    GetTrainingPodSpecResponse,
+    ListTrainingContainersRequest,
+    ListTrainingUsersResponse,
+    ListTrainingContainersResponse,
+    Pod,
+)
+
+_BASE = "/model-training-platform/api/v1"
+
+
+class ModelTrainingPlatformService:
+
+    def __init__(self, http: httpx.Client):
+        self._training = _Training(http)
+
+    def create_training(self, payload: CreateTrainingRequest) -> int:
+        """创建训练任务
+
+        Args:
+            payload (CreateTrainingRequest): 创建训练任务参数
+
+        Returns:
+            int: 训练任务ID
+        """
+        return self._training.create(payload)
+
+    def list_trainings(self, payload: ListTrainingsRequest) -> ListTrainingsResponse:
+        return self._training.list(payload)
+
+    def get_training(self, training_id: int) -> Training:
+        """获取训练任务详情
+
+        Args:
+            training_id (int): 训练任务ID
+        Returns:
+            Training
+
+        """
+        return self._training.get(training_id)
+
+    def stop_training(self, training_id: int) -> None:
+        """停止训练瑞文
+
+        Args:
+            training_id (int): 训练任务ID
+
+        Returns:
+
+        """
+        self._training.stop(training_id)
+
+    def list_training_pods(
+        self, training_id: int, payload: ListTrainingPodsRequest
+    ) -> ListTrainingPodsResponse:
+        return self._training.list_training_pods(training_id, payload)
+
+    def get_training_pod(self, training_id: int, pod_id: int) -> Pod:
+        return self._training.get_training_pod(training_id, pod_id)
+
+    def get_pod_logs_new(
+        self, training_id: int, pod_id: int
+    ) -> GetTrainingPodLogsNewResponse:
+        return self._training.get_training_logs_new(training_id, pod_id)
+
+    def get_pod_spec(self, training_id: int, pod_id: int) -> GetTrainingPodSpecResponse:
+        return self._training.get_training_spec(training_id, pod_id)
+
+    def get_pod_events(
+        self, training_id: int, pod_id: int
+    ) -> GetTrainingPodEventsResponse:
+        return self._training.get_training_events(training_id, pod_id)
+
+    def list_training_users(
+        self, payload: ListTrainingUsersRequest
+    ) -> ListTrainingUsersResponse:
+        return self._training.list_training_users(payload)
+
+    def list_training_containers(
+        self, payload: ListTrainingContainersRequest
+    ) -> ListTrainingContainersResponse:
+        return self._training.list_training_containers(payload)
+
+    @property
+    def training(self) -> _Training:
+        return self._training
+
+
+class _Training:
+
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def create(self, payload: CreateTrainingRequest) -> int:
+        resp = self._http.post(
+            f"{_BASE}/trainings",
+            json=payload.model_dump(by_alias=True, exclude_none=True),
+        )
+        wrapper = APIWrapper[CreateTrainingResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data.id
+
+    def list(self, payload: ListTrainingsRequest) -> ListTrainingsResponse:
+        resp = self._http.get(
+            f"{_BASE}/trainings",
+            params=payload.model_dump(by_alias=True, exclude_none=True),
+        )
+        wrapper = APIWrapper[ListTrainingsResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def get(self, training_id: int) -> Training:
+        resp = self._http.get(f"{_BASE}/trainings/{training_id}")
+        wrapper = APIWrapper[Training].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def stop(self, training_id: int) -> None:
+        resp = self._http.post(f"{_BASE}/trainings/{training_id}/stop")
+        wrapper = APIWrapper[dict].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+
+    def list_training_pods(
+        self, training_id: int, payload: ListTrainingPodsRequest
+    ) -> ListTrainingPodsResponse:
+        resp = self._http.get(
+            f"{_BASE}/trainings/{training_id}/pods",
+            params=payload.model_dump(by_alias=True, exclude_none=True),
+        )
+        wrapper = APIWrapper[ListTrainingPodsResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def get_training_pod(self, training_id: int, pod_id: int) -> Pod:
+        resp = self._http.get(f"{_BASE}/trainings/{training_id}/pods/{pod_id}")
+        wrapper = APIWrapper[Pod].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def get_training_logs_new(
+        self, training_id: int, pod_id: int
+    ) -> GetTrainingPodLogsNewResponse:
+        resp = self._http.get(f"{_BASE}/trainings/{training_id}/pods/{pod_id}/logs/new")
+        wrapper = APIWrapper[GetTrainingPodLogsNewResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def get_training_spec(
+        self, training_id: int, pod_id: int
+    ) -> GetTrainingPodSpecResponse:
+        resp = self._http.get(f"{_BASE}/trainings/{training_id}/pods/{pod_id}/spec")
+        wrapper = APIWrapper[GetTrainingPodSpecResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def get_training_events(
+        self, training_id: int, pod_id: int
+    ) -> GetTrainingPodEventsResponse:
+        resp = self._http.get(f"{_BASE}/trainings/{training_id}/pods/{pod_id}/events")
+        wrapper = APIWrapper[GetTrainingPodEventsResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def list_training_users(
+        self, payload: ListTrainingUsersRequest
+    ) -> ListTrainingUsersResponse:
+        resp = self._http.get(
+            f"{_BASE}/training-users", params=payload.model_dump(by_alias=True)
+        )
+        wrapper = APIWrapper[ListTrainingUsersResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def list_training_containers(
+        self, payload: ListTrainingContainersRequest
+    ) -> ListTrainingContainersResponse:
+        resp = self._http.get(
+            f"{_BASE}/training-containers", params=payload.model_dump(by_alias=True)
+        )
+        wrapper = APIWrapper[ListTrainingContainersResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data