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

aihub/models/model_center.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from enum import IntEnum
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class User(BaseModel):
+    """用户"""
+    id: int = Field(description="用户ID")
+    name: str = Field(description="用户名")
+
+
+class ModelType(BaseModel):
+    """模型类型"""
+    id: int = Field(description="类型ID")
+    name: str = Field(description="类型名称")
+
+
+class DeployPlatform(BaseModel):
+    """部署平台"""
+    id: int = Field(description="部署平台ID")
+    name: str = Field(description="部署平台名称")
+
+
+class QuantLevel(BaseModel):
+    """量化等级"""
+    id: int = Field(description="量化等级ID")
+    name: str = Field(description="量化等级名称")
+
+
+class ModelStatus(IntEnum):
+    """模型状态"""
+    Waiting = 1
+    Creating = 2
+    Success = 3
+    Fail = 4
+
+
+class Model(BaseModel):
+    """模型详情"""
+    id: int = Field(description="模型ID")
+    name: str = Field(description="模型名称")
+    description: str = Field(description="描述")
+    model_type: ModelType = Field(alias="model_type", description="模型类型")
+    model_path: str = Field(alias="model_path", description="模型路径")
+    deploy_platform: DeployPlatform = Field(alias="deploy_platform", description="部署平台")
+    param_cnt: str = Field(alias="param_cnt", description="参数量")
+    quant_level: QuantLevel = Field(alias="quant_level", description="量化等级")
+    creator: User = Field(description="创建人")
+    status: ModelStatus = Field(description="模型状态")
+    created_at: int = Field(alias="created_at", description="创建时间戳 (ms)")
+    updated_at: int = Field(alias="updated_at", description="更新时间戳 (ms)")
+
+    model_config = {"use_enum_values": True}
+
+
+class ListModelTypesRequest(BaseModel):
+    """查询模型类型列表请求"""
+    page_size: int = Field(999, alias="page_size", description="每页数量")
+    page_num: int = Field(1, alias="page_num", description="当前页码")
+
+
+class ListModelTypesResponse(BaseModel):
+    """查询模型类型列表返回"""
+    total: int = Field(description="总条数")
+    page_size: int = Field(alias="page_size", description="每页数量")
+    page_num: int = Field(alias="page_num", description="当前页码")
+    data: List[ModelType] = Field(default_factory=list, description="类型列表")
+
+
+class ListDeployPlatformsRequest(BaseModel):
+    """查询部署平台列表请求"""
+    page_size: int = Field(999, alias="page_size", description="每页数量")
+    page_num: int = Field(1, alias="page_num", description="当前页码")
+
+
+class ListDeployPlatformsResponse(BaseModel):
+    """查询部署平台列表返回"""
+    total: int = Field(description="总条数")
+    page_size: int = Field(alias="page_size", description="每页数量")
+    page_num: int = Field(alias="page_num", description="当前页码")
+    data: List[DeployPlatform] = Field(default_factory=list, description="平台列表")
+
+
+class ListQuantLevelsRequest(BaseModel):
+    """查询量化等级列表请求"""
+    page_size: int = Field(999, alias="page_size", description="每页数量")
+    page_num: int = Field(1, alias="page_num", description="当前页码")
+
+
+class ListQuantLevelsResponse(BaseModel):
+    """查询量化等级列表返回"""
+    total: int = Field(description="总条数")
+    page_size: int = Field(alias="page_size", description="每页数量")
+    page_num: int = Field(alias="page_num", description="当前页码")
+    data: List[QuantLevel] = Field(default_factory=list, description="量化等级列表")
+
+
+class ListModelsRequest(BaseModel):
+    """查询模型列表请求"""
+    page_size: int = Field(20, alias="page_size", description="每页数量")
+    page_num: int = Field(1, alias="page_num", description="当前页码")
+    name: Optional[str] = Field(None, description="名称过滤")
+
+
+class ListModelsResponse(BaseModel):
+    """查询模型列表返回"""
+    total: int = Field(description="总条数")
+    page_size: int = Field(alias="page_size", description="每页数量")
+    page_num: int = Field(alias="page_num", description="当前页码")
+    data: List[Model] = Field(default_factory=list, description="模型列表")
+
+
+class CreateModelRequest(BaseModel):
+    """创建模型请求"""
+    name: str = Field(description="模型名称")
+    description: Optional[str] = Field(None, description="描述")
+    model_type_id: int = Field(alias="model_type_id", description="模型类型ID")
+    model_path: Optional[str] = Field(None, alias="model_path", description="模型路径")
+    deploy_platform_id: Optional[int] = Field(None, alias="deploy_platform_id", description="部署平台ID")
+    param_cnt: Optional[str] = Field(None, alias="param_cnt", description="参数量")
+    quant_level_id: Optional[int] = Field(None, alias="quant_level_id", description="量化等级ID")
+
+
+class CreateModelResponse(BaseModel):
+    """创建模型返回"""
+    id: int = Field(description="模型ID")
+
+
+class EditModelRequest(BaseModel):
+    """编辑模型请求"""
+    id: int = Field(description="模型ID")
+    name: str = Field(description="模型名称")
+    description: Optional[str] = Field(None, description="描述")
+    model_type_id: int = Field(alias="model_type_id", description="模型类型ID")
+    model_path: Optional[str] = Field(None, alias="model_path", description="模型路径")
+    deploy_platform_id: Optional[int] = Field(None, alias="deploy_platform_id", description="部署平台ID")
+    param_cnt: Optional[str] = Field(None, alias="param_cnt", description="参数量")
+    quant_level_id: Optional[int] = Field(None, alias="quant_level_id", description="量化等级ID")

aihub/models/model_training_platform.py

@@ -1,234 +1,268 @@
 from __future__ import annotations
 
+from enum import IntEnum
 from typing import List, Optional
 
 from pydantic import BaseModel, Field
 
 
 class Env(BaseModel):
-    key: str
-    value: str
-
-
-class Mount(BaseModel):
-    name: str
-    path: str
+    """环境变量"""
+    key: str = Field(description="变量名")
+    value: str = Field(description="变量值")
 
 
 class Sku(BaseModel):
-    cpu: int
-    gpu: int
-    memory: int
+    """sku"""
+    cpu: int = Field(description="CPU数")
+    gpu: int = Field(description="GPU数")
+    memory: int = Field(description="内存GiB")
 
 
 class VirtualCluster(BaseModel):
-    id: int
-    name: str
-    gpu_type: str = Field(alias="gpu_type")
-    label: str
-    sku: Sku
+    """虚拟集群"""
+    id: int = Field(description="ID")
+    name: str = Field(description="名称")
+    gpu_type: str = Field(alias="gpu_type", description="GPU类型，A800/4090/3090/2080")
+    label: str = Field(description="标签")
+    sku: Sku = Field(description="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
+    """存储"""
+    id: int = Field(description="存储ID")
+    name: str = Field(description="存储名称")
+    path: str = Field(description="挂载路径")
+    server_path: str = Field(alias="server_path", description="服务器路径")
+    server_host: str = Field(alias="server_host", description="服务器地址")
+    server_type: str = Field(alias="server_type", description="服务器类型")
+    permission: str = Field(description="权限")
+    description: str = Field(description="说明")
 
 
 class Category(BaseModel):
-    id: int
-    name: str
+    """分类"""
+    id: int = Field(description="分类ID")
+    name: str = Field(description="分类名称")
 
 
 class Project(BaseModel):
-    id: int
-    name: str
-    description: str
+    """项目"""
+    id: int = Field(description="项目ID")
+    name: str = Field(description="项目名称")
+    description: str = Field(description="项目描述")
 
 
 class User(BaseModel):
-    id: int
-    name: str
+    """用户"""
+    id: int = Field(description="用户ID")
+    name: str = Field(description="用户名")
+
 
+class TrainingStatus(IntEnum):
+    """训练任务状态"""
+    Waiting = 1
+    Running = 2
+    Success = 3
+    Failed = 4
+    Stopped = 5
 
-class Department(BaseModel):
-    id: int
-    name: str
+
+class PreemptPolicy(IntEnum):
+    """抢占策略"""
+    WAIT = 1  # 等待任务完成
+    STOP = 2  # 停止任务
 
 
 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")
+    """创建训练任务请求"""
+    framework: str = Field(description="训练框架，如PyTorchJob/MpiJobMpiRun/MpiJobDeepspeed")
+    name: str = Field(description="任务名称")
+    description: Optional[str] = Field(None, description="描述")
+    command: Optional[str] = Field(None, description="命令")
+    image: str = Field(description="镜像")
+    virtual_cluster_id: int = Field(alias="virtual_cluster_id", description="虚拟集群ID")
+    sku_cnt: int = Field(alias="sku_cnt", description="sku数量")
+    enable_ssh: Optional[bool] = Field(False, alias="enable_ssh", description="是否开启SSH")
+    envs: Optional[List[Env]] = Field(default_factory=list, description="环境变量")
+    storage_ids: Optional[List[int]] = Field(default_factory=list, alias="storage_ids", description="挂载存储ID列表")
+    instances: int = Field(description="实例数")
+    use_ib_network: Optional[bool] = Field(False, alias="use_ib_network", description="是否使用IB网络")
+    always_pull_image: Optional[bool] = Field(False, alias="always_pull_image", description="每次强制拉镜像")
+    shm: Optional[int] = Field(None, description="共享内存")
+    category_id: int = Field(alias="category_id", description="分类ID")
+    project_id: int = Field(alias="project_id", description="项目ID")
+    estimate_run_time: Optional[int] = Field(None, alias="estimate_run_time", description="预计运行时长(s)")
+    is_vip: Optional[bool] = Field(False, alias="is_vip", description="是否vip任务")
+    preempt_policy: Optional[PreemptPolicy] = Field(None, alias="preempt_policy", description="抢占策略")
+    vip_node_names: Optional[List[str]] = Field(default_factory=list, alias="vip_node_names", description="VIP节点名字")
+    is_quota_schedule: Optional[bool] = Field(False, alias="is_quota_schedule", description="是否配额调度")
+
+    model_config = {"use_enum_values": True}
 
 
 class CreateTrainingResponse(BaseModel):
-    id: int
+    """创建训练任务返回"""
+    id: int = Field(description="训练任务ID")
 
 
 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")
+    """查询训练任务列表请求"""
+    page_size: int = Field(20, alias="page_size", description="每页数量")
+    page_num: int = Field(1, alias="page_num", description="当前页码")
+    user_id: Optional[int] = Field(None, alias="user_id", description="用户过滤")
+    name: Optional[str] = Field(None, description="名称过滤")
+    virtual_cluster_id: Optional[int] = Field(None, alias="virtual_cluster_id", description="虚拟集群过滤")
+    status: Optional[TrainingStatus] = Field(None, description="状态过滤")
+    category_id: Optional[int] = Field(None, alias="category_id", description="分类过滤")
+    project_id: Optional[int] = Field(None, alias="project_id", description="项目过滤")
+    is_quota_schedule: Optional[bool] = Field(None, alias="is_quota_schedule", description="配额调度过滤")
+
+    model_config = {"use_enum_values": True}
 
 
 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")
+    """训练任务详情"""
+    id: int = Field(description="任务ID")
+    framework: str = Field(description="训练框架")
+    name: str = Field(description="名称")
+    description: str = Field(description="描述")
+    command: str = Field(description="命令")
+    image: str = Field(description="镜像")
+    virtual_cluster: VirtualCluster = Field(alias="virtual_cluster", description="虚拟集群")
+    sku_cnt: int = Field(alias="sku_cnt", description="SKU数量")
+    enable_ssh: bool = Field(alias="enable_ssh", description="SSH开启")
+    envs: Optional[List[Env]] = Field(None, description="环境变量")
+    storages: Optional[List[Storage]] = Field(None, description="挂载存储")
+    instances: int = Field(description="实例数")
+    created_at: int = Field(alias="created_at", description="创建时间")
+    username: str = Field(description="提交人")
+    user_id: int = Field(alias="user_id", description="提交人ID")
+    namespace: str = Field(description="K8s Namespace")
+    res_name: str = Field(alias="res_name", description="K8s资源名")
+    status: TrainingStatus = Field(description="状态")
+    use_ib_network: bool = Field(alias="use_ib_network", description="IB网络")
+    always_pull_image: bool = Field(alias="always_pull_image", description="每次拉镜像")
+    shm: int = Field(description="共享内存")
+    category: Category = Field(description="分类")
+    project: Project = Field(description="项目")
+    avg_gpu_util: float = Field(alias="avg_gpu_util", description="平均GPU利用率")
+    finished_at: int = Field(alias="finished_at", description="结束时间")
+    started_at: int = Field(alias="started_at", description="开始时间")
+    estimate_run_time: int = Field(alias="estimate_run_time", description="预计运行时长")
+    is_vip: bool = Field(alias="is_vip", description="是否VIP")
+    cluster_partition: str = Field(alias="cluster_partition", description="集群分区")
+    preempt_policy: PreemptPolicy = Field(alias="preempt_policy", description="抢占策略")
+    vip_node_names: List[str] = Field(alias="vip_node_names", description="VIP节点")
+    stop_op_user: Optional[User] = Field(None, alias="stop_op_user", description="停止操作人")
+    use_new_log: bool = Field(alias="use_new_log", description="是否使用新版日志")
+    is_quota_schedule: bool = Field(alias="is_quota_schedule", description="配额调度")
+
+    model_config = {"use_enum_values": True}
 
 
 class ListTrainingsResponse(BaseModel):
-    total: int
-    page_size: int = Field(alias="page_size")
-    page_num: int = Field(alias="page_num")
-    data: List[Training]
+    """查询训练任务列表返回"""
+    total: int = Field(description="总条数")
+    page_size: int = Field(alias="page_size", description="每页数量")
+    page_num: int = Field(alias="page_num", description="当前页码")
+    data: List[Training] = Field(description="训练列表")
 
 
 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")
+    """训练任务Pod"""
+    id: int = Field(description="ID")
+    namespace: str = Field(description="Namespace")
+    name: str = Field(description="名称")
+    status: str = Field(description="状态")
+    created_at: int = Field(alias="created_at", description="创建时间")
+    started_at: int = Field(alias="started_at", description="启动时间")
+    finished_at: int = Field(alias="finished_at", description="结束时间")
+    host_ip: str = Field(alias="host_ip", description="宿主机IP")
+    node_name: str = Field(alias="node_name", description="节点名")
+    ssh_port: int = Field(alias="ssh_port", description="SSH端口")
+    ssh_info: str = Field(alias="ssh_info", description="SSH连接信息")
+    use_new_log: bool = Field(alias="use_new_log", description="是否使用新版日志")
 
 
 class ListTrainingPodsRequest(BaseModel):
-    page_size: int = Field(20, alias="page_size")
-    page_num: int = Field(1, alias="page_num")
+    """查询训练任务Pod列表请求"""
+    page_size: int = Field(20, alias="page_size", description="每页数量")
+    page_num: int = Field(1, alias="page_num", description="当前页码")
 
 
 class ListTrainingPodsResponse(BaseModel):
-    total: int
-    page_size: int = Field(alias="page_size")
-    page_num: int = Field(alias="page_num")
-    data: List[Pod]
+    """查询训练任务Pod列表返回"""
+    total: int = Field(description="总条数")
+    page_size: int = Field(alias="page_size", description="每页数量")
+    page_num: int = Field(alias="page_num", description="当前页码")
+    data: List[Pod] = Field(description="Pod列表")
 
 
 class PodLogInfo(BaseModel):
-    name: str
-    url: str
+    """pod日志信息"""
+    name: str = Field(description="日志名称")
+    url: str = Field(description="URL")
 
 
 class GetTrainingPodLogsNewResponse(BaseModel):
-    logs: List[PodLogInfo]
+    """查询训练任务Pod日志返回"""
+    logs: List[PodLogInfo] = Field(description="日志列表")
 
 
 class GetTrainingPodSpecResponse(BaseModel):
-    spec: str
+    """查询训练任务Pod Spec返回"""
+    spec: str = Field(description="Pod YAML Spec")
 
 
 class GetTrainingPodEventsResponse(BaseModel):
-    events: str
+    """查询训练任务Pod Event返回"""
+    events: str = Field(description="事件内容")
 
 
 class TrainingUser(BaseModel):
-    user_id: int = Field(alias="user_id")
-    username: str
+    """训练任务用户"""
+    user_id: int = Field(alias="user_id", description="用户ID")
+    username: str = Field(description="用户名")
 
 
 class ListTrainingUsersRequest(BaseModel):
-    page_size: int = Field(20, alias="page_size")
-    page_num: int = Field(1, alias="page_num")
+    """查询训练任务用户列表请求"""
+    page_size: int = Field(20, alias="page_size", description="每页数量")
+    page_num: int = Field(1, alias="page_num", description="当前页码")
 
 
 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)
+    """查询训练任务用户列表返回"""
+    total: int = Field(description="总条数")
+    page_size: int = Field(alias="page_size", description="每页数量")
+    page_num: int = Field(alias="page_num", description="当前页码")
+    data: Optional[List[TrainingUser]] = Field(default_factory=list, description="用户列表")
 
 
 class Container(BaseModel):
-    namespace: str
-    pod_name: str = Field(alias="pod_name")
-    container_name: str = Field(alias="container_name")
+    """容器"""
+    namespace: str = Field(description="Namespace")
+    pod_name: str = Field(alias="pod_name", description="Pod名")
+    container_name: str = Field(alias="container_name", description="容器名")
 
 
 class TrainingContainer(BaseModel):
-    training_id: int = Field(alias="training_id")
-    training_name: str = Field(alias="training_name")
-    containers: List[Container]
+    """训练容器"""
+    training_id: int = Field(alias="training_id", description="训练ID")
+    training_name: str = Field(alias="training_name", description="训练名")
+    containers: List[Container] = Field(description="容器列表")
 
 
 class ListTrainingContainersRequest(BaseModel):
-    page_size: int = Field(20, alias="page_size")
-    page_num: int = Field(1, alias="page_num")
+    """查询训练任务容器列表请求"""
+    page_size: int = Field(20, alias="page_size", description="每页数量")
+    page_num: int = Field(1, alias="page_num", description="当前页码")
 
 
 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)
+    """查询训练任务容器列表返回"""
+    total: int = Field(description="总条数")
+    page_size: int = Field(alias="page_size", description="每页数量")
+    page_num: int = Field(alias="page_num", description="当前页码")
+    data: Optional[List[TrainingContainer]] = Field(default_factory=list, description="训练容器列表")