从 0 到 1 搭建 AI 训练平台：架构设计与技术选型
企业级 AI 训练平台需要解决资源调度、任务管理、监控告警等问题。本文提供完整的技术方案，涵盖架构设计、技术选型和落地实践。
一、需求分析
核心功能
用户角色
| 角色 | 需求 | 权限 |
|---|---|---|
| 管理员 | 资源管理、用户管理 | 全部权限 |
| 算法工程师 | 任务提交、环境配置 | 个人资源 |
| 数据工程师 | 数据管理、预处理 | 数据权限 |
二、架构设计
整体架构
┌─────────────────────────────────────────┐
│ 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 调度
- 丰富的扩展能力
yaml
# 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
- 队列管理完善
yaml
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）
用途：模型、数据集、日志存储
部署：
bash
docker run -p 9000:9000 minio/minio server /data#### 分布式文件系统（NFS/Ceph）
用途：共享数据集
4. 监控方案
#### Prometheus + Grafana
监控指标：
- GPU 利用率、显存使用
- CPU、内存使用
- 网络 IO、磁盘 IO
- 任务状态、队列长度
yaml
# prometheus.yml
scrape_configs:
- job_name: 'gpu-nodes'
static_configs:
- targets: ['node1:9100', 'node2:9100']
metrics_path: /metrics四、核心功能实现
1. 任务提交 API
python
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. 资源调度策略
python
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_node3. 监控数据采集
python
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. 单节点部署（测试）
bash
# 安装 Docker
curl -fsSL https://get.docker.com | sh
# 安装 K3s（轻量级 K8s）
curl -sfL https://get.k3s.io | sh -
# 部署平台
kubectl apply -f deployment.yaml2. 多节点部署（生产）
bash
# 主节点
kubeadm init --pod-network-cidr=10.244.0.0/16
# 工作节点
kubeadm join : --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 |
实施建议
核心观点：AI 训练平台建设是系统工程，需要平衡功能、成本、易用性。建议基于开源生态构建，避免重复造轮子。
_欢迎交流讨论，分享你的实践经验。_