https://github.com/data-infra/cube-studio
Cube Studio 是一个开源的一站式云原生机器学习/深度学习/大模型训练与推理平台。
主要特点:
1. 功能模块
- 数据管理:支持数据集管理、特征工程
- 模型开发:Jupyter Notebook、VSCode 在线开发环境
- 模型训练:支持分布式训练(TensorFlow、PyTorch、MXNet 等)
- 模型推理:支持模型部署、AB测试、流量管理
- Pipeline 编排:基于 Argo Workflow 的 DAG 任务编排
2. 技术架构
- 基于 Kubernetes 云原生架构
- 使用 Kubeflow 组件
- 支持 MLflow 进行实验追踪
- 集成 Prometheus + Grafana 监控
3. 支持的框架
- TensorFlow
- PyTorch
- PaddlePaddle
- MindSpore
- 以及各种大模型训练框架
4. 适用场景
- 企业级 MLOps 平台建设
- AI 中台搭建
- 机器学习团队协作
Cube Studio 架构与技术栈分析
整体架构
Cube Studio 采用云原生微服务架构,主要分为以下几层:
复制代码
┌─────────────────────────────────────────────────────────────────┐
│ 前端层 (Web UI) │
│ React/Vue + Ant Design Pro │
├─────────────────────────────────────────────────────────────────┤
│ 网关层 (Gateway) │
│ Nginx / Kubernetes Ingress │
├─────────────────────────────────────────────────────────────────┤
│ 服务层 (Backend) │
│ Flask/Python + RESTful API + Celery │
├─────────────────────────────────────────────────────────────────┤
│ 调度编排层 (Orchestration) │
│ Argo Workflow + Kubeflow Pipelines + Airflow │
├─────────────────────────────────────────────────────────────────┤
│ 计算引擎层 (Compute) │
│ TFJob / PyTorchJob / Spark / Ray / Volcano / MPI Operator │
├─────────────────────────────────────────────────────────────────┤
│ 基础设施层 (Infrastructure) │
│ Kubernetes + Docker + GPU调度 + 存储(NFS/Ceph) │
└─────────────────────────────────────────────────────────────────┘
核心技术栈
1. 前端技术
| 技术 |
用途 |
| React |
前端框架 |
| Ant Design Pro |
UI 组件库 |
| ECharts |
数据可视化 |
| Monaco Editor |
代码编辑器 |
2. 后端技术
| 技术 |
用途 |
| Python |
主要开发语言 |
| Flask |
Web 框架 |
| SQLAlchemy |
ORM |
| Celery |
异步任务队列 |
| Redis |
缓存、消息队列 |
| MySQL |
元数据存储 |
3. 云原生与容器
| 技术 |
用途 |
| Kubernetes |
容器编排平台 |
| Docker |
容器运行时 |
| Helm |
K8s 包管理 |
| Kubernetes Ingress |
流量入口 |
4. ML/AI 工作流
| 技术 |
用途 |
| Argo Workflow |
DAG 工作流引擎 |
| Kubeflow |
ML 工具套件 |
| MLflow |
实验追踪、模型注册 |
| Jupyter Hub |
在线开发环境 |
5. 分布式训练 Operators
| Operator |
用途 |
| TF-Operator |
TensorFlow 分布式训练 |
| PyTorch-Operator |
PyTorch 分布式训练 |
| MPI-Operator |
MPI 分布式训练 (Horovod) |
| Paddle-Operator |
PaddlePaddle 分布式训练 |
| Spark-Operator |
Spark 任务 |
| Volcano |
批调度、Gang Scheduling |
| Ray |
分布式计算框架 |
6. 模型推理服务
| 技术 |
用途 |
| KFServing / KServe |
模型服务 |
| Triton Inference Server |
高性能推理 |
| TensorFlow Serving |
TF 模型服务 |
| TorchServe |
PyTorch 模型服务 |
| Seldon Core |
模型部署 |
7. 监控与日志
| 技术 |
用途 |
| Prometheus |
指标采集 |
| Grafana |
监控可视化 |
| ELK Stack |
日志收集分析 |
| Jaeger |
链路追踪 |
8. 存储
| 技术 |
用途 |
| NFS |
共享文件存储 |
| Ceph |
分布式存储 |
| MinIO |
对象存储 (S3兼容) |
| HDFS |
大数据存储 |
核心模块架构
复制代码
cube-studio/
├── myapp/ # 后端主应用
│ ├── views/ # API 视图层
│ ├── models/ # 数据模型
│ ├── tasks/ # Celery 异步任务
│ └── utils/ # 工具类
├── install/ # 部署脚本
│ ├── kubernetes/ # K8s 部署配置
│ └── docker/ # Docker 镜像构建
├── job-template/ # 任务模板
│ ├── tf-distributed/ # TensorFlow 分布式
│ ├── pytorch-distributed/ # PyTorch 分布式
│ ├── spark/ # Spark 任务
│ └── ...
├── images/ # 各组件 Docker 镜像
└── frontend/ # 前端代码
关键能力
Pipeline 工作流
- 基于 Argo Workflow 实现 DAG 编排
- 支持拖拽式可视化编排
- 任务依赖、并行、条件分支
GPU 调度
- 支持 NVIDIA GPU 调度
- 支持 GPU 共享 (GPU Share)
- 集成 Volcano 实现 Gang Scheduling
多租户
- 基于 Kubernetes Namespace 隔离
- RBAC 权限控制
- 资源配额管理
Cube Studio Pipeline 编排功能详解
什么是 Pipeline 编排?
Pipeline 编排是将机器学习的各个步骤(数据处理、特征工程、模型训练、模型评估、模型部署等)串联成一个有向无环图(DAG),实现自动化、可重复、可追溯的端到端工作流。
复制代码
┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐
│ 数据拉取 │───▶│ 数据清洗 │───▶│ 特征工程 │───▶│ 模型训练 │───▶│ 模型部署 │
└──────────┘ └──────────┘ └──────────┘ └──────────┘ └──────────┘
│
▼
┌──────────┐
│ 模型评估 │
└──────────┘
Cube Studio Pipeline 的核心特点
1. 可视化拖拽编排
- 通过 Web UI 拖拽组件构建 Pipeline
- 无需编写复杂的 YAML 配置
- 直观展示任务依赖关系
2. 丰富的任务模板
预置了大量开箱即用的任务模板:
| 类别 |
模板示例 |
| 数据处理 |
SQL 查询、Spark ETL、数据校验 |
| 特征工程 |
特征提取、特征存储 |
| 模型训练 |
TensorFlow、PyTorch、XGBoost 分布式训练 |
| 模型评估 |
指标计算、A/B 测试 |
| 模型部署 |
KServe、Triton 部署 |
| 通用任务 |
Python 脚本、Shell 脚本、HTTP 请求 |
3. 底层技术
- 基于 Argo Workflow 引擎
- 每个任务节点运行在独立的 Pod 中
- 支持任务重试、超时、条件分支
实际案例:推荐系统日更新 Pipeline
业务场景
电商推荐系统需要每天:
- 从数据仓库拉取用户行为数据
- 生成用户和商品特征
- 训练推荐模型
- 评估模型效果
- 模型达标则自动上线
Pipeline DAG 设计
复制代码
┌─────────────────┐
│ 定时触发 │
│ (每天凌晨2点) │
└────────┬────────┘
│
▼
┌─────────────────┐
│ 数据拉取 │
│ (Hive SQL) │
└────────┬────────┘
│
┌──────────────┼──────────────┐
▼ ▼ ▼
┌──────────────┐ ┌──────────────┐ ┌──────────────┐
│ 用户特征生成 │ │ 商品特征生成 │ │ 交互特征生成 │
│ (Spark) │ │ (Spark) │ │ (Spark) │
└──────┬───────┘ └──────┬───────┘ └──────┬───────┘
│ │ │
└────────────────┼────────────────┘
│
▼
┌─────────────────┐
│ 特征合并 │
│ (Spark) │
└────────┬────────┘
│
▼
┌─────────────────┐
│ 模型训练 │
│ (PyTorch 分布式) │
│ 4 GPU 节点 │
└────────┬────────┘
│
▼
┌─────────────────┐
│ 模型评估 │
│ (AUC、召回率) │
└────────┬────────┘
│
┌──────┴──────┐
▼ ▼
┌─────────────┐ ┌─────────────┐
│ AUC > 0.75 │ │ AUC ≤ 0.75 │
│ 模型部署 │ │ 告警通知 │
│ (KServe) │ │ (钉钉) │
└─────────────┘ └─────────────┘
在 Cube Studio 中的配置示例
yaml
复制代码
# Pipeline 定义(简化示例)
pipeline:
name: recommend-daily-update
schedule: "0 2 * * *" # 每天凌晨2点
tasks:
- name: data-extract
template: sql-query
params:
sql: "SELECT * FROM user_behavior WHERE dt='{{ds}}'"
output: /data/raw/{{ds}}/
- name: user-feature
template: spark-job
depends: [data-extract]
params:
script: user_feature.py
resources:
executor: 10
memory: 8Gi
- name: item-feature
template: spark-job
depends: [data-extract]
params:
script: item_feature.py
- name: model-training
template: pytorch-distributed
depends: [user-feature, item-feature]
params:
script: train.py
resources:
workers: 4
gpu: 1
memory: 32Gi
- name: model-evaluate
template: python-script
depends: [model-training]
params:
script: evaluate.py
outputs:
- auc
- name: model-deploy
template: kserve-deploy
depends: [model-evaluate]
condition: "{{tasks.model-evaluate.outputs.auc}} > 0.75"
params:
model_path: /models/recommend/{{ds}}
replicas: 3
- name: alert
template: dingtalk-notify
depends: [model-evaluate]
condition: "{{tasks.model-evaluate.outputs.auc}} <= 0.75"
params:
message: "模型AUC不达标: {{tasks.model-evaluate.outputs.auc}}"
Pipeline 编排的核心价值
1. 自动化 → 降本增效
| 对比项 |
手动方式 |
Pipeline 自动化 |
| 执行方式 |
人工依次执行脚本 |
自动调度运行 |
| 耗时 |
需要人工值守 |
无人值守 |
| 出错处理 |
人工排查、重跑 |
自动重试、断点续跑 |
| 人力成本 |
高 |
低 |
案例量化:
某团队之前每天需要 1 名工程师花 4 小时人工跑批,使用 Pipeline 后完全自动化,每年节省 1400+ 人时。
2. 可重复性 → 实验管理
复制代码
实验1: learning_rate=0.01, batch_size=32 → AUC=0.72
实验2: learning_rate=0.001, batch_size=64 → AUC=0.75
实验3: learning_rate=0.001, batch_size=128 → AUC=0.78 ✓
- 每次 Pipeline 运行的参数、代码版本、数据版本都被记录
- 可随时回溯对比不同实验
- 便于复现最佳结果
3. 可追溯性 → 问题定位
复制代码
Pipeline 运行 #1024 失败
├── data-extract ✅ 成功 (耗时: 5min)
├── user-feature ✅ 成功 (耗时: 20min)
├── item-feature ✅ 成功 (耗时: 15min)
├── model-training ❌ 失败 (耗时: 45min) ← 点击查看日志
│ └── Error: CUDA out of memory
├── model-evaluate ⏸️ 跳过
└── model-deploy ⏸️ 跳过
- 清晰的任务状态可视化
- 每个节点独立日志
- 快速定位失败原因
4. 资源弹性 → 成本优化
复制代码
┌─────────────────────────────────────────────────────────┐
│ 资源使用时间线 │
├─────────────────────────────────────────────────────────┤
│ │
│ Spark 集群 ████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ │
│ (50 CPU) 用完即释放 │
│ │
│ GPU 集群 ░░░░░░░░░░████████████░░░░░░░░░░░░░░░░░ │
│ (4 x V100) 训练时才申请 │
│ │
│ 推理服务 ░░░░░░░░░░░░░░░░░░░░░░████████████████ │
│ (8 CPU) 部署后常驻 │
│ │
└─────────────────────────────────────────────────────────┘
0h 2h 4h 6h 8h
- 按需申请资源,任务完成自动释放
- 避免资源闲置浪费
- GPU 等昂贵资源利用率大幅提升
5. 协作标准化 → 团队效率
传统方式的问题:
复制代码
小A: "你那个训练脚本怎么跑的?"
小B: "先跑 prepare.sh,然后 train.py,参数记得改..."
小A: "我跑出来结果不一样啊"
小B: "你用的哪个版本的数据?"
Pipeline 方式:
复制代码
- 所有步骤在 Pipeline 中标准化定义
- 新成员直接复用现有 Pipeline
- 修改参数即可发起新实验
- 知识沉淀在平台中,不依赖个人
总结
| 价值维度 |
具体收益 |
| 效率提升 |
端到端自动化,人力节省 80%+ |
| 质量保障 |
可重复、可追溯、可审计 |
| 成本优化 |
资源按需使用,利用率提升 |
| 协作提升 |
标准化流程,知识沉淀 |
| 快速迭代 |
从天级到小时级的实验周期 |
Pipeline 编排是 MLOps 的核心能力,Cube Studio 通过可视化 + 模板化的方式大大降低了使用门槛,让算法工程师可以专注于算法本身,而不是繁琐的工程细节。