Kubernetes 워크로드를 위한 GPU 리소스 동적 할당

Kubernetes 워크로드를 위한 GPU 리소스 동적 할당

현재 Kubernetes(k8s)에서 GPU Pod를 스케줄링하기 위해 Device Plugin, Extended Resource, scheduler extender, scheduler framework, 또는 새로운 스케줄러 개발 등 다양한 확장 솔루션이 사용되고 있습니다. 이러한 솔루션은 일반적으로 nodeSelector나 nodeAffinity를 사용하여 특정 GPU 모델을 제공하는 노드를 선택합니다. 또한 각 노드의 GPU 모델을 자동으로 탐지하려면 NFD(Node Feature Discovery) 또는 GFD(GPU Feature Discovery)를 배포해야 합니다.

이후, 서드파티 디바이스를 컨테이너에 표준화된 방식으로 노출하기 위해 새로운 런타임 사양인 CDI(Container Device Interface)가 도입되었습니다. 물론 CDI는 CRI 측면에서만 솔루션을 제공합니다.

다행히 k8s v1.26부터 Dynamic Resource Allocation(DRA)이 지원되어 스케줄링 측면에서 리소스 요청 및 할당을 위한 일관된 솔루션을 제공합니다. 그러나 초기 DRA(Control Plane Controller 방식)에는 몇 가지 문제가 있었으며, k8s v1.30에서 Structured Parameters를 사용하는 DRA가 도입되면서 개선되었습니다. 이것이 이 글의 핵심 주제입니다.

k8s 1.31에서는 이를 기반으로 추가 변경이 이루어졌으며, 이 글에서는 k8s 1.31 버전을 기준으로 Structured Parameters를 사용하는 DRA를 소개합니다.

(이전 확장 솔루션에 대해서는 이 글에서 자세히 다루지 않습니다. 시간이 허락한다면 나중에 별도로 소개하는 글을 작성하겠습니다.)

이제 바로 이 글의 주요 주제인 Structured Parameters를 사용하는 DRA 로 들어가 보겠습니다.

Structured Parameters를 사용하는 DRA

DRA는 주로 Pod 간 또는 단일 Pod 내 여러 컨테이너 간의 리소스 요청 및 공유를 위해 사용되며, GPU 리소스뿐만 아니라 다양한 리소스를 포괄합니다. 이전 Device Plugin 방식보다 훨씬 복잡하지만, PVC/PV/SC와 유사한 ResourceClaimTemplate/ResourceClaim/DeviceClass의 세 가지 요소는 유연성과 포괄적인 리소스 정의 세트를 제공하여 다양한 이기종 리소스를 공식적으로 수용할 수 있습니다.

물론 사용자는 여전히 이러한 리소스의 준비 및 유지 관리를 처리할 자체 리소스 드라이버를 개발해야 합니다.

주요 개념

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는 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"`
}

사용자가 ResourceClaim 또는 ResourceClaimTemplate을 생성할 때는 deviceClassName을 통해 해당 DeviceClass를 지정해야 합니다. 스케줄러는 DeviceClass의 selector를 사용하여 장치를 일치시켜 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에서는 Structured Parameters를 사용하는 DRA의 주요 개선 기능으로 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")
}

리소스 드라이버(DRA 드라이버)

리소스 드라이버(일명 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는 아직 성숙되지 않았으며 지속적으로 개선되고 있습니다. 커뮤니티의 최신 개발 상황을 계속 주시하는 것이 좋습니다.