Alocação dinâmica de recursos de GPU para cargas de trabalho Kubernetes

Alocação dinâmica de recursos de GPU para cargas de trabalho Kubernetes

Atualmente, para agendar Pods de GPU no Kubernetes (k8s), várias soluções de extensão são colocadas em ação, incluindo Device Plugin, Extended Resource, scheduler extender, scheduler framework, ou desenvolvendo um novo scheduler. Essas soluções geralmente dependem de nodeSelector ou nodeAffinity para selecionar nós que forneçam modelos específicos de GPU. Além disso, para descobrir automaticamente os modelos de GPU em cada nó, é necessário implantar o Node Feature Discovery (NFD) ou o GPU Feature Discovery (GFD).

Posteriormente, para padronizar a exposição de dispositivos de terceiros aos contêineres, foi introduzida a nova especificação de runtime CDI (Container Device Interface). Claro, o CDI fornece apenas uma solução no lado do CRI.

Felizmente, a partir do k8s v1.26, o Dynamic Resource Allocation (DRA) é suportado, oferecendo uma solução consistente para solicitações e alocação de recursos no lado do agendamento. No entanto, a abordagem inicial do DRA com Control Plane Controller tinha algumas falhas, que foram melhoradas no k8s v1.30 com a introdução do DRA com Structured Parameters, que é o foco deste artigo.

O k8s 1.31 fez novas alterações com base nisso, e este artigo apresentará o DRA com Structured Parameters baseado na versão k8s 1.31.

(Para soluções de extensão anteriores, este artigo não entrará em detalhes; se o tempo permitir, escreverei outro artigo para apresentá-las posteriormente.)

Agora, vamos direto ao tópico principal deste artigo, DRA com Structured Parameters.

DRA com Structured Parameters

O DRA é utilizado principalmente para a solicitação e compartilhamento de recursos entre Pods ou entre múltiplos contêineres dentro de um único Pod, abrangendo mais do que apenas recursos de GPU. Embora seja significativamente mais complexo do que a abordagem anterior do Device Plugin, o trio ResourceClaimTemplate/ResourceClaim/DeviceClass, análogo a PVC/PV/SC, oferece flexibilidade e um conjunto abrangente de definições de recursos que podem acomodar oficialmente uma variedade de recursos heterogêneos.

Naturalmente, os usuários ainda são obrigados a desenvolver seus próprios drivers de recurso para lidar com a preparação e manutenção desses recursos.

Conceitos-chave

ResourceClaim

Um ResourceClaim é usado para solicitações de acesso a recursos. O Scheduler seleciona nós disponíveis com base no ResourceClaim configurado no Pod.

---
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 é usado para criar ResourceClaims. Ele é gerenciado pelo controller resourceclaim do kube-controller-manager. Aqui está o trecho de código para criá-lo:

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)
    }
    

Por que introduzir o ResourceClaimTemplate?

Considere o seguinte exemplo dentro de um Pod Spec: o campo resourceClaims permite fazer referência a um ResourceClaim ou a um ResourceClaimTemplate, mas não ambos – eles são mutuamente exclusivos. Quando você faz referência a um ResourceClaim via resourceClaimName, todos os Pods que utilizam este Pod Spec (como aqueles em um Deployment ou StatefulSet) compartilham a mesma instância de ResourceClaim. Esta instância corresponde ao ResourceClaim com o nome especificado por resourceClaimName dentro do mesmo namespace que o Pod.

Por outro lado, quando você faz referência a um ResourceClaimTemplate via resourceClaimTemplateName, cada Pod obtém sua própria instância distinta de ResourceClaim. Isso porque o controller resourceclaim do kube-controller-manager gera automaticamente um ResourceClaim separado para cada Pod.

Aqui está um exemplo simplificado de 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

O DeviceClass engloba a configuração e o seletor para um Device, conforme definido abaixo.

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

Quando os usuários criam um ResourceClaim ou um ResourceClaimTemplate, eles precisam especificar o DeviceClass correspondente através do deviceClassName. O Scheduler atribui nós aos Pods combinando dispositivos usando o seletor do DeviceClass e, por fim, preenche a configuração do Device do DeviceClass no Status da instância do ResourceClaim.


    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,
                    })
                }
            }
        }

Caminho do código: 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

O Kubernetes 1.30 introduz o ResourceSlice como uma grande melhoria para o DRA com parâmetros estruturados. A criação e manutenção de instâncias ResourceSlice são determinadas pelo kubelet e pelos drivers de recurso fornecidos pelos fornecedores. Além disso, o Scheduler determina o nó de agendamento com base no ResourceClaim e ResourceSlice durante o agendamento do Pod, assumindo efetivamente a responsabilidade pela alocação de recursos.

O kubelet é o principal responsável por limpar automaticamente o ResourceSlice no nó durante a inicialização, reinicialização ou quando um plugin kubelet é removido.

// 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, também conhecido como DRA driver

O Resource Driver, também conhecido como DRA driver, precisa ser desenvolvido pelos fornecedores ou pelos próprios usuários. Ele é o principal responsável por publicar recursos ResourceSlice e, como um plugin kubelet, lida com solicitações de preparação de recursos acionadas pelo kubelet. Por fim, os recursos são utilizados através do CDI.

Abaixo está o código relevante para a inicialização do kubelet nas solicitações de preparação de recursos.

// 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})
......
}

Falhas na abordagem DRA com Control Plane Controller

A desvantagem da abordagem DRA com Control Plane Controller é que, durante o agendamento do Pod, o scheduler cria um PodSchedulingContext e o resource driver aloca recursos. Como o agendamento de Pods é sequencial, isso inevitavelmente aumenta o atraso para outros Pods aguardando na fila para agendamento.

Para otimizar isso, foram introduzidos parâmetros estruturados, permitindo que o resource driver relate proativamente ResourceSlices, permitindo que o scheduler tome decisões de agendamento diretamente dos ResourceSlices durante o agendamento do Pod.

Resumo

Atualmente, o treinamento e a inferência de IA podem ser agendados pelo Kubernetes como contêineres de CPU comuns, mas isso é alcançado por meio de mecanismos de extensão existentes do Kubernetes, sem uma solução unificada. Além disso, quanto mais complexos os cenários de alocação de recursos de GPU, mais complexas são as funções que os usuários precisam implementar por conta própria. O DRA do Kubernetes oferece uma solução unificada, onde fornecedores ou usuários só precisam implementar o resource driver e fornecer configurações de instância DriverClass, e tem suporte oficial, portanto não há necessidade de se preocupar com a interrupção da manutenção de uma solução de código aberto para agendamento de GPU.

O DRA ainda não está maduro e está sendo continuamente melhorado. Vamos ficar de olho nos últimos desenvolvimentos da comunidade.