从 0 到 1 搭建 AI 训练平台：架构设计与技术选型

企业级 AI 训练平台需要解决资源调度、任务管理、监控告警等问题。本文提供完整的技术方案，涵盖架构设计、技术选型和落地实践。

---

一、需求分析

核心功能

1. 资源管理：GPU 资源池化、动态分配
2. 任务调度：训练任务提交、队列管理
3. 环境管理：镜像管理、依赖配置
4. 监控告警：资源监控、任务状态、异常告警
5. 数据管理：数据集管理、版本控制

用户角色

角色	需求	权限
管理员	资源管理、用户管理	全部权限
算法工程师	任务提交、环境配置	个人资源
数据工程师	数据管理、预处理	数据权限

---

二、架构设计

整体架构

┌─────────────────────────────────────────┐
│              Web UI                      │
│    (任务提交、监控、管理界面)            │
└─────────────────┬───────────────────────┘
                  │
┌─────────────────▼───────────────────────┐
│              API Gateway                 │
│    (认证、限流、路由)                    │
└─────────────────┬───────────────────────┘
                  │
┌─────────────────▼───────────────────────┐
│           Control Plane                  │
│  ┌─────────┐  ┌─────────┐  ┌─────────┐  │
│  │Scheduler│  │ResourceManager│ │Monitor│  │
│  └─────────┘  └─────────┘  └─────────┘  │
└─────────────────┬───────────────────────┘
                  │
┌─────────────────▼───────────────────────┐
│            Data Plane                    │
│  ┌─────────┐  ┌─────────┐  ┌─────────┐  │
│  │ Node 1  │  │ Node 2  │  │ Node N  │  │
│  │(GPU×8)  │  │(GPU×8)  │  │(GPU×8)  │  │
│  └─────────┘  └─────────┘  └─────────┘  │
└─────────────────────────────────────────┘

核心组件

组件	职责	技术选型
Web UI	用户界面	React + Ant Design
API Gateway	接口网关	Kong/Nginx
Scheduler	任务调度	Kubernetes + Volcano
ResourceManager	资源管理	Kubernetes
Monitor	监控告警	Prometheus + Grafana
Storage	数据存储	MinIO + MySQL

---

三、技术选型

1. 容器编排

Kubernetes（推荐）

优势：

• 生态完善，社区活跃
• 支持 GPU 调度
• 丰富的扩展能力

关键配置：

# GPU 资源定义
apiVersion: v1
kind: Pod
metadata:
  name: gpu-training
spec:
  containers:
  - name: training
    image: pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
    resources:
      limits:
        nvidia.com/gpu: 4  # 申请 4 张 GPU
      requests:
        nvidia.com/gpu: 4

Docker Swarm（备选）

适用场景：小规模集群（<10 节点）

2. 任务调度

Volcano（推荐）

优势：

• 专为 AI 训练设计
• 支持 Gang Scheduling
• 队列管理完善

配置示例：

apiVersion: batch.volcano.sh/v1alpha1
kind: Job
metadata:
  name: distributed-training
spec:
  schedulerName: volcano
  queue: high-priority
  tasks:
    - replicas: 8
      template:
        spec:
          containers:
          - name: worker
            image: pytorch/pytorch:2.0.1
            command: ["python", "train.py"]

KubeFlow

适用场景：需要完整 ML 流水线

3. 存储方案

对象存储（MinIO）

用途：模型、数据集、日志存储

部署：

docker run -p 9000:9000 minio/minio server /data

分布式文件系统（NFS/Ceph）

用途：共享数据集

4. 监控方案

Prometheus + Grafana

监控指标：

• GPU 利用率、显存使用
• CPU、内存使用
• 网络 IO、磁盘 IO
• 任务状态、队列长度

配置示例：

# prometheus.yml
scrape_configs:
  - job_name: 'gpu-nodes'
    static_configs:
      - targets: ['node1:9100', 'node2:9100']
    metrics_path: /metrics

---

四、核心功能实现

1. 任务提交 API

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import kubernetes

app = FastAPI()

class TrainingJob(BaseModel):
    name: str
    image: str
    gpu_count: int
    command: list[str]
    env: dict = {}

@app.post("/api/jobs")
async def create_job(job: TrainingJob):
    # 创建 Kubernetes Job
    job_manifest = {
        "apiVersion": "batch/v1",
        "kind": "Job",
        "metadata": {"name": job.name},
        "spec": {
            "template": {
                "spec": {
                    "containers": [{
                        "name": "training",
                        "image": job.image,
                        "command": job.command,
                        "resources": {
                            "limits": {"nvidia.com/gpu": job.gpu_count}
                        },
                        "env": [{"name": k, "value": v} for k, v in job.env.items()]
                    }]
                }
            }
        }
    }
    
    kubernetes.client.BatchV1Api().create_namespaced_job(
        namespace="default",
        body=job_manifest
    )
    
    return {"status": "created", "job_name": job.name}

2. 资源调度策略

class GPUScheduler:
    def __init__(self):
        self.node_resources = {}  # 节点资源信息
    
    def schedule(self, gpu_count: int, memory_gb: int) -> str:
        """选择最优节点"""
        candidates = []
        
        for node, resources in self.node_resources.items():
            if (resources['available_gpus'] >= gpu_count and 
                resources['available_memory'] >= memory_gb):
                candidates.append(node)
        
        if not candidates:
            raise Exception("资源不足")
        
        # 选择资源最紧张的节点（碎片整理）
        best_node = min(
            candidates,
            key=lambda n: self.node_resources[n]['available_gpus']
        )
        
        return best_node

3. 监控数据采集

from prometheus_client import Counter, Gauge, start_http_server
import pynvml
import time

# 定义指标
gpu_utilization = Gauge('gpu_utilization_percent', 'GPU 利用率', ['node', 'gpu_id'])
gpu_memory = Gauge('gpu_memory_used_bytes', 'GPU 显存使用', ['node', 'gpu_id'])
training_jobs = Counter('training_jobs_total', '训练任务总数', ['status'])

def collect_gpu_metrics():
    pynvml.nvmlInit()
    device_count = pynvml.nvmlDeviceGetCount()
    
    for i in range(device_count):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        
        # GPU 利用率
        util = pynvml.nvmlDeviceGetUtilizationRates(handle)
        gpu_utilization.labels(node='node1', gpu_id=i).set(util.gpu)
        
        # 显存使用
        memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
        gpu_memory.labels(node='node1', gpu_id=i).set(memory.used)
    
    pynvml.nvmlShutdown()

# 启动监控
start_http_server(8000)
while True:
    collect_gpu_metrics()
    time.sleep(15)

---

五、部署方案

1. 单节点部署（测试）

# 安装 Docker
curl -fsSL https://get.docker.com | sh

# 安装 K3s（轻量级 K8s）
curl -sfL https://get.k3s.io | sh -

# 部署平台
kubectl apply -f deployment.yaml

2. 多节点部署（生产）

# 主节点
kubeadm init --pod-network-cidr=10.244.0.0/16

# 工作节点
kubeadm join <master-ip>:<port> --token <token>

# 安装网络插件
kubectl apply -f https://raw.githubusercontent.com/flannel-io/flannel/master/Documentation/kube-flannel.yml

# 安装 GPU 驱动
kubectl apply -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/v0.14.0/nvidia-device-plugin.yml

---

六、成本估算

硬件成本（8 卡服务器×5 台）

项目	单价	数量	总价
GPU 服务器	80 万	5	400 万
网络设备	20 万	1	20 万
存储设备	30 万	1	30 万
合计	-	-	450 万

软件成本

项目	费用
Kubernetes	免费
监控系统	免费
定制开发	50-100 万

运维成本

项目	年费用
电费	50 万
机房	30 万
运维人员	100 万
合计	180 万/年

---

七、最佳实践

1. 资源隔离

• 使用 Namespace 隔离不同团队
• 设置 ResourceQuota 限制资源使用
• 使用 NetworkPolicy 控制网络访问

2. 高可用

• Control Plane 多副本部署
• 使用负载均衡
• 定期备份 etcd 数据

3. 安全加固

• 启用 RBAC 权限控制
• 使用私有镜像仓库
• 定期更新安全补丁

---

八、总结

技术栈总览

层级	技术
前端	React + Ant Design
后端	FastAPI + Kubernetes
调度	Volcano
存储	MinIO + MySQL
监控	Prometheus + Grafana
部署	K3s/Kubernetes

实施建议

1. 从小规模开始：先搭建单节点测试环境
2. 逐步扩展：根据需求增加节点
3. 重视监控：上线前完善监控体系
4. 文档先行：编写完善的用户文档

---

核心观点：AI 训练平台建设是系统工程，需要平衡功能、成本、易用性。建议基于开源生态构建，避免重复造轮子。

欢迎交流讨论，分享你的实践经验。