From eda98fb070246fbdb9eb208dfc28eda0a2c28a04 Mon Sep 17 00:00:00 2001
From: Preetha Appan <preetha@hashicorp.com>
Date: Thu, 8 Nov 2018 17:01:58 -0600
Subject: [PATCH] Normalize scores correctly

---
 scheduler/device.go      |  5 +++--
 scheduler/device_test.go |  1 -
 scheduler/rank.go        | 20 +++++++++++---------
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/scheduler/device.go b/scheduler/device.go
index 20063a30e..2efcdc4e9 100644
--- a/scheduler/device.go
+++ b/scheduler/device.go
@@ -41,7 +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
 	for id, devInst := range d.Devices {
@@ -86,6 +86,7 @@ func (d *deviceAllocator) AssignDevice(ask *structs.RequestedDevice) (out *struc
 				}
 				choiceScore += a.Weight
 			}
+			matchedWeights += choiceScore
 
 			// normalize
 			choiceScore /= totalWeight
@@ -124,5 +125,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
 }
diff --git a/scheduler/device_test.go b/scheduler/device_test.go
index 36643555f..167339fb7 100644
--- a/scheduler/device_test.go
+++ b/scheduler/device_test.go
@@ -302,7 +302,6 @@ func TestDeviceAllocator_Allocate_Affinities(t *testing.T) {
 				},
 			},
 			ExpectedDevice: nvidia0,
-			ZeroScore:      true,
 		},
 		{
 			Name: "nvidia/gpu",
diff --git a/scheduler/rank.go b/scheduler/rank.go
index 1674b2f76..7d2b4ad4e 100644
--- a/scheduler/rank.go
+++ b/scheduler/rank.go
@@ -196,8 +196,8 @@ OUTER:
 		devAllocator.AddAllocs(proposed)
 
 		// Track the affinities of the devices
-		devicesWithAffinities := 0.0
-		deviceAffinityScore := 0.0
+		totalDeviceAffinityWeight := 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
@@ -295,8 +295,10 @@ OUTER:
 
 				// Add the scores
 				if len(req.Affinities) != 0 {
-					devicesWithAffinities++
-					deviceAffinityScore += score
+					for _, a := range req.Affinities {
+						totalDeviceAffinityWeight += math.Abs(a.Weight)
+					}
+					sumMatchingAffinities += sumAffinities
 				}
 			}
 
@@ -350,10 +352,10 @@ OUTER:
 		iter.ctx.Metrics().ScoreNode(option.Node, "binpack", normalizedFit)
 
 		// Score the device affinity
-		if devicesWithAffinities != 0 {
-			deviceAffinityScore /= float64(devicesWithAffinities)
-			option.Scores = append(option.Scores, deviceAffinityScore)
-			iter.ctx.Metrics().ScoreNode(option.Node, "devices", deviceAffinityScore)
+		if totalDeviceAffinityWeight != 0 {
+			sumMatchingAffinities /= totalDeviceAffinityWeight
+			option.Scores = append(option.Scores, sumMatchingAffinities)
+			iter.ctx.Metrics().ScoreNode(option.Node, "devices", sumMatchingAffinities)
 		}
 
 		return option