intellif-aihub 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

aihub/services/model_training_platform.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import httpx
+
+from ..exceptions import APIError
+from ..models.common import APIWrapper
+from ..models.model_training_platform import *
+
+_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:
+        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:
+        return self._training.get(training_id)
+
+    def stop_training(self, training_id: int) -> None:
+        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

aihub/services/quota_schedule_management.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+_ENV_KEY = "PRE_STOP_SENTINEL_FILE"
+_DEFAULT_SENTINEL = "/tmp/pre_stop_sentinel_file"
+_SENTINEL_PATH = Path(os.getenv(_ENV_KEY, _DEFAULT_SENTINEL))
+
+
+def is_pre_stopped() -> bool:
+    return _SENTINEL_PATH.exists()
+
+
+class PreStopService:
+    @staticmethod
+    def is_pre_stopped() -> bool:
+        return is_pre_stopped()

aihub/services/reporter.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import json
+import os
+
+
+class ReportService:
+    @staticmethod
+    def report_envs():
+        pod_namespace = os.environ.get('POD_NAMESPACE', '')
+        pod_name = os.environ.get('POD_NAME', '')
+        if not pod_namespace or not pod_name:
+            return
+        env_file_path = f'/var/mtp/log/{pod_namespace}/{pod_name}/env'
+        try:
+            os.makedirs(os.path.dirname(env_file_path), exist_ok=True)
+            with open(env_file_path, 'w', encoding='utf-8') as f:
+                f.write(json.dumps(dict(os.environ)))
+        except Exception as e:
+            print(f"[report_envs] 写入失败: {e}", flush=True)

aihub/services/tag_resource_management.py

@@ -0,0 +1,55 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+from __future__ import annotations
+
+import httpx
+
+from ..exceptions import APIError
+from ..models.common import APIWrapper
+from ..models.tag_resource_management import *
+
+_BASE = "/tag-resource-management/api/v1"
+
+
+class TagResourceManagementService:
+    def __init__(self, http: httpx.Client):
+        self._project = _Project(http)
+        self._virtual_cluster = _VirtualCluster(http)
+
+    def select_projects(self) -> List[Project]:
+        return self._project.select_projects()
+
+    def select_virtual_clusters(self, payload: SelectVirtualClustersRequest) -> List[VirtualClusterBrief]:
+        return self._virtual_cluster.select(payload).data
+
+    @property
+    def project(self) -> _Project:
+        return self._project
+
+    @property
+    def virtual_cluster(self) -> _VirtualCluster:
+        return self._virtual_cluster
+
+
+class _Project:
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def select_projects(self) -> List[Project]:
+        resp = self._http.get(f"{_BASE}/select-projects")
+        wrapper = APIWrapper[ProjectListData].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data.data
+
+
+class _VirtualCluster:
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def select(self, payload: SelectVirtualClustersRequest) -> SelectVirtualClustersResponse:
+        resp = self._http.get(f"{_BASE}/select-clusters", params=payload.model_dump(by_alias=True, exclude_none=True))
+        wrapper = APIWrapper[SelectVirtualClustersResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data

aihub/services/task_center.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import datetime
+
+import httpx
+from loguru import logger
+
+from .tag_resource_management import TagResourceManagementService
+from .user_system import UserSystemService
+from ..exceptions import APIError
+from ..models.common import APIWrapper
+from ..models.task_center import (
+    CreateTaskReq,
+    CreateTaskResp,
+    CreateTaskOtherInfo,
+    LabelProjectTypeEnum,
+    TaskCenterPriorityEnum,
+    LabelTaskDetail,
+)
+from ..models.user_system import SearchUsersRequest
+
+_BASE = "/task-center/api/v1"
+
+
+def date_str_to_timestamp(date_str: str) -> int:
+    date_obj = datetime.datetime.strptime(date_str, "%Y-%m-%d")
+    return int(date_obj.timestamp())
+
+
+class TaskCenterService:
+
+    def __init__(self, http: httpx.Client):
+        self._TaskCenter = _TaskCenter(http)
+        self._http = http
+
+    def create(self, payload: CreateTaskReq) -> int:
+        return self._TaskCenter.create(payload)
+
+    def get(self, task_id: int) -> LabelTaskDetail:
+        """获取任务信息
+
+        Args:
+            task_id (int): 任务ID
+
+
+        Returns:
+            LabelTaskDetail: 任务信息
+
+        """
+
+        return self._TaskCenter.get(task_id)
+
+    # 如果想要访问子对象，也保留属性
+    @property
+    def TaskCenter(self) -> _TaskCenter:
+        return self._TaskCenter
+
+    def create_label_task(
+            self,
+            name: str,
+            dataset_version_name: str,
+            feishu_doc_name: str,
+            task_receiver_name: str,
+            estimated_delivery_at: str,
+            project_name: str,
+            label_type: LabelProjectTypeEnum = LabelProjectTypeEnum.IMAGE_CLASSIFICATION,
+            description: str = "",
+            task_priority: TaskCenterPriorityEnum = TaskCenterPriorityEnum.low,
+    ) -> int:
+        """创建标注任务
+
+        Examples:
+            >>> from aihub.client import Client
+            >>> client = Client(base_url="xxx", token="xxxx")
+            >>> task_id = client.task_center.create_label_task( \
+name="test_task", dataset_version_name="re/V1",\
+feishu_doc_name="人脸质量人脸照片分类", task_receiver_name="hyc", \
+project_name="hycpro", label_type=LabelProjectTypeEnum.IMAGE_CLASSIFICATION, description="test_description", \
+task_priority="low", estimated_delivery_at= "2025-08-01")
+            1
+
+        Args:
+            name (str): 任务名称
+            dataset_version_name (str): 数据集版本名称
+            feishu_doc_name (str): 飞书文档名称
+            task_receiver_name (str): 任务接收者名称
+            estimated_delivery_at (str): 预计交付时间，格式为 "YYYY-MM-DD"
+            project_name (str): 项目名称
+            label_type (LabelProjectTypeEnum): 标注项目类型，默认为图像分类
+            description (str): 任务描述，默认为空
+            task_priority (TaskCenterPriorityEnum): 任务优先级，默认为低优先级
+        Returns:
+            任务ID
+        """
+        # 获取接收者ID
+        user_service = UserSystemService(self._http)
+        task_receiver_id = user_service.search_one(
+            payload=SearchUsersRequest(nickname=task_receiver_name)
+        )
+
+        # 获取项目ID
+        tag_service = TagResourceManagementService(self._http)
+        projects = tag_service.select_projects()
+        project_id = None
+        for project in projects:
+            if project.name == project_name:
+                project_id = project.id
+                break
+
+        if project_id is None:
+            raise APIError(f"未找到项目: {project_name}")
+
+        # 获取数据集ID
+        from .dataset_management import DatasetManagementService
+
+        dataset_service = DatasetManagementService(self._http)
+        dataset_version = dataset_service.get_dataset_version_by_name(
+            version_name=dataset_version_name
+        )
+
+        dataset_id = dataset_version.dataset_id
+        dataset_version_id = dataset_version.id
+
+        # 获取文档ID
+        from .document_center import DocumentCenterService
+
+        doc_service = DocumentCenterService(self._http)
+        docs = doc_service.get_documents(name=feishu_doc_name)
+
+        if not docs:
+            raise APIError(f"未找到文档: {feishu_doc_name}")
+
+        doc_id = docs[0].id
+
+        # 创建任务
+        other_info = CreateTaskOtherInfo(
+            label_project_type=label_type,
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_id,
+            doc_id=doc_id,
+            doc_type="doc_center",
+        )
+        estimated_delivery_at_timestamp = date_str_to_timestamp(estimated_delivery_at)
+        task_req = CreateTaskReq(
+            name=name,
+            description=description,
+            task_priority=task_priority,
+            type="label",
+            receiver_id=task_receiver_id,
+            other_info=other_info,
+            project_id=project_id,
+            estimated_delivery_at=estimated_delivery_at_timestamp,
+        )
+
+        return self.create(task_req)
+
+
+class _TaskCenter:
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def create(self, payload: CreateTaskReq) -> int:
+        """创建任务"""
+        logger.debug(f"create task: {payload}")
+        resp = self._http.post(
+            f"{_BASE}/tasks",
+            json=payload.model_dump(),
+        )
+        logger.debug(f"create task response: {resp.text}")
+        if resp.status_code != 200:
+            raise APIError(
+                message="API Error", status=resp.status_code, detail=resp.json()
+            )
+        wrapper = APIWrapper[CreateTaskResp].model_validate(resp.json())
+
+        return wrapper.data.id
+
+    def get(self, task_id: int) -> LabelTaskDetail:
+        """获取任务"""
+        logger.debug(f"get task: {task_id}")
+        resp = self._http.get(
+            f"{_BASE}/tasks/{task_id}",
+        )
+        logger.debug(f"get task response: {resp.text}")
+        if resp.status_code != 200:
+            raise APIError(
+                message="API Error", status=resp.status_code, detail=resp.json()
+            )
+        wrapper = APIWrapper[LabelTaskDetail].model_validate(resp.json())
+        return wrapper.data