Kubernetesワークロード向けGPUリソースの動的割り当て

Kubernetesワークロード向けGPUリソースの動的割り当て

現在、Kubernetes (k8s) で GPU Pod をスケジューリングするために、Device Plugin、Extended Resource、scheduler extender、scheduler framework、または新しいスケジューラの開発など、さまざまな拡張ソリューションが導入されています。これらのソリューションは、特定のGPUモデルを提供するノードを選択するために、nodeSelector や nodeAffinity に依存することがよくあります。さらに、各ノードのGPUモデルを自動的に検出するには、Node Feature Discovery (NFD) や GPU Feature Discovery (GFD) をデプロイする必要があります。

その後、サードパーティ製デバイスのコンテナへの公開を標準化するために、新しいランタイム仕様である CDI (Container Device Interface) が導入されました。もちろん、CDI は CRI 側でのみソリューションを提供します。

幸いなことに、k8s v1.26 から Dynamic Resource Allocation (DRA) がサポートされ、スケジューリング側でリソース要求と割り当てのための一貫したソリューションが提供されています。しかし、初期の Control Plane Controller を使用した DRA アプローチにはいくつかの欠点があり、k8s v1.30 で Structured Parameters を使用した DRA が導入され改善されました。これが本記事の焦点です。

k8s 1.31 ではさらに変更が加えられており、本記事では k8s 1.31 バージョンに基づいて Structured Parameters を使用した DRA を紹介します。

(以前の拡張ソリューションについては、本記事では詳細に立ち入りません。時間があれば、後日別の記事で紹介する予定です。)

それでは、本記事のメインテーマである Structured Parameters を使用した DRA に直接入りましょう。

DRA with Structured Parameters

DRA は主に、Pod 間、または単一 Pod 内の複数のコンテナ間でのリソースの要求と共有に使用され、GPU リソースだけでなく、より広範囲をカバーします。以前の Device Plugin アプローチよりもはるかに複雑ですが、ResourceClaimTemplate/ResourceClaim/DeviceClass の3つ組は、PVC/PV/SC と同様に、柔軟性と包括的なリソース定義のセットを提供し、公式にさまざまな異種リソースに対応できます。

当然ながら、ユーザーはこれらのリソースの準備とメンテナンスを処理するために、独自のリソースドライバを開発する必要があります。

主要な概念

ResourceClaim

ResourceClaim は、リソースアクセス要求に使用されます。スケジューラは、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-manager の resourceclaim コントローラによって管理されます。以下は、その作成のためのコードスニペットです。

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 と同じネームスペース内で resourceClaimName で指定された名前の ResourceClaim に対応します。

逆に、resourceClaimTemplateName を介して ResourceClaimTemplate を参照する場合、各 Pod は独自の個別の ResourceClaim インスタンスを取得します。これは、kube-controller-manager の resourceclaim コントローラが各 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 は、デバイスの構成とセレクタを含みます。以下に定義します。

// +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"`
}

ユーザーが ResourceClaim または ResourceClaimTemplate を作成するとき、deviceClassName を介して対応する DeviceClass を指定する必要があります。スケジューラは、DeviceClass のセレクタを使用してデバイスをマッチングすることにより、Pod にノードを割り当て、最終的に DeviceClass からデバイス構成を 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 では、DRA with structured parameters の主要な改善機能として ResourceSlice が導入されました。ResourceSlice インスタンスの作成と保守は、kubelet およびベンダー提供のリソースドライバによって決定されます。さらに、スケジューラは Pod スケジューリング中に ResourceClaim と ResourceSlice に基づいてスケジューリングノードを決定し、実質的にリソース割り当ての責任を負います。

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 ドライバ)

Resource Driver(DRA ドライバとも呼ばれる)は、ベンダーまたはユーザー自身が開発する必要があります。主に 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})
......
}

Control Plane Controller を使用した DRA アプローチの欠点

Control Plane Controller を使用した DRA アプローチの欠点は、Pod スケジューリング中にスケジューラが PodSchedulingContext を作成し、リソースドライバがリソースを割り当てることです。Pod スケジューリングは逐次的であるため、スケジューリングを待っている他の Pod の遅延が避けられません。これを最適化するために structured parameters が導入され、リソースドライバが ResourceSlice をプロアクティブに報告できるようになり、スケジューラが Pod スケジューリング中に ResourceSlice から直接スケジューリング判断を下せるようになりました。

まとめ

現在、AI トレーニングと推論は、通常の CPU コンテナと同様に Kubernetes でスケジューリングできますが、これは既存の Kubernetes 拡張メカニズムを通じて実現されており、統一されたソリューションはありません。さらに、GPU リソース割り当てのシナリオが複雑になるほど、ユーザーが自分で実装する必要がある機能も複雑になります。Kubernetes DRA は統一されたソリューションを提供し、ベンダーまたはユーザーはリソースドライバを実装し、DriverClass インスタンス構成を提供するだけで済みます。これは公式にサポートされているため、GPU スケジューリングのオープンソースソリューションのメンテナンスが停止することを心配する必要はありません。DRA はまだ成熟しておらず、継続的に改善されています。最新の動向に注目しましょう。