fix(awsneuron): allocate by device index, not slice position

pkg/device/awsneuron/device.go

@@ -17,6 +17,7 @@ limitations under the License.
 package awsneuron
 
 import (
+	"cmp"
 	"fmt"
 	"math"
 	"slices"
@@ -389,6 +390,10 @@ func addCoreUsage(prev map[string]any, require int) map[string]any {
 	}
 	return res
 }
+
+// continuousDeviceAvailable reports the device indexes of a free run of count
+// devices starting at slice position start. It returns indexes rather than
+// positions because Fit matches its result against DeviceUsage.Index.
 func continuousDeviceAvailable(devices []*device.DeviceUsage, start int, count int) []int {
 	if len(devices) < start+count {
 		return []int{}
@@ -399,14 +404,29 @@ func continuousDeviceAvailable(devices []*device.DeviceUsage, start int, count i
 		if devices[iterator].Used > 0 || !devices[iterator].Health {
 			return []int{}
 		}
-		res = append(res, iterator)
+		res = append(res, int(devices[iterator].Index))
 		iterator++
 	}
 	return res
 }
 
 func graphSelect(devices []*device.DeviceUsage, count int) []int {
-	if len(devices) == 0 || devices[0].CustomInfo == nil || devices[0].CustomInfo[AWSNodeType] == nil {
+	if len(devices) == 0 {
+		return []int{}
+	}
+	// The ring and power of two groupings below read adjacency from slice
+	// positions, but the scheduler sorts devices by score before calling Fit.
+	// Restore index order on a copy so position math matches the physical
+	// layout, the same way kunlun's topology selector does. Ordering first
+	// also puts the device carrying the node type back at position zero.
+	sorted := make([]*device.DeviceUsage, len(devices))
+	copy(sorted, devices)
+	slices.SortFunc(sorted, func(a, b *device.DeviceUsage) int {
+		return cmp.Compare(a.Index, b.Index)
+	})
+	devices = sorted
+
+	if devices[0].CustomInfo == nil || devices[0].CustomInfo[AWSNodeType] == nil {
 		return []int{}
 	}
 	AWSNodetype := ""

pkg/device/awsneuron/device_test.go

@@ -1102,3 +1102,115 @@ func TestDevices_Fit(t *testing.T) {
 		})
 	}
 }
+
+// Test_graphSelect_ScoreSortedInput covers the order the scheduler actually
+// hands to Fit. DeviceUsageList sorts by score, so a busy device can sit at
+// position zero while its Index is elsewhere. graphSelect must reason about
+// physical adjacency and report indexes, because Fit matches its result
+// against DeviceUsage.Index. Every device carries the node type here, as
+// GetNodeDevices sets CustomInfo on all of them.
+func Test_graphSelect_ScoreSortedInput(t *testing.T) {
+	devices := makeDeviceUsages("trn", map[int]int32{5: 1}, true)
+	for _, d := range devices {
+		d.CustomInfo = map[string]any{AWSNodeType: "trn"}
+	}
+
+	// Move the busy device to the front, as a score sort would.
+	scoreOrdered := []*device.DeviceUsage{devices[5]}
+	for i, d := range devices {
+		if i != 5 {
+			scoreOrdered = append(scoreOrdered, d)
+		}
+	}
+
+	got := graphSelect(scoreOrdered, 4)
+	assert.DeepEqual(t, got, []int{0, 1, 2, 3})
+
+	// Every returned value has to name a free device.
+	byIndex := map[uint]*device.DeviceUsage{}
+	for _, d := range devices {
+		byIndex[d.Index] = d
+	}
+	for _, idx := range got {
+		d, ok := byIndex[uint(idx)]
+		assert.Assert(t, ok, "graphSelect returned %d, which is not a device index", idx)
+		assert.Assert(t, d.Used == 0, "graphSelect returned busy device index %d", idx)
+	}
+}
+
+// makeOffsetDeviceUsages builds devices whose Index deliberately does not equal
+// its slice position, so a result made of positions is distinguishable from one
+// made of indexes.
+func makeOffsetDeviceUsages(nodeType string, base uint, n int, used map[uint]int32) []*device.DeviceUsage {
+	devices := make([]*device.DeviceUsage, n)
+	for i := range n {
+		idx := base + uint(i)
+		du := &device.DeviceUsage{
+			Index:      idx,
+			Health:     true,
+			CustomInfo: map[string]any{AWSNodeType: nodeType},
+		}
+		if u, ok := used[idx]; ok {
+			du.Used = u
+		}
+		devices[i] = du
+	}
+	return devices
+}
+
+// Test_graphSelect_ReturnsIndexesNotPositions pins the contract Fit relies on.
+// Fit matches the returned values against DeviceUsage.Index, so graphSelect has
+// to report indexes. With an index set that does not start at zero, a result
+// built from slice positions names devices that do not exist, and Fit's lookup
+// loop then appends nothing while still reporting a successful fit.
+func Test_graphSelect_ReturnsIndexesNotPositions(t *testing.T) {
+	devices := makeOffsetDeviceUsages("trn", 100, 16, nil)
+
+	got := graphSelect(devices, 4)
+	assert.DeepEqual(t, got, []int{100, 101, 102, 103})
+
+	known := map[int]bool{}
+	for _, d := range devices {
+		known[int(d.Index)] = true
+	}
+	for _, idx := range got {
+		assert.Assert(t, known[idx], "graphSelect returned %d, which is not a device index", idx)
+	}
+}
+
+// Test_graphSelect_BusyDeviceWithOffsetIndexes checks the same contract when a
+// device is taken, which is the case that used to hand out an in-use device.
+func Test_graphSelect_BusyDeviceWithOffsetIndexes(t *testing.T) {
+	devices := makeOffsetDeviceUsages("trn", 100, 16, map[uint]int32{101: 1})
+
+	got := graphSelect(devices, 4)
+	assert.DeepEqual(t, got, []int{104, 105, 106, 107})
+
+	byIndex := map[uint]*device.DeviceUsage{}
+	for _, d := range devices {
+		byIndex[d.Index] = d
+	}
+	for _, idx := range got {
+		d, ok := byIndex[uint(idx)]
+		assert.Assert(t, ok, "graphSelect returned %d, which is not a device index", idx)
+		assert.Assert(t, d.Used == 0, "graphSelect returned busy device index %d", idx)
+	}
+}
+
+// Test_graphSelect_NodeTypeReadFromLowestIndex covers the node type lookup.
+// GetNodeDevices hands the same CustomInfo map to every device, but the check
+// only inspects the first element, so it has to run after the slice is ordered
+// by index rather than by score.
+func Test_graphSelect_NodeTypeReadFromLowestIndex(t *testing.T) {
+	devices := makeDeviceUsages("trn", nil, true)
+	// Only the lowest-index device carries the node type, and the score sort
+	// has moved a different device to the front.
+	scoreOrdered := []*device.DeviceUsage{devices[7]}
+	for i, d := range devices {
+		if i != 7 {
+			scoreOrdered = append(scoreOrdered, d)
+		}
+	}
+
+	assert.DeepEqual(t, graphSelect(scoreOrdered, 4), []int{0, 1, 2, 3})
+}