為 Kubernetes 工作負載動態分配 GPU 資源

為 Kubernetes 工作負載動態分配 GPU 資源

目前,要在 Kubernetes (k8s) 中排程 GPU Pod,已實作多種擴充方案,包括 Device Plugin、Extended Resource、scheduler extender、scheduler framework,或是開發新的排程器。這些方案通常依賴 nodeSelectornodeAffinity 來選擇提供特定 GPU 型號的節點。此外,為了自動發現每個節點上的 GPU 型號,必須部署 Node Feature Discovery (NFD) 或 GPU Feature Discovery (GFD)。

後來,為了標準化將第三方裝置暴露給容器的方式,引入了新的執行時期規格 CDI (Container Device Interface)。當然,CDI 只提供了 CRI 端的解決方案。

幸運的是,從 k8s v1.26 開始支援 Dynamic Resource Allocation (DRA),為排程端的資源請求和分配提供了一致的解決方案。然而,初始的 DRA with Control Plane Controller 方式存在一些缺陷,這些缺陷在 k8s v1.30 中透過引入 DRA with Structured Parameters 獲得改善,而這正是本文的重點。

k8s 1.31 在此基礎上做了進一步的變更,本文將基於 k8s 1.31 版本介紹 DRA with Structured Parameters。

(關於先前的擴充方案,本文不會深入探討;如果時間允許,我之後會另寫一篇文章介紹。)

現在,讓我們直接進入本文的主題:DRA with Structured Parameters。

DRA with Structured Parameters

DRA 主要用於 Pod 之間或單個 Pod 內多個容器之間的資源請求和共享,涵蓋的不只 GPU 資源。雖然它比之前的 Device Plugin 方式複雜得多,但 ResourceClaimTemplate / ResourceClaim / DeviceClass 這三件組,類似於 PVC / PV / SC,提供了彈性且完整的資源定義,能夠正式支援各種異質資源。

當然,使用者仍然需要自行開發資源驅動程式來處理這些資源的準備和維護。

關鍵概念

ResourceClaim

ResourceClaim 用於資源存取請求。Scheduler 會根據 Pod 中設定的 ResourceClaim 來選擇可用的節點。

---
apiVersion: resource.k8s.io/v1alpha3
kind: ResourceClaim
metadata:
  name: small-white-cat-claim
spec:
  devices:
    requests:
    - name: req-0
      deviceClassName: resource.example.com
      selectors:
      - cel:
         expression: |-
            device.attributes["resource-driver.example.com"].color == "white" &&
            device.attributes["resource-driver.example.com"].size == "small"

ResourceClaimTemplate

ResourceClaimTemplate 用於建立 ResourceClaim。它由 kube-controller-managerresourceclaim controller 管理。以下是建立它的程式碼片段:

func (ec *Controller) handleClaim(ctx context.Context, pod *v1.Pod, podClaim v1.PodResourceClaim, newPodClaims *map[string]string) error {
    claimName, mustCheckOwner, err := resourceclaim.Name(pod, &podClaim)
    switch {
    case errors.Is(err, resourceclaim.ErrClaimNotFound):
    ......
    }
    // Before we create a new ResourceClaim, check if there is an orphaned one.
    // This covers the case that the controller has created it, but then fails
    // before it can update the pod status.
    claim, err := ec.findPodResourceClaim(pod, podClaim)
    if err != nil {
        return fmt.Errorf("finding ResourceClaim for claim %s in pod %s/%s failed: %v", podClaim.Name, pod.Namespace, pod.Name, err)
    }

    if claim == nil {
        template, err := ec.templateLister.ResourceClaimTemplates(pod.Namespace).Get(*templateName)
        if err != nil {
            return fmt.Errorf("resource claim template %q: %v", *templateName, err)
        }

        // Create the ResourceClaim with pod as owner, with a generated name that uses
        // <pod>-<claim name> as base.
        isTrue := true
        annotations := template.Spec.ObjectMeta.Annotations
        if annotations == nil {
            annotations = make(map[string]string)
        }
        annotations[podResourceClaimAnnotation] = podClaim.Name
        generateName := pod.Name + "-" + podClaim.Name + "-"
        maxBaseLen := 57 // Leave space for hyphen and 5 random characters in a name with 63 characters.
        if len(generateName) > maxBaseLen {
            // We could leave truncation to the apiserver, but as
            // it removes at the end, we would loose everything
            // from the pod claim name when the pod name is long.
            // We can do better and truncate both strings,
            // proportional to their length.
            generateName = pod.Name[0:len(pod.Name)*maxBaseLen/len(generateName)] +
                "-" +
                podClaim.Name[0:len(podClaim.Name)*maxBaseLen/len(generateName)]
        }
        claim = &resourceapi.ResourceClaim{
            ObjectMeta: metav1.ObjectMeta{
                GenerateName: generateName,
                OwnerReferences: []metav1.OwnerReference{
                    {
                        APIVersion:         "v1",
                        Kind:               "Pod",
                        Name:               pod.Name,
                        UID:                pod.UID,
                        Controller:         &isTrue,
                        BlockOwnerDeletion: &isTrue,
                    },
                },
                Annotations: annotations,
                Labels:      template.Spec.ObjectMeta.Labels,
            },
            Spec: template.Spec.Spec,
        }
        metrics.ResourceClaimCreateAttempts.Inc()
        claimName := claim.Name
        claim, err = ec.kubeClient.ResourceV1alpha3().ResourceClaims(pod.Namespace).Create(ctx, claim, metav1.CreateOptions{})
        if err != nil {
            metrics.ResourceClaimCreateFailures.Inc()
            return fmt.Errorf("create ResourceClaim %s: %v", claimName, err)
        }
        logger.V(4).Info("Created ResourceClaim", "claim", klog.KObj(claim), "pod", klog.KObj(pod))
        ec.claimCache.Mutation(claim)
    }
    

為什麼要引入 ResourceClaimTemplate?

考慮以下 Pod Spec 中的範例:resourceClaims 欄位允許你引用一個 ResourceClaim 或一個 ResourceClaimTemplate,但不能同時引用兩者——它們是互斥的。當你透過 resourceClaimName 引用一個 ResourceClaim 時,所有使用此 Pod Spec 的 Pod(例如 Deployment 或 StatefulSet 中的 Pod)會共用同一個 ResourceClaim 實例。該實例對應於與 Pod 相同命名空間中名稱為 resourceClaimNameResourceClaim

相反地,當你透過 resourceClaimTemplateName 引用一個 ResourceClaimTemplate 時,每個 Pod 會獲得自己的獨立 ResourceClaim 實例。這是因為 kube-controller-managerresourceclaim controller 會自動為每個 Pod 產生一個單獨的 ResourceClaim

以下是一個簡化的 Pod Spec 範例:

–--
apiVersion: v1
kind: Pod
metadata:
  name: pod-with-cats
spec:
  containers:
  - name: container0
    image: ubuntu:20.04
    command: ["sleep", "9999"]
    resources:
      claims:
      - name: cat-0
  - name: container1
    image: ubuntu:20.04
    command: ["sleep", "9999"]
    resources:
      claims:
      - name: cat-1
  resourceClaims:
  - name: cat-0
    resourceClaimName: small-white-cat-claim
  - name: cat-1
    resourceClaimTemplateName: large-black-cat-claim-template

DeviceClass

DeviceClass 包含 Device 的配置和選擇器,定義如下。

// +genclient
// +genclient:nonNamespaced
// +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object
// +k8s:prerelease-lifecycle-gen:introduced=1.31

// DeviceClass is a vendor- or admin-provided resource that contains
// device configuration and selectors. It can be referenced in
// the device requests of a claim to apply these presets.
// Cluster scoped.
//
// This is an alpha type and requires enabling the DynamicResourceAllocation
// feature gate.
type DeviceClass struct {
    metav1.TypeMeta `json:",inline"`
    // Standard object metadata
    // +optional
    metav1.ObjectMeta `json:"metadata,omitempty" protobuf:"bytes,1,opt,name=metadata"`

    // Spec defines what can be allocated and how to configure it.
    //
    // This is mutable. Consumers have to be prepared for classes changing
    // at any time, either because they get updated or replaced. Claim
    // allocations are done once based on whatever was set in classes at
    // the time of allocation.
    //
    // Changing the spec automatically increments the metadata.generation number.
    Spec DeviceClassSpec `json:"spec" protobuf:"bytes,2,name=spec"`
}

// DeviceClassSpec is used in a [DeviceClass] to define what can be allocated
// and how to configure it.
type DeviceClassSpec struct {
    // Each selector must be satisfied by a device which is claimed via this class.
    //
    // +optional
    // +listType=atomic
    Selectors []DeviceSelector `json:"selectors,omitempty" protobuf:"bytes,1,opt,name=selectors"`

    // Config defines configuration parameters that apply to each device that is claimed via this class.
    // Some classses may potentially be satisfied by multiple drivers, so each instance of a vendor
    // configuration applies to exactly one driver.
    //
    // They are passed to the driver, but are not considered while allocating the claim.
    //
    // +optional
    // +listType=atomic
    Config []DeviceClassConfiguration `json:"config,omitempty" protobuf:"bytes,2,opt,name=config"`

    // Only nodes matching the selector will be considered by the scheduler
    // when trying to find a Node that fits a Pod when that Pod uses
    // a claim that has not been allocated yet *and* that claim
    // gets allocated through a control plane controller. It is ignored
    // when the claim does not use a control plane controller
    // for allocation.
    //
    // Setting this field is optional. If unset, all Nodes are candidates.
    //
    // This is an alpha field and requires enabling the DRAControlPlaneController
    // feature gate.
    //
    // +optional
    // +featureGate=DRAControlPlaneController
    SuitableNodes *v1.NodeSelector `json:"suitableNodes,omitempty" protobuf:"bytes,3,opt,name=suitableNodes"`
}

使用者在建立 ResourceClaimResourceClaimTemplate 時,需要透過 deviceClassName 指定對應的 DeviceClassScheduler 透過 DeviceClassselector 來匹配裝置,從而為 Pod 分配節點,最後將 DeviceClass 中的 Device 配置填入 ResourceClaim 實例的 Status 中。


    if requestData.class != nil {
        match, err := alloc.selectorsMatch(r, device, deviceID, requestData.class, requestData.class.Spec.Selectors)
        if err != nil {
            return false, err
        }
        if !match {
            alloc.deviceMatchesRequest[matchKey] = false
            return false, nil
        }
    }

    request := &alloc.claimsToAllocate[r.claimIndex].Spec.Devices.Requests[r.requestIndex]
    match, err := alloc.selectorsMatch(r, device, deviceID, nil, request.Selectors)
        

    for claimIndex, allocationResult := range alloc.result {
        claim := alloc.claimsToAllocate[claimIndex]

        // Populate configs.
        for requestIndex := range claim.Spec.Devices.Requests {
            class := alloc.requestData[requestIndices{claimIndex: claimIndex, requestIndex: requestIndex}].class
            if class != nil {
                for _, config := range class.Spec.Config {
                    allocationResult.Devices.Config = append(allocationResult.Devices.Config, resourceapi.DeviceAllocationConfiguration{
                        Source:              resourceapi.AllocationConfigSourceClass,
                        Requests:            nil, // All of them...
                        DeviceConfiguration: config.DeviceConfiguration,
                    })
                }
            }
        }

程式碼路徑:*pkg\\scheduler\\framework\\plugins\\dynamicresources\\dynamicresources.go*

func (pl *dynamicResources) bindClaim(ctx context.Context, state *stateData, index int, pod *v1.Pod, nodeName string) (patchedClaim *resourceapi.ResourceClaim, finalErr error) {
        if allocation != nil {
            if claim.Status.Allocation != nil {
                return fmt.Errorf("claim %s got allocated elsewhere in the meantime", klog.KObj(claim))
            }
......
            claim.Status.Allocation = allocation
        }
        claim.Status.ReservedFor = append(claim.Status.ReservedFor, resourceapi.ResourceClaimConsumerReference{Resource: "pods", Name: pod.Name, UID: pod.UID})
        updatedClaim, err := pl.clientset.ResourceV1alpha3().ResourceClaims(claim.Namespace).UpdateStatus(ctx, claim, metav1.UpdateOptions{})

ResourceSlice

Kubernetes 1.30 引入了 ResourceSlice 作為 DRA with structured parameters 的一項重大改進功能。ResourceSlice 實例的建立和維護由 kubelet 和供應商提供的資源驅動程式決定。此外,Scheduler 在 Pod 排程期間根據 ResourceClaimResourceSlice 決定排程節點,實際上承擔了資源分配的責任。

kubelet 主要負責在啟動、重新啟動或移除 kubelet 外掛程式時自動清除節點上的 ResourceSlice

// NewPluginHandler returns new registration handler.
//
// Must only be called once per process because it manages global state.
// If a kubeClient is provided, then it synchronizes ResourceSlices
// with the resource information provided by plugins.
func NewRegistrationHandler(kubeClient kubernetes.Interface, getNode func() (*v1.Node, error)) *RegistrationHandler {
    handler := &RegistrationHandler{
        // The context and thus logger should come from the caller.
        backgroundCtx: klog.NewContext(context.TODO(), klog.LoggerWithName(klog.TODO(), "DRA registration handler")),
        kubeClient:    kubeClient,
        getNode:       getNode,
    }

    // When kubelet starts up, no DRA driver has registered yet. None of
    // the drivers are usable until they come back, which might not happen
    // at all. Therefore it is better to not advertise any local resources
    // because pods could get stuck on the node waiting for the driver
    // to start up.
    //
    // This has to run in the background.
    go handler.wipeResourceSlices("")

    return handler
}

// DeRegisterPlugin is called when a plugin has removed its socket,
// signaling it is no longer available.
func (h *RegistrationHandler) DeRegisterPlugin(pluginName string) {
    if p := draPlugins.delete(pluginName); p != nil {
        logger := klog.FromContext(p.backgroundCtx)
        logger.V(3).Info("Deregister DRA plugin", "endpoint", p.endpoint)

        // Clean up the ResourceSlices for the deleted Plugin since it
        // may have died without doing so itself and might never come
        // back.
        go h.wipeResourceSlices(pluginName)

        return
    }

    logger := klog.FromContext(h.backgroundCtx)
    logger.V(3).Info("Deregister DRA plugin not necessary, was already removed")
}

Resource Driver,又稱 DRA driver

資源驅動程式(Resource Driver),也稱為 DRA driver,需要由供應商或使用者自行開發。它主要負責發布 ResourceSlice 資源,並作為 kubelet 外掛程式處理 kubelet 觸發的資源準備請求。最終,資源透過 CDI 來使用。

以下是 kubelet 發起資源準備請求的相關程式碼。

// PrepareResources attempts to prepare all of the required resources
// for the input container, issue NodePrepareResources rpc requests
// for each new resource requirement, process their responses and update the cached
// containerResources on success.
func (m *ManagerImpl) PrepareResources(pod *v1.Pod) error {
    batches := make(map[string][]*drapb.Claim)
    resourceClaims := make(map[types.UID]*resourceapi.ResourceClaim)
    for i := range pod.Spec.ResourceClaims {
        podClaim := &pod.Spec.ResourceClaims[i]
        klog.V(3).InfoS("Processing resource", "pod", klog.KObj(pod), "podClaim", podClaim.Name)
        claimName, mustCheckOwner, err := resourceclaim.Name(pod, podClaim)
......       
        resourceClaim, err := m.kubeClient.ResourceV1alpha3().ResourceClaims(pod.Namespace).Get(
            context.TODO(),
            *claimName,
            metav1.GetOptions{})
......
        // Atomically perform some operations on the claimInfo cache.
        err = m.cache.withLock(func() error {
            claimInfo, exists := m.cache.get(resourceClaim.Name, resourceClaim.Namespace)
......
            // Loop through all drivers and prepare for calling NodePrepareResources.
            claim := &drapb.Claim{
                Namespace: claimInfo.Namespace,
                UID:       string(claimInfo.ClaimUID),
                Name:      claimInfo.ClaimName,
            }
            for driverName := range claimInfo.DriverState {
                batches[driverName] = append(batches[driverName], claim)
            }

            return nil
        })
        if err != nil {
            return fmt.Errorf("locked cache operation: %w", err)
        }
    }

    // Call NodePrepareResources for all claims in each batch.
    // If there is any error, processing gets aborted.
    // We could try to continue, but that would make the code more complex.
    for driverName, claims := range batches {
        // Call NodePrepareResources RPC for all resource handles.
        client, err := dra.NewDRAPluginClient(driverName)
        if err != nil {
            return fmt.Errorf("failed to get gRPC client for driver %s: %w", driverName, err)
        }
        response, err := client.NodePrepareResources(context.Background(), &drapb.NodePrepareResourcesRequest{Claims: claims})
......
}

DRA with Control Plane Controller 方式的缺陷

DRA with Control Plane Controller 方式的缺點在於,Pod 排程期間,排程器會建立一個 PodSchedulingContext,而資源驅動程式則負責分配資源。由於 Pod 排程是順序進行的,這無可避免地會增加排程佇列中其他等待 Pod 的延遲。

為了最佳化這個問題,引入了結構化參數(structured parameters),讓資源驅動程式能夠主動報告 ResourceSlice,從而使排程器在 Pod 排程時能直接從 ResourceSlice 中做出排程決策。

總結

目前,AI 訓練和推論可以像一般 CPU 容器一樣由 Kubernetes 排程,但這是透過現有的 Kubernetes 擴充機制實現的,缺乏統一的解決方案。而且,GPU 資源分配的場景越複雜,使用者需要自行實作的功能也越複雜。Kubernetes DRA 提供了一個統一的解決方案,供應商或使用者只需實作資源驅動程式並提供 DriverClass 實例設定,而且由官方支援,因此無需擔心 GPU 排程開源方案停止維護的問題。

DRA 尚未成熟,仍在持續改進中。讓我們持續關注社群的最新發展。