Skip to content

Commit

Permalink
Clean up nodedb tests (#2633)
Browse files Browse the repository at this point in the history
  • Loading branch information
@zuqq
zuqq authored Jul 3, 2023
1 parent 24fb052 commit ee1455f
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 93 deletions.
2 changes: 1 addition & 1 deletion internal/scheduler/nodedb/nodedb.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ type Node struct {

TotalResources schedulerobjects.ResourceList

// This field is set when inserting the Node into the node database.
// This field is set when inserting the Node into a NodeDb.
Keys [][]byte

NodeTypeId uint64
Expand Down
16 changes: 8 additions & 8 deletions internal/scheduler/nodedb/nodedb_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ func TestNodeDbSchema(t *testing.T) {

// Test the accounting of total resources across all nodes.
func TestTotalResources(t *testing.T) {
nodeDb, err := createNodeDb([]*schedulerobjects.Node{})
nodeDb, err := newNodeDbWithNodes([]*schedulerobjects.Node{})
require.NoError(t, err)

expected := schedulerobjects.ResourceList{Resources: make(map[string]resource.Quantity)}
Expand Down Expand Up @@ -65,7 +65,7 @@ func TestSelectNodeForPod_NodeIdLabel_Success(t *testing.T) {
nodes := testfixtures.N32CpuNodes(2, testfixtures.TestPriorities)
nodeId := nodes[1].Id
require.NotEmpty(t, nodeId)
db, err := createNodeDb(nodes)
db, err := newNodeDbWithNodes(nodes)
require.NoError(t, err)
jobs := testfixtures.WithNodeSelectorJobs(
map[string]string{schedulerconfig.NodeIdLabel: nodeId},
Expand Down Expand Up @@ -94,7 +94,7 @@ func TestSelectNodeForPod_NodeIdLabel_Failure(t *testing.T) {
nodes := testfixtures.N32CpuNodes(1, testfixtures.TestPriorities)
nodeId := nodes[0].Id
require.NotEmpty(t, nodeId)
db, err := createNodeDb(nodes)
db, err := newNodeDbWithNodes(nodes)
require.NoError(t, err)
jobs := testfixtures.WithNodeSelectorJobs(
map[string]string{schedulerconfig.NodeIdLabel: "this node does not exist"},
Expand All @@ -119,7 +119,7 @@ func TestSelectNodeForPod_NodeIdLabel_Failure(t *testing.T) {

func TestNodeBindingEvictionUnbinding(t *testing.T) {
node := testfixtures.Test8GpuNode(testfixtures.TestPriorities)
nodeDb, err := createNodeDb([]*schedulerobjects.Node{node})
nodeDb, err := newNodeDbWithNodes([]*schedulerobjects.Node{node})
require.NoError(t, err)
entry, err := nodeDb.GetNode(node.Id)
require.NoError(t, err)
Expand Down Expand Up @@ -276,7 +276,7 @@ func TestEviction(t *testing.T) {
}
for name, tc := range tests {
t.Run(name, func(t *testing.T) {
nodeDb, err := createNodeDb([]*schedulerobjects.Node{})
nodeDb, err := newNodeDbWithNodes([]*schedulerobjects.Node{})
require.NoError(t, err)
txn := nodeDb.Txn(true)
jobs := []*jobdb.Job{
Expand Down Expand Up @@ -432,7 +432,7 @@ func TestScheduleIndividually(t *testing.T) {
}
for name, tc := range tests {
t.Run(name, func(t *testing.T) {
nodeDb, err := createNodeDb(tc.Nodes)
nodeDb, err := newNodeDbWithNodes(tc.Nodes)
require.NoError(t, err)

jctxs := schedulercontext.JobSchedulingContextsFromJobs(testfixtures.TestPriorityClasses, tc.Jobs)
Expand Down Expand Up @@ -516,7 +516,7 @@ func TestScheduleMany(t *testing.T) {
}
for name, tc := range tests {
t.Run(name, func(t *testing.T) {
nodeDb, err := createNodeDb(tc.Nodes)
nodeDb, err := newNodeDbWithNodes(tc.Nodes)
require.NoError(t, err)
for i, jobs := range tc.Jobs {
jctxs := schedulercontext.JobSchedulingContextsFromJobs(testfixtures.TestPriorityClasses, jobs)
Expand Down Expand Up @@ -696,7 +696,7 @@ func BenchmarkScheduleManyResourceConstrained(b *testing.B) {
)
}

func createNodeDb(nodes []*schedulerobjects.Node) (*NodeDb, error) {
func newNodeDbWithNodes(nodes []*schedulerobjects.Node) (*NodeDb, error) {
nodeDb, err := NewNodeDb(
testfixtures.TestPriorityClasses,
testfixtures.TestMaxExtraNodesToConsider,
Expand Down
129 changes: 45 additions & 84 deletions internal/scheduler/nodedb/nodeiteration_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,6 @@ import (
"fmt"
"testing"

"github.com/hashicorp/go-memdb"
"github.com/pkg/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/exp/maps"
Expand Down Expand Up @@ -37,7 +35,7 @@ func TestNodesIterator(t *testing.T) {
for i, node := range tc.Nodes {
indexById[node.Id] = i
}
nodeDb, err := createNodeDb(tc.Nodes)
nodeDb, err := newNodeDbWithNodes(tc.Nodes)
if !assert.NoError(t, err) {
return
}
Expand All @@ -64,30 +62,21 @@ func TestNodesIterator(t *testing.T) {
}

func TestNodePairIterator(t *testing.T) {
nodeDb, err := createNodeDb(nil)
require.NoError(t, err)

nodes := testfixtures.TestCluster()
for i, c := range []string{"A", "B", "C"} {
nodes[i].Id = c
for i, nodeId := range []string{"A", "B", "C"} {
nodes[i].Id = nodeId
}
nodeDb, err := newNodeDbWithNodes(nodes)
require.NoError(t, err)
entries := make([]*Node, len(nodes))
for i, node := range nodes {
entry, err := nodeDb.create(node)
entry, err := nodeDb.GetNode(node.Id)
require.NoError(t, err)
entries[i] = entry
}

for _, node := range entries {
node.Keys = make([][]byte, len(nodeDb.prioritiesToTryAssigningAt))
for i, p := range nodeDb.prioritiesToTryAssigningAt {
node.Keys[i] = nodeDb.nodeDbKey(node.Keys[i], node.NodeTypeId, node.AllocatableByPriority[p])
}
}

txn := nodeDb.Txn(true)
require.NoError(t, txn.Insert("nodes", entries[0]))
require.NoError(t, txn.Insert("nodes", entries[1]))
require.NoError(t, txn.Delete("nodes", entries[2]))
txn.Commit()
txnA := nodeDb.Txn(false)

Expand Down Expand Up @@ -420,37 +409,24 @@ func TestNodeTypeIterator(t *testing.T) {
}
for name, tc := range tests {
t.Run(name, func(t *testing.T) {
nodeDb, err := createNodeDb(nil)
nodeDb, err := newNodeDbWithNodes(nil)
require.NoError(t, err)

// Set monotonically increaseing node ids to ensure nodes appear in predictable order.
entries := make([]*Node, len(tc.nodes))
for i, node := range tc.nodes {
// Set monotonically increasing node IDs to ensure nodes appear in predictable order.
node.Id = fmt.Sprintf("%d", i)
}
indexByNodeId := make(map[string]int)
for i, node := range tc.nodes {
indexByNodeId[node.Id] = i
}

entries := make([]*Node, len(tc.nodes))
for i, node := range tc.nodes {
entry, err := nodeDb.create(node)
require.NoError(t, err)

// We can safely override NodeTypeId, because Keys is recomputed upon insertion.
entry.NodeTypeId = node.NodeTypeId
entries[i] = entry
}

// Compute the keys necessary to efficiently iterate over nodes
// and populate the database. We do this manually instead of using nodeDb.Upsert to control the nodeTypeId.
for _, node := range entries {
node.Keys = make([][]byte, len(nodeDb.prioritiesToTryAssigningAt))
for i, p := range nodeDb.prioritiesToTryAssigningAt {
node.Keys[i] = nodeDb.nodeDbKey(node.Keys[i], node.NodeTypeId, node.AllocatableByPriority[p])
}
entries[i] = entry
}
require.NoError(t, populateDatabase(nodeDb.db, entries))
require.NoError(t, nodeDb.UpsertMany(entries))

// Create iterator.
indexedResourceRequests := make([]resource.Quantity, len(testfixtures.TestResources))
for i, t := range nodeDb.indexedResources {
indexedResourceRequests[i] = tc.resourceRequests.Get(t)
Expand All @@ -462,22 +438,31 @@ func TestNodeTypeIterator(t *testing.T) {
}
}
require.NotEqual(t, -1, keyIndex)
it, err := NewNodeTypeIterator(nodeDb.Txn(false), tc.nodeTypeId, nodeIndexName(keyIndex), tc.priority, testfixtures.TestResourceNames, indexedResourceRequests, testfixtures.TestIndexedResourceResolutionMillis)
it, err := NewNodeTypeIterator(
nodeDb.Txn(false),
tc.nodeTypeId,
nodeIndexName(keyIndex),
tc.priority,
testfixtures.TestResourceNames,
indexedResourceRequests,
testfixtures.TestIndexedResourceResolutionMillis,
)
require.NoError(t, err)

// Compare actual with expected order.
actual := make([]int, 0)
expected := make([]string, len(tc.expected))
for i, nodeId := range tc.expected {
expected[i] = fmt.Sprintf("%d", nodeId)
}
actual := make([]string, 0)
for {
node, err := it.NextNode()
require.NoError(t, err)
if node == nil {
break
}
i, ok := indexByNodeId[node.Id]
require.True(t, ok)
actual = append(actual, i)
actual = append(actual, node.Id)
}
assert.Equal(t, tc.expected, actual)
assert.Equal(t, expected, actual)

// Calling next should always return nil from now on.
for i := 0; i < 100; i++ {
Expand Down Expand Up @@ -814,35 +799,23 @@ func TestNodeTypesIterator(t *testing.T) {
}
for name, tc := range tests {
t.Run(name, func(t *testing.T) {
nodeDb, err := createNodeDb(nil)
nodeDb, err := newNodeDbWithNodes(nil)
require.NoError(t, err)

// Set monotonically increaseing node ids to ensure nodes appear in predictable order.
entries := make([]*Node, len(tc.nodes))
for i, node := range tc.nodes {
// Set monotonically increasing node IDs to ensure nodes appear in predictable order.
node.Id = fmt.Sprintf("%d", i)
}
indexByNodeId := make(map[string]int)
for i, node := range tc.nodes {
indexByNodeId[node.Id] = i
}

entries := make([]*Node, len(tc.nodes))
for i, node := range tc.nodes {
entry, err := nodeDb.create(node)
require.NoError(t, err)

// We can safely override NodeTypeId, because Keys is recomputed upon insertion.
entry.NodeTypeId = node.NodeTypeId
entries[i] = entry
}

// Compute the keys necessary to efficiently iterate over nodes
// and populate the database. We do this manually instead of using nodeDb.Upsert to control the nodeTypeId.
for _, node := range entries {
node.Keys = make([][]byte, len(nodeDb.prioritiesToTryAssigningAt))
for i, p := range nodeDb.prioritiesToTryAssigningAt {
node.Keys[i] = nodeDb.nodeDbKey(node.Keys[i], node.NodeTypeId, node.AllocatableByPriority[p])
}
entries[i] = entry
}
require.NoError(t, populateDatabase(nodeDb.db, entries))
require.NoError(t, nodeDb.UpsertMany(entries))

indexedResourceRequests := make([]resource.Quantity, len(testfixtures.TestResources))
for i, t := range testfixtures.TestResourceNames {
Expand All @@ -859,19 +832,20 @@ func TestNodeTypesIterator(t *testing.T) {
)
require.NoError(t, err)

// Compare actual with expected order.
actual := make([]int, 0)
expected := make([]string, len(tc.expected))
for i, nodeId := range tc.expected {
expected[i] = fmt.Sprintf("%d", nodeId)
}
actual := make([]string, 0)
for {
node, err := it.NextNode()
require.NoError(t, err)
if node == nil {
break
}
i, ok := indexByNodeId[node.Id]
require.True(t, ok)
actual = append(actual, i)
actual = append(actual, node.Id)
}
assert.Equal(t, tc.expected, actual)
assert.Equal(t, expected, actual)

// Calling next again should still return nil.
node, err := it.NextNode()
Expand All @@ -881,19 +855,6 @@ func TestNodeTypesIterator(t *testing.T) {
}
}

func populateDatabase(db *memdb.MemDB, nodes []*Node) error {
txn := db.Txn(true)
defer txn.Abort()
for _, node := range nodes {
err := txn.Insert("nodes", node)
if err != nil {
return errors.WithStack(err)
}
}
txn.Commit()
return nil
}

func BenchmarkNodeTypeIterator(b *testing.B) {
// Create nodes with varying amounts of CPU available.
numNodes := 1000
Expand All @@ -912,7 +873,7 @@ func BenchmarkNodeTypeIterator(b *testing.B) {
[]*schedulerobjects.Node{node},
)
}
nodeDb, err := createNodeDb(nodes)
nodeDb, err := newNodeDbWithNodes(nodes)
require.NoError(b, err)

// Create iterator for 0 CPU required and an unfeasible memory request,
Expand Down

0 comments on commit ee1455f

Please sign in to comment.