作者：AI Platform Team | 发布日期：2025-10-01 | 标签：AI Platform、MLOps、平台架构、算法工程
一、背景与定位
1.1 为什么要搭建 AI Infra Platform
算法工程师直接登录服务器，用Jupyter跑实验，训练完手动上传模型到生产环境。
听起来很原始，但确实够用——因为量小。
但当团队扩张到20人、50人、100人的时候，问题就暴露了：
┌────────────────────────────────────────────────────────────────────────────┐
│ 痛点一：实验管理混乱 │
├────────────────────────────────────────────────────────────────────────────┤
│ │
│ 算法工程师A: "我上周跑的那个模型在哪来着？"... │
│ 算法工程师B: "这个参数配置是谁改的？结果对不上啊" │
│ 算法工程师C: "我本地训练的结果和服务器不一样，GPU型号差异？" │
│ │
│ 实际情况： │
│ - 代码版本混乱（V1/V2/V3...最终版/最终版2/最终版final） │
│ - 数据集分散（/data/similar_name/ 目录下有5个版本） │
│ - 实验记录靠Excel（人工记录，日期写错、参数遗漏） │
│ - 模型artifact找不着（不知道哪个对应哪个实验） │
│ │
└────────────────────────────────────────────────────────────────────────────┘
┌────────────────────────────────────────────────────────────────────────────┐
│ 痛点二：资源利用率低 │
├────────────────────────────────────────────────────────────────────────────┤
│ │
│ 现象： │
│ - GPU服务器负载图像心电图（有训练时100%，空闲时5%） │
│ - 资源分配靠"申请-审批-手动配置"的低效流程 │
│ - 排队现象严重，紧急任务无法优先 │
│ - 资源共享困难（A的卡被B的任务占着） │
│ │
│ 根因： │
│ - 没有统一的资源调度系统 │
│ - 任务优先级不明确 │
│ - 缺乏弹性扩缩容能力 │
│ │
└────────────────────────────────────────────────────────────────────────────┘
┌────────────────────────────────────────────────────────────────────────────┐
│ 痛点三：上线流程冗长 │
├────────────────────────────────────────────────────────────────────────────┤
│ │
│ 从实验到生产的路径： │
│ │
│ 算法工程师 ──> 训练完成 ──> 手动导出模型 ──> 上传S3 │
│ │ │ │
│ │ 2-3天后... ▼ │
│ │ 开发同事接收 │
│ │ │ │
│ │ 写推理服务代码 │
│ │ │ │
│ │ 测试环境部署 │
│ │ │ │
│ │ 修复Bug... │
│ │ │ │
│ │ 预发布环境 │
│ │ │ │
│ └──────────────> 生产环境部署 <────────────┘ │
│ │
│ 典型耗时：算法训练2天，上线等待2周 │
│ │
└────────────────────────────────────────────────────────────────────────────┘1.2 AI Infra Platform 的定位
基于上述痛点，我们定义了 AI Infra Platform 的核心定位：
┌─────────────────────────────────────────────────────────────────────────────┐
│ AI Infra Platform 定位 │
├─────────────────────────────────────────────────────────────────────────────┤
│ │
│ ┌─────────────────────────┐ │
│ │ 业务团队/产品 │ │
│ │ "我要用AI能力解决问题" │ │
│ └───────────┬─────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────┐ │
│ │ AI Infra Platform │ │
│ │ │ │
│ │ • 模型托管与推理服务 │ │
│ │ • 数据管理与特征工程 │ │
│ │ • 实验追踪与版本控制 │ │
│ │ • 资源调度与成本优化 │ │
│ │ • 监控告警与日志分析 │ │
│ └───────────┬─────────────┘ │
│ │ │
│ ┌───────────────┼───────────────┐ │
│ ▼ ▼ ▼ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ 算法团队 │ │ 数据团队 │ │ 工程团队 │ │
│ │ │ │ │ │ │ │
│ │ • 训练 │ │ • 数据 │ │ • 服务 │ │
│ │ • 实验 │ │ • 特征 │ │ • 部署 │ │
│ │ • 调优 │ │ • 质量 │ │ • 集成 │ │
│ └──────────┘ └──────────┘ └──────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────┘核心价值：
对算法工程师：专注算法研发，不用关心基础设施
对数据工程师：统一数据管理，保障数据质量
对工程团队：标准化部署流程，提升协作效率
对管理层：资源可见可控，成本可量化
二、整体架构设计
2.1 架构图（多层架构）
╔═══════════════════════════════════════════════════════════════════════════════════════════╗
║ AI Infra Platform ║
╠═══════════════════════════════════════════════════════════════════════════════════════════╣
║ ║
║ ┌─────────────────────────────────────────────────────────────────────────────────────┐ ║
║ │ 应用层 (Application Layer) │ ║
║ │ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ ┌────────────┐ │ ║
║ │ │ Web UI │ │ Python SDK │ │ CLI Tool │ │ RESTful API│ │ gRPC │ │ ║
║ │ │ (前端门户) │ │ (客户端) │ │ (命令行) │ │ (网关) │ │ (内部) │ │ ║
║ │ └─────────────┘ └─────────────┘ └─────────────┘ └─────────────┘ └────────────┘ │ ║
║ └─────────────────────────────────────────────────────────────────────────────────────┘ ║
║ │ ║
║ ▼ ║
║ ┌─────────────────────────────────────────────────────────────────────────────────────┐ ║
║ │ 平台层 (Platform Layer) │ ║
║ │ │ ║
║ │ ┌─────────────────────────────────────────────────────────────────────────────────┐ │ ║
║ │ │ 核心服务 (Core Services) │ │ ║
║ │ │ │ │ ║
║ │ │ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ │ │ ║
║ │ │ │ Model │ │ Training │ │ Inference │ │ Data │ │ Feature │ │ │ ║
║ │ │ │ Registry │ │ Pipeline │ │ Service │ │ Catalog │ │ Store │ │ │ ║
║ │ │ │ (模型注册) │ │ (训练流程) │ │ (推理服务) │ │ (数据目录) │ │ (特征存储) │ │ │ ║
║ │ │ └────────────┘ └────────────┘ └────────────┘ └────────────┘ └────────────┘ │ │ ║
║ │ │ │ │ ║
║ │ │ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ │ │ ║
║ │ │ │ Experiment │ │ Artifact │ │ Metric │ │ Alert │ │ Permission │ │ │ ║
║ │ │ │ Tracking │ │ Store │ │ Monitor │ │ Manager │ │ & Auth │ │ │ ║
║ │ │ │ (实验追踪) │ │ (产物存储) │ │ (指标监控) │ │ (告警管理) │ │ (权限认证) │ │ │ ║
║ │ │ └────────────┘ └────────────┘ └────────────┘ └────────────┘ └────────────┘ │ │ ║
║ │ └─────────────────────────────────────────────────────────────────────────────────┘ │ ║
║ │ │ ║
║ │ ┌─────────────────────────────────────────────────────────────────────────────────┐ │ ║
║ │ │ 消息队列 (Message Queue) │ │ ║
║ │ │ RabbitMQ / Redis Streams / Kafka (根据场景选择) │ │ ║
║ │ └─────────────────────────────────────────────────────────────────────────────────┘ │ ║
║ └─────────────────────────────────────────────────────────────────────────────────────┘ ║
║ │ ║
║ ▼ ║
║ ┌─────────────────────────────────────────────────────────────────────────────────────┐ ║
║ │ 基础设施层 (Infrastructure Layer) │ ║
║ │ │ ║
║ │ ┌────────────────┐ ┌────────────────┐ ┌────────────────┐ │ ║
║ │ │ Compute │ │ Storage │ │ Network │ │ ║
║ │ │ ┌──────────┐ │ │ ┌──────────┐ │ │ ┌──────────┐ │ │ ║
║ │ │ │Kubernetes │ │ │ │ MinIO/S3 │ │ │ │ VPC │ │ │ ║
║ │ │ │ Cluster │ │ │ │ (对象存储)│ │ │ │ (私有网络)│ │ │ ║
║ │ │ ├──────────┤ │ │ ├──────────┤ │ │ ├──────────┤ │ │ ║
║ │ │ │ GPU Pool │ │ │ │ NAS/EFS │ │ │ │ Load │ │ │ ║
║ │ │ │ (GPU节点) │ │ │ │ (文件存储)│ │ │ │ Balancer │ │ │ ║
║ │ │ ├──────────┤ │ │ ├──────────┤ │ │ ├──────────┤ │ │ ║
║ │ │ │ CPU Pool │ │ │ │ Redis │ │ │ │ DNS │ │ │ ║
║ │ │ │ (CPU节点) │ │ │ │ (缓存) │ │ │ │ (内部) │ │ │ ║
║ │ │ └──────────┘ │ │ └──────────┘ │ │ └──────────┘ │ │ ║
║ │ └────────────────┘ └────────────────┘ └────────────────┘ │ ║
║ │ │ ║
║ │ ┌────────────────┐ ┌────────────────┐ ┌────────────────┐ │ ║
║ │ │ Monitoring │ │ Logging │ │ Database │ │ ║
║ │ │ ┌──────────┐ │ │ ┌──────────┐ │ │ ┌──────────┐ │ │ ║
║ │ │ │Prometheus│ │ │ │ Loki │ │ │ │PostgreSQL│ │ │ ║
║ │ │ │Grafana │ │ │ │Promtail │ │ │ │ (主库) │ │ │ ║
║ │ │ └──────────┘ │ │ └──────────┘ │ │ ├──────────┤ │ │ ║
║ │ └────────────────┘ └────────────────┘ │ │ Milvus │ │ │ ║
║ │ │ │(向量库) │ │ │ ║
║ │ │ └──────────┘ │ │ ║
║ │ └────────────────┘ │ ║
║ └─────────────────────────────────────────────────────────────────────────────────────┘ ║
║ ║
╚═══════════════════════════════════════════════════════════════════════════════════════════╝2.2 模块交互图
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│ 模块交互流程 │
├─────────────────────────────────────────────────────────────────────────────────────────┤
│ │
│ 用户/Client 外部系统 │
│ │ │ │
│ │ create_experiment() │ │
│ │ │ │
│ ▼ ▼ │
│ ┌──────────────────────────────────────────────────────────────────────────────┐ │
│ │ API Gateway / Web Server │ │
│ │ (Django + DRF / FastAPI) │ │
│ └──────────────────────────────────────────────────────────────────────────────┘ │
│ │ │
│ │ submit_training_job() │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────────────────────────┐ │
│ │ Training Pipeline Service │ │
│ │ │ │
│ │ 1. 创建实验记录 → ExperimentTracker │ │
│ │ 2. 分配资源 → ResourceScheduler (K8s API) │ │
│ │ 3. 下发任务 → MessageQueue (Kafka/RabbitMQ) │ │
│ │ 4. 更新状态 → StatusTracker │ │
│ │ │ │
│ └──────────────────────────────────────────────────────────────────────────────┘ │
│ │ │
│ │ consume_training_task │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────────────────────────┐ │
│ │ Message Queue │ │
│ │ ┌─────────────────┐ │ │
│ │ │ training_tasks │ ← 训练任务队列 │ │
│ │ │ inference_tasks │ ← 推理任务队列 │ │
│ │ │ notification_q │ ← 通知队列 │ │
│ │ └─────────────────┘ │ │
│ └──────────────────────────────────────────────────────────────────────────────┘ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌────────────┐ ┌────────────┐ ┌────────────┐ │
│ │ Training │ │ Training │ │ Monitoring │ │
│ │ Worker #1 │ │ Worker #2 │ │ Worker │ │
│ │ (GPU Pod) │ │ (GPU Pod) │ │ │ │
│ └────────────┘ └────────────┘ └────────────┘ │
│ │ │
│ │ log_metrics() upload_artifact() save_model() │
│ ▼ │ │ │
│ ┌──────────────────────────────────────────────────────────────────────────────┐ │
│ │ Storage Layer │ │
│ │ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ │ │
│ │ │ MinIO │ │ Redis │ │ PostgreSQL │ │ Loki │ │ │
│ │ │ (Model/ │ │ (Cache/ │ │ (Metadata) │ │ (Logs) │ │ │
│ │ │ Dataset) │ │ Queue) │ │ │ │ │ │ │
│ │ └────────────┘ └────────────┘ └────────────┘ └────────────┘ │ │
│ └──────────────────────────────────────────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────────────────┘三、核心模块设计
3.1 模型管理 (Model Registry)
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│ Model Registry │
├─────────────────────────────────────────────────────────────────────────────────────────┤
│ │
│ 模型注册流程: │
│ │
│ 算法工程师 Platform Model Registry │
│ │ │ │ │
│ │ 1. 上传模型文件 │ │ │
│ │───────────────────────────>│ │ │
│ │ │ 2. 保存元信息 │ │
│ │ │────────────────────────────────>│ │
│ │ │ │ │
│ │ │ 3. 生成模型版本号 │ │
│ │ │<────────────────────────────────│ │
│ │ │ │ │
│ │ 4. 返回模型ID/版本 │ │ │
│ │<───────────────────────────│ │ │
│ │ │ │ │
│ │
│ 模型元信息结构: │
│ ┌─────────────────────────────────────────────────────────────────────────────┐ │
│ │ ModelMetadata { │ │
│ │ model_id: "uuid", // 唯一标识 │ │
│ │ name: "推荐模型", // 模型名称 │ │
│ │ version: "v1.2.3", // 版本号 │ │
│ │ framework: "pytorch", // 训练框架 │ │
│ │ framework_version: "2.0.1", // 框架版本 │ │
│ │ task_type: "recommendation", // 任务类型 │ │
│ │ metrics: { // 评估指标 │ │
│ │ accuracy: 0.95, // │ │
│ │ latency_ms: 12 // │ │
│ │ }, │ │
│ │ artifacts: [ // 模型产物 │ │
│ │ { path: "model.pt", size: "150MB", checksum: "md5:..." } │ │
│ │ ], │ │
│ │ created_by: "user_id", // 创建者 │ │
│ │ created_at: "2025-10-01", // 创建时间 │ │
│ │ tags: ["production", "stable"], // 标签 │ │
│ │ stage: "staging" // 阶段: development/staging/production │ │
│ │ } │ │
│ └─────────────────────────────────────────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────────────────┘3.2 训练管理 (Training Pipeline)
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│ Training Pipeline │
├─────────────────────────────────────────────────────────────────────────────────────────┤
│ │
│ 训练任务生命周期: │
│ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ Pending │───>│ Preparing│───>│ Running │───>│ Completed│ │ Failed │ │
│ │ (等待) │ │ (准备) │ │ (运行) │ │ (完成) │───>│ (失败) │ │
│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ └──────────┘ │
│ │ │ │ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ 资源分配失败 拉取镜像/数据 中间保存 │
│ │ │ │ │
│ │ ▼ ▼ │
│ │ ┌──────────┐ ┌──────────┐ │
│ └────────>│ Cancelled│ │ Checkpoint│ │
│ │ (取消) │ │ (保存点) │ │
│ └──────────┘ └──────────┘ │
│ │
│ 分布式训练支持: │
│ │
│ ┌───────────────────────────────────────────────────────────────────────────────┐ │
│ │ 单节点多卡 / 多节点多卡 │ │
│ │ │ │
│ │ Node 1 Node 2 │ │
│ │ ┌─────────────┐ ┌─────────────┐ │ │
│ │ │ Worker 0 │◄─────────►│ Worker 2 │ │ │
│ │ │ (GPU 0,1) │ NCCL │ (GPU 0,1) │ │ │
│ │ ├─────────────┤ ├─────────────┤ │ │
│ │ │ Worker 1 │◄─────────►│ Worker 3 │ │ │
│ │ │ (GPU 2,3) │ NCCL │ (GPU 2,3) │ │ │
│ │ └─────────────┘ └─────────────┘ │ │
│ │ │ │ │ │
│ │ └──────────┬──────────────┘ │ │
│ │ │ │ │
│ │ ┌─────┴─────┐ │ │
│ │ │ Parameter │ │ │
│ │ │ Server │ │ │
│ │ └───────────┘ │ │
│ └───────────────────────────────────────────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────────────────┘3.3 推理服务 (Inference Service)
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│ Inference Service │
├─────────────────────────────────────────────────────────────────────────────────────────┤
│ │
│ 推理服务架构: │
│ │
│ ┌────────────────┐ │
│ │ Load Balancer │ │
│ │ (Nginx/Envoy)│ │
│ └───────┬────────┘ │
│ │ │
│ ┌────────────────────────────┼────────────────────────────┐ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌────────────────┐ ┌────────────────┐ ┌────────────────┐ │
│ │ Inference Pod │ │ Inference Pod │ │ Inference Pod │ │
│ │ # 1 │ │ # 2 │ │ # 3 │ │
│ │ ┌────────────┐ │ │ ┌────────────┐ │ │ ┌────────────┐ │ │
│ │ │ Model │ │ │ │ Model │ │ │ │ Model │ │ │
│ │ │ Loader │ │ │ │ Loader │ │ │ │ Loader │ │ │
│ │ └─────┬──────┘ │ │ └─────┬──────┘ │ │ └─────┬──────┘ │ │
│ │ │ │ │ │ │ │ │ │ │
│ │ ▼ │ │ ▼ │ │ ▼ │ │
│ │ ┌────────────┐ │ │ ┌────────────┐ │ │ ┌────────────┐ │ │
│ │ │ Preprocess│ │ │ │ Preprocess│ │ │ │ Preprocess│ │ │
│ │ └─────┬──────┘ │ │ └─────┬──────┘ │ │ └─────┬──────┘ │ │
│ │ │ │ │ │ │ │ │ │ │
│ │ ▼ │ │ ▼ │ │ ▼ │ │
│ │ ┌────────────┐ │ │ ┌────────────┐ │ │ ┌────────────┐ │ │
│ │ │ Inference │ │ │ │ Inference │ │ │ │ Inference │ │ │
│ │ │ (GPU) │ │ │ │ (GPU) │ │ │ │ (GPU) │ │ │
│ │ └─────┬──────┘ │ │ └─────┬──────┘ │ │ └─────┬──────┘ │ │
│ │ │ │ │ │ │ │ │ │ │
│ │ ▼ │ │ ▼ │ │ ▼ │ │
│ │ ┌────────────┐ │ │ ┌────────────┐ │ │ ┌────────────┐ │ │
│ │ │ Postproc │ │ │ │ Postproc │ │ │ │ Postproc │ │ │
│ │ └────────────┘ │ │ └────────────┘ │ │ └────────────┘ │ │
│ └────────────────┘ └────────────────┘ └────────────────┘ │
│ │
│ 服务特性: │
│ ✓ 自动扩缩容 (HPA/KEDA) ✓ 模型热更新 ✓ 灰度发布 │
│ ✓ 多版本并存 ✓ A/B Testing ✓ 流量分配 │
│ ✓ 请求限流 ✓ 熔断降级 ✓ 指标采集 │
│ │
└─────────────────────────────────────────────────────────────────────────────────────────┘四、技术选型决策
4.1 为什么选 K8s 而不是纯 VM
这是我们讨论最多的一个问题。让我直接说结论：
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│ K8s vs 纯VM 对比 │
├───────────────────────┬─────────────────────────────┬─────────────────────────────────┤
│ 维度 │ 纯VM方案 │ Kubernetes方案 │
├───────────────────────┼─────────────────────────────┼─────────────────────────────────┤
│ 资源利用率 │ 30-40% │ 60-70% │
│ 扩缩容速度 │ 分钟级 │ 秒级 │
│ 部署复杂度 │ 简单 │ 复杂 │
│ 学习成本 │ 低 │ 高 │
│ 运维成本 │ 中等 │ 中等(有门槛) │
│ GPU支持 │ 需要手动管理 │ 原生支持 │
│ 服务发现 │ 需要额外工具 │ 内置 │
│ 灰度发布 │ 难以实现 │ 原生支持 │
│ 成本(机器>20台) │ 线性增长 │ 节省20-30% │
└───────────────────────┴─────────────────────────────┴─────────────────────────────────┘我们的判断：
团队规模 > 10人，GPU服务器 > 5台 → K8s 更划算
算法工程师需要频繁启停训练任务 → K8s 的弹性更合适
推理服务需要快速迭代 → K8s 的滚动更新更方便
选K8s的原因：
弹性伸缩：训练任务潮汐现象明显，K8s可以根据负载自动扩缩
资源隔离：多用户共享GPU，避免相互干扰
声明式部署：配置文件即代码，便于版本管理
生态丰富：Prometheus/Grafana/Loki等监控组件开箱即用
4.2 为什么选 Python 而不是 Go
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│ Python vs Go 对比 │
├───────────────────────┬─────────────────────────────┬─────────────────────────────────┤
│ 维度 │ Go方案 │ Python方案 │
├───────────────────────┼─────────────────────────────┼─────────────────────────────────┤
│ 性能 │ 高 │ 中等 │
│ 并发 │ 原生协程 │ 需要asyncio │
│ 算法团队友好度 │ 低 │ 高 │
│ 机器学习库支持 │ 一般 │ 丰富 │
│ 团队技能储备 │ 2人熟悉 │ 15人熟悉 │
│ 开发效率 │ 中等 │ 高 │
│ 部署复杂度 │ 单二进制 │ 需要环境管理 │
│ 长期维护 │ 好 │ 需要注意依赖管理 │
└───────────────────────┴─────────────────────────────┴─────────────────────────────────┘我们的选择：
平台核心服务（API网关、资源调度）→ Go
业务逻辑层（训练管理、模型服务）→ Python
CLI工具 → Go（最终用户使用，体验好）
理由：
算法工程师100%会用Python，降低学习成本
ML生态（PyTorch/TensorFlow）都是Python优先
Go用于基础设施层（CLI、高并发网关），发挥其性能优势
五、API 设计
5.1 RESTful + gRPC 双协议支持
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│ API 协议设计 │
├─────────────────────────────────────────────────────────────────────────────────────────┤
│ │
│ 协议选择策略: │
│ │
│ ┌─────────────────────────────────────────────────────────────────────────────┐ │
│ │ │ │
│ │ RESTful API (HTTP/JSON) gRPC (HTTP/2 + Protocol Buffers) │ │
│ │ ────────────────────────────── ─────────────────────────────── │ │
│ │ │ │
│ │ 适用场景: 适用场景: │ │
│ │ • 浏览器端调用 • 服务间内部调用 │ │
│ │ • 简单CRUD操作 • 高性能场景 │ │
│ │ • 调试/测试方便 • 流式处理 │ │
│ │ • 外部API开放 • 低延迟要求 │ │
│ │ │ │
│ │ 优点: 优点: │ │
│ │ • 通用性强 • 性能高 │ │
│ │ • 易于调试 • 类型安全 │ │
│ │ • 文档完善 • 代码生成 │ │
│ │ │ │
│ └─────────────────────────────────────────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────────────────┘5.2 API 端点设计
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│ API 端点设计 │
├─────────────────────────────────────────────────────────────────────────────────────────┤
│ │
│ /api/v1 │
│ ├── experiments # 实验管理 │
│ │ ├── GET / # 列出实验 │
│ │ ├── POST / # 创建实验 │
│ │ ├── GET /{id} # 获取实验详情 │
│ │ ├── PATCH /{id} # 更新实验 │
│ │ ├── DELETE /{id} # 删除实验 │
│ │ └── POST /{id}/cancel # 取消实验 │
│ │ │
│ ├── training # 训练任务 │
│ │ ├── POST /jobs # 提交训练任务 │
│ │ ├── GET /jobs/{id} # 获取任务状态 │
│ │ ├── GET /jobs/{id}/logs # 获取训练日志 │
│ │ ├── POST /jobs/{id}/stop # 停止训练 │
│ │ └── GET /jobs # 列出训练任务 │
│ │ │
│ ├── models # 模型管理 │
│ │ ├── GET / # 列出模型 │
│ │ ├── POST / # 注册模型 │
│ │ ├── GET /{id} # 获取模型详情 │
│ │ ├── POST /{id}/deploy # 部署模型 │
│ │ ├── POST /{id}/stage # 更新模型阶段 │
│ │ └── GET /{id}/versions # 获取版本列表 │
│ │ │
│ ├── inference # 推理服务 │
│ │ ├── POST /predict # 同步推理 │
│ │ ├── POST /batch_predict # 批量推理 │
│ │ └── GET /services # 列出推理服务 │
│ │ │
│ ├── datasets # 数据集管理 │
│ │ ├── GET / # 列出数据集 │
│ │ ├── POST / # 创建数据集 │
│ │ ├── POST /{id}/upload # 上传数据 │
│ │ └── GET /{id}/versions # 获取版本 │
│ │ │
│ └── features # 特征管理 │
│ ├── GET / # 列出特征 │
│ ├── POST / # 注册特征 │
│ └── GET /{name}/values # 获取特征值 │
│ │
└─────────────────────────────────────────────────────────────────────────────────────────┘六、项目结构与代码实现
6.1 Django 项目结构
ai_platform/
├── config/ # 配置目录
│ ├── settings/
│ │ ├── __init__.py
│ │ ├── base.py # 基础配置
│ │ ├── development.py # 开发环境
│ │ └── production.py # 生产环境
│ ├── urls.py # URL路由
│ ├── wsgi.py # WSGI配置
│ └── asgi.py # ASGI配置（支持WebSocket）
│
├── apps/ # 应用目录
│ ├── core/ # 核心应用（公共组件）
│ │ ├── __init__.py
│ │ ├── models.py # 公共模型（BaseModel等）
│ │ ├── serializers.py # 公共序列化器
│ │ ├── permissions.py # 权限类
│ │ ├── throttling.py # 限流类
│ │ └── middleware.py # 中间件
│ │
│ ├── experiments/ # 实验管理
│ │ ├── __init__.py
│ │ ├── models.py
│ │ ├── views.py
│ │ ├── serializers.py
│ │ ├── services/
│ │ │ ├── experiment_tracker.py
│ │ │ └── metric_collector.py
│ │ └── urls.py
│ │
│ ├── training/ # 训练管理
│ │ ├── __init__.py
│ │ ├── models.py
│ │ ├── views.py
│ │ ├── serializers.py
│ │ ├── services/
│ │ │ ├── job_scheduler.py # 任务调度
│ │ │ ├── resource_allocator.py # 资源分配
│ │ │ └── training_operator.py # K8s Operator
│ │ ├── tasks/
│ │ │ ├── __init__.py
│ │ │ └── training_tasks.py # Celery任务
│ │ └── urls.py
│ │
│ ├── models/ # 模型管理
│ │ ├── __init__.py
│ │ ├── models.py
│ │ ├── views.py
│ │ ├── serializers.py
│ │ ├── services/
│ │ │ ├── model_registry.py # 模型注册
│ │ │ ├── model_store.py # 模型存储
│ │ │ └── model_versioning.py # 版本管理
│ │ └── urls.py
│ │
│ ├── inference/ # 推理服务
│ │ ├── __init__.py
│ │ ├── models.py
│ │ ├── views.py
│ │ ├── serializers.py
│ │ ├── services/
│ │ │ ├── serving_manager.py # 服务管理
│ │ │ ├── predictor.py # 预测器
│ │ │ └── model_loader.py # 模型加载
│ │ ├── handlers/
│ │ │ └── grpc_handler.py # gRPC处理器
│ │ └── urls.py
│ │
│ └── datasets/ # 数据管理
│ ├── __init__.py
│ ├── models.py
│ ├── views.py
│ ├── services/
│ │ ├── dataset_manager.py
│ │ └── data_processor.py
│ └── urls.py
│
├── common/ # 公共工具
│ ├── __init__.py
│ ├── utils/
│ │ ├── k8s_client.py # K8s客户端封装
│ │ ├── storage_client.py # 对象存储客户端
│ │ └── metrics_client.py # 监控客户端
│ ├── exceptions.py # 自定义异常
│ └── responses.py # 统一响应格式
│
├── worker/ # Worker服务
│ ├── __init__.py
│ ├── training_worker.py # 训练Worker
│ ├── inference_worker.py # 推理Worker
│ └── celery_app.py # Celery配置
│
├── scripts/ # 脚本
│ ├── init_k8s_resources.py # 初始化K8s资源
│ └── migrate_data.py # 数据迁移
│
├── tests/ # 测试
│ ├── unit/
│ ├── integration/
│ └── fixtures/
│
├── docker/ # Docker配置
│ ├── Dockerfile.api
│ ├── Dockerfile.worker
│ └── docker-compose.yml
│
├── k8s/ # K8s配置
│ ├── base/
│ │ ├── deployment.yaml
│ │ ├── service.yaml
│ │ └── configmap.yaml
│ ├── overlays/
│ │ ├── development/
│ │ └── production/
│ └── operators/
│ └── training_operator.py
│
├── requirements.txt
├── requirements-dev.txt
├── manage.py
└── README.md6.2 核心代码实现
apps/core/models.py - 基础模型
"""
AI Platform 核心模型
提供基础模型类和公共字段
Author: AI Platform Team
Date: 2025-10-01
"""
import uuid
from django.db import models
from django.utils import timezone
class BaseModel(models.Model):
"""
基础模型类
提供：
- UUID主键
- 创建/更新时间
- 软删除支持
- 公共字段
"""
id = models.UUIDField(
primary_key=True,
default=uuid.uuid4,
editable=False,
help_text="唯一标识"
)
created_at = models.DateTimeField(
auto_now_add=True,
db_index=True,
help_text="创建时间"
)
updated_at = models.DateTimeField(
auto_now=True,
help_text="更新时间"
)
created_by = models.CharField(
max_length=128,
null=True,
blank=True,
db_index=True,
help_text="创建者"
)
is_deleted = models.BooleanField(
default=False,
db_index=True,
help_text="软删除标记"
)
class Meta:
abstract = True
def soft_delete(self):
"""软删除"""
self.is_deleted = True
self.save(update_fields=['is_deleted', 'updated_at'])
def restore(self):
"""恢复删除"""
self.is_deleted = False
self.save(update_fields=['is_deleted', 'updated_at'])
class TenantModel(BaseModel):
"""
租户模型基类
支持多租户隔离
"""
tenant_id = models.UUIDField(
null=True,
blank=True,
db_index=True,
help_text="租户ID"
)
class Meta:
abstract = True
class Experiment(BaseModel):
"""
实验记录
记录每次实验的配置、参数、结果
"""
class Status(models.TextChoices):
PENDING = 'pending', '等待中'
RUNNING = 'running', '运行中'
SUCCEEDED = 'succeeded', '成功'
FAILED = 'failed', '失败'
CANCELLED = 'cancelled', '已取消'
name = models.CharField(
max_length=255,
help_text="实验名称"
)
description = models.TextField(
null=True,
blank=True,
help_text="实验描述"
)
project = models.CharField(
max_length=128,
db_index=True,
help_text="所属项目"
)
status = models.CharField(
max_length=32,
choices=Status.choices,
default=Status.PENDING,
db_index=True,
help_text="实验状态"
)
# 实验配置（JSON格式存储）
config = models.JSONField(
default=dict,
help_text="实验配置"
)
# 参数
parameters = models.JSONField(
default=dict,
help_text="超参数"
)
# 数据集引用
dataset_id = models.UUIDField(
null=True,
blank=True,
help_text="数据集ID"
)
# 模型引用
base_model_id = models.UUIDField(
null=True,
blank=True,
help_text="基础模型ID"
)
# 计算资源
resources = models.JSONField(
default=dict,
help_text="资源配置"
)
# 代码引用
code_version = models.CharField(
max_length=64,
null=True,
blank=True,
help_text="代码版本/Git Commit"
)
# Git信息
git_repo = models.URLField(
null=True,
blank=True,
help_text="Git仓库地址"
)
git_branch = models.CharField(
max_length=128,
null=True,
blank=True,
help_text="Git分支"
)
git_commit = models.CharField(
max_length=64,
null=True,
blank=True,
help_text="Git提交hash"
)
# 执行信息
started_at = models.DateTimeField(
null=True,
blank=True,
help_text="开始时间"
)
finished_at = models.DateTimeField(
null=True,
blank=True,
help_text="结束时间"
)
duration_seconds = models.IntegerField(
null=True,
blank=True,
help_text="运行时长(秒)"
)
# 失败信息
error_message = models.TextField(
null=True,
blank=True,
help_text="错误信息"
)
# 标签
tags = models.JSONField(
default=list,
help_text="标签列表"
)
# 元数据
metadata = models.JSONField(
default=dict,
help_text="额外元信息"
)
class Meta:
db_table = 'experiments'
ordering = ['-created_at']
indexes = [
models.Index(fields=['project', 'status']),
models.Index(fields=['created_by', '-created_at']),
models.Index(fields=['-created_at']),
]
def __str__(self):
return f"{self.name} ({self.status})"
def mark_started(self):
"""标记为开始"""
self.status = self.Status.RUNNING
self.started_at = timezone.now()
self.save(update_fields=['status', 'started_at', 'updated_at'])
def mark_completed(self, success: bool = True):
"""标记为完成"""
self.status = self.Status.SUCCEEDED if success else self.Status.FAILED
self.finished_at = timezone.now()
if self.started_at:
self.duration_seconds = int(
(self.finished_at - self.started_at).total_seconds()
)
self.save(update_fields=[
'status', 'finished_at', 'duration_seconds', 'updated_at'
])
def mark_cancelled(self, reason: str = None):
"""标记为取消"""
self.status = self.Status.CANCELLED
self.finished_at = timezone.now()
if reason:
self.error_message = reason
if self.started_at:
self.duration_seconds = int(
(self.finished_at - self.started_at).total_seconds()
)
self.save(update_fields=[
'status', 'finished_at', 'duration_seconds',
'error_message', 'updated_at'
])
class TrainingJob(BaseModel):
"""
训练任务
对应一个实际的K8s Job/ChoriaJob
"""
class Status(models.TextChoices):
PENDING = 'pending', '等待中'
PREPARING = 'preparing', '准备中'
RUNNING = 'running', '运行中'
SUCCEEDED = 'succeeded', '成功'
FAILED = 'failed', '失败'
CANCELLED = 'cancelled', '已取消'
experiment = models.ForeignKey(
Experiment,
on_delete=models.CASCADE,
related_name='jobs',
help_text="关联实验"
)
job_name = models.CharField(
max_length=255,
unique=True,
help_text="K8s Job名称"
)
status = models.CharField(
max_length=32,
choices=Status.choices,
default=Status.PENDING,
db_index=True,
help_text="任务状态"
)
# 资源请求
gpu_request = models.IntegerField(
default=0,
help_text="请求GPU数量"
)
cpu_request = models.CharField(
max_length=32,
default='2',
help_text="请求CPU核数"
)
memory_request = models.CharField(
max_length=32,
default='8Gi',
help_text="请求内存大小"
)
# K8s信息
namespace = models.CharField(
max_length=64,
default='ai-platform',
help_text="K8s命名空间"
)
pod_name = models.CharField(
max_length=255,
null=True,
blank=True,
help_text="Pod名称"
)
node_name = models.CharField(
max_length=255,
null=True,
blank=True,
help_text="调度到的节点"
)
# 镜像
image = models.CharField(
max_length=512,
help_text="训练镜像"
)
# 启动命令
command = models.JSONField(
default=list,
help_text="启动命令"
)
# 环境变量
env_vars = models.JSONField(
default=dict,
help_text="环境变量"
)
# 挂载配置
volume_mounts = models.JSONField(
default=list,
help_text="存储挂载"
)
# 分布式配置
replicas = models.IntegerField(
default=1,
help_text="副本数(分布式训练)"
)
rank = models.IntegerField(
default=0,
help_text="当前副本编号"
)
# 执行时间
started_at = models.DateTimeField(
null=True,
blank=True,
help_text="开始时间"
)
finished_at = models.DateTimeField(
null=True,
blank=True,
help_text="结束时间"
)
# 日志
log_url = models.URLField(
null=True,
blank=True,
help_text="日志访问URL"
)
# 输出
output_path = models.CharField(
max_length=512,
null=True,
blank=True,
help_text="输出路径"
)
# Checkpoint
checkpoint_path = models.CharField(
max_length=512,
null=True,
blank=True,
help_text="Checkpoint路径"
)
# 指标（最新值）
latest_metrics = models.JSONField(
default=dict,
help_text="最新指标值"
)
class Meta:
db_table = 'training_jobs'
ordering = ['-created_at']
indexes = [
models.Index(fields=['experiment', 'status']),
models.Index(fields=['job_name']),
]
def __str__(self):
return f"{self.job_name} ({self.status})"
class ModelRegistry(BaseModel):
"""
模型注册表
管理所有已注册的模型
"""
class Framework(models.TextChoices):
PYTORCH = 'pytorch', 'PyTorch'
TENSORFLOW = 'tensorflow', 'TensorFlow'
JAX = 'jax', 'JAX'
SKLEARN = 'sklearn', 'Scikit-learn'
XGBOOST = 'xgboost', 'XGBoost'
ONNX = 'onnx', 'ONNX'
OTHER = 'other', '其他'
class TaskType(models.TextChoices):
CLASSIFICATION = 'classification', '分类'
REGRESSION = 'regression', '回归'
NLP = 'nlp', 'NLP'
CV = 'cv', '计算机视觉'
RECOMMENDATION = 'recommendation', '推荐'
GENERATION = 'generation', '生成'
OTHER = 'other', '其他'
class Stage(models.TextChoices):
DEVELOPMENT = 'development', '开发'
STAGING = 'staging', '预发布'
PRODUCTION = 'production', '生产'
ARCHIVED = 'archived', '归档'
name = models.CharField(
max_length=255,
help_text="模型名称"
)
version = models.CharField(
max_length=64,
help_text="模型版本"
)
# 唯一性约束：name + version
description = models.TextField(
null=True,
blank=True,
help_text="模型描述"
)
framework = models.CharField(
max_length=32,
choices=Framework.choices,
help_text="训练框架"
)
framework_version = models.CharField(
max_length=32,
null=True,
blank=True,
help_text="框架版本"
)
task_type = models.CharField(
max_length=32,
choices=TaskType.choices,
help_text="任务类型"
)
# 模型文件
artifact_path = models.CharField(
max_length=512,
help_text="模型文件路径"
)
artifact_size = models.BigIntegerField(
null=True,
blank=True,
help_text="模型大小(字节)"
)
artifact_checksum = models.CharField(
max_length=64,
null=True,
blank=True,
help_text="文件校验和"
)
# 训练信息
experiment = models.ForeignKey(
Experiment,
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name='produced_models',
help_text="来源实验"
)
training_job = models.ForeignKey(
TrainingJob,
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name='output_models',
help_text="来源训练任务"
)
# 指标
metrics = models.JSONField(
default=dict,
help_text="评估指标"
)
# 阶段
stage = models.CharField(
max_length=32,
choices=Stage.choices,
default=Stage.DEVELOPMENT,
db_index=True,
help_text="模型阶段"
)
# 签名（用于推理）
signature = models.JSONField(
default=dict,
help_text="模型签名"
)
# 元数据
input_shape = models.JSONField(
default=dict,
help_text="输入形状"
)
output_shape = models.JSONField(
default=dict,
help_text="输出形状"
)
labels = models.JSONField(
default=list,
help_text="标签列表"
)
metadata = models.JSONField(
default=dict,
help_text="额外元信息"
)
# 服务信息
current_deployment = models.ForeignKey(
'InferenceDeployment',
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name='deployed_models',
help_text="当前部署"
)
class Meta:
db_table = 'model_registry'
ordering = ['-created_at']
unique_together = ['name', 'version']
indexes = [
models.Index(fields=['name', 'stage']),
models.Index(fields=['framework', 'task_type']),
models.Index(fields=['experiment']),
]
def __str__(self):
return f"{self.name}:{self.version}"
class InferenceDeployment(BaseModel):
"""
推理服务部署
管理模型的在线服务
"""
class Status(models.TextChoices):
PENDING = 'pending', '部署中'
RUNNING = 'running', '运行中'
FAILED = 'failed', '失败'
STOPPED = 'stopped', '已停止'
UPDATING = 'updating', '更新中'
class ServiceType(models.TextChoices):
REALTIME = 'realtime', '实时推理'
BATCH = 'batch', '批量推理'
STREAMING = 'streaming', '流式推理'
name = models.CharField(
max_length=255,
unique=True,
help_text="服务名称"
)
model = models.ForeignKey(
ModelRegistry,
on_delete=models.CASCADE,
related_name='deployments',
help_text="部署的模型"
)
status = models.CharField(
max_length=32,
choices=Status.choices,
default=Status.PENDING,
db_index=True,
help_text="服务状态"
)
service_type = models.CharField(
max_length=32,
choices=ServiceType.choices,
default=ServiceType.REALTIME,
help_text="服务类型"
)
# 副本配置
replicas = models.IntegerField(
default=1,
help_text="副本数"
)
min_replicas = models.IntegerField(
default=1,
help_text="最小副本数"
)
max_replicas = models.IntegerField(
default=10,
help_text="最大副本数"
)
# 资源
gpu_per_replica = models.FloatField(
default=0,
help_text="每副本GPU数"
)
cpu_per_replica = models.CharField(
max_length=32,
default='1',
help_text="每副本CPU核数"
)
memory_per_replica = models.CharField(
max_length=32,
default='2Gi',
help_text="每副本内存"
)
# K8s配置
namespace = models.CharField(
max_length=64,
default='ai-inference',
help_text="命名空间"
)
service_name = models.CharField(
max_length=255,
null=True,
blank=True,
help_text="K8s Service名称"
)
deployment_name = models.CharField(
max_length=255,
null=True,
blank=True,
help_text="K8s Deployment名称"
)
# 访问配置
endpoint = models.URLField(
null=True,
blank=True,
help_text="访问端点"
)
port = models.IntegerField(
default=8000,
help_text="服务端口"
)
# 流量配置
traffic_weight = models.IntegerField(
default=100,
help_text="流量权重(%)"
)
# 指标
qps = models.FloatField(
default=0,
help_text="当前QPS"
)
latency_p99_ms = models.FloatField(
default=0,
help_text="P99延迟(ms)"
)
# 健康检查
health_check_url = models.URLField(
null=True,
blank=True,
help_text="健康检查URL"
)
last_health_check = models.DateTimeField(
null=True,
blank=True,
help_text="上次健康检查时间"
)
is_healthy = models.BooleanField(
default=True,
help_text="是否健康"
)
# 配置
config = models.JSONField(
default=dict,
help_text="服务配置"
)
class Meta:
db_table = 'inference_deployments'
ordering = ['-created_at']
indexes = [
models.Index(fields=['status', 'service_type']),
models.Index(fields=['model', 'stage']),
]
def __str__(self):
return f"{self.name} ({self.status})"6.3 训练任务调度服务
apps/training/services/job_scheduler.py - 任务调度器
"""
训练任务调度器
负责将训练任务调度到K8s集群
核心功能：
1. 任务提交与状态跟踪
2. 资源分配与调度
3. 日志收集与指标上报
Author: AI Platform Team
Date: 2025-10-01
"""
import time
import json
import logging
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field
from datetime import datetime
import uuid
from django.conf import settings
from kubernetes import client, config
from kubernetes.client.rest import ApiException
logger = logging.getLogger(__name__)
@dataclass
class ResourceRequest:
"""资源请求"""
cpu: str = "2"
memory: str = "8Gi"
gpu: int = 0
gpu_type: str = "nvidia" # nvidia, amd
storage: str = "50Gi"
timeout: int = 86400 # 秒
@dataclass
class TrainingJobSpec:
"""训练任务规格"""
job_name: str
image: str
command: List[str]
env_vars: Dict[str, str] = field(default_factory=dict)
resources: ResourceRequest = field(default_factory=ResourceRequest)
working_dir: str = "/workspace"
volume_mounts: List[Dict] = field(default_factory=list)
node_selector: Dict[str, str] = field(default_factory=dict)
tolerations: List[Dict] = field(default_factory=list)
replicas: int = 1 # 分布式训练副本数
chief_port: int = 2222 # Chief节点端口
labels: Dict[str, str] = field(default_factory=dict)
def to_dict(self) -> Dict:
"""转换为字典"""
return {
"job_name": self.job_name,
"image": self.image,
"command": self.command,
"env_vars": self.env_vars,
"resources": {
"cpu": self.resources.cpu,
"memory": self.resources.memory,
"gpu": self.resources.gpu,
"gpu_type": self.resources.gpu_type
},
"working_dir": self.working_dir,
"replicas": self.replicas,
"chief_port": self.chief_port
}
class JobScheduler:
"""
训练任务调度器
封装K8s API，实现：
- 单节点训练任务提交
- 分布式训练任务提交（PyTorchJob/TFJob）
- 任务状态跟踪
- 日志收集
"""
def __init__(self, namespace: str = "ai-platform"):
"""
初始化调度器
Args:
namespace: K8s命名空间
"""
self.namespace = namespace
# 初始化K8s客户端
try:
# 尝试加载in-cluster配置（Pod内运行）
config.load_incluster_config()
except config.ConfigException:
# 回退到本地配置（开发环境）
try:
config.load_kube_config()
except Exception as e:
logger.warning(f"无法加载K8s配置: {e}")
self.in_cluster = False
return
self.in_cluster = True
self.batch_api = client.BatchV1Api()
self.core_api = client.CoreV1Api()
self.custom_api = client.CustomObjectsApi()
# 默认配置
self.default_image = settings.TRAINING.get(
"default_image",
"pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime"
)
self.default_storage_class = settings.TRAINING.get(
"default_storage_class",
"nfs-storage"
)
logger.info(f"JobScheduler初始化完成，namespace: {namespace}")
def submit_job(self, spec: TrainingJobSpec) -> Dict:
"""
提交训练任务
Args:
spec: 任务规格
Returns:
提交结果，包含job_name等信息
"""
try:
if spec.replicas > 1:
# 分布式训练
return self._submit_distributed_job(spec)
else:
# 单节点训练
return self._submit_single_job(spec)
except ApiException as e:
logger.error(f"K8s API错误: {e}")
raise TrainingJobError(f"提交任务失败: {e}")
def _submit_single_job(self, spec: TrainingJobSpec) -> Dict:
"""提交单节点训练任务"""
# 生成唯一Job名称
job_name = spec.job_name or f"train-{uuid.uuid4().hex[:8]}"
# 构建容器
container = client.V1Container(
name="training",
image=spec.image or self.default_image,
command=spec.command,
env=self._build_env_vars(spec.env_vars),
resources=self._build_resources(spec.resources),
working_dir=spec.working_dir,
volume_mounts=spec.volume_mounts
)
# 构建Pod模板
pod_template = client.V1PodTemplateSpec(
metadata=client.V1ObjectMeta(
labels={
"app": "ai-training",
"job-name": job_name,
**spec.labels
},
annotations={
"sidecar.istio.io/inject": "false" # 禁用Istio注入
}
),
spec=client.V1PodSpec(
restart_policy="Never",
containers=[container],
node_selector=spec.node_selector,
tolerations=spec.tolerations if spec.tolerations else [
# 默认允许调度到任意节点
client.V1Toleration(
key="node.kubernetes.io/not-ready",
operator="Exists",
effect="NoSchedule",
toleration_seconds=300
)
],
# 挂载存储
volumes=self._build_volumes(job_name)
)
)
# 构建Job
job = client.V1Job(
api_version="batch/v1",
kind="Job",
metadata=client.V1ObjectMeta(
name=job_name,
namespace=self.namespace,
labels={
"app": "ai-training",
**spec.labels
}
),
spec=client.V1JobSpec(
backoff_limit=0, # 不重试
ttl_seconds_after_finished=3600, # 完成后1小时清理
template=pod_template
)
)
# 提交Job
self.batch_api.create_namespaced_job(
namespace=self.namespace,
body=job
)
logger.info(f"训练任务已提交: {job_name}")
return {
"job_name": job_name,
"namespace": self.namespace,
"status": "pending",
"replicas": 1
}
def _submit_distributed_job(self, spec: TrainingJobSpec) -> Dict:
"""
提交分布式训练任务
使用PyTorchJob CRD（需要安装PyTorch Operator）
"""
job_name = spec.job_name or f"dist-train-{uuid.uuid4().hex[:8]}"
# 构建PyTorch Job spec
from kubernetes.client import V1ResourceRequirements
# Worker规格
worker_spec = client.V1PodSpec(
containers=[
client.V1Container(
name="pytorch",
image=spec.image or self.default_image,
command=[
"/bin/bash", "-c",
f"""
torchrun \
--nnodes={spec.replicas} \
--nproc_per_node=1 \
--master_addr=$MASTER_ADDR \
--master_port={spec.chief_port} \
{' '.join(spec.command)}
"""
],
env=self._build_env_vars({
**spec.env_vars,
"WORLD_SIZE": str(spec.replicas),
"RANK": "$(POD_INDEX)",
"MASTER_ADDR": "pytorch-master",
"MASTER_PORT": str(spec.chief_port)
}),
resources=V1ResourceRequirements(
requests={
"cpu": spec.resources.cpu,
"memory": spec.resources.memory,
"nvidia.com/gpu": str(spec.resources.gpu)
},
limits={
"nvidia.com/gpu": str(spec.resources.gpu)
}
),
working_dir=spec.working_dir
)
],
restart_policy="Never",
node_selector=spec.node_selector,
tolerations=spec.tolerations
)
# 构建分布式Job
torch_job = {
"apiVersion": "kubeflow.org/v1",
"kind": "PyTorchJob",
"metadata": {
"name": job_name,
"namespace": self.namespace,
"labels": spec.labels
},
"spec": {
"backoffLimit": 0,
"pytorchReplicaSpecs": {
"Master": {
"replicas": 1,
"restartPolicy": "Never",
"template": {
"metadata": {
"labels": {"app": "ai-training"}
},
"spec": {
**worker_spec,
"containers": [
client.V1Container(
name="pytorch",
image=spec.image or self.default_image,
command=["/bin/bash", "-c",
f"""
torchrun \
--nnodes={spec.replicas} \
--nproc_per_node=1 \
--master_addr=$MASTER_ADDR \
--master_port={spec.chief_port} \
{' '.join(spec.command)}
"""
],
env=self._build_env_vars({
**spec.env_vars,
"WORLD_SIZE": str(spec.replicas),
"RANK": "0",
"MASTER_ADDR": "localhost",
"MASTER_PORT": str(spec.chief_port)
}),
resources=V1ResourceRequirements(
requests={
"cpu": spec.resources.cpu,
"memory": spec.resources.memory,
"nvidia.com/gpu": str(spec.resources.gpu)
}
),
working_dir=spec.working_dir
)
]
}
}
},
"Worker": {
"replicas": spec.replicas - 1,
"restartPolicy": "Never",
"template": {
"metadata": {
"labels": {"app": "ai-training"}
},
"spec": worker_spec
}
}
}
}
}
# 提交Job
self.custom_api.create_namespaced_custom_object(
group="kubeflow.org",
version="v1",
namespace=self.namespace,
plural="pytorchjobs",
body=torch_job
)
logger.info(f"分布式训练任务已提交: {job_name}, replicas: {spec.replicas}")
return {
"job_name": job_name,
"namespace": self.namespace,
"status": "pending",
"replicas": spec.replicas
}
def get_job_status(self, job_name: str) -> Dict:
"""
获取任务状态
Returns:
状态信息
"""
try:
job = self.batch_api.read_namespaced_job(
name=job_name,
namespace=self.namespace
)
status = job.status
# 解析状态
if status.active and status.active > 0:
state = "running"
elif status.succeeded and status.succeeded > 0:
state = "succeeded"
elif status.failed and status.failed > 0:
state = "failed"
else:
state = "pending"
# 获取Pod信息
pods = self.core_api.list_namespaced_pod(
namespace=self.namespace,
label_selector=f"job-name={job_name}"
)
pod_info = []
for pod in pods.items:
pod_info.append({
"name": pod.metadata.name,
"phase": pod.status.phase if pod.status else "Unknown",
"node": pod.spec.node_name if pod.spec else None,
"start_time": pod.status.start_time.isoformat() if pod.status and pod.status.start_time else None
})
return {
"job_name": job_name,
"state": state,
"active": status.active or 0,
"succeeded": status.succeeded or 0,
"failed": status.failed or 0,
"conditions": [
{"type": c.type, "status": c.status, "reason": c.reason}
for c in (status.conditions or [])
],
"pods": pod_info
}
except ApiException as e:
if e.status == 404:
return {
"job_name": job_name,
"state": "not_found",
"error": f"Job不存在: {job_name}"
}
raise TrainingJobError(f"获取状态失败: {e}")
def cancel_job(self, job_name: str) -> bool:
"""
取消任务
Args:
job_name: 任务名称
Returns:
是否成功
"""
try:
self.batch_api.delete_namespaced_job(
name=job_name,
namespace=self.namespace,
body=client.V1DeleteOptions(
propagation_policy="Foreground"
)
)
logger.info(f"训练任务已取消: {job_name}")
return True
except ApiException as e:
if e.status == 404:
logger.warning(f"任务不存在，无法取消: {job_name}")
return False
raise TrainingJobError(f"取消任务失败: {e}")
def get_job_logs(self, job_name: str, tail_lines: int = 100) -> Dict:
"""
获取任务日志
Args:
job_name: 任务名称
tail_lines: 返回最近N行
Returns:
日志内容
"""
try:
# 获取Pod
pods = self.core_api.list_namespaced_pod(
namespace=self.namespace,
label_selector=f"job-name={job_name}"
)
if not pods.items:
return {"error": "未找到Pod"}
# 获取主Pod日志
pod = pods.items[0]
logs = self.core_api.read_namespaced_pod_log(
name=pod.metadata.name,
namespace=self.namespace,
tail_lines=tail_lines,
follow=False
)
return {
"job_name": job_name,
"pod_name": pod.metadata.name,
"logs": logs
}
except ApiException as e:
raise TrainingJobError(f"获取日志失败: {e}")
def _build_env_vars(self, env_vars: Dict[str, str]) -> List[client.V1EnvVar]:
"""构建环境变量"""
result = []
for key, value in env_vars.items():
result.append(client.V1EnvVar(
name=key,
value=str(value)
))
# 添加默认环境变量
result.extend([
client.V1EnvVar(
name="PYTHONUNBUFFERED",
value="1"
),
client.V1EnvVar(
name="NVIDIA_VISIBLE_DEVICES",
value="all"
)
])
return result
def _build_resources(self, resources: ResourceRequest) -> client.V1ResourceRequirements:
"""构建资源请求"""
requests = {
"cpu": resources.cpu,
"memory": resources.memory
}
limits = {
"cpu": resources.cpu,
"memory": resources.memory
}
if resources.gpu > 0:
requests[f"nvidia.com/gpu"] = str(resources.gpu)
limits[f"nvidia.com/gpu"] = str(resources.gpu)
return client.V1ResourceRequirements(
requests=requests,
limits=limits
)
def _build_volumes(self, job_name: str) -> List[client.V1Volume]:
"""构建存储卷"""
volumes = [
# 工作目录PVC
client.V1Volume(
name="workspace",
persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource(
claim_name=f"workspace-{job_name}",
default_capacity=resources.storage if 'resources' in locals() else "50Gi"
)
)
]
return volumes
class TrainingJobError(Exception):
"""训练任务异常"""
pass6.4 关键 API 接口代码
apps/training/views.py - 训练任务 API
"""
训练任务 API 接口
Author: AI Platform Team
Date: 2025-10-01
"""
import time
from typing import Optional
from dataclasses import dataclass
from rest_framework import viewsets, status
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework.views import APIView
from rest_framework.permissions import IsAuthenticated
from django_filters.rest_framework import DjangoFilterBackend
from .models import TrainingJob, Experiment
from .serializers import (
TrainingJobSerializer,
TrainingJobCreateSerializer,
TrainingJobUpdateSerializer
)
from .services.job_scheduler import JobScheduler, TrainingJobSpec, ResourceRequest
from common.permissions import TeamMemberPermission, ProjectAdminPermission
from common.pagination import StandardPagination
import logging
logger = logging.getLogger(__name__)
@dataclass
class SubmitJobRequest:
"""提交任务请求"""
experiment_id: str
image: str
command: list
env_vars: dict
resources: dict
working_dir: str = "/workspace"
node_selector: dict = None
replicas: int = 1
@classmethod
def from_request(cls, data: dict) -> "SubmitJobRequest":
"""从请求数据创建"""
return cls(
experiment_id=data["experiment_id"],
image=data.get("image", "pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime"),
command=data["command"],
env_vars=data.get("env_vars", {}),
resources=data.get("resources", {
"cpu": "2",
"memory": "8Gi",
"gpu": 1
}),
working_dir=data.get("working_dir", "/workspace"),
node_selector=data.get("node_selector"),
replicas=data.get("replicas", 1)
)
class TrainingJobViewSet(viewsets.ModelViewSet):
"""
训练任务视图集
提供完整的CRUD操作和自定义动作
"""
queryset = TrainingJob.objects.filter(is_deleted=False)
serializer_class = TrainingJobSerializer
pagination_class = StandardPagination
filter_backends = [DjangoFilterBackend]
filterset_fields = ['status', 'experiment', 'created_by']
def get_permissions(self):
"""权限配置"""
if self.action in ['create', 'list', 'retrieve']:
return [IsAuthenticated()]
elif self.action in ['update', 'partial_update', 'destroy']:
return [IsAuthenticated(), ProjectAdminPermission()]
else:
return [IsAuthenticated()]
def get_queryset(self):
"""过滤数据集"""
queryset = super().get_queryset()
# 按项目过滤
project = self.request.query_params.get('project')
if project:
queryset = queryset.filter(experiment__project=project)
# 按状态过滤
status_filter = self.request.query_params.get('status')
if status_filter:
queryset = queryset.filter(status=status_filter)
return queryset.select_related('experiment')
def create(self, request, *args, **kwargs):
"""
提交训练任务
POST /api/v1/training/jobs
Request Body:
{
"experiment_id": "uuid",
"image": "pytorch/pytorch:2.0.1",
"command": ["python", "train.py", "--epochs", "100"],
"env_vars": {
"DATA_PATH": "/data/train.csv",
"MODEL_DIR": "/workspace/output"
},
"resources": {
"cpu": "4",
"memory": "16Gi",
"gpu": 2
},
"replicas": 1
}
"""
# 1. 验证请求
serializer = TrainingJobCreateSerializer(data=request.data)
serializer.is_valid(raise_exception=True)
# 2. 获取实验
experiment_id = serializer.validated_data['experiment_id']
try:
experiment = Experiment.objects.get(id=experiment_id)
except Experiment.DoesNotExist:
return Response(
{"error": f"实验不存在: {experiment_id}"},
status=status.HTTP_404_NOT_FOUND
)
# 3. 构建任务规格
spec = SubmitJobRequest.from_request(serializer.validated_data)
job_spec = TrainingJobSpec(
job_name=f"train-{experiment.name[:20]}-{int(time.time())}",
image=spec.image,
command=spec.command,
env_vars={
**spec.env_vars,
"EXPERIMENT_ID": str(experiment.id),
"PROJECT": experiment.project
},
resources=ResourceRequest(
cpu=spec.resources.get("cpu", "2"),
memory=spec.resources.get("memory", "8Gi"),
gpu=spec.resources.get("gpu", 0)
),
working_dir=spec.working_dir,
replicas=spec.replicas,
labels={
"experiment_id": str(experiment.id),
"project": experiment.project,
"created_by": request.user.username
}
)
# 4. 提交到调度器
scheduler = JobScheduler()
try:
result = scheduler.submit_job(job_spec)
except Exception as e:
logger.error(f"提交训练任务失败: {e}")
return Response(
{"error": f"提交任务失败: {str(e)}"},
status=status.HTTP_500_INTERNAL_SERVER_ERROR
)
# 5. 创建数据库记录
job = TrainingJob.objects.create(
experiment=experiment,
job_name=result["job_name"],
status=TrainingJob.Status.PENDING,
image=spec.image,
command=spec.command,
env_vars=spec.env_vars,
gpu_request=spec.resources.get("gpu", 0),
cpu_request=spec.resources.get("cpu", "2"),
memory_request=spec.resources.get("memory", "8Gi"),
replicas=spec.replicas,
namespace=result["namespace"],
created_by=request.user.username
)
# 6. 更新实验状态
experiment.mark_started()
# 7. 返回结果
output_serializer = TrainingJobSerializer(job)
return Response(
{
"code": 0,
"message": "任务提交成功",
"data": output_serializer.data
},
status=status.HTTP_201_CREATED
)
@action(detail=True, methods=['get'])
def status(self, request, pk=None):
"""
获取任务状态
GET /api/v1/training/jobs/{id}/status
"""
job = self.get_object()
# 从K8s获取最新状态
scheduler = JobScheduler()
try:
k8s_status = scheduler.get_job_status(job.job_name)
except Exception as e:
logger.warning(f"获取K8s状态失败: {e}")
k8s_status = {"state": "unknown"}
# 更新数据库状态
if k8s_status.get("state"):
status_map = {
"running": TrainingJob.Status.RUNNING,
"succeeded": TrainingJob.Status.SUCCEEDED,
"failed": TrainingJob.Status.FAILED,
"pending": TrainingJob.Status.PENDING,
"not_found": TrainingJob.Status.FAILED
}
new_status = status_map.get(k8s_status["state"])
if new_status and job.status != new_status:
job.status = new_status
if new_status in [TrainingJob.Status.SUCCEEDED, TrainingJob.Status.FAILED]:
job.finished_at = timezone.now()
job.save(update_fields=['status', 'finished_at', 'updated_at'])
return Response({
"code": 0,
"data": {
"job_id": str(job.id),
"job_name": job.job_name,
"status": job.status,
"k8s_status": k8s_status
}
})
@action(detail=True, methods=['get'])
def logs(self, request, pk=None):
"""
获取训练日志
GET /api/v1/training/jobs/{id}/logs?tail=100
"""
job = self.get_object()
tail_lines = int(request.query_params.get('tail', 100))
scheduler = JobScheduler()
try:
logs = scheduler.get_job_logs(job.job_name, tail_lines=tail_lines)
except Exception as e:
return Response(
{"error": f"获取日志失败: {str(e)}"},
status=status.HTTP_500_INTERNAL_SERVER_ERROR
)
return Response({
"code": 0,
"data": logs
})
@action(detail=True, methods=['post'])
def stop(self, request, pk=None):
"""
停止训练任务
POST /api/v1/training/jobs/{id}/stop
"""
job = self.get_object()
if job.status in [TrainingJob.Status.SUCCEEDED, TrainingJob.Status.FAILED]:
return Response(
{"error": "任务已完成，无法停止"},
status=status.HTTP_400_BAD_REQUEST
)
scheduler = JobScheduler()
try:
scheduler.cancel_job(job.job_name)
except Exception as e:
return Response(
{"error": f"停止任务失败: {str(e)}"},
status=status.HTTP_500_INTERNAL_SERVER_ERROR
)
# 更新状态
job.status = TrainingJob.Status.CANCELLED
job.finished_at = timezone.now()
job.save(update_fields=['status', 'finished_at', 'updated_at'])
# 更新实验
job.experiment.mark_cancelled("用户手动停止")
return Response({
"code": 0,
"message": "任务已停止"
})
@action(detail=True, methods=['post'])
def restart(self, request, pk=None):
"""
重新提交任务（基于之前的配置）
POST /api/v1/training/jobs/{id}/restart
"""
job = self.get_object()
# 创建新任务
serializer = TrainingJobCreateSerializer(data={
"experiment_id": str(job.experiment_id),
"image": job.image,
"command": job.command,
"env_vars": job.env_vars,
"resources": {
"cpu": job.cpu_request,
"memory": job.memory_request,
"gpu": job.gpu_request
},
"replicas": job.replicas
})
serializer.is_valid(raise_exception=True)
# 复用创建逻辑
request._full_data = serializer.validated_data
return self.create(request)
class TrainingMetricsView(APIView):
"""
训练指标 API
提供实时指标查询
"""
permission_classes = [IsAuthenticated]
def get(self, request, job_id):
"""
获取训练指标
GET /api/v1/training/jobs/{id}/metrics
"""
try:
job = TrainingJob.objects.get(id=job_id)
except TrainingJob.DoesNotExist:
return Response(
{"error": "任务不存在"},
status=status.HTTP_404_NOT_FOUND
)
# 从Redis/时序数据库获取指标
# 这里简化处理
metrics = {
"job_id": str(job.id),
"job_name": job.job_name,
"status": job.status,
"latest_metrics": job.latest_metrics or {},
"resources": {
"gpu_utilization": [], # 从Prometheus获取
"memory_usage": [],
"cpu_usage": [],
"step_time": []
},
"timestamps": [] # 对应的时间戳
}
return Response({
"code": 0,
"data": metrics
})七、总结与展望
7.1 架构亮点
7.2 踩过的坑
过早优化：初期引入了过多组件，导致维护成本高
权限过粗：初期只有项目级权限，后改为资源级
存储规划不足：早期没有考虑数据增长，频繁扩容
API版本管理：RESTful API没有版本控制，后续兼容困难
7.3 未来规划
平台智能化
AutoML集成
超参数自动优化
资源智能调度
开发者体验
JupyterLab集成
VSCode插件
本地调试能力
成本控制
资源利用率分析
成本分摊报表
闲时资源调度
📚 相关资源
架构设计文档：[内部Wiki链接]
技术交流群：AI Platform Team
版本：v1.0.0
评论区