运行 PyTorchJob

使用 Kueue 调度 PyTorchJob

本页面展示了在运行 Trainer PyTorchJob 时,如何利用 Kueue 的调度和资源管理能力。

本指南适用于对 Kueue 有基本了解的批处理用户。 欲了解更多详情,请参见 Kueue 概述

开始之前

查阅管理集群配额, 了解初始集群设置的详细信息。

查阅 Trainer 安装指南

请注意,Trainer 的最低版本要求是 v1.7.0。

你可以修改已安装版本的 Kueue 配置, 以将 PyTorchJob 包括为允许的工作负载。

PyTorchJob 定义

a. 队列选择

目标本地队列应当在 PyTorchJob 配置的 metadata.labels 部分指定。

metadata:
  labels:
    kueue.x-k8s.io/queue-name: user-queue

b. 可选择在 PyTorchJob 中设置 Suspend 字段

spec:
  runPolicy:
    suspend: true

默认情况下,Kueue 会通过 Webhook 将 suspend 设置为 true, 并在 PyTorchJob 被接受时取消挂起。

PyTorchJob 示例

此示例基于 https://github.com/kubeflow/trainer/blob/855e0960668b34992ba4e1fd5914a08a3362cfb1/examples/pytorch/simple.yaml

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: pytorch-simple
  namespace: default
  labels:
    kueue.x-k8s.io/queue-name: user-queue
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: docker.io/kubeflowkatib/pytorch-mnist-cpu:v1beta1-21320b6
#              If you have gpu, pytorch-mnist-gpu would be helpful. pytorch-mnist-gpu is approximately 22GB
#              image: docker.io/kubeflowkatib/pytorch-mnist-cpu:latest
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"
              resources:
                requests:
                  cpu: 1
                  memory: "200Mi"
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: docker.io/kubeflowkatib/pytorch-mnist-cpu:v1beta1-21320b6
#              If you have gpu, pytorch-mnist-gpu would be helpful. pytorch-mnist-gpu is approximately 22GB
#              image: docker.io/kubeflowkatib/pytorch-mnist-cpu:latest
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"
              resources:
                requests:
                  cpu: 1
                  memory: "200Mi"

最后修改 July 24, 2025: Localize pytorchjobs.md (#6171) (c7bdb15f)