Get allocated nvidia gpus from pod namespace

Signed-off-by: Cyclinder Kuo <[email protected]>

Author: cyclinder

go.mod

@@ -45,7 +45,7 @@ require (
 	go.uber.org/zap v1.27.0
 	golang.org/x/net v0.35.0
 	golang.org/x/sync v0.11.0 // indirect
-	golang.org/x/sys v0.30.0
+	golang.org/x/sys v0.33.0
 	golang.org/x/tools v0.29.0
 	gopkg.in/natefinch/lumberjack.v2 v2.2.1
 	gopkg.in/yaml.v2 v2.4.0

go.sum

@@ -805,8 +805,8 @@ golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
-golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
+golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210615171337-6886f2dfbf5b/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=

pkg/dra/config.go

@@ -3,9 +3,6 @@ package dra
 import (
 	"encoding/json"
 	"fmt"
-	"regexp"
-	"strings"
-	"unicode"
 
 	resourcev1beta1 "k8s.io/api/resource/v1beta1"
 )
@@ -111,27 +108,3 @@ func (d *NetworkConfig) GetResourceNames() []string {
 func MultusAnnotationValue(namespace, name string) string {
 	return fmt.Sprintf("%s/%s", namespace, name)
 }
-
-// NormalizedDNS1123Label normalizes the interface name to a valid DNS1123 label
-func NormalizedDNS1123Label(iface string) string {
-	// Convert to lowercase
-	normalized := strings.ToLower(iface)
-	// Replace invalid chars with hyphen
-	reg := regexp.MustCompile("[^a-z0-9-]")
-	normalized = reg.ReplaceAllString(normalized, "-")
-	// Remove leading and trailing hyphens
-	normalized = strings.Trim(normalized, "-")
-	// Replace multiple consecutive hyphens with a single one
-	reg = regexp.MustCompile("-+")
-	normalized = reg.ReplaceAllString(normalized, "-")
-
-	// If the string is empty after normalization, use a default name
-	if normalized == "" {
-		normalized = "iface"
-	}
-	// If it starts with a number, prefix it
-	if unicode.IsDigit(rune(normalized[0])) {
-		normalized = "iface-" + normalized
-	}
-	return normalized
-}

pkg/dra/driver.go

@@ -161,10 +161,11 @@ func (d *Driver) prepareMultusConfigs(pod *corev1.Pod, configs []resourcev1beta1
 
 // PublishResources periodically publishes the available SR-IOV resources
 func (d *Driver) PublishResources(ctx context.Context) {
-	d.logger.Info("Starting to publish resources")
 	devices := d.state.GetNetDevices()
 	if err := d.draPlugin.PublishResources(ctx, kubeletplugin.Resources{Devices: devices}); err != nil {
 		d.logger.Error("failed to publish resources", zap.Error(err))
+	} else {
+		d.logger.Info("Published DRA resources")
 	}
 
 	ticker := time.NewTicker(time.Minute)
@@ -183,8 +184,6 @@ func (d *Driver) PublishResources(ctx context.Context) {
 			devices := d.state.GetNetDevices()
 			if err := d.draPlugin.PublishResources(ctx, kubeletplugin.Resources{Devices: devices}); err != nil {
 				d.logger.Error("failed to publish resources", zap.Error(err))
-			} else {
-				d.logger.Info("published resources")
 			}
 		}
 	}

pkg/dra/nri/gpu.go

@@ -2,43 +2,59 @@ package nri
 
 import (
 	"context"
-	"strings"
 
 	"github.com/containerd/nri/pkg/api"
 	"go.uber.org/zap"
 	podresourcesapi "k8s.io/kubelet/pkg/apis/podresources/v1"
 )
 
 const (
+	NvidiaGPU = iota
+	// More GPU vendor
+)
+
+var (
 	NvidiaGPUResourceName = "nvidia.com/gpu"
+	NvidiaDriverGPUPath   = "/proc/driver/nvidia/gpus"
 )
 
 func (n *nriPlugin) getAllocatedGpusForPodSandbox(ctx context.Context, pod *api.PodSandbox) (gpus []string, err error) {
 	n.logger.Info("Getting allocated GPUs for pod", zap.String("podID", pod.GetId()))
 
-	podResources, err := n.kubeletClient.Get(ctx, &podresourcesapi.GetPodResourcesRequest{
-		PodName:      pod.GetName(),
-		PodNamespace: pod.GetNamespace(),
-	})
+	// It shoule be better use Get function here, but we should enable the kubelet feature-gate
+	// "KubeletPodResourcesGetAllocatable"(alpha in 1.27).
+	// podResources, err := n.kubeletClient.Get(ctx, &podresourcesapi.GetPodResourcesRequest{
+	// 	PodName:      pod.GetName(),
+	// 	PodNamespace: pod.GetNamespace(),
+	// })
+
+	resp, err := n.kubeletClient.List(ctx, &podresourcesapi.ListPodResourcesRequest{})
 	if err != nil {
 		n.logger.Error("Failed to get pod resource map", zap.Error(err))
 		return
 	}
 
-	return GetPodAllocatedGpuResources(NvidiaGPUResourceName, podResources), nil
-}
+	for _, r := range resp.PodResources {
+		if r.Name == pod.Name && r.Namespace == pod.Namespace {
+			return n.getPodAllocatedGpuResources(pod, r)
+		}
+	}
 
-func (n *nriPlugin) getAllocatedGpusForContainer(contaienr *api.Container) (gpus []string, err error) {
-	n.logger.Info("Debug GPU in CreateContainer", zap.Any("container", contaienr))
-	return nil, nil
+	// return if no any resources allocated
+	return
 }
 
-// 从设备路径中提取设备ID
-func extractDeviceIDFromPath(path string) string {
-	// 假设路径格式为 /dev/nvidia0
-	parts := strings.Split(path, "nvidia")
-	if len(parts) < 2 {
-		return ""
+func (n *nriPlugin) getPodAllocatedGpuResources(sandbox *api.PodSandbox, PodResources *podresourcesapi.PodResources) ([]string, error) {
+	var gpuDevices []string
+	for _, c := range PodResources.Containers {
+		// device plugin resouresouces
+		for _, dev := range c.Devices {
+			if dev.ResourceName == NvidiaGPUResourceName {
+				// get allocated nvidia gpus from the pod
+				return GetAllocatedNvidiaGpusFromPodNamespace(int(sandbox.Pid), NvidiaDriverGPUPath)
+			}
+		}
 	}
-	return parts[1]
+
+	return gpuDevices, nil
 }

pkg/dra/nri/kubelet_resources.go

@@ -22,15 +22,17 @@ var (
 )
 
 type nriPlugin struct {
-	logger        *zap.Logger
-	nri           stub.Stub
-	kubeletClient podresourcesapi.PodResourcesListerClient
-	conn          *grpc.ClientConn
+	gpuResourceNames map[string]struct{}
+	logger           *zap.Logger
+	nri              stub.Stub
+	kubeletClient    podresourcesapi.PodResourcesListerClient
+	conn             *grpc.ClientConn
 }
 
 func Run(ctx context.Context) error {
 	n := &nriPlugin{
-		logger: logutils.Logger.Named("nri"),
+		logger:           logutils.Logger.Named("nri"),
+		gpuResourceNames: make(map[string]struct{}),
 	}
 	// register the NRI plugin
 	nriOpts := []stub.Option{
@@ -43,13 +45,13 @@ func Run(ctx context.Context) error {
 	}
 	n.nri = stub
 
-	kubeletClient, conn, err := GetKubeletResourceClient()
+	n.kubeletClient, n.conn, err = GetKubeletResourceClient()
 	if err != nil {
 		return err
 	}
 
-	n.kubeletClient = kubeletClient
-	n.conn = conn
+	// TODO: make it configuiretable
+	n.gpuResourceNames[NvidiaGPUResourceName] = struct{}{}
 
 	go func() {
 		if err = n.nri.Run(ctx); err != nil {
@@ -64,10 +66,11 @@ func Run(ctx context.Context) error {
 
 func (n *nriPlugin) RunPodSandbox(ctx context.Context, pod *api.PodSandbox) error {
 	n.logger.Info("RunPodSandbox is called", zap.Any("pod", pod))
-	// 1. get pod net namespace
-	// 2. get multus config
-	// 3. get
-	gpus, _ := n.getAllocatedGpusForPodSandbox(ctx, pod)
+	gpus, err := n.getAllocatedGpusForPodSandbox(ctx, pod)
+	if err != nil {
+		n.logger.Error(err.Error())
+		return err
+	}
 	n.logger.Info("Allocated GPUs for pod", zap.Strings("gpus", gpus))
 	return nil
 }
@@ -81,8 +84,7 @@ func (n *nriPlugin) Configure(ctx context.Context, config, runtime, version stri
 	return api.EventMask(
 		api.Event_RUN_POD_SANDBOX |
 			api.Event_STOP_POD_SANDBOX |
-			api.Event_REMOVE_POD_SANDBOX |
-			api.Event_CREATE_CONTAINER), nil
+			api.Event_REMOVE_POD_SANDBOX), nil
 }
 
 func (n *nriPlugin) StopPodSandbox(ctx context.Context, pod *api.PodSandbox) error {
@@ -91,12 +93,12 @@ func (n *nriPlugin) StopPodSandbox(ctx context.Context, pod *api.PodSandbox) err
 }
 
 func (n *nriPlugin) CreateContainer(ctx context.Context, pod *api.PodSandbox, container *api.Container) (*api.ContainerAdjustment, []*api.ContainerUpdate, error) {
-	n.getAllocatedGpusForContainer(container)
 	n.logger.Info("CreateContainer is called", zap.Any("container", container))
 	return nil, nil, nil
 }
 
 func (n *nriPlugin) RemovePodSandbox(ctx context.Context, pod *api.PodSandbox) error {
+	n.logger.Info("RemovePodSandbox is called", zap.Any("pod", pod))
 	return nil
 }