构建云原生机器学习平台的架构与实践总结

摘要

在入职该汽车公司的7个月时间内,我的重点工作职责之一就是与基础设施团队一起研发适合内部使用的机器学习平台,经过几个月的长期奋战,目前该机器学习平台已经初步可以服务内部的算法工程师进行模型训推、模型推理以及模型注册等功能。

我的主要职责是MLOps基础设施架构师后端软件架构师,利用云原生技术栈(特别是Kubernetes、Volcano、PyTorch分布式框架)为算法工程师打造一个服务于模型训练、推理与注册全生命周期的、高效、稳定且易用的平台。


1. 核心挑战:从“AI作坊”到“AI工厂”的转型之路

在平台构建之前,面临的算法模型研发呈现出典型的“作坊式”特征,具体痛点如下:

  • 环境不一致与脆弱 (Inconsistent Environments): 算法工程师需耗费大量精力处理 requirements.txt 冲突和 CUDA/cuDNN 版本地狱,导致“在我机器上能跑”成为常态,严重影响协作与复现。
  • 数据与模型的孤岛 (Data & Model Silos): 训练数据与模型产物散落在个人目录,版本管理混乱,实验结果难以追溯与公平比较。
  • 手动且低效的流程 (Manual & Inefficient Processes): 整个ML生命周期高度依赖手动执行脚本和线下沟通,缺乏自动化与标准化。
  • 资源竞争与浪费 (Resource Contention & Waste): 昂贵的GPU资源通过SSH“抢占”模式使用,缺乏统一调度,导致资源死锁和利用率低下(普遍低于30%)。

因此:核心目标是构建一个抽象化、标准化、自动化的平台,让算法工程师能聚焦于模型算法本身的创新。


2. 平台总体架构:以云原生技术为基石

为了应对上述挑战,我们设计的机器学习平台在技术选型上全面拥抱云原生,其核心架构组件包括:

  • 计算与调度层 (Compute & Scheduling):

    • Kubernetes (K8s): 作为容器编排的基座。
    • Volcano: 作为批处理调度引擎,替代K8s默认调度器,解决AI/大数据工作负载的核心调度难题。
    • NVIDIA Device Plugin: 实现对GPU等异构资源的精细化管理。
  • MLOps工作流与追踪层 (Workflow & Tracking):

    • Argo Workflows: 作为核心的ML流水线执行引擎。
    • MLflow: 作为实验跟踪、模型注册与版本管理的“记录引擎”。
  • 平台抽象与服务层 (Abstraction & Service):

    • 自定义Kubernetes Operator: 平台的核心,负责将用户的简单意图翻译为底层的复杂K8s资源操作。
    • 后端服务 (Go + Gin): 提供稳定的RESTful API,作为平台门户与底层K8s交互的桥梁。
    • 前端门户 (React + Ant Design): 为算法工程师提供图形化的、一站式的操作界面。

3. 核心经验一:作为MLOps基础设施架构师

在这一角色中,我专注于为算法工程师构建稳定、高效且资源利用率最大化的底层基础设施。

3.1 以Volcano实现高级批处理调度

K8s默认调度器无法满足大规模分布式训练的需求,核心痛点在于缺乏“成组调度”(Gang Scheduling)。一个分布式任务的部分Pod启动而另一部分因资源不足失败,会导致已启动的Pod空占资源,造成集群死锁。

解决方案是引入并深度整合Volcano调度器。

  • 成组调度 (Gang Scheduling): Volcano确保一个分布式作业(在Volcano中抽象为JobPodGroup)所需的所有资源(如1个PS、8个Worker)在集群中被一次性满足后,才统一创建所有Pod。这从根本上解决了资源死锁问题,是保障大规模训练任务成功率的关键。

  • 弹性分层队列系统: 为不同业务团队和任务类型(如生产级ADAS训练、算法研究)建立了具有不同权重和资源配额的队列。这实现了:

    • 资源保障: 核心业务的队列拥有更高的资源保障和优先级。
    • 公平性与弹性: 通过DRF(Dominant Resource Fairness)算法和资源“借用”机制,在保障公平的同时,让空闲队列的资源能被繁忙队列临时使用,将集群整体GPU利用率从30%提升至75%
    # 示例:为ADAS训练和算法研究划分不同优先级的队列
    apiVersion: scheduling.volcano.sh/v1beta1
    kind: Queue
    metadata:
      name: adas-training-prod
    spec:
      weight: 100 # 高权重,高优先级
      capability:
        nvidia.com/gpu: "64"
    ---
    apiVersion: scheduling.volcano.sh/v1beta1
    kind: Queue
    metadata:
      name: algo-research-dev
    spec:
      weight: 20 # 低权重
      reclaimable: true # 允许资源被高优队列抢占
      capability:
        nvidia.com/gpu: "16"
    
  • 拓扑感知调度: 在百卡以上的大规模训练中,节点间通信是瓶颈。Volcano的拓扑感知能力,能将需要频繁通信的Pod(如PyTorch DDP中的workers)调度到同一机架或交换机下,极大降低网络延迟,提升训练效率约30%

3.2 PyTorch大规模分布式训练的协调与优化

为了服务算法工程师进行超大规模模型训练,我们不仅提供了调度能力,还在平台层面封装了复杂的分布式训练配置。

  • 硬件与网络基座: 我们定义了标准的计算节点配置,采用胖树网络拓扑200Gb/s InfiniBand高速网络,为上层通信优化提供物理保障。这部分主要是基础设施平台团队提供支持。

  • 通信后端优化 (NCCL): 平台在启动训练任务时,会根据节点硬件和网络环境,自动注入最优的NCCL环境变量。这避免了算法工程师手动进行复杂的网络调优。

    # 平台内部根据环境自动生成的NCCL优化配置
    def setup_nccl_env():
        os.environ['NCCL_SOCKET_IFNAME'] = 'ib0'      # 强制使用InfiniBand
        os.environ['NCCL_IB_DISABLE'] = '0'           # 启用InfiniBand
        os.environ['NCCL_ALGO'] = 'Tree'              # 针对胖树拓扑优化通信算法
        os.environ['NCCL_P2P_DISABLE'] = '0'          # 启用GPU间直接点对点通信
    
  • 训练范式抽象 (FSDP): 针对大模型训练,平台原生支持并自动化配置FSDP (Fully Sharded Data Parallel)。算法工程师只需在界面上选择“FSDP训练模式”,平台即可自动生成如下复杂的配置,并应用到训练任务中。

    # 平台为用户自动生成的FSDP配置片段
    config["fsdp"] = {
        "enable_fsdp": True,
        "sharding_strategy": "HYBRID_SHARD",      # 混合分片策略,兼顾计算和通信效率
        "mixed_precision": "bf16",                # 自动启用bf16混合精度
        "forward_prefetch": True,                 # 启用前向预取,隐藏数据加载延迟
        "auto_wrap_policy": "TRANSFORMER_BASED_WRAP", # 自动包装Transformer层
        "transformer_layer_cls": "LlamaDecoderLayer"
    }
    config["training"]["enable_gradient_checkpointing"] = True # 自动启用梯度检查点,节省显存
    

    通过这种方式,将复杂的分布式训练技术细节对用户透明化,使其可以轻松驾驭百亿甚至千亿参数级别的大模型训练。


4. 核心经验二:作为后端软件架构师

在这一角色中,我专注于平台自身的产品化和工程卓越,核心是设计一套优雅的抽象层,将底层基础设施的复杂性完美封装

4.1 架构核心:基于“CRD + Operator”的声明式API

这是整个平台能够成功的技术基石。其实算法工程师只关心几个核心问题:“代码在哪?用什么环境跑?跑什么命令?需要多少资源?超参数是什么?”

基于此,我设计了CustomTrainJob这个自定义资源(CRD),并为其开发了一个专属的Kubernetes Operator

  • CRD - 简单的用户意图: 算法工程师提交的CustomTrainJob资源非常简洁,只包含他们关心的字段。

  • Operator - 复杂的“翻译官”与“执行者”:

    1. 监听与翻译: Operator监听到CustomTrainJob的创建后,会读取其spec,并使用内置的Go模板,将其“翻译”成一个包含上百行配置、功能完备的Argo Workflow YAML。这个翻译过程会注入前述的FSDP配置、NCCL优化、MLflow集成等所有“最佳实践”。
    2. 创建与关联: Operator创建Argo Workflow,并设置OwnerReference,实现CRD与Workflow的父子绑定,从而做到级联删除和状态同步。
    3. 状态同步: Operator持续监听Argo Workflow的状态,并将Running, Succeeded等信息写回CustomTrainJobstatus字段,让用户能通过简单接口实时了解任务进展。

这个模式实现了彻底的关注点分离:算法工程师面向领域API编程,平台工程师则在Operator中迭代底层技术,双方互不干扰。

4.2 工程实践:可维护、可扩展的后端系统

为了支撑平台的稳定运行和快速迭代,我设计了清晰的工程规范和代码结构。

  • 技术栈: 前端采用 React + TypeScript + Ant Design,后端采用 Go + Gin,API规范遵循 OpenAPI (Swagger)

  • Monorepo与CI/CD: 所有代码在Monorepo中管理,并通过GitHub Actions建立了包含代码检查、单元测试、镜像构建和自动部署的完整CI/CD流程。

  • 后端分层架构: 后端代码严格遵循三层架构(API-Service-Repository),确保高内聚、低耦合。

    /backend
    ├── cmd/server/main.go         # 程序入口:初始化所有服务
    ├── internal/
    │   ├── api/                   # API层 (Gin Handlers),负责HTTP请求响应
    │   │   └── v1/training_handler.go
    │   ├── service/               # 业务逻辑层,处理核心业务
    │   │   └── training_service.go  # 核心逻辑:接收请求,调用Repository创建CRD
    │   └── repository/            # 数据与基础设施交互层
    │       └── k8s/
    │           └── train_job_client.go # 封装client-go与K8s API交互的逻辑
    └── Dockerfile
    

    我编写了repository/k8s/中的核心代码,使用client-go与Kubernetes API Server进行交互,并主导开发了上述的Operator。


5. 统一工作流:为算法工程师提供无缝体验

我设计的双重角色最终在平台统一的工作流中融合,为算法工程师提供了“一键式”的无缝体验:

  1. 用户交互: 算法工程师在Web门户上填写训练任务表单(Git仓库、资源、超参数等)。
  2. API请求: 前端将表单数据发送到后端Go服务。
  3. CRD创建: 后端服务根据请求,在K8s集群中创建一个CustomTrainJob资源。
  4. 自动化编排: 开发的Operator监听到新CRD,自动“翻译”并创建包含Argo WorkflowMLflow集成的复杂工作流。
  5. 执行与追踪: Argo执行训练流程,训练脚本中的指标和模型产物被自动记录到MLflow。
  6. 状态反馈: Operator将Argo Workflow的状态实时同步回CRD,前端轮询该状态并展示给用户。

整个过程对用户完全透明,模型开发部署周期从数月缩短至数周


6. 结论与展望

通过深度融合Kubernetes、Volcano、PyTorch分布式等云原生技术,并设计以“CRD+Operator”为核心的优雅抽象层,我们成功构建了一个高效、稳定、自动化的机器学习平台。作为架构师,我主要为团队提供了云原生的核心设计理念及落地、规范了工程化的平台代码开发流程等,最终实现了为算法工程师服务的核心目标。

未来,规划平台将向以下方向演进:

  • 引入特征商店 (Feature Store): 解决特征共用和线上线下一致性问题。
  • 支持多集群调度: 利用Karmada与Volcano Global,实现跨地域、跨云的统一资源调度。
  • 完善AIOps闭环: 建立更完善的模型性能监控、数据漂移检测和自动化重训练机制。
  • 引入微服务治理与Nacos 的二次开发:实现模型研发过程中的灰度验证与发布及配置的灰度变更。