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

aihub/__init__.py

@@ -1 +1 @@
-__version__ = "0.1.2"
+__version__ = "0.1.4"

aihub/client.py

@@ -9,14 +9,20 @@ from .services import artifact
 from .services import dataset_management
 from .services import document_center
 from .services import labelfree
-from .services import tag_management
+from .services import model_training_platform
+from .services import quota_schedule_management
+from .services import tag_resource_management
 from .services import task_center
+from .services import user_system
 from .services.artifact import ArtifactService
 from .services.dataset_management import DatasetManagementService
 from .services.document_center import DocumentCenterService
 from .services.labelfree import LabelfreeService
-from .services.tag_management import TagManagementService
+from .services.model_training_platform import ModelTrainingPlatformService
+from .services.quota_schedule_management import QuotaScheduleManagementService
+from .services.tag_resource_management import TagResourceManagementService
 from .services.task_center import TaskCenterService
+from .services.user_system import UserSystemService
 
 
 class Client:
@@ -31,15 +37,16 @@ class Client:
     """
 
     labelfree: LabelfreeService = None
-    tag_management: TagManagementService = None
+    tag_resource_management: TagResourceManagementService = None
     document_center: DocumentCenterService = None
     task_center: TaskCenterService = None
     dataset_management: DatasetManagementService = None
     artifact: ArtifactService = None
+    user_system: UserSystemService = None
+    model_training_platform: ModelTrainingPlatformService = None
+    quota_schedule_management: QuotaScheduleManagementService = None
 
-    def __init__(
-        self, *, base_url: str, token: str | None = None, timeout: float = 60.0
-    ):
+    def __init__(self, *, base_url: str, token: str | None = None, timeout: float = 60.0):
         """AI-HUB python SDK 客户端
 
         Args:
@@ -54,7 +61,7 @@ class Client:
         if not base_url:
             raise ValueError("base_url必须填写")
 
-        token = token or os.getenv("AI_HUB_TOKEN")
+        token = os.getenv("AI_HUB_TOKEN") or token
         if not token:
             raise ValueError("缺少token：请显式传入，或在环境变量AI_HUB_TOKEN中设置")
 
@@ -64,14 +71,15 @@ class Client:
             headers={"Authorization": f"Bearer {token}"},
             # event_hooks={"response": [self._raise_for_status]},
         )
-        self.dataset_management = dataset_management.DatasetManagementService(
-            self._http
-        )
+        self.dataset_management = dataset_management.DatasetManagementService(self._http)
         self.labelfree = labelfree.LabelfreeService(self._http)
-        self.tag_management = tag_management.TagManagementService(self._http)
+        self.tag_resource_management = tag_resource_management.TagResourceManagementService(self._http)
         self.document_center = document_center.DocumentCenterService(self._http)
         self.task_center = task_center.TaskCenterService(self._http)
         self.artifact = artifact.ArtifactService(self._http)
+        self.user_system = user_system.UserSystemService(self._http)
+        self.model_training_platform = model_training_platform.ModelTrainingPlatformService(self._http)
+        self.quota_schedule_management = quota_schedule_management.QuotaScheduleManagementService(self._http)
 
     @staticmethod
     def _raise_for_status(r: httpx.Response):

aihub/models/eval.py

@@ -0,0 +1,17 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+from pydantic import BaseModel
+
+
+class CreatEvalReq(BaseModel):
+    dataset_id: int
+    dataset_version_id: int
+    prediction_artifact_path: str
+    evaled_artifact_path: str
+    run_id: str
+    user_id: int = 0
+    report: dict = {}
+
+
+class CreatEvalResp(BaseModel):
+    id: int

aihub/models/labelfree.py

@@ -8,14 +8,26 @@ from pydantic import BaseModel, Field
 class Stats(BaseModel):
     """标注统计信息"""
 
-    total_annotations: int = Field(alias="total_annotations")
-    labeled_annotations: int = Field(alias="labeled_annotations")
-    total_labels: int = Field(alias="total_labels")
-    total_reviews: Optional[int] = Field(None, alias="total_reviews")
-    unlabeled_reviews: Optional[int] = Field(None, alias="unlabeled_reviews")
-    labeled_reviews: Optional[int] = Field(None, alias="labeled_reviews")
-    accepted_count: Optional[int] = Field(None, alias="accepted_count")
-    rejected_count: Optional[int] = Field(None, alias="rejected_count")
+    total_annotations: int = Field(alias="total_annotations", description="总数据量")
+    labeled_annotations: int = Field(
+        alias="labeled_annotations", description="已标注数据量"
+    )
+    total_labels: int = Field(alias="total_labels", description="总标签量")
+    total_reviews: Optional[int] = Field(
+        None, alias="total_reviews", description="总质检量"
+    )
+    unlabeled_reviews: Optional[int] = Field(
+        None, alias="unlabeled_reviews", description="未质检量"
+    )
+    labeled_reviews: Optional[int] = Field(
+        None, alias="labeled_reviews", description="已质检量"
+    )
+    accepted_count: Optional[int] = Field(
+        None, alias="accepted_count", description="质检通过量"
+    )
+    rejected_count: Optional[int] = Field(
+        None, alias="rejected_count", description="质检打回量"
+    )
 
 
 class GetGlobalStatsResponse(BaseModel):
@@ -24,8 +36,18 @@ class GetGlobalStatsResponse(BaseModel):
     """
 
     global_stats: Stats = Field(alias="global_stats")
-    valid_ten_percent: bool = Field(alias="valid_ten_percent")
-    valid_fifty_percent: bool = Field(alias="valid_fifty_percent")
-    valid_hundred_percent: bool = Field(alias="valid_hundred_percent")
-    data_exported_count: int = Field(alias="data_exported_count")
-    exported_dataset_name: str = Field(alias="exported_dataset_name")
+    valid_ten_percent: bool = Field(
+        alias="valid_ten_percent", description="是否完成验收10%"
+    )
+    valid_fifty_percent: bool = Field(
+        alias="valid_fifty_percent", description="是否完成验收50%"
+    )
+    valid_hundred_percent: bool = Field(
+        alias="valid_hundred_percent", description="是否完成验收100%"
+    )
+    data_exported_count: int = Field(
+        alias="data_exported_count", description="已导出数据次数"
+    )
+    exported_dataset_name: str = Field(
+        alias="exported_dataset_name", description="最新数据集名称"
+    )

aihub/models/model_training_platform.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class Env(BaseModel):
+    key: str
+    value: str
+
+
+class Mount(BaseModel):
+    name: str
+    path: str
+
+
+class Sku(BaseModel):
+    cpu: int
+    gpu: int
+    memory: int
+
+
+class VirtualCluster(BaseModel):
+    id: int
+    name: str
+    gpu_type: str = Field(alias="gpu_type")
+    label: str
+    sku: Sku
+
+
+class Storage(BaseModel):
+    id: int
+    name: str
+    path: str
+    server_path: str = Field(alias="server_path")
+    server_host: str = Field(alias="server_host")
+    server_type: str = Field(alias="server_type")
+    permission: str
+    description: str
+
+
+class Category(BaseModel):
+    id: int
+    name: str
+
+
+class Project(BaseModel):
+    id: int
+    name: str
+    description: str
+
+
+class User(BaseModel):
+    id: int
+    name: str
+
+
+class Department(BaseModel):
+    id: int
+    name: str
+
+
+class CreateTrainingRequest(BaseModel):
+    framework: str
+    name: str
+    description: Optional[str] = None
+    command: Optional[str] = None
+    image: str
+    virtual_cluster_id: int = Field(alias="virtual_cluster_id")
+    sku_cnt: int = Field(alias="sku_cnt")
+    enable_ssh: Optional[bool] = Field(False, alias="enable_ssh")
+    envs: Optional[List[Env]] = Field(default_factory=list, alias="envs")
+    storage_ids: Optional[List[int]] = Field(default_factory=list, alias="storage_ids")
+    instances: int
+    use_ib_network: Optional[bool] = Field(False, alias="use_ib_network")
+    always_pull_image: Optional[bool] = Field(False, alias="always_pull_image")
+    shm: Optional[int] = None
+    category_id: int = Field(alias="category_id")
+    project_id: int = Field(alias="project_id")
+    estimate_run_time: Optional[int] = Field(None, alias="estimate_run_time")
+    is_vip: Optional[bool] = Field(False, alias="is_vip")
+    preempt_policy: Optional[int] = Field(None, alias="preempt_policy")
+    vip_node_names: Optional[List[str]] = Field(default_factory=list, alias="vip_node_names")
+    is_quota_schedule: Optional[bool] = Field(False, alias="is_quota_schedule")
+
+
+class CreateTrainingResponse(BaseModel):
+    id: int
+
+
+class ListTrainingsRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+    user_id: Optional[int] = Field(None, alias="user_id")
+    name: Optional[str] = None
+    virtual_cluster_id: Optional[int] = Field(None, alias="virtual_cluster_id")
+    status: Optional[int] = None
+    category_id: Optional[int] = Field(None, alias="category_id")
+    project_id: Optional[int] = Field(None, alias="project_id")
+    is_quota_schedule: Optional[bool] = Field(None, alias="is_quota_schedule")
+
+
+class Training(BaseModel):
+    """训练任务"""
+
+    id: int
+    framework: str
+    name: str
+    description: str
+    command: str
+    image: str
+    virtual_cluster: VirtualCluster = Field(alias="virtual_cluster")
+    sku_cnt: int = Field(alias="sku_cnt")
+    enable_ssh: bool = Field(alias="enable_ssh")
+    envs: Optional[List[Env]] = Field(None, alias="envs")
+    storages: Optional[List[Storage]] = Field(None, alias="storages")
+    instances: int
+    created_at: int = Field(alias="created_at")
+    username: str
+    user_id: int = Field(alias="user_id")
+    namespace: str
+    res_name: str = Field(alias="res_name")
+    status: int
+    use_ib_network: bool = Field(alias="use_ib_network")
+    always_pull_image: bool = Field(alias="always_pull_image")
+    shm: int
+    category: Category
+    project: Project
+    avg_gpu_util: float = Field(alias="avg_gpu_util")
+    finished_at: int = Field(alias="finished_at")
+    started_at: int = Field(alias="started_at")
+    estimate_run_time: int = Field(alias="estimate_run_time")
+    is_vip: bool = Field(alias="is_vip")
+    cluster_partition: str = Field(alias="cluster_partition")
+    preempt_policy: int = Field(alias="preempt_policy")
+    vip_node_names: List[str] = Field(alias="vip_node_names")
+    stop_op_user: Optional[User] = Field(None, alias="stop_op_user")
+    use_new_log: bool = Field(alias="use_new_log")
+    is_quota_schedule: bool = Field(alias="is_quota_schedule")
+
+
+class ListTrainingsResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[Training]
+
+
+class Pod(BaseModel):
+    """训练任务容器"""
+
+    id: int
+    namespace: str
+    name: str
+    status: str
+    created_at: int = Field(alias="created_at")
+    started_at: int = Field(alias="started_at")
+    finished_at: int = Field(alias="finished_at")
+    host_ip: str = Field(alias="host_ip")
+    node_name: str = Field(alias="node_name")
+    ssh_port: int = Field(alias="ssh_port")
+    ssh_info: str = Field(alias="ssh_info")
+    use_new_log: bool = Field(alias="use_new_log")
+
+
+class ListTrainingPodsRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+
+
+class ListTrainingPodsResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[Pod]
+
+
+class PodLogInfo(BaseModel):
+    name: str
+    url: str
+
+
+class GetTrainingPodLogsNewResponse(BaseModel):
+    logs: List[PodLogInfo]
+
+
+class GetTrainingPodSpecResponse(BaseModel):
+    spec: str
+
+
+class GetTrainingPodEventsResponse(BaseModel):
+    events: str
+
+
+class TrainingUser(BaseModel):
+    user_id: int = Field(alias="user_id")
+    username: str
+
+
+class ListTrainingUsersRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+
+
+class ListTrainingUsersResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: Optional[List[TrainingUser]] = Field(default_factory=list)
+
+
+class Container(BaseModel):
+    namespace: str
+    pod_name: str = Field(alias="pod_name")
+    container_name: str = Field(alias="container_name")
+
+
+class TrainingContainer(BaseModel):
+    training_id: int = Field(alias="training_id")
+    training_name: str = Field(alias="training_name")
+    containers: List[Container]
+
+
+class ListTrainingContainersRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+
+
+class ListTrainingContainersResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: Optional[List[TrainingContainer]] = Field(default_factory=list)

aihub/models/quota_schedule_management.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class Env(BaseModel):
+    key: str
+    value: str
+
+
+class Sku(BaseModel):
+    cpu: int
+    gpu: int
+    memory: int
+
+
+class VirtualCluster(BaseModel):
+    id: int
+    name: str
+    gpu_type: str = Field(alias="gpu_type")
+    label: str
+    sku: Sku
+
+
+class Storage(BaseModel):
+    id: int
+    name: str
+    path: str
+    server_path: str = Field(alias="server_path")
+    server_host: str = Field(alias="server_host")
+    server_type: str = Field(alias="server_type")
+    permission: str
+    description: str
+
+
+class Category(BaseModel):
+    id: int
+    name: str
+
+
+class Project(BaseModel):
+    id: int
+    name: str
+    description: str
+
+
+class User(BaseModel):
+    id: int
+    name: str
+
+
+class SourceTask(BaseModel):
+    id: int
+    name: str
+
+
+class CreateTaskRequest(BaseModel):
+    priority: int
+    framework: str
+    name: str
+    description: Optional[str] = None
+    command: Optional[str] = None
+    image: str
+    virtual_cluster_id: int = Field(alias="virtual_cluster_id")
+    sku_cnt: int = Field(alias="sku_cnt")
+    enable_ssh: Optional[bool] = Field(False, alias="enable_ssh")
+    envs: Optional[List[Env]] = Field(default_factory=list, alias="envs")
+    storage_ids: Optional[List[int]] = Field(default_factory=list, alias="storage_ids")
+    instances: int
+    use_ib_network: Optional[bool] = Field(False, alias="use_ib_network")
+    always_pull_image: Optional[bool] = Field(False, alias="always_pull_image")
+    shm: Optional[int] = None
+    category_id: int = Field(alias="category_id")
+    project_id: int = Field(alias="project_id")
+    estimate_run_time: Optional[int] = Field(None, alias="estimate_run_time")
+    is_vip: Optional[bool] = Field(False, alias="is_vip")
+    preempt_policy: Optional[int] = Field(None, alias="preempt_policy")
+    vip_node_names: Optional[List[str]] = Field(default_factory=list, alias="vip_node_names")
+    enable_reschedule: Optional[bool] = Field(False, alias="enable_reschedule")
+
+
+class CreateTaskResponse(BaseModel):
+    id: int
+
+
+class Task(BaseModel):
+    id: int
+    priority: int
+    mtp_id: int = Field(alias="mtp_id")
+    framework: str
+    name: str
+    description: str
+    command: str
+    image: str
+    virtual_cluster: VirtualCluster = Field(alias="virtual_cluster")
+    sku_cnt: int = Field(alias="sku_cnt")
+    enable_ssh: bool = Field(alias="enable_ssh")
+    envs: Optional[List[Env]] = Field(default_factory=list, alias="envs")
+    storages: Optional[List[Storage]] = Field(default_factory=list, alias="storages")
+    instances: int
+    created_at: int = Field(alias="created_at")
+    username: str
+    user_id: int = Field(alias="user_id")
+    namespace: str
+    res_name: str = Field(alias="res_name")
+    status: int
+    use_ib_network: bool = Field(alias="use_ib_network")
+    always_pull_image: bool = Field(alias="always_pull_image")
+    shm: int
+    category: Category
+    project: Project
+    avg_gpu_util: float = Field(alias="avg_gpu_util")
+    finished_at: int = Field(alias="finished_at")
+    started_at: int = Field(alias="started_at")
+    estimate_run_time: int = Field(alias="estimate_run_time")
+    is_vip: bool = Field(alias="is_vip")
+    cluster_partition: str = Field(alias="cluster_partition")
+    preempt_policy: int = Field(alias="preempt_policy")
+    vip_node_names: Optional[List[str]] = Field(None, alias="vip_node_names")
+    stop_op_user: Optional[User] = Field(None, alias="stop_op_user")
+    use_new_log: bool = Field(alias="use_new_log")
+    is_quota_schedule: bool = Field(alias="is_quota_schedule")
+    enable_reschedule: bool = Field(alias="enable_reschedule")
+    remain_schedule_cnt: int = Field(alias="remain_schedule_cnt")
+    source_task: Optional[SourceTask] = Field(None, alias="source_task")
+
+
+class ListTasksRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+    user_id: Optional[int] = Field(None, alias="user_id")
+    name: Optional[str] = None
+    virtual_cluster_id: Optional[int] = Field(None, alias="virtual_cluster_id")
+    status: Optional[int] = None
+    category_id: Optional[int] = Field(None, alias="category_id")
+    project_id: Optional[int] = Field(None, alias="project_id")
+    priority: Optional[int] = None
+
+
+class ListTasksResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[Task] = Field(default_factory=list)
+
+
+class Pod(BaseModel):
+    id: int
+    namespace: str
+    name: str
+    status: str
+    created_at: int = Field(alias="created_at")
+    started_at: int = Field(alias="started_at")
+    finished_at: int = Field(alias="finished_at")
+    host_ip: str = Field(alias="host_ip")
+    node_name: str = Field(alias="node_name")
+    ssh_port: int = Field(alias="ssh_port")
+    ssh_info: str = Field(alias="ssh_info")
+    use_new_log: bool = Field(alias="use_new_log")
+
+
+class ListTaskPodsRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+
+
+class ListTaskPodsResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[Pod] = Field(default_factory=list)
+
+
+class PodLogInfo(BaseModel):
+    name: str
+    url: str
+
+
+class GetTaskPodLogsNewResponse(BaseModel):
+    logs: List[PodLogInfo]
+
+
+class GetTaskPodSpecResponse(BaseModel):
+    spec: str
+
+
+class GetTaskPodEventsResponse(BaseModel):
+    events: str
+
+
+class MachineOverview(BaseModel):
+    high: int
+    low: int
+    free: int
+
+
+class HighPrioritySummary(BaseModel):
+    group_id: int = Field(alias="group_id")
+    group_name: str = Field(alias="group_name")
+    used: int
+    total: int
+
+
+class MetricsOverview(BaseModel):
+    vc_id: int = Field(alias="vc_id")
+    vc_name: str = Field(alias="vc_name")
+    machine: MachineOverview
+    high_priority: List[HighPrioritySummary] = Field(alias="high_priority")
+
+
+class GetMetricsOverviewRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+
+
+class GetMetricsOverviewResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[MetricsOverview] = Field(default_factory=list)
+
+
+class TaskUser(BaseModel):
+    user_id: int = Field(alias="user_id")
+    username: str
+
+
+class ListTaskUsersRequest(BaseModel):
+    page_size: int = Field(20, alias="page_size")
+    page_num: int = Field(1, alias="page_num")
+
+
+class ListTaskUsersResponse(BaseModel):
+    total: int
+    page_size: int = Field(alias="page_size")
+    page_num: int = Field(alias="page_num")
+    data: List[TaskUser] = Field(default_factory=list)

aihub/models/tag_resource_management.py

@@ -0,0 +1,50 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class Project(BaseModel):
+    id: int
+    name: str
+
+
+class ProjectListData(BaseModel):
+    data: List[Project]
+
+
+class SelectProjectsResponse(BaseModel):
+    data: List[Project]
+
+
+class SkuBrief(BaseModel):
+    id: int
+    description: str
+    cpu: int
+    memory: int
+    gpu_type: int = Field(alias="gpu_type")
+    gpu_memory: int = Field(alias="gpu_memory")
+    network: int
+    created_at: int = Field(alias="created_at")
+
+
+class VirtualClusterBrief(BaseModel):
+    id: int
+    name: str
+    uuid: str
+    sku: Optional[SkuBrief] = None
+    created_at: int = Field(alias="created_at")
+
+
+class SelectVirtualClustersRequest(BaseModel):
+    user_id: int = Field(alias="user_id")
+    module_type: Optional[int] = Field(None, alias="module_type")
+    new_module_type: Optional[str] = Field(None, alias="new_module_type")
+
+
+class SelectVirtualClustersResponse(BaseModel):
+    data: Optional[List[VirtualClusterBrief]] = Field(default_factory=list)