Kubernetes与机器学习工作负载集成最佳实践

引言

随着机器学习和人工智能的快速发展，在Kubernetes上部署和管理ML工作负载已经成为趋势。本文将深入探讨如何在Kubernetes环境中高效运行机器学习训练和推理工作负载。

一、ML工作负载架构

1.1 典型ML工作流

┌─────────────────────────────────────────────────────────────┐
│                   ML工作负载架构                            │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────┐   │
│  │   数据准备   │───▶│   模型训练   │───▶│   模型部署   │   │
│  │   (Data)     │    │   (Train)    │    │   (Serve)    │   │
│  └──────────────┘    └──────────────┘    └──────────────┘   │
│         │                   │                   │           │
│         ▼                   ▼                   ▼           │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────┐   │
│  │  数据存储    │    │  GPU资源     │    │  推理服务    │   │
│  │  (S3/NFS)   │    │  (NVIDIA)    │    │  (Triton)   │   │
│  └──────────────┘    └──────────────┘    └──────────────┘   │
│                                                             │
└─────────────────────────────────────────────────────────────┘

1.2 ML工作负载类型

类型 特点 资源需求
数据预处理 CPU密集型 高CPU、内存
模型训练 GPU密集型 GPU、高内存
模型推理 低延迟 GPU/CPU、网络
超参数调优 并行任务 多GPU
批量预测 高吞吐量 多节点

二、GPU资源管理

2.1 NVIDIA设备插件部署

apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: nvidia-device-plugin-daemonset
  namespace: kube-system
spec:
  selector:
    matchLabels:
      name: nvidia-device-plugin-ds
  template:
    metadata:
      labels:
        name: nvidia-device-plugin-ds
    spec:
      tolerations:
      - key: nvidia.com/gpu
        operator: Exists
        effect: NoSchedule
      containers:
      - name: nvidia-device-plugin
        image: nvidia/k8s-device-plugin:v0.14.0
        securityContext:
          allowPrivilegeEscalation: false
          capabilities:
            drop: ["ALL"]
        volumeMounts:
        - name: device-plugin
          mountPath: /var/lib/kubelet/device-plugins
      volumes:
      - name: device-plugin
        hostPath:
          path: /var/lib/kubelet/device-plugins

2.2 GPU资源请求

apiVersion: v1
kind: Pod
metadata:
  name: ml-training-pod
spec:
  containers:
  - name: trainer
    image: pytorch/pytorch:latest
    command: ["python", "train.py"]
    resources:
      requests:
        cpu: "4"
        memory: "32Gi"
        nvidia.com/gpu: "2"
      limits:
        cpu: "8"
        memory: "64Gi"
        nvidia.com/gpu: "2"
    volumeMounts:
    - name: data
      mountPath: /data
    - name: models
      mountPath: /models
  volumes:
  - name: data
    persistentVolumeClaim:
      claimName: ml-data-pvc
  - name: models
    persistentVolumeClaim:
      claimName: ml-models-pvc

2.3 GPU节点亲和性

apiVersion: apps/v1
kind: Job
metadata:
  name: ml-training-job
spec:
  template:
    spec:
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
            - matchExpressions:
              - key: nvidia.com/gpu.count
                operator: GreaterThanOrEqual
                values:
                - "2"
              - key: nvidia.com/gpu.product
                operator: In
                values:
                - NVIDIA-A100-SXM4-80GB
      containers:
      - name: trainer
        image: pytorch/pytorch:latest
        resources:
          requests:
            nvidia.com/gpu: "2"

三、模型训练工作负载

3.1 Training Job配置

apiVersion: batch/v1
kind: Job
metadata:
  name: mnist-training
spec:
  parallelism: 1
  completions: 1
  template:
    spec:
      restartPolicy: OnFailure
      containers:
      - name: trainer
        image: tensorflow/tensorflow:latest-gpu
        command: ["python", "/app/train.py"]
        args:
        - "--epochs=10"
        - "--batch-size=64"
        - "--learning-rate=0.001"
        resources:
          requests:
            cpu: "4"
            memory: "16Gi"
            nvidia.com/gpu: "1"
        volumeMounts:
        - name: training-data
          mountPath: /data
        - name: output
          mountPath: /output
      volumes:
      - name: training-data
        persistentVolumeClaim:
          claimName: training-data-pvc
      - name: output
        persistentVolumeClaim:
          claimName: model-output-pvc

3.2 分布式训练

apiVersion: kubeflow.org/v1
kind: TFJob
metadata:
  name: distributed-training
spec:
  tfReplicaSpecs:
    Chief:
      replicas: 1
      template:
        spec:
          containers:
          - name: tensorflow
            image: tensorflow/tensorflow:latest-gpu
            command:
            - python
            - /app/train.py
            - --tf-config={{.TF_CONFIG}}
            resources:
              requests:
                cpu: "2"
                memory: "8Gi"
                nvidia.com/gpu: "1"
    Worker:
      replicas: 3
      template:
        spec:
          containers:
          - name: tensorflow
            image: tensorflow/tensorflow:latest-gpu
            command:
            - python
            - /app/train.py
            - --tf-config={{.TF_CONFIG}}
            resources:
              requests:
                cpu: "4"
                memory: "16Gi"
                nvidia.com/gpu: "1"

四、模型推理服务

4.1 Triton Inference Server

apiVersion: v1
kind: Service
metadata:
  name: triton-server
spec:
  selector:
    app: triton-server
  ports:
  - name: http
    port: 8000
  - name: grpc
    port: 8001
  - name: metrics
    port: 8002

---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: triton-server
spec:
  replicas: 2
  selector:
    matchLabels:
      app: triton-server
  template:
    metadata:
      labels:
        app: triton-server
    spec:
      containers:
      - name: triton
        image: nvcr.io/nvidia/tritonserver:23.04-py3
        command: ["tritonserver"]
        args:
        - "--model-repository=/models"
        - "--http-port=8000"
        - "--grpc-port=8001"
        - "--metrics-port=8002"
        resources:
          requests:
            cpu: "2"
            memory: "8Gi"
            nvidia.com/gpu: "1"
        volumeMounts:
        - name: models
          mountPath: /models
      volumes:
      - name: models
        persistentVolumeClaim:
          claimName: models-pvc

4.2 KFServing部署

apiVersion: serving.kubeflow.org/v1beta1
kind: InferenceService
metadata:
  name: mnist-service
spec:
  predictor:
    tensorflow:
      storageUri: "gs://my-bucket/models/mnist"
      resources:
        requests:
          cpu: "1"
          memory: "4Gi"
        limits:
          nvidia.com/gpu: "1"

4.3 推理客户端

import tritonclient.http as httpclient

client = httpclient.InferenceServerClient("triton-server.default.svc.cluster.local:8000")

inputs = httpclient.InferInput("input", [1, 28, 28, 1], "FP32")
inputs.set_data_from_numpy(image_data)

outputs = httpclient.InferRequestedOutput("output")

response = client.infer("mnist_model", inputs=[inputs], outputs=[outputs])
result = response.as_numpy("output")

五、数据管理

5.1 数据存储配置

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: ml-data-pvc
spec:
  accessModes:
    - ReadWriteMany
  resources:
    requests:
      storage: 100Gi
  storageClassName: nfs-storage

5.2 数据加载策略

apiVersion: v1
kind: Pod
metadata:
  name: data-loader
spec:
  initContainers:
  - name: data-downloader
    image: busybox:1.35
    command: ["sh", "-c", "wget -P /data https://example.com/dataset.tar && tar -xzf /data/dataset.tar -C /data"]
    volumeMounts:
    - name: data
      mountPath: /data
  containers:
  - name: trainer
    image: pytorch/pytorch:latest
    volumeMounts:
    - name: data
      mountPath: /data
  volumes:
  - name: data
    persistentVolumeClaim:
      claimName: ml-data-pvc

六、监控与日志

6.1 ML指标监控

apiVersion: monitoring.coreos.com/v1
kind: ServiceMonitor
metadata:
  name: triton-monitor
  namespace: monitoring
spec:
  selector:
    matchLabels:
      app: triton-server
  endpoints:
  - port: metrics
    interval: 15s

6.2 训练指标收集

from prometheus_client import Summary, Counter, push_to_gateway

# 定义指标
training_duration = Summary('training_duration_seconds', 'Training duration')
batch_counter = Counter('training_batches_total', 'Number of batches processed')
loss_summary = Summary('training_loss', 'Training loss')

@training_duration.time()
def train_epoch(model, data_loader):
    for batch in data_loader:
        batch_counter.inc()
        loss = model.train_step(batch)
        loss_summary.observe(loss)

七、工作流编排

7.1 Kubeflow Pipeline

apiVersion: pipelines.kubeflow.org/v1beta1
kind: Pipeline
metadata:
  name: ml-pipeline
spec:
  tasks:
    - name: data-prep
      taskRef:
        name: data-prep-task
    - name: train
      taskRef:
        name: train-task
      dependencies:
        - data-prep
    - name: evaluate
      taskRef:
        name: evaluate-task
      dependencies:
        - train
    - name: deploy
      taskRef:
        name: deploy-task
      dependencies:
        - evaluate

7.2 Airflow集成

from airflow import DAG
from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator

with DAG('ml_workflow', schedule_interval='@daily') as dag:
    data_prep = KubernetesPodOperator(
        task_id='data_prep',
        name='data-prep',
        image='data-prep:latest',
        cmds=['python', 'prep.py']
    )
    
    train = KubernetesPodOperator(
        task_id='train',
        name='train',
        image='trainer:latest',
        cmds=['python', 'train.py'],
        resources={'request_gpu': '1'}
    )
    
    data_prep >> train

八、最佳实践

8.1 资源配置建议

工作负载类型 CPU 内存 GPU
数据预处理 4-8 16-32Gi 0
小型模型训练 4-8 32-64Gi 1-2
大型模型训练 8-16 64-128Gi 4-8
推理服务 2-4 8-16Gi 1

8.2 模型版本管理

# 使用Git管理模型配置
git clone https://github.com/example/ml-models.git

# 使用DVC管理数据和模型
dvc init
dvc add data/training_data
dvc push

# 使用OCI镜像存储模型
docker build -t myregistry/model:v1.0 .
docker push myregistry/model:v1.0

8.3 成本优化

# 使用Spot实例
apiVersion: v1
kind: Pod
metadata:
  name: training-spot
spec:
  affinity:
    nodeAffinity:
      preferredDuringSchedulingIgnoredDuringExecution:
      - weight: 100
        preference:
          matchExpressions:
          - key: cloud.google.com/gke-spot
            operator: In
            values: ["true"]

结论

在Kubernetes上运行机器学习工作负载需要综合考虑资源管理、数据处理、模型部署和监控等多个方面。通过合理配置GPU资源、使用专业的推理框架和建立完整的工作流，可以高效地运行ML训练和推理任务。

# 云原生 # kubernetes
Logo
https://edu.csdn.net/learn/39067/627173?utm_source=2019755004

汇聚全球AI编程工具，助力开发者即刻编程。

更多推荐

cover

一个 API 入口调用多个大模型：AiiOnly客户端 + CC Switch 打通 Codex AI编程流程

avatar AI编程社区
cover

VSCode Claude Code + MiMo 中转免登录完整配置教程（解决原生登录弹窗 + 模型报错）

avatar AI编程社区
cover

从零开始做一个属于自己的 Skill（保姆级教程）

avatar AI编程社区