intellif-aihub 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

aihub/services/model_training_platform.py

@@ -1,3 +1,16 @@
+# !/usr/bin/env python
+# -*- coding:utf-8 -*-
+"""模型训练平台服务模块
+
+对接 **model‑training‑platform** 后端，以 **训练任务** 为核心提供常用业务能力，包括：
+
+- **创建训练任务**
+- **分页查询 / 获取训练详情**
+- **停止训练任务**
+- **查询训练 Pod 及其 Logs / Events / Spec**
+- **查询训练相关用户、容器等辅助信息**
+"""
+
 from __future__ import annotations
 
 import httpx
@@ -20,15 +33,18 @@ from ..models.model_training_platform import (
     ListTrainingUsersResponse,
     ListTrainingContainersResponse,
     Pod,
+    ListStoragesResponse,
 )
 
 _BASE = "/model-training-platform/api/v1"
 
 
 class ModelTrainingPlatformService:
+    """训练任务业务封装"""
 
     def __init__(self, http: httpx.Client):
         self._training = _Training(http)
+        self._storage = _Storage(http)
 
     def create_training(self, payload: CreateTrainingRequest) -> int:
         """创建训练任务
@@ -42,65 +58,133 @@ class ModelTrainingPlatformService:
         return self._training.create(payload)
 
     def list_trainings(self, payload: ListTrainingsRequest) -> ListTrainingsResponse:
+        """分页查询训练任务
+
+        Args:
+            payload: 分页与过滤条件，详见 ``ListTrainingsRequest``
+
+        Returns:
+            ListTrainingsResponse: 训练任务分页结果
+        """
         return self._training.list(payload)
 
     def get_training(self, training_id: int) -> Training:
         """获取训练任务详情
 
         Args:
-            training_id (int): 训练任务ID
-        Returns:
-            Training
+            training_id: 训练任务 ID
 
+        Returns:
+            Training: 训练任务完整信息
         """
         return self._training.get(training_id)
 
     def stop_training(self, training_id: int) -> None:
-        """停止训练瑞文
+        """停止训练任务
 
         Args:
-            training_id (int): 训练任务ID
-
-        Returns:
-
+            training_id: 训练任务 ID
         """
         self._training.stop(training_id)
 
-    def list_training_pods(
-        self, training_id: int, payload: ListTrainingPodsRequest
-    ) -> ListTrainingPodsResponse:
+    def list_training_pods(self, training_id: int, payload: ListTrainingPodsRequest) -> ListTrainingPodsResponse:
+        """查询训练任务的 Pod 列表
+
+        Args:
+            training_id: 训练任务 ID
+            payload: 分页参数，详见 ``ListTrainingPodsRequest``
+
+        Returns:
+            ListTrainingPodsResponse: Pod 分页结果
+        """
         return self._training.list_training_pods(training_id, payload)
 
     def get_training_pod(self, training_id: int, pod_id: int) -> Pod:
+        """获取单个 Pod 详情
+
+        Args:
+            training_id: 训练任务 ID
+            pod_id: Pod ID
+
+        Returns:
+            Pod: Pod 详细信息
+        """
         return self._training.get_training_pod(training_id, pod_id)
 
-    def get_pod_logs_new(
-        self, training_id: int, pod_id: int
-    ) -> GetTrainingPodLogsNewResponse:
+    def get_pod_logs_new(self, training_id: int, pod_id: int) -> GetTrainingPodLogsNewResponse:
+        """获取 Pod 日志
+
+        Args:
+            training_id: 训练任务 ID
+            pod_id: Pod ID
+
+        Returns:
+            GetTrainingPodLogsNewResponse: 日志文件信息（名称 / URL 列表）
+        """
         return self._training.get_training_logs_new(training_id, pod_id)
 
     def get_pod_spec(self, training_id: int, pod_id: int) -> GetTrainingPodSpecResponse:
+        """获取 Pod Spec
+
+        Args:
+            training_id: 训练任务 ID
+            pod_id: Pod ID
+
+        Returns:
+            GetTrainingPodSpecResponse: Pod Spec 字符串
+        """
         return self._training.get_training_spec(training_id, pod_id)
 
-    def get_pod_events(
-        self, training_id: int, pod_id: int
-    ) -> GetTrainingPodEventsResponse:
+    def get_pod_events(self, training_id: int, pod_id: int) -> GetTrainingPodEventsResponse:
+        """获取 Pod Events
+
+        Args:
+            training_id: 训练任务 ID
+            pod_id: Pod ID
+
+        Returns:
+            GetTrainingPodEventsResponse: 事件文本
+        """
         return self._training.get_training_events(training_id, pod_id)
 
-    def list_training_users(
-        self, payload: ListTrainingUsersRequest
-    ) -> ListTrainingUsersResponse:
+    def list_training_users(self, payload: ListTrainingUsersRequest) -> ListTrainingUsersResponse:
+        """查询训练任务的用户列表
+
+        Args:
+            payload: 分页参数 ``ListTrainingUsersRequest``
+
+        Returns:
+            ListTrainingUsersResponse: 用户分页结果
+        """
         return self._training.list_training_users(payload)
 
-    def list_training_containers(
-        self, payload: ListTrainingContainersRequest
-    ) -> ListTrainingContainersResponse:
+    def list_training_containers(self, payload: ListTrainingContainersRequest) -> ListTrainingContainersResponse:
+        """查询训练容器信息
+
+        Args:
+            payload: 分页参数 ``ListTrainingContainersRequest``
+
+        Returns:
+            ListTrainingContainersResponse: 训练容器列表
+        """
         return self._training.list_training_containers(payload)
 
+    def list_storages(self) -> ListStoragesResponse:
+        """获取存储列表
+
+        Returns:
+            ListStoragesResponse: 存储列表
+        """
+        return self._storage.list_storages()
+
     @property
     def training(self) -> _Training:
         return self._training
 
+    @property
+    def storage(self) -> _Storage:
+        return self._storage
+
 
 class _Training:
 
@@ -108,20 +192,14 @@ class _Training:
         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),
-        )
+        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),
-        )
+        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}")
@@ -140,13 +218,9 @@ class _Training:
         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),
-        )
+    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}")
@@ -159,51 +233,50 @@ class _Training:
             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:
+    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:
+    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:
+    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)
-        )
+    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)
-        )
+    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
+
+
+class _Storage:
+
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def list_storages(self) -> ListStoragesResponse:
+        resp = self._http.get(f"{_BASE}/storages")
+        wrapper = APIWrapper[ListStoragesResponse].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

@@ -1,13 +1,43 @@
+# !/usr/bin/env python
+# -*- coding:utf-8 -*-
+"""配额调度服务模块
+
+封装 **quota‑schedule‑management** 后端接口，围绕 **配额调度任务** 提供常用能力：
+
+- **创建 / 查询 / 停止调度任务**
+- **Pod 维度信息**（列表 / Logs / Spec / Events）
+- **GPU 资源概览、活跃用户、容器信息** 等辅助查询
+- **Pre‑Stop** 机制：进程内检测平台下发的优雅终止信号
+"""
+
 from __future__ import annotations
 
 import os
 from pathlib import Path
+from typing import List
 
 import httpx
 
 from ..exceptions import APIError
 from ..models.common import APIWrapper
-from ..models.quota_schedule_management import *
+from ..models.quota_schedule_management import (
+    CreateTaskRequest,
+    CreateTaskResponse,
+    ListTasksRequest,
+    ListTasksResponse,
+    Task,
+    ListTaskPodsRequest,
+    ListTaskPodsResponse,
+    Pod,
+    PodLogInfo,
+    GetTaskPodLogsNewResponse,
+    GetTaskPodSpecResponse,
+    GetTaskPodEventsResponse,
+    ListTaskUsersRequest,
+    ListTaskUsersResponse,
+    GetMetricsOverviewRequest,
+    GetMetricsOverviewResponse,
+)
 
 _ENV_KEY = "PRE_STOP_SENTINEL_FILE"
 _DEFAULT_SENTINEL = "/tmp/pre_stop_sentinel_file"
@@ -16,6 +46,14 @@ _BASE = "/quota-schedule-management/api/v1"
 
 
 def is_pre_stopped() -> bool:
+    """判断当前进程是否已收到 *pre‑stop* 信号
+
+    该信号由调度平台通过创建哨兵文件的方式下发。文件路径可通过环境变量 ``PRE_STOP_SENTINEL_FILE`` 覆盖，
+    默认为``/tmp/pre_stop_sentinel_file``。
+
+    Returns:
+        bool:  若哨兵文件存在，返回 ``True``，表示应尽快停止任务。
+    """
     return _SENTINEL_PATH.exists()
 
 
@@ -26,6 +64,7 @@ class PreStopService:
 
 
 class QuotaScheduleManagementService:
+    """配额调度任务业务封装"""
 
     def __init__(self, http: httpx.Client):
         self._task = _Task(http)
@@ -34,7 +73,7 @@ class QuotaScheduleManagementService:
         """创建调度任务
 
         Args:
-            payload: 创建任务的各项参数
+            payload: 创建任务的各项参数，详见 ``CreateTaskRequest``
 
         Returns:
             int: 后端生成的任务 ID
@@ -45,10 +84,10 @@ class QuotaScheduleManagementService:
         """分页查询任务列表
 
         Args:
-            payload: 查询条件（分页 / 过滤）
+            payload: 分页 / 过滤参数，详见 ``ListTasksRequest``
 
         Returns:
-            ListTasksResponse: 列表及分页信息
+            ListTasksResponse: 任务分页结果
         """
         return self._task.list(payload)
 
@@ -68,31 +107,89 @@ class QuotaScheduleManagementService:
 
         Args:
             task_id: 任务 ID
-
-        Returns:
-            None
         """
         self._task.stop(task_id)
 
     def list_task_pods(self, task_id: int, payload: ListTaskPodsRequest) -> ListTaskPodsResponse:
+        """分页查询任务的 Pod 列表
+
+        Args:
+            task_id: 调度任务 ID
+            payload: 分页参数，见 :class:`ListTaskPodsRequest`
+
+        Returns:
+            ListTaskPodsResponse: Pod 列表及分页信息
+        """
         return self._task.list_pods(task_id, payload)
 
     def get_task_pod(self, task_id: int, pod_id: int) -> Pod:
+        """获取单个 Pod 详情
+
+        Args:
+            task_id: 调度任务 ID
+            pod_id:  Pod ID（数据库主键）
+
+        Returns:
+            Pod: Pod 详细信息
+        """
         return self._task.get_pod(task_id, pod_id)
 
     def get_pod_logs_new(self, task_id: int, pod_id: int) -> List[PodLogInfo]:
+        """获取新版日志
+
+        Args:
+            task_id: 调度任务 ID
+            pod_id:  Pod ID
+
+        Returns:
+            list[PodLogInfo]: 每条日志文件的名称与下载地址
+        """
         return self._task.get_logs_new(task_id, pod_id).logs
 
     def get_pod_spec(self, task_id: int, pod_id: int) -> str:
+        """获取 Pod 运行时 Spec（YAML）
+
+        Args:
+            task_id: 调度任务 ID
+            pod_id:  Pod ID
+
+        Returns:
+            str: Pod 描述 YAML 字符串
+        """
         return self._task.get_spec(task_id, pod_id).spec
 
     def get_pod_events(self, task_id: int, pod_id: int) -> str:
+        """获取 Pod 事件
+
+        Args:
+            task_id: 调度任务 ID
+            pod_id:  Pod ID
+
+        Returns:
+            str: 事件文本（kubectl describe events）
+        """
         return self._task.get_events(task_id, pod_id).events
 
     def list_task_users(self, payload: ListTaskUsersRequest) -> ListTaskUsersResponse:
+        """查询任务的用户列表
+
+        Args:
+            payload: 分页参数，见 :class:`ListTaskUsersRequest`
+
+        Returns:
+            ListTaskUsersResponse: 用户列表及分页信息
+        """
         return self._task.list_users(payload)
 
     def get_metrics_overview(self, payload: GetMetricsOverviewRequest) -> GetMetricsOverviewResponse:
+        """获取虚拟集群维度的 GPU 使用概览
+
+        Args:
+            payload: 分页参数，见 :class:`GetMetricsOverviewRequest`
+
+        Returns:
+            GetMetricsOverviewResponse: 各虚拟集群资源统计
+        """
         return self._task.get_metrics_overview(payload)
 
     @property

aihub/services/tag_resource_management.py

@@ -1,25 +1,56 @@
 # !/usr/bin/env python
-# -*-coding:utf-8 -*-
+# -*- coding:utf-8 -*-
+"""标签资源管理服务模块
+
+封装 **tag‑resource‑management** 相关接口，提供两组下拉选项：
+
+- **选择项目（Project）下拉列表**
+- **选择虚拟集群（Virtual‑Cluster）下拉列表**
+"""
+
 from __future__ import annotations
 
+from typing import List
+
 import httpx
 
 from ..exceptions import APIError
 from ..models.common import APIWrapper
-from ..models.tag_resource_management import *
+from ..models.tag_resource_management import (
+    Project,
+    ProjectListData,
+    SelectVirtualClustersRequest,
+    SelectVirtualClustersResponse,
+    VirtualClusterBrief,
+)
 
 _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]:
+        """获取全部项目下拉选项
+
+        Returns:
+            list[Project]: 项目列表，每项仅包含 `id` / `name`
+        """
         return self._project.select_projects()
 
     def select_virtual_clusters(self, payload: SelectVirtualClustersRequest) -> List[VirtualClusterBrief]:
+        """获取虚拟集群下拉选项
+
+        Args:
+            payload: 查询参数，包含用户 ID / 模块类型等过滤条件，详见 :class:`SelectVirtualClustersRequest`
+
+        Returns:
+            list[VirtualClusterBrief]: 过滤后的虚拟集群简要列表
+        """
         return self._virtual_cluster.select(payload).data
 
     @property

aihub/services/task_center.py

@@ -1,3 +1,13 @@
+# !/usr/bin/env python
+# -*- coding:utf-8 -*-
+"""任务中心服务模块
+
+封装 **task‑center** 常用能力：
+
+- **创建 / 查询任务**
+- **一键创建标注任务**（高阶封装）
+"""
+
 from __future__ import annotations
 
 import datetime
@@ -28,31 +38,37 @@ def date_str_to_timestamp(date_str: str) -> int:
 
 
 class TaskCenterService:
+    """任务中心服务"""
 
     def __init__(self, http: httpx.Client):
         self._TaskCenter = _TaskCenter(http)
         self._http = http
 
     def create(self, payload: CreateTaskReq) -> int:
+        """创建任务
+
+        Args:
+            payload: 任务创建请求体，详见 :class:`CreateTaskReq`
+
+        Returns:
+            int: 后端生成的任务 ID
+        """
         return self._TaskCenter.create(payload)
 
     def get(self, task_id: int) -> LabelTaskDetail:
-        """获取任务信息
+        """根据任务 ID 获取详情
 
         Args:
-            task_id (int): 任务ID
-
+            task_id: 任务 ID
 
         Returns:
-            LabelTaskDetail: 任务信息
-
+            LabelTaskDetail: 任务完整信息
         """
-
         return self._TaskCenter.get(task_id)
 
     # 如果想要访问子对象，也保留属性
     @property
-    def TaskCenter(self) -> _TaskCenter:
+    def task_center(self) -> _TaskCenter:
         return self._TaskCenter
 
     def create_label_task(
@@ -160,7 +176,6 @@ class _TaskCenter:
         self._http = http
 
     def create(self, payload: CreateTaskReq) -> int:
-        """创建任务"""
         logger.debug(f"create task: {payload}")
         resp = self._http.post(
             f"{_BASE}/tasks",
@@ -176,7 +191,6 @@ class _TaskCenter:
         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}",