Skip to main content

概述​

Kubeflow Trainer提供了三个核心的CRD（Custom Resource Definition）来支持分布式训练任务的定义和执行：

CRD名称作用域主要作用
TrainJob命名空间定义具体的训练任务，包括训练的配置参数、资源需求、数据集和模型初始化等
TrainingRuntime命名空间命名空间级别的训练运行时模板，定义了特定框架（如 PyTorch、DeepSpeed）的执行环境，只能被同一命名空间的TrainJob引用
ClusterTrainingRuntime集群集群级别的训练运行时模板，可以被任何命名空间的 TrainJob引用，适用于跨团队共享的标准运行时配置

本文将详细介绍这三个CRD的完整配置格式，并对每个配置项进行说明。

TrainJob​

TrainJob是用户创建分布式训练任务的核心CRD，它定义了训练任务的所有配置参数。

完整模板示例​

apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  # TrainJob 的名称，必须符合 RFC 1035 DNS 标签格式（小写字母、数字、连字符，最多63个字符）
  name: pytorch-distributed-training
  # TrainJob 所在的命名空间
  namespace: kubeflow
  # 可选的标签
  labels:
    app: training
    framework: pytorch
  # 可选的注解
  annotations:
    description: "PyTorch distributed training example"
spec:
  # 引用的训练运行时配置（必填）
  runtimeRef:
    # 运行时的名称（必填，最小长度为1）
    name: torch-distributed
    # 运行时的 API 组，默认为 trainer.kubeflow.org
    apiGroup: trainer.kubeflow.org
    # 运行时的类型，可选值：TrainingRuntime 或 ClusterTrainingRuntime
    # 默认为 ClusterTrainingRuntime
    kind: ClusterTrainingRuntime

  # 初始化器配置，用于数据集和模型的初始化（可选）
  initializer:
    # 数据集初始化配置
    dataset:
      # 数据集存储的 URI，支持各种存储协议（如 s3://、gs://、pvc:// 等）
      storageUri: s3://my-bucket/datasets/fashion-mnist
      # 环境变量列表，用于数据集初始化容器
      env:
        - name: AWS_REGION
          value: us-west-2
        - name: DATASET_FORMAT
          value: pytorch
      # Secret 引用，包含访问数据集的凭证
      # Secret 必须在 TrainJob 所在的命名空间中创建
      secretRef:
        name: s3-credentials

    # 预训练模型初始化配置
    model:
      # 预训练模型存储的 URI
      storageUri: s3://my-bucket/models/pretrained-resnet
      # 环境变量列表，用于模型初始化容器
      env:
        - name: MODEL_FORMAT
          value: pytorch
        - name: MODEL_VERSION
          value: v1.0
      # Secret 引用，包含访问模型的凭证
      secretRef:
        name: s3-credentials

  # 训练器配置（可选）
  trainer:
    # 训练容器的镜像
    image: pytorch/pytorch:2.7.1-cuda12.8-cudnn9-runtime
    # 容器的启动命令
    command:
      - python
      - /workspace/train.py
    # 容器的参数
    args:
      - --epochs=10
      - --batch-size=64
      - --learning-rate=0.001
    # 环境变量列表
    env:
      - name: NCCL_DEBUG
        value: INFO
      - name: PYTORCH_CUDA_ALLOC_CONF
        value: max_split_size_mb:512
    # 训练节点数量
    numNodes: 4
    # 每个节点的资源配置
    resourcesPerNode:
      requests:
        cpu: "4"
        memory: 16Gi
        nvidia.com/gpu: "1"
      limits:
        cpu: "8"
        memory: 32Gi
        nvidia.com/gpu: "1"
    # 每个节点的进程/worker 数量
    # 对于 PyTorch：可以设置为 auto、cpu、gpu 或整数值
    # 对于 MPI：只能设置为整数值
    numProcPerNode: auto

  # 应用到衍生 JobSet 和 Jobs 的标签（可选）
  # 这些标签会与 TrainingRuntime 中的标签合并
  # 注意：当存在相同键时，TrainJob 中的标签值会覆盖 TrainingRuntime 中的标签值
  labels:
    team: ml-team
    project: image-classification

  # 应用到衍生 JobSet 和 Jobs 的注解（可选）
  # 这些注解会与 TrainingRuntime 中的注解合并
  # 注意：当存在相同键时，TrainJob 中的注解值会覆盖 TrainingRuntime 中的注解值
  annotations:
    owner: john@example.com
    cost-center: "1234"

  # Pod 模板覆盖配置，用于自定义特定 Job 的 Pod 配置（可选）
  podTemplateOverrides:
    # 可以配置多个覆盖规则，后面的规则会覆盖前面的字段值
    - targetJobs:
        # 目标 Job 名称列表
        - name: node
      # Pod 元数据覆盖
      metadata:
        labels:
          custom-label: custom-value
        annotations:
          custom-annotation: custom-value
      # Pod Spec 覆盖
      spec:
        # 服务账号名称
        serviceAccountName: training-sa
        # 节点选择器
        nodeSelector:
          node-type: gpu-node
          gpu-type: nvidia-a100
        # 亲和性配置
        affinity:
          nodeAffinity:
            requiredDuringSchedulingIgnoredDuringExecution:
              nodeSelectorTerms:
                - matchExpressions:
                    - key: gpu-type
                      operator: In
                      values:
                        - nvidia-a100
                        - nvidia-v100
          podAntiAffinity:
            preferredDuringSchedulingIgnoredDuringExecution:
              - weight: 100
                podAffinityTerm:
                  labelSelector:
                    matchExpressions:
                      - key: app
                        operator: In
                        values:
                          - training
                  topologyKey: kubernetes.io/hostname
        # 容忍度配置
        tolerations:
          - key: nvidia.com/gpu
            operator: Exists
            effect: NoSchedule
          - key: training-workload
            operator: Equal
            value: "true"
            effect: NoSchedule
        # 卷配置
        volumes:
          - name: workspace
            persistentVolumeClaim:
              claimName: training-workspace-pvc
          - name: shared-memory
            emptyDir:
              medium: Memory
              sizeLimit: 8Gi
        # 初始化容器覆盖
        initContainers:
          - name: setup
            env:
              - name: SETUP_MODE
                value: distributed
            volumeMounts:
              - name: workspace
                mountPath: /workspace
        # 容器覆盖
        containers:
          - name: node
            env:
              - name: CUSTOM_ENV
                value: custom-value
            volumeMounts:
              - name: workspace
                mountPath: /workspace
              - name: shared-memory
                mountPath: /dev/shm
        # 调度门控（用于 Pod 调度就绪性控制）
        schedulingGates:
          - name: custom-gate
        # 镜像拉取 Secret
        imagePullSecrets:
          - name: docker-registry-secret

  # 是否暂停运行中的 TrainJob（可选）
  # 默认为 false
  suspend: false

  # 管理者标识，指示由哪个控制器或实体管理此 TrainJob（可选）
  # 可选值：
  # - trainer.kubeflow.org/trainjob-controller（默认，由内置控制器管理）
  # - kueue.x-k8s.io/multikueue（委托给 Kueue 管理）
  # 此字段不可变
  managedBy: trainer.kubeflow.org/trainjob-controller

重要字段说明​

runtimeRef​

runtimeRef 是必填字段，用于引用训练运行时配置：

字段类型必填默认值说明
namestring是-运行时名称，最小长度为1
apiGroupstring否trainer.kubeflow.org运行时的API组
kindstring否ClusterTrainingRuntime运行时类型，可选 TrainingRuntime 或 ClusterTrainingRuntime

numProcPerNode​

每个节点的进程/worker数量配置：

框架支持的值说明
PyTorchauto、cpu、gpu、整数auto表示自动检测，cpu/gpu表示使用所有 CPU/GPU，整数表示具体进程数
MPI整数只能设置具体的整数值
DeepSpeed整数只能设置具体的整数值

managedBy​

控制TrainJob的管理者：

值说明
trainer.kubeflow.org/trainjob-controller由内置的 TrainJob 控制器管理（默认）
kueue.x-k8s.io/multikueue委托给 Kueue 进行管理，支持多集群调度

initializer（初始化器）​

初始化器用于在训练开始前自动完成数据集和预训练模型的准备工作，极大简化了训练任务的配置和管理。

dataset（数据集初始化器）​

数据集初始化器的作用：

  • 自动下载数据集：从各种存储源（S3、GCS、Hugging Face、PVC 等）自动下载训练所需的数据集
  • 数据预处理：可以在初始化阶段对数据进行预处理和格式转换
  • 数据持久化：将下载的数据集存储到共享卷中，供所有训练节点访问
  • 凭证管理：通过secretRef安全地管理访问私有数据集所需的凭证

配置示例：

initializer:
  dataset:
    # 数据集 URI，支持多种协议：
    # - s3://bucket/path（AWS S3）
    # - gs://bucket/path（Google Cloud Storage）
    # - hf://dataset-name（Hugging Face）
    # - pvc://pvc-name/path（Kubernetes PVC）
    storageUri: s3://my-bucket/datasets/imagenet
    env:
      - name: AWS_REGION
        value: us-west-2
    secretRef:
      name: s3-credentials  # 包含 AWS_ACCESS_KEY_ID 和 AWS_SECRET_ACCESS_KEY

数据集初始化完成后，数据会被挂载到训练容器的/workspace/dataset路径。

model（模型初始化器）​

模型初始化器的作用：

  • 自动下载预训练模型：从模型仓库下载预训练模型权重
  • 模型格式转换：支持不同框架间的模型格式转换
  • 模型持久化：将模型文件存储到共享卷中
  • 版本管理：可以通过环境变量指定具体的模型版本

配置示例：

initializer:
  model:
    # 预训练模型 URI
    storageUri: hf://meta-llama/Llama-2-7b-hf
    env:
      - name: MODEL_VERSION
        value: v1.0
      - name: HUGGING_FACE_TOKEN
        valueFrom:
          secretKeyRef:
            name: hf-token
            key: token
    secretRef:
      name: hf-credentials

模型初始化完成后，模型文件会被挂载到训练容器的/workspace/model路径。

初始化器的执行流程​
  1. dataset-initializer Job首先启动，下载并准备数据集
  2. model-initializer Job随后启动，下载并准备预训练模型
  3. 两个初始化Job都成功完成后，trainer Job才会启动
  4. 训练容器可以直接访问/workspace/dataset和/workspace/model路径下的数据和模型

这种机制确保了训练任务启动时所有必需的资源都已就绪，避免了训练过程中的资源加载失败。

Status Conditions​

TypeReason说明
SuspendedSuspendedTrainJob 被暂停
SuspendedResumedTrainJob 从暂停状态恢复
FailedTrainingRuntimeNotSupported引用的TrainingRuntime不受支持
CompleteJobsCompletedTrainJob 成功完成

TrainingRuntime​

TrainingRuntime是命名空间级别的训练运行时模板，定义了特定ML框架的执行环境。它只能被同一命名空间中的TrainJob引用。

完整模板示例​

apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainingRuntime
metadata:
  # TrainingRuntime 的名称
  name: pytorch-distributed
  # TrainingRuntime 所在的命名空间
  namespace: kubeflow
  # 标签，通常标识框架类型
  labels:
    trainer.kubeflow.org/framework: pytorch
    environment: production
spec:
  # ML 策略配置，提供ML特定的参数（可选）
  mlPolicy:
    # 训练节点数量，默认为1
    numNodes: 2
    
    # PyTorch运行时配置
    torch:
      # 每个节点的进程数量
      # 支持的值：auto（自动检测）、cpu（CPU数量）、gpu（GPU数量）或整数值
      # 默认为 auto
      numProcPerNode: auto
      
      # PyTorch 弹性训练策略（可选）
      # 注意：如果配置了 elasticPolicy，则不能同时设置 mlPolicy.numNodes
      elasticPolicy:
        # 训练任务可以重启的最大次数
        # 此值会插入到 torchrun 的 --max-restarts 参数中
        maxRestarts: 3
        # 可以缩减到的最小节点数
        minNodes: 1
        # 可以扩展到的最大节点数
        maxNodes: 4
        # 用于计算期望节点数的指标
        # 将创建 HPA 来执行自动缩放
        metrics:
          - type: Resource
            resource:
              name: cpu
              target:
                type: Utilization
                averageUtilization: 80

  # PodGroup 策略配置，用于启用 gang-scheduling（可选）
  podGroupPolicy:
    # 使用 Kubernetes scheduler-plugins 的 coscheduling 插件
    coscheduling:
      # gang-scheduling 的最大调度超时时间（秒）
      # 如果为 0，则使用默认值
      # 默认为 60 秒
      scheduleTimeoutSeconds: 120
    
    # 或者使用 Volcano gang-scheduler
    # 注意：coscheduling 和 volcano 只能配置其中一个
    # volcano:
    #   # 网络拓扑配置，与网络拓扑特性和 hyperNode CRD 配合使用
    #   networkTopology:
    #     # 配置具体的网络拓扑策略
    #     # 详见 Volcano 文档

  # JobSet 模板配置
  template:
    # JobSet 的元数据
    metadata:
      labels:
        runtime: pytorch
        version: v2.7
      annotations:
        description: PyTorch distributed training runtime
    
    # JobSet 的规范
    spec:
      # 复制的 Job 列表
      replicatedJobs:
        # 主要的训练节点 Job
        - name: node
          # 该 Job 的副本数量（节点数量）
          # 注意：实际的副本数会被 mlPolicy.numNodes 或 TrainJob.trainer.numNodes 覆盖
          replicas: 1
          
          # Job 模板
          template:
            metadata:
              labels:
                # 此标签标识这是训练器步骤
                trainer.kubeflow.org/trainjob-ancestor-step: trainer
            spec:
              # Job 完成模式
              # Indexed 表示每个 Pod 都有唯一的索引
              completionMode: Indexed
              
              # Job 的 Pod 模板
              template:
                spec:
                  # 添加 initContainers
                  initContainers:
                    # 示例：环境准备
                    - name: setup-env
                      image: busybox:latest
                      command:
                        - sh
                        - -c
                        - |
                          echo "Setting up environment..."
                          mkdir -p /workspace/logs
                          echo "Setup complete"
                      volumeMounts:
                        - name: workspace
                          mountPath: /workspace

                  # 容器列表
                  containers:
                    - name: node
                      # 默认镜像，会被 TrainJob.trainer.image 覆盖
                      image: pytorch/pytorch:2.7.1-cuda12.8-cudnn9-runtime
                      # 默认命令，会被 TrainJob.trainer.command 覆盖
                      command:
                        - torchrun
                      # 默认参数
                      args:
                        - --nnodes=$(TRAINER_NNODES)
                        - --nproc_per_node=$(TRAINER_NPROC_PER_NODE)
                        - --node_rank=$(JOB_COMPLETION_INDEX)
                        - --rdzv_backend=c10d
                        - --rdzv_endpoint=$(TRAINER_NODE_0_HOSTNAME):29400
                        - --rdzv_id=$(TRAINER_JOB_ID)
                        - /workspace/train.py
                      # 环境变量
                      env:
                        - name: LOGLEVEL
                          value: INFO
                        - name: NCCL_DEBUG
                          value: INFO
                      # 资源配置
                      resources:
                        requests:
                          cpu: "2"
                          memory: 8Gi
                        limits:
                          cpu: "4"
                          memory: 16Gi
                      # 卷挂载
                      volumeMounts:
                        - name: workspace
                          mountPath: /workspace
                  
                  # 卷定义
                  volumes:
                    - name: workspace
                      emptyDir: {}
                  
                  # 重启策略
                  restartPolicy: OnFailure
      
      # 失败策略
      failurePolicy:
        # 允许的最大重启次数
        maxRestarts: 3
      
      # 成功策略
      successPolicy:
        # 操作符，表示如何判断成功
        # All 表示所有目标 Job 都必须成功
        operator: All
        # 目标 Job 列表
        targetReplicatedJobs:
          - node
      
      # 网络配置
      network:
        # 是否发布未就绪的地址
        # 对于分布式训练，通常设置为 true
        publishNotReadyAddresses: true
        # 是否启用 DNS 主机名
        enableDNSHostnames: true

重要字段说明​

mlPolicy​

ML策略配置的互斥规则：

配置互斥规则说明
numNodes不能与 torch.elasticPolicy 同时使用弹性训练使用 minNodes/maxNodes 代替固定节点数
torch不能与 mpi 同时使用只能配置一种运行时策略
mpi不能与 torch 同时使用只能配置一种运行时策略

PyTorch 弹性训练​

启用PyTorch弹性训练时的关键配置：

字段说明
maxRestarts插入到 torchrun 的 --max-restarts 参数和Job的 .spec.failurePolicy.maxRestarts
minNodes插入到 torchrun 的 --nnodes 参数的最小值
maxNodes插入到 torchrun 的 --nnodes 参数的最大值
metrics用于创建HPA（Horizontal Pod Autoscaler）进行自动缩放

MPI 实现类型​

支持的MPI实现：

类型说明
OpenMPIOpen MPI 实现（默认）
IntelIntel MPI 实现
MPICHMPICH 实现

特殊标签和注解​

Kubeflow Trainer使用一些特殊的标签和注解来控制训练任务的行为和标识资源关系。

系统预留标签​

标签键可能的值用途应用位置
trainer.kubeflow.org/trainjob-ancestor-stepdataset-initializer
model-initializer
trainer		标识Pod模板在训练流程中的角色，用于关联TrainJob的不同配置项与Runtime中的Job模板TrainingRuntime/ClusterTrainingRuntime的replicatedJobs[].template.labels
trainer.kubeflow.org/frameworktorch
deepspeed
mpi
mlx
torchtune		标识 Runtime 支持的训练框架类型TrainingRuntime/ClusterTrainingRuntime 的 metadata.labels
trainer.kubeflow.org/supportdeprecated标识Runtime的支持状态，当值为 deprecated 时会在创建 TrainJob时发出警告TrainingRuntime/ClusterTrainingRuntime 的 metadata.labels

标签使用说明​

trainer.kubeflow.org/trainjob-ancestor-step​

这是最重要的系统标签，用于建立TrainJob配置与Runtime中Job模板之间的映射关系：

  • dataset-initializer：带有此标签的Job模板会被TrainJob.spec.initializer.dataset配置覆盖
  • model-initializer：带有此标签的Job模板会被TrainJob.spec.initializer.model配置覆盖
  • trainer：带有此标签的Job模板会被TrainJob.spec.trainer配置覆盖

是否必须配置：

该标签在TrainingRuntime/ClusterTrainingRuntime的replicatedJobs中不是强制要求的，但通常应该要配置，但具有以下重要作用：

  1. 不配置该标签：

    • 该Job模板将不会被TrainJob的任何配置项覆盖
    • 适用于辅助性的Job（如日志收集、监控等），这些Job不需要TrainJob级别的定制
    • Job将完全按照Runtime中定义的配置运行

  2. 配置 trainer 值：

    • 强烈推荐为主训练节点配置此标签
    • 允许用户通过TrainJob.spec.trainer定制训练镜像、命令、参数、资源等
    • 提供了训练任务级别的灵活性，是Kubeflow Trainer的核心功能之一
    • 不配置此标签意味着放弃了TrainJob级别的训练配置能力

  3. 配置初始化器标签：

    • 如果需要数据集/模型初始化功能，必须为对应的Job配置相应的标签
    • 不配置则无法使用TrainJob.spec.initializer功能

最佳实践示例：

# 推荐的配置方式
spec:
  template:
    spec:
      replicatedJobs:
        # 可选：数据集初始化 Job
        - name: dataset-initializer
          template:
            metadata:
              labels:
                trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
        
        # 可选：模型初始化 Job
        - name: model-initializer
          template:
            metadata:
              labels:
                trainer.kubeflow.org/trainjob-ancestor-step: model-initializer
        
        # 强烈推荐：主训练节点必须配置此标签
        - name: node
          template:
            metadata:
              labels:
                trainer.kubeflow.org/trainjob-ancestor-step: trainer  # 推荐配置
        
        # 可选：辅助 Job 可以不配置此标签
        - name: monitoring
          template:
            metadata:
              labels:
                app: monitoring  # 不配置 trainjob-ancestor-step

示例：

# 在 TrainingRuntime 中定义
spec:
  template:
    spec:
      replicatedJobs:
        - name: dataset-initializer
          template:
            metadata:
              labels:
                # 此标签标识这是数据集初始化步骤
                trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
        - name: node
          template:
            metadata:
              labels:
                # 此标签标识这是训练器步骤
                trainer.kubeflow.org/trainjob-ancestor-step: trainer
trainer.kubeflow.org/framework​

用于标识Runtime的框架类型，便于用户筛选和管理。

标签作用：

  1. 资源分类和管理：

    • 帮助平台管理员对不同框架的Runtime进行分类
    • 便于通过标签选择器查询特定框架的Runtime
    • 支持在UI或CLI工具中按框架类型过滤和展示

  2. 用户体验优化：

    • 用户可以通过 kubectl get clustertrainingruntimes -l trainer.kubeflow.org/framework=torch 快速查找PyTorch训练任务模板
    • 在多租户环境中，可以基于框架类型实现RBAC权限控制
    • 便于监控和审计特定框架的使用情况

  3. 文档和自动化：

    • 作为Runtime的元数据，帮助生成文档和使用指南
    • 支持自动化工具根据框架类型进行配置推荐
    • 便于集成到CI/CD流程中进行框架相关的校验

是否必须配置：

该标签不是强制要求的，但强烈推荐配置：

  • ✅ 推荐配置：可以提供更好的用户体验和管理便利性
  • ⚠️ 不配置影响：Runtime功能完全正常，但会失去上述分类和筛选能力
  • 📝 命名约定：建议使用官方支持的框架名称（torch、deepspeed、mpi、mlx、torchtune）

支持的框架类型：

框架值对应框架说明
torchPyTorch标准 PyTorch 分布式训练
deepspeedDeepSpeedMicrosoft DeepSpeed 框架
mpiMPI基于 MPI 的分布式训练
mlxMLXApple MLX 框架
torchtuneTorchTunePyTorch 微调框架

使用示例：

# 示例1：PyTorch Runtime
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-distributed
  labels:
    trainer.kubeflow.org/framework: torch  # 标识为 PyTorch 运行时

# 示例2：DeepSpeed Runtime
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  name: deepspeed-distributed
  labels:
    trainer.kubeflow.org/framework: deepspeed  # 标识为 DeepSpeed 运行时

# 示例3：查询特定框架的 Runtime
# kubectl get clustertrainingruntimes -l trainer.kubeflow.org/framework=torch

最佳实践：

  1. 为所有生产环境的Runtime配置此标签
  2. 保持标签值的一致性，避免使用自定义值
  3. 结合其他标签（如版本、环境）实现更细粒度的管理：
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-distributed
  labels:
    trainer.kubeflow.org/framework: torch  # 标识这是 PyTorch 运行时
trainer.kubeflow.org/support​

用于标识Runtime的废弃状态：

apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  name: legacy-runtime
  labels:
    trainer.kubeflow.org/support: deprecated  # 标识此运行时已废弃

当引用带有 deprecated 标签的Runtime时，TrainJob创建时会收到警告，提示用户该Runtime将在未来版本中移除。详见运行时废弃策略。

用户自定义标签和注解​

除了系统预留标签外，用户可以在以下位置添加自定义标签和注解：

  1. TrainJob.metadata.labels/annotations：应用TrainJob资源本身
  2. TrainJob.spec.labels/annotations：应用于衍生的JobSet和Jobs，会与Runtime中的标签/注解合并
  3. TrainJob.spec.podTemplateOverrides[].metadata.labels/annotations：应用于特定Job的Pod模板

标签和注解的合并规则：

  • TrainJob中定义的标签/注解会与Runtime中的标签/注解合并
  • 当存在相同的键时，TrainJob中的值会覆盖Runtime中的值
  • 这种设计允许用户在不修改Runtime的情况下为特定训练任务自定义标签和注解

ClusterTrainingRuntime​

ClusterTrainingRuntime是集群级别的训练运行时模板，可以被任何命名空间中的TrainJob引用。其配置格式与TrainingRuntime完全相同，唯一区别是资源作用域。

完整模板示例​

apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  # ClusterTrainingRuntime 的名称（集群唯一）
  name: torch-distributed
  # 标签，标识框架类型
  labels:
    trainer.kubeflow.org/framework: torch
    version: v2.7
spec:
  # ML 策略配置（与 TrainingRuntime 相同）
  mlPolicy:
    numNodes: 1
    torch:
      numProcPerNode: auto

  # PodGroup 策略配置（可选）
  podGroupPolicy:
    # Volcano gang-scheduler 配置示例
    volcano:
      # 网络拓扑配置，用于网络感知调度
      networkTopology:
        # 网络拓扑键，用于标识网络域
        # 例如：topology.kubernetes.io/zone
        topologyKey: topology.kubernetes.io/zone
        # 网络策略，定义 Pod 应该如何分布
        # 可选值：Best-effort（尽力而为）、Strict（严格模式）
        # Best-effort：尽量将 Pod 调度到同一网络域，但不强制
        # Strict：必须将所有 Pod 调度到同一网络域，否则调度失败
        networkPolicy: Best-effort

  # JobSet 模板配置
  template:
    metadata:
      labels:
        runtime-type: cluster-wide
        framework: pytorch
    
    spec:
      # 复制的 Job 列表
      replicatedJobs:
        - name: node
          replicas: 1
          template:
            metadata:
              labels:
                # 标识这是训练器步骤
                trainer.kubeflow.org/trainjob-ancestor-step: trainer
            spec:
              # Job 的 Pod 模板
              template:
                spec:
                  containers:
                    - name: node
                      # 默认 PyTorch 镜像
                      image: pytorch/pytorch:2.7.1-cuda12.8-cudnn9-runtime
                      # torchrun 命令配置
                      command:
                        - torchrun
                      args:
                        - --nnodes=$(TRAINER_NNODES)
                        - --nproc_per_node=$(TRAINER_NPROC_PER_NODE)
                        - --node_rank=$(JOB_COMPLETION_INDEX)
                        - --rdzv_backend=c10d
                        - --rdzv_endpoint=$(TRAINER_NODE_0_HOSTNAME):29400
                        - --rdzv_id=$(TRAINER_JOB_ID)
                      # 环境变量
                      env:
                        - name: LOGLEVEL
                          value: INFO
                        - name: PYTHONUNBUFFERED
                          value: "1"
                        - name: NCCL_DEBUG
                          value: INFO
                        - name: TORCH_DISTRIBUTED_DEBUG
                          value: INFO
                      # 资源请求和限制
                      resources:
                        requests:
                          cpu: "4"
                          memory: 16Gi
                        limits:
                          cpu: "8"
                          memory: 32Gi
                  
                  # 重启策略
                  restartPolicy: OnFailure
      
      # 失败策略
      failurePolicy:
        maxRestarts: 3
      
      # 成功策略
      successPolicy:
        operator: All
        targetReplicatedJobs:
          - node
      
      # 网络配置
      network:
        publishNotReadyAddresses: true
        enableDNSHostnames: true

集群级数据初始化运行时示例​

apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-distributed-with-initializer
  labels:
    trainer.kubeflow.org/framework: torch
spec:
  mlPolicy:
    numNodes: 1
    torch:
      numProcPerNode: auto

  template:
    spec:
      replicatedJobs:
        # 数据集初始化 Job
        - name: dataset-initializer
          replicas: 1
          template:
            metadata:
              labels:
                trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
            spec:
              template:
                spec:
                  containers:
                    - name: dataset-initializer
                      # 数据集初始化镜像
                      image: ghcr.io/kubeflow/trainer/storage-initializer
                      # 环境变量会被 TrainJob.initializer.dataset.env 合并
                      env:
                        - name: STORAGE_TYPE
                          value: s3
                      # 卷挂载
                      volumeMounts:
                        - name: dataset
                          mountPath: /mnt/dataset
                  
                  # 卷定义
                  volumes:
                    - name: dataset
                      emptyDir: {}
                  
                  restartPolicy: OnFailure
        
        # 模型初始化 Job
        - name: model-initializer
          replicas: 1
          template:
            metadata:
              labels:
                trainer.kubeflow.org/trainjob-ancestor-step: model-initializer
            spec:
              template:
                spec:
                  containers:
                    - name: model-initializer
                      image: ghcr.io/kubeflow/trainer/storage-initializer
                      env:
                        - name: STORAGE_TYPE
                          value: s3
                      volumeMounts:
                        - name: pretrained-model
                          mountPath: /mnt/model
                  
                  volumes:
                    - name: pretrained-model
                      emptyDir: {}
                  
                  restartPolicy: OnFailure
        
        # 训练节点 Job
        - name: node
          replicas: 1
          template:
            metadata:
              labels:
                trainer.kubeflow.org/trainjob-ancestor-step: trainer
            spec:
              template:
                spec:
                  containers:
                    - name: node
                      image: pytorch/pytorch:2.7.1-cuda12.8-cudnn9-runtime
                      command:
                        - torchrun
                      args:
                        - --nnodes=$(TRAINER_NNODES)
                        - --nproc_per_node=$(TRAINER_NPROC_PER_NODE)
                        - --node_rank=$(JOB_COMPLETION_INDEX)
                        - --rdzv_backend=c10d
                        - --rdzv_endpoint=$(TRAINER_NODE_0_HOSTNAME):29400
                        - --rdzv_id=$(TRAINER_JOB_ID)
                        - /workspace/train.py
                      # 挂载初始化的数据集和模型
                      volumeMounts:
                        - name: dataset
                          mountPath: /mnt/dataset
                        - name: pretrained-model
                          mountPath: /mnt/model
                  
                  volumes:
                    - name: dataset
                      emptyDir: {}
                    - name: pretrained-model
                      emptyDir: {}
                  
                  restartPolicy: OnFailure
      
      # 成功策略：只要训练节点成功即可
      successPolicy:
        operator: All
        targetReplicatedJobs:
          - node
      
      network:
        publishNotReadyAddresses: true

TrainingRuntime vs ClusterTrainingRuntime​

两种运行时的对比：

特性TrainingRuntimeClusterTrainingRuntime
作用域命名空间级别集群级别
可见性只能被同一命名空间的TrainJob引用可以被任何命名空间的TrainJob引用
使用场景团队或项目特定的运行时配置跨团队共享的标准运行时配置
权限要求命名空间级别的权限集群级别的权限
配置格式完全相同完全相同

环境变量注入​

Kubeflow Trainer会根据使用的训练框架，为训练容器自动注入相应的环境变量。不同框架使用的环境变量不同。

PyTorch框架环境变量​

使用PyTorch（通过torchrun启动）时，Torch Plugin会自动注入以下环境变量：

PET_*系列（PyTorch Elastic Training）​

这些是torchrun使用的标准环境变量，由Kubeflow Trainer自动注入：

环境变量说明示例值注入来源
PET_NNODES训练节点总数2spec.trainer.numNodes
PET_NPROC_PER_NODE每节点进程数4spec.mlPolicy.torch.numProcPerNode
PET_NODE_RANK当前节点编号0-1JOB_COMPLETION_INDEX
PET_MASTER_ADDR主节点地址myjob-node-0-0.myjob{trainjob-name}-node-0-0.{trainjob-name}
PET_MASTER_PORT主节点端口29400固定值

说明：

  • PET代表PyTorch Elastic Training，是torchrun的标准环境变量前缀
  • 这些变量在TrainingRuntime的command中使用，如：torchrun --nproc_per_node=$(PET_NPROC_PER_NODE) ...
  • PET_NODE_RANK通过Kubernetes的fieldRef机制自动从JOB_COMPLETION_INDEX获取

标准分布式训练环境变量​

torchrun会根据PET_*变量，为每个进程设置以下标准环境变量：

环境变量说明计算方式示例值
WORLD_SIZE总进程数PET_NNODES × PET_NPROC_PER_NODE8
RANK全局进程编号自动计算0-7
LOCAL_RANK本地进程编号节点内进程索引0-3
MASTER_ADDR主节点地址继承PET_MASTER_ADDRmyjob-node-0-0.myjob
MASTER_PORT主节点端口继承PET_MASTER_PORT29400

说明：

  • 这些变量由torchrun自动设置，无需在TrainingRuntime中配置
  • 训练代码中可直接使用torch.distributed.get_rank()等API，或通过os.environ访问

OpenMPI框架环境变量​

使用OpenMPI（如DeepSpeed）时，会注入以下环境变量：

环境变量说明示例值
OMPI_MCA_orte_default_hostfile主机列表文件路径/etc/mpi/hostfile
OMPI_MCA_plm_rsh_agentSSH代理（禁用）/usr/bin/false
OMPI_MCA_orte_keep_fqdn_hostnames保持完整主机名true

说明：

  • 这些变量用于配置OpenMPI运行时行为
  • 通常由ClusterTrainingRuntime的容器模板配置

通用环境变量​

所有框架都可以使用以下Kubernetes原生环境变量：

环境变量说明示例值
JOB_COMPLETION_INDEX节点索引（0开始）0-1

说明：

  • 此变量来自Kubernetes Job的索引机制，表示当前Pod在Job中的序号
  • 可用于获取节点编号或配置节点特定行为

使用示例​

PyTorch训练脚本示例​

import os
import torch.distributed as dist

# 初始化分布式进程组（torchrun会自动设置所需环境变量）
dist.init_process_group(backend="nccl")

# 获取分布式信息
rank = dist.get_rank()           # 等价于 int(os.environ['RANK'])
world_size = dist.get_world_size()  # 等价于 int(os.environ['WORLD_SIZE'])
local_rank = int(os.environ['LOCAL_RANK'])
node_rank = int(os.environ.get('JOB_COMPLETION_INDEX', 0))

print(f"Node Rank: {node_rank}")
print(f"Global Rank: {rank}, World Size: {world_size}")
print(f"Local Rank: {local_rank}")

TrainingRuntime command配置示例​

command:
  - torchrun
  - --nproc_per_node=$(PET_NPROC_PER_NODE) 
  - --nnodes=$(PET_NNODES)
  - --node_rank=$(PET_NODE_RANK)
  - --master_addr=$(PET_MASTER_ADDR)
  - --master_port=$(PET_MASTER_PORT)
  - /workspace/train.py

最佳实践​

  1. 资源配置

    • 为训练任务设置合理的资源请求（requests）和限制（limits）
    • GPU资源通常设置相同的requests和limits
    • 为需要大量内存的操作（如数据加载）预留足够的内存

  2. 存储配置

    • 使用PVC或对象存储（S3、GCS）来持久化数据集和模型
    • 对于需要高I/O性能的场景，考虑使用hostPath或local volume
    • 使用emptyDir（medium: Memory）作为共享内存，避免/dev/shm不足

  3. 网络配置

    • 启用publishNotReadyAddresses: true确保分布式训练的网络连通性
    • 对于大规模分布式训练，考虑使用高速网络（如InfiniBand）
    • 配置适当的网络插件和CNI

  4. 调度策略

    • 使用gang-scheduling（coscheduling或Volcano）确保所有训练节点同时启动
    • 使用节点亲和性和反亲和性控制Pod分布
    • 对于GPU训练，使用节点选择器确保Pod调度到GPU节点

  5. 容错配置

    • 设置合理的maxRestarts值，避免无限重试
    • 使用checkpoint机制定期保存训练状态
    • 配置健康检查（readinessProbe、livenessProbe）

  6. 运行时选择

    • 使用ClusterTrainingRuntime定义组织级别的标准运行时
    • 使用TrainingRuntime为特定团队或项目定制运行时
    • 保持运行时配置的版本化管理

参考资料​