Merge pull request #4855 from hashicorp/b-device-scheduling-fixtest

Fixes device scheduling unit tests

scheduler/device.go

@@ -3,6 +3,8 @@ package scheduler
 import (
 	"fmt"
 
+	"math"
+
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -39,6 +41,7 @@ func (d *deviceAllocator) AssignDevice(ask *structs.RequestedDevice) (out *struc
 	// Hold the current best offer
 	var offer *structs.AllocatedDeviceResource
 	var offerScore float64
+	var matchedWeights float64
 
 	// Determine the devices that are feasible based on availability and
 	// constraints
@@ -63,6 +66,10 @@ func (d *deviceAllocator) AssignDevice(ask *structs.RequestedDevice) (out *struc
 
 		// Score the choice
 		var choiceScore float64
+
+		// Track the sum of matched affinity weights in a separate variable
+		// We return this if this device had the best score compared to other devices considered
+		var sumMatchedWeights float64
 		if l := len(ask.Affinities); l != 0 {
 			totalWeight := 0.0
 			for _, a := range ask.Affinities {
@@ -76,13 +83,14 @@ func (d *deviceAllocator) AssignDevice(ask *structs.RequestedDevice) (out *struc
 					continue
 				}
 
-				totalWeight += a.Weight
+				totalWeight += math.Abs(a.Weight)
 
 				// Check if satisfied
 				if !checkAttributeAffinity(d.ctx, a.Operand, lVal, rVal) {
 					continue
 				}
 				choiceScore += a.Weight
+				sumMatchedWeights += a.Weight
 			}
 
 			// normalize
@@ -97,6 +105,9 @@ func (d *deviceAllocator) AssignDevice(ask *structs.RequestedDevice) (out *struc
 		// Set the new highest score
 		offerScore = choiceScore
 
+		// Set the new sum of matching affinity weights
+		matchedWeights = sumMatchedWeights
+
 		// Build the choice
 		offer = &structs.AllocatedDeviceResource{
 			Vendor:    id.Vendor,
@@ -122,5 +133,5 @@ func (d *deviceAllocator) AssignDevice(ask *structs.RequestedDevice) (out *struc
 		return nil, 0.0, fmt.Errorf("no devices match request")
 	}
 
-	return offer, offerScore, nil
+	return offer, matchedWeights, nil
 }

scheduler/device_test.go

@@ -301,8 +301,8 @@ func TestDeviceAllocator_Allocate_Affinities(t *testing.T) {
 					Weight:  -0.2,
 				},
 			},
-			ExpectedDevice: nvidia0,
 			ZeroScore:      true,
+			ExpectedDevice: nvidia0,
 		},
 		{
 			Name: "nvidia/gpu",

scheduler/rank.go

@@ -197,7 +197,7 @@ OUTER:
 
 		// Track the affinities of the devices
 		totalDeviceAffinityWeight := 0.0
-		deviceAffinityScore := 0.0
+		sumMatchingAffinities := 0.0
 
 		// Assign the resources for each task
 		total := &structs.AllocatedResources{
@@ -283,7 +283,7 @@ OUTER:
 
 			// Check if we need to assign devices
 			for _, req := range task.Resources.Devices {
-				offer, score, err := devAllocator.AssignDevice(req)
+				offer, sumAffinities, err := devAllocator.AssignDevice(req)
 				if offer == nil {
 					iter.ctx.Metrics().ExhaustedNode(option.Node, fmt.Sprintf("devices: %s", err))
 					continue OUTER
@@ -296,9 +296,9 @@ OUTER:
 				// Add the scores
 				if len(req.Affinities) != 0 {
 					for _, a := range req.Affinities {
-						totalDeviceAffinityWeight += a.Weight
+						totalDeviceAffinityWeight += math.Abs(a.Weight)
 					}
-					deviceAffinityScore += score
+					sumMatchingAffinities += sumAffinities
 				}
 			}
 
@@ -353,9 +353,9 @@ OUTER:
 
 		// Score the device affinity
 		if totalDeviceAffinityWeight != 0 {
-			deviceAffinityScore /= totalDeviceAffinityWeight
-			option.Scores = append(option.Scores, deviceAffinityScore)
-			iter.ctx.Metrics().ScoreNode(option.Node, "devices", deviceAffinityScore)
+			sumMatchingAffinities /= totalDeviceAffinityWeight
+			option.Scores = append(option.Scores, sumMatchingAffinities)
+			iter.ctx.Metrics().ScoreNode(option.Node, "devices", sumMatchingAffinities)
 		}
 
 		return option

scheduler/rank_test.go

@@ -607,7 +607,7 @@ func TestBinPackIterator_Devices(t *testing.T) {
 					},
 				},
 			},
-			DeviceScore: 0.9,
+			DeviceScore: 1.0,
 		},
 		{
 			Name: "single request over count, no match",
@@ -784,7 +784,7 @@ func TestBinPackIterator_Devices(t *testing.T) {
 			// Check potential affinity scores
 			if c.DeviceScore != 0.0 {
 				require.Len(out.Scores, 2)
-				require.Equal(out.Scores[1], c.DeviceScore)
+				require.Equal(c.DeviceScore, out.Scores[1])
 			}
 		})
 	}