fix(kunlun): honour use-gpuuuid and nouse-gpuuuid annotations

pkg/device/kunlun/device.go

@@ -184,10 +184,29 @@ func (kl *KunlunDevices) Fit(devices []*device.DeviceUsage, request device.Conta
 	tmpDevs := make(map[string]device.ContainerDevices)
 	reason := make(map[string]int)
 
-	alloc := graghSelect(devices, request, FitXPU)
+	// graghSelect decides topology from the position of a device in the slice,
+	// so the uuid constraint has to be applied through fitFn rather than by
+	// filtering the slice first.
+	uuidMismatches := make(map[string]bool)
+	fitFn := func(d *device.DeviceUsage, r device.ContainerDeviceRequest) bool {
+		if !device.CheckUUID(pod.GetAnnotations(), d.ID, UseUUIDAnno, NoUseUUIDAnno, kl.CommonWord()) ||
+			!device.CheckUUID(pod.GetAnnotations(), d.ID, KunlunUseUUID, KunlunNoUseUUID, kl.CommonWord()) {
+			uuidMismatches[d.ID] = true
+			klog.V(5).InfoS(common.CardUUIDMismatch, "pod", klog.KObj(pod), "device", d.ID)
+			return false
+		}
+		return FitXPU(d, r)
+	}
+
+	alloc := graghSelect(devices, request, fitFn)
 	if len(alloc) == 0 {
-		reason[common.NumaNotFit]++
-		klog.V(5).InfoS(common.NumaNotFit, "pod", klog.KObj(pod), "device", devices, "request nums", request.Nums, "numa")
+		uuidMismatch := len(uuidMismatches)
+		if uuidMismatch > 0 {
+			reason[common.CardUUIDMismatch] += uuidMismatch
+		} else {
+			reason[common.NumaNotFit]++
+			klog.V(5).InfoS(common.NumaNotFit, "pod", klog.KObj(pod), "device", devices, "request nums", request.Nums, "numa")
+		}
 		return false, tmpDevs, common.GenReason(reason, len(devices))
 	}
 

pkg/device/kunlun/device_test.go

@@ -17,6 +17,8 @@ limitations under the License.
 package kunlun
 
 import (
+	"fmt"
+	"strings"
 	"testing"
 
 	"gotest.tools/v3/assert"
@@ -511,3 +513,139 @@ func Test_ScoreNode(t *testing.T) {
 		})
 	}
 }
+
+// hami.io/use-kunlun-uuid and nouse-kunlun-uuid were defined but never consulted, so
+// a pod asking for a specific XPU was scheduled onto any of them.
+func TestKunlunDevices_Fit_UseUUID(t *testing.T) {
+	dev := &KunlunDevices{}
+	devices := make([]*device.DeviceUsage, 8)
+	for i := range devices {
+		devices[i] = &device.DeviceUsage{
+			Index:     uint(i),
+			ID:        fmt.Sprintf("xpu-%d", i),
+			Count:     1,
+			Totalmem:  KunlunMaxMemory,
+			Totalcore: 100,
+			Health:    true,
+		}
+	}
+	req := device.ContainerDeviceRequest{Nums: 1, Type: KunlunGPUDevice}
+
+	podWith := func(annos map[string]string) *corev1.Pod {
+		return &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Annotations: annos}}
+	}
+
+	// Test legacy annotation
+	fit, res, _ := dev.Fit(devices, req, podWith(map[string]string{
+		KunlunUseUUID: "xpu-5",
+	}), &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, true)
+	assert.Equal(t, len(res[KunlunGPUDevice]), 1)
+	assert.Equal(t, res[KunlunGPUDevice][0].UUID, "xpu-5")
+
+	// Test new standard annotation
+	fit, res, _ = dev.Fit(devices, req, podWith(map[string]string{
+		UseUUIDAnno: "xpu-3",
+	}), &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, true)
+	assert.Equal(t, len(res[KunlunGPUDevice]), 1)
+	assert.Equal(t, res[KunlunGPUDevice][0].UUID, "xpu-3")
+
+	// the literal key, so a rename of the constant cannot silently stop the
+	// physical path from reading what vdevice.go documents.
+	fit, res, _ = dev.Fit(devices, req, podWith(map[string]string{
+		"hami.io/use-xpu-uuid": "xpu-6",
+	}), &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, true)
+	assert.Equal(t, len(res[KunlunGPUDevice]), 1)
+	assert.Equal(t, res[KunlunGPUDevice][0].UUID, "xpu-6")
+
+	// asking for a card that is not on the node must not fall back to another
+	fit, _, reason := dev.Fit(devices, req, podWith(map[string]string{
+		UseUUIDAnno: "xpu-99",
+	}), &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, false)
+	assert.Assert(t, strings.Contains(reason, common.CardUUIDMismatch))
+}
+
+func TestKunlunDevices_Fit_NoUseUUID(t *testing.T) {
+	dev := &KunlunDevices{}
+	devices := make([]*device.DeviceUsage, 8)
+	for i := range devices {
+		devices[i] = &device.DeviceUsage{
+			Index:     uint(i),
+			ID:        fmt.Sprintf("xpu-%d", i),
+			Count:     1,
+			Totalmem:  KunlunMaxMemory,
+			Totalcore: 100,
+			Health:    true,
+		}
+	}
+	req := device.ContainerDeviceRequest{Nums: 1, Type: KunlunGPUDevice}
+
+	// exclude every card and nothing should be allocatable (using legacy annotation)
+	all := make([]string, 0, 8)
+	for i := range devices {
+		all = append(all, fmt.Sprintf("xpu-%d", i))
+	}
+	fit, _, reason := dev.Fit(devices, req, &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
+			KunlunNoUseUUID: strings.Join(all, ","),
+		}},
+	}, &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, false)
+	assert.Assert(t, strings.Contains(reason, common.CardUUIDMismatch))
+
+	// exclude every card and nothing should be allocatable (using new standard annotation)
+	fit, _, reason = dev.Fit(devices, req, &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
+			NoUseUUIDAnno: strings.Join(all, ","),
+		}},
+	}, &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, false)
+	assert.Assert(t, strings.Contains(reason, common.CardUUIDMismatch))
+}
+
+func TestKunlunVDevices_Fit_UUIDAnnotations(t *testing.T) {
+	dev := &KunlunVDevices{}
+	devices := make([]*device.DeviceUsage, 8)
+	for i := range devices {
+		devices[i] = &device.DeviceUsage{
+			Index:     uint(i),
+			ID:        fmt.Sprintf("xpu-%d", i),
+			Count:     10,
+			Totalmem:  KunlunMaxMemory,
+			Totalcore: 100,
+			Health:    true,
+		}
+	}
+	req := device.ContainerDeviceRequest{Nums: 1, Type: XPUDevice, Memreq: KunlunMaxMemory}
+
+	// Test standard hami annotation
+	fit, res, _ := dev.Fit(devices, req, &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{UseUUIDAnno: "xpu-3"}},
+	}, &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, true)
+	assert.Equal(t, res[XPUDevice][0].UUID, "xpu-3")
+
+	// Test legacy baidu annotation (backward compatibility)
+	fit, res, _ = dev.Fit(devices, req, &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{KunlunUseUUID: "xpu-2"}},
+	}, &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, true)
+	assert.Equal(t, res[XPUDevice][0].UUID, "xpu-2")
+
+	// Test standard hami nouse annotation
+	fit, _, reason := dev.Fit(devices, req, &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{NoUseUUIDAnno: "xpu-0,xpu-1,xpu-2,xpu-3,xpu-4,xpu-5,xpu-6,xpu-7"}},
+	}, &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, false)
+	assert.Assert(t, strings.Contains(reason, common.CardUUIDMismatch))
+
+	// Test legacy baidu nouse annotation
+	fit, _, reason = dev.Fit(devices, req, &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{KunlunNoUseUUID: "xpu-0,xpu-1,xpu-2,xpu-3,xpu-4,xpu-5,xpu-6,xpu-7"}},
+	}, &device.NodeInfo{}, &device.PodDevices{})
+	assert.Equal(t, fit, false)
+	assert.Assert(t, strings.Contains(reason, common.CardUUIDMismatch))
+}

pkg/device/kunlun/vdevice.go

@@ -234,13 +234,26 @@ func (dev *KunlunVDevices) Fit(devices []*device.DeviceUsage, request device.Con
 	reason := make(map[string]int)
 
 	isMutex := util.PolicyContains(util.GetGPUSchedulerPolicyByPod(device.GPUSchedulerPolicy, pod), util.GPUSchedulerPolicyMutex)
-	fitFn := FitFn(FitVXPU)
+	base := FitFn(FitVXPU)
 	if isMutex {
 		// mutex: only idle devices are eligible, no sharing onto a used device.
-		fitFn = func(d *device.DeviceUsage, r device.ContainerDeviceRequest) bool {
+		base = func(d *device.DeviceUsage, r device.ContainerDeviceRequest) bool {
 			return d.Used == 0 && FitVXPU(d, r)
 		}
 	}
+	// graghSelect decides topology from the position of a device in the slice,
+	// so the uuid constraint has to be applied through fitFn rather than by
+	// filtering the slice first.
+	uuidMismatches := make(map[string]bool)
+	fitFn := func(d *device.DeviceUsage, r device.ContainerDeviceRequest) bool {
+		if !device.CheckUUID(pod.GetAnnotations(), d.ID, UseUUIDAnno, NoUseUUIDAnno, dev.CommonWord()) ||
+			!device.CheckUUID(pod.GetAnnotations(), d.ID, KunlunUseUUID, KunlunNoUseUUID, dev.CommonWord()) {
+			uuidMismatches[d.ID] = true
+			klog.V(5).InfoS(common.CardUUIDMismatch, "pod", klog.KObj(pod), "device", d.ID)
+			return false
+		}
+		return base(d, r)
+	}
 	alloc := graghSelect(devices, request, fitFn)
 	if len(alloc) == 0 {
 		if isMutex {
@@ -251,6 +264,10 @@ func (dev *KunlunVDevices) Fit(devices []*device.DeviceUsage, request device.Con
 				}
 			}
 		}
+		uuidMismatch := len(uuidMismatches)
+		if len(reason) == 0 && uuidMismatch > 0 {
+			reason[common.CardUUIDMismatch] += uuidMismatch
+		}
 		if len(reason) == 0 {
 			reason[common.NumaNotFit]++
 			klog.V(5).InfoS(common.NumaNotFit, "pod", klog.KObj(pod), "device", devices, "request nums", request.Nums)