From ec58c3162c1ac218b7ea19ab9d031f76d485c62f Mon Sep 17 00:00:00 2001
From: Moshe Levi <moshele@nvidia.com>
Date: Thu, 29 Jun 2023 10:22:32 +0300
Subject: [PATCH] add support for Dynamic Resource Allocation

Signed-off-by: Moshe Levi <moshele@nvidia.com>
---
 pkg/kubeletclient/kubeletclient.go      | 19 +++++++
 pkg/kubeletclient/kubeletclient_test.go | 76 ++++++++++++++++++++++---
 2 files changed, 87 insertions(+), 8 deletions(-)

diff --git a/pkg/kubeletclient/kubeletclient.go b/pkg/kubeletclient/kubeletclient.go
index d0438f012..457d99ac9 100644
--- a/pkg/kubeletclient/kubeletclient.go
+++ b/pkg/kubeletclient/kubeletclient.go
@@ -21,6 +21,7 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"strings"
 	"time"
 
 	"golang.org/x/net/context"
@@ -144,6 +145,24 @@ func (rc *kubeletClient) GetPodResourceMap(pod *v1.Pod) (map[string]*types.Resou
 						resourceMap[dev.ResourceName] = &types.ResourceInfo{DeviceIDs: dev.DeviceIds}
 					}
 				}
+				for _, dynamicResource := range cnt.DynamicResources {
+					var deviceIDs []string
+					for _, claimResource := range dynamicResource.ClaimResources {
+						for _, cdiDevice := range claimResource.CDIDevices {
+							res := strings.Split(cdiDevice.Name, "=")
+							if len(res) == 2 {
+								deviceIDs = append(deviceIDs, res[1])
+							} else {
+								logging.Errorf("GetPodResourceMap: Invalid CDI format")
+							}
+						}
+					}
+					if rInfo, ok := resourceMap[dynamicResource.ClassName]; ok {
+						rInfo.DeviceIDs = append(rInfo.DeviceIDs, deviceIDs...)
+					} else {
+						resourceMap[dynamicResource.ClassName] = &types.ResourceInfo{DeviceIDs: deviceIDs}
+					}
+				}
 			}
 		}
 	}
diff --git a/pkg/kubeletclient/kubeletclient_test.go b/pkg/kubeletclient/kubeletclient_test.go
index 5e0e43b82..620b86aaa 100644
--- a/pkg/kubeletclient/kubeletclient_test.go
+++ b/pkg/kubeletclient/kubeletclient_test.go
@@ -52,11 +52,12 @@ func (m *fakeResourceServer) GetAllocatableResources(_ context.Context, _ *podre
 	return &podresourcesapi.AllocatableResourcesResponse{}, nil
 }
 
-func (m *fakeResourceServer) List(_ context.Context, _ *podresourcesapi.ListPodResourcesRequest) (*podresourcesapi.ListPodResourcesResponse, error) {
-	podName := "pod-name"
-	podNamespace := "pod-namespace"
-	containerName := "container-name"
+//In K8s 1.27 Get metod was added
+func (m *fakeResourceServer) Get(_ context.Context, _ *podresourcesapi.GetPodResourcesRequest) (*podresourcesapi.GetPodResourcesResponse, error) {
+	return &podresourcesapi.GetPodResourcesResponse{}, nil
+}
 
+func (m *fakeResourceServer) List(_ context.Context, _ *podresourcesapi.ListPodResourcesRequest) (*podresourcesapi.ListPodResourcesResponse, error) {
 	devs := []*podresourcesapi.ContainerDevices{
 		{
 			ResourceName: "resource",
@@ -64,18 +65,49 @@ func (m *fakeResourceServer) List(_ context.Context, _ *podresourcesapi.ListPodR
 		},
 	}
 
+	cdiDevices := []*podresourcesapi.CDIDevice{
+		{
+			Name: "cdi-kind=cdi-resource",
+		},
+	}
+
+	claimsResource := []*podresourcesapi.ClaimResource{
+		{
+			CDIDevices: cdiDevices,
+		},
+	}
+
+	dynamicResources := []*podresourcesapi.DynamicResource{
+		{
+			ClassName:      "resource-class",
+			ClaimName:      "resource-claim",
+			ClaimNamespace: "dynamic-resource-pod-namespace",
+			ClaimResources: claimsResource,
+		},
+	}
+
 	resp := &podresourcesapi.ListPodResourcesResponse{
 		PodResources: []*podresourcesapi.PodResources{
 			{
-				Name:      podName,
-				Namespace: podNamespace,
+				Name:      "pod-name",
+				Namespace: "pod-namespace",
 				Containers: []*podresourcesapi.ContainerResources{
 					{
-						Name:    containerName,
+						Name:    "container-name",
 						Devices: devs,
 					},
 				},
 			},
+			{
+				Name:      "dynamic-resource-pod-name",
+				Namespace: "dynamic-resource-pod-namespace",
+				Containers: []*podresourcesapi.ContainerResources{
+					{
+						Name:             "dynamic-resource-container-name",
+						DynamicResources: dynamicResources,
+					},
+				},
+			},
 		},
 	}
 	return resp, nil
@@ -181,7 +213,7 @@ var _ = Describe("Kubelet resource endpoint data read operations", func() {
 		})
 	})
 	Context("GetPodResourceMap() with valid pod name and namespace", func() {
-		It("should return no error", func() {
+		It("should return no error with device plugin resource", func() {
 			podUID := k8sTypes.UID("970a395d-bb3b-11e8-89df-408d5c537d23")
 			fakePod := &v1.Pod{
 				ObjectMeta: metav1.ObjectMeta{
@@ -209,6 +241,34 @@ var _ = Describe("Kubelet resource endpoint data read operations", func() {
 			Expect(resourceMap).To(Equal(outputRMap))
 		})
 
+		It("should return no error with dynamic resource", func() {
+			podUID := k8sTypes.UID("9f94e27b-4233-43d6-bd10-f73b4de6f456")
+			fakePod := &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "dynamic-resource-pod-name",
+					Namespace: "dynamic-resource-pod-namespace",
+					UID:       podUID,
+				},
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{
+						{
+							Name: "dynamic-resource-container-name",
+						},
+					},
+				},
+			}
+			client, err := getKubeletClient(testKubeletSocket)
+			Expect(err).NotTo(HaveOccurred())
+
+			outputRMap := map[string]*mtypes.ResourceInfo{
+				"resource-class": {DeviceIDs: []string{"cdi-resource"}},
+			}
+			resourceMap, err := client.GetPodResourceMap(fakePod)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resourceMap).ShouldNot(BeNil())
+			Expect(resourceMap).To(Equal(outputRMap))
+		})
+
 		It("should return an error with garbage socket value", func() {
 			u, err := url.Parse("/badfilepath!?//")
 			Expect(err).NotTo(HaveOccurred())