Asignación dinámica de recursos GPU para cargas de trabajo en Kubernetes

Asignación dinámica de recursos GPU para cargas de trabajo en Kubernetes

Actualmente, para programar Pods de GPU en Kubernetes (k8s), se implementan varias soluciones de extensión, como Device Plugin, Extended Resource, scheduler extender, scheduler framework o el desarrollo de un nuevo scheduler. Estas soluciones a menudo dependen de nodeSelector o nodeAffinity para seleccionar nodos que proporcionen modelos de GPU específicos. Además, para descubrir automáticamente los modelos de GPU en cada nodo, es necesario implementar Node Feature Discovery (NFD) o GPU Feature Discovery (GFD).

Más tarde, para estandarizar la exposición de dispositivos de terceros a los contenedores, se introdujo la nueva especificación de runtime CDI (Container Device Interface). Por supuesto, CDI solo proporciona una solución en el lado de CRI.

Afortunadamente, a partir de k8s v1.26, se admite la asignación dinámica de recursos (DRA), que ofrece una solución coherente para las solicitudes y asignación de recursos en el lado de la programación. Sin embargo, el enfoque inicial de DRA con Control Plane Controller tenía algunos defectos, que se mejoraron en k8s v1.30 con la introducción de DRA con parámetros estructurados, que es el tema central de este artículo.

k8s 1.31 introdujo más cambios sobre esta base, y este artículo presentará DRA con parámetros estructurados basado en la versión k8s 1.31.

(No profundizaré en las soluciones de extensión anteriores; si el tiempo lo permite, escribiré otro artículo para presentarlas más adelante.)

Ahora, vayamos directamente al tema principal de este artículo: DRA con parámetros estructurados.

DRA con parámetros estructurados

DRA se utiliza principalmente para la solicitud y el intercambio de recursos entre Pods o entre múltiples contenedores dentro de un solo Pod, abarcando más que solo recursos de GPU. Si bien es significativamente más complejo que el enfoque anterior de Device Plugin, el trío ResourceClaimTemplate/ResourceClaim/DeviceClass, similar a PVC/PV/SC, ofrece flexibilidad y un conjunto completo de definiciones de recursos que pueden admitir oficialmente una variedad de recursos heterogéneos.

Naturalmente, los usuarios aún deben desarrollar sus propios controladores de recursos para manejar la preparación y el mantenimiento de estos recursos.

Conceptos clave

ResourceClaim

Un ResourceClaim se utiliza para solicitudes de acceso a recursos. El Scheduler selecciona nodos disponibles según el ResourceClaim configurado en el 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 se utiliza para crear ResourceClaims. Es gestionado por el controlador resourceclaim del kube-controller-manager. Aquí está el fragmento de código para crearlo:

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 qué introducir ResourceClaimTemplate?

Considere el siguiente ejemplo dentro de un Pod Spec: el campo resourceClaims le permite hacer referencia a un ResourceClaim o a un ResourceClaimTemplate, pero no a ambos; son mutuamente excluyentes. Cuando hace referencia a un ResourceClaim mediante resourceClaimName, todos los Pods que usan este Pod Spec (como los de un Deployment o StatefulSet) comparten la misma instancia de ResourceClaim. Esta instancia corresponde al ResourceClaim con el nombre especificado por resourceClaimName dentro del mismo namespace que el Pod.

Por el contrario, cuando hace referencia a un ResourceClaimTemplate mediante resourceClaimTemplateName, cada Pod obtiene su propia instancia de ResourceClaim distinta. Esto se debe a que el controlador resourceclaim del kube-controller-manager genera automáticamente un ResourceClaim separado para cada Pod.

Aquí hay un ejemplo simplificado de un 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 abarca la configuración y el selector de un Device, como se define a continuación.

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

Cuando los usuarios crean un ResourceClaim o un ResourceClaimTemplate, deben especificar el DeviceClass correspondiente mediante deviceClassName. El Scheduler asigna nodos a los Pods haciendo coincidir dispositivos utilizando el selector del DeviceClass y, finalmente, completa la configuración del Device desde el DeviceClass en el Status de la instancia de 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,
                    })
                }
            }
        }

Ruta de 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

Kubernetes 1.30 introduce ResourceSlice como una característica de mejora importante para DRA con parámetros estructurados. La creación y el mantenimiento de las instancias de ResourceSlice los determinan kubelet y los controladores de recursos proporcionados por los proveedores. Además, el Scheduler determina el nodo de programación basándose en ResourceClaim y ResourceSlice durante la programación de Pods, asumiendo efectivamente la responsabilidad de la asignación de recursos.

kubelet es principalmente responsable de limpiar automáticamente el ResourceSlice en el nodo durante el inicio, reinicio o cuando se elimina un plugin de kubelet.

// 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 (controlador de recursos), también conocido como DRA driver

El controlador de recursos, también conocido como DRA driver, debe ser desarrollado por los proveedores o los propios usuarios. Es principalmente responsable de publicar recursos ResourceSlice y, como plugin de kubelet, maneja las solicitudes de preparación de recursos iniciadas por kubelet. Finalmente, los recursos se utilizan a través de CDI.

A continuación se muestra el código relevante para que kubelet inicie solicitudes de preparación 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})
......
}

Defectos del enfoque DRA con Control Plane Controller

El inconveniente del enfoque DRA con Control Plane Controller es que durante la programación de Pods, el scheduler crea un PodSchedulingContext y el controlador de recursos asigna recursos. Dado que la programación de Pods es secuencial, inevitablemente aumenta la demora para otros Pods que esperan en la cola para ser programados.

Para optimizar esto, se introdujeron parámetros estructurados, lo que permite que el controlador de recursos informe proactivamente ResourceSlices, permitiendo que el scheduler tome decisiones de programación directamente desde los ResourceSlices al programar los Pods.

Resumen

Actualmente, el entrenamiento y la inferencia de IA pueden ser programados por Kubernetes como contenedores de CPU normales, pero esto se logra a través de los mecanismos de extensión existentes de Kubernetes sin una solución unificada. Además, cuanto más complejos son los escenarios de asignación de recursos de GPU, más complejas son las funciones que los usuarios deben implementar por sí mismos. Kubernetes DRA proporciona una solución unificada donde los proveedores o usuarios solo necesitan implementar el controlador de recursos y proporcionar configuraciones de instancias de DriverClass, y tiene soporte oficial, por lo que no hay necesidad de preocuparse de que una solución de código abierto para la programación de GPU deje de recibir mantenimiento.

DRA aún no está maduro y se sigue mejorando continuamente. Mantengámonos atentos a las últimas novedades en la comunidad.