session_plugins_test.go

/*
Copyright 2023 NVIDIA CORPORATION
SPDX-License-Identifier: Apache-2.0
*/

package framework

import (
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	kaiv1 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/kai/v1"
	"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api"
	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/node_info"
	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/podgroup_info"
	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/podgroup_info/subgroup_info"
	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/resource_info"
	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/conf"
)

type testScenarioGeneratorContext struct {
	action ActionType
}

func (ctx testScenarioGeneratorContext) Action() ActionType {
	return ctx.action
}

type testScenarioGenerator struct {
	name string
}

func (generator *testScenarioGenerator) Name() string {
	return generator.name
}

func (generator *testScenarioGenerator) Next() api.ScenarioInfo {
	return nil
}

func newTestScenarioGenerator(name string) ScenarioGeneratorFactory {
	return func(_ ScenarioGeneratorContext) ScenarioGenerator {
		return &testScenarioGenerator{name: name}
	}
}

func TestLogNodeSetsPluginResultDoesNotAllocateWhenVerboseLoggingIsDisabled(t *testing.T) {
	podGroup := podgroup_info.NewPodGroupInfo("pod-group")
	nodeSets := []node_info.NodeSet{{{Name: "node-a"}}}

	allocations := testing.AllocsPerRun(100, func() {
		logNodeSetsPluginResult(nil, podGroup, nodeSets)
	})

	assert.Zero(t, allocations)
}

func TestMutateBindRequestAnnotations(t *testing.T) {
	tests := []struct {
		name                string
		mutateFns           []api.BindRequestMutateFn
		expectedAnnotations map[string]string
	}{
		{
			name:                "no mutate functions",
			mutateFns:           []api.BindRequestMutateFn{},
			expectedAnnotations: map[string]string{},
		},
		{
			name: "single mutate function",
			mutateFns: []api.BindRequestMutateFn{
				func(pod *pod_info.PodInfo, nodeName string) map[string]string {
					return map[string]string{"key1": "value1"}
				},
			},
			expectedAnnotations: map[string]string{"key1": "value1"},
		},
		{
			name: "multiple mutate functions with different keys",
			mutateFns: []api.BindRequestMutateFn{
				func(pod *pod_info.PodInfo, nodeName string) map[string]string {
					return map[string]string{"key1": "value1"}
				},
				func(pod *pod_info.PodInfo, nodeName string) map[string]string {
					return map[string]string{"key2": "value2"}
				},
			},
			expectedAnnotations: map[string]string{"key1": "value1", "key2": "value2"},
		},
		{
			name: "multiple mutate functions with overlapping keys - later should override",
			mutateFns: []api.BindRequestMutateFn{
				func(pod *pod_info.PodInfo, nodeName string) map[string]string {
					return map[string]string{"key1": "value1", "common": "first"}
				},
				func(pod *pod_info.PodInfo, nodeName string) map[string]string {
					return map[string]string{"key2": "value2", "common": "second"}
				},
			},
			expectedAnnotations: map[string]string{"key1": "value1", "key2": "value2", "common": "second"},
		},
		{
			name: "mutate function returns nil map",
			mutateFns: []api.BindRequestMutateFn{
				func(pod *pod_info.PodInfo, nodeName string) map[string]string {
					return map[string]string{"key1": "value1"}
				},
				func(pod *pod_info.PodInfo, nodeName string) map[string]string {
					return nil
				},
			},
			expectedAnnotations: map[string]string{"key1": "value1"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ssn := &Session{
				BindRequestMutateFns: tt.mutateFns,
			}
			pod := &pod_info.PodInfo{
				Name: "test-pod",
			}
			nodeName := "test-node"
			annotations := ssn.MutateBindRequestAnnotations(pod, nodeName)
			assert.Equal(t, tt.expectedAnnotations, annotations)
		})
	}
}

func TestPartitionMultiImplementation(t *testing.T) {
	nodes := []*node_info.NodeInfo{
		{
			Name: "cluster1rack0-1",
		},
		{
			Name: "cluster0rack0",
		},
		{
			Name: "cluster1rack1-1",
		},
		{
			Name: "cluster0rack1",
		},
		{
			Name: "cluster1rack0-2",
		},
		{
			Name: "cluster1rack1-2",
		},
	}

	shardClusterSubseting := func(_ *podgroup_info.PodGroupInfo, _ *subgroup_info.SubGroupInfo, _ map[string]*subgroup_info.PodSet, _ []*pod_info.PodInfo, nodeset node_info.NodeSet) ([]node_info.NodeSet, error) {
		var subset1 []*node_info.NodeInfo
		var subset2 []*node_info.NodeInfo
		for _, node := range nodeset {
			if strings.Contains(node.Name, "cluster0") {
				subset1 = append(subset1, node)
			} else {
				subset2 = append(subset2, node)
			}
		}
		return []node_info.NodeSet{subset1, subset2}, nil
	}

	topologySubseting := func(_ *podgroup_info.PodGroupInfo, _ *subgroup_info.SubGroupInfo, _ map[string]*subgroup_info.PodSet, _ []*pod_info.PodInfo, nodeset node_info.NodeSet) ([]node_info.NodeSet, error) {
		var subset1 []*node_info.NodeInfo
		var subset2 []*node_info.NodeInfo
		for _, node := range nodeset {
			if strings.Contains(node.Name, "rack0") {
				subset1 = append(subset1, node)
			} else {
				subset2 = append(subset2, node)
			}
		}
		return []node_info.NodeSet{subset1, subset2}, nil
	}

	ssn := &Session{}

	ssn.AddSubsetNodesFn(shardClusterSubseting)
	ssn.AddSubsetNodesFn(topologySubseting)

	partitions, _ := ssn.SubsetNodesFn(podgroup_info.NewPodGroupInfo("a"), nil, nil, nil, nodes)

	assert.Equal(t, 4, len(partitions))

	assert.Equal(t, len(partitions[0]), 1)
	assert.Equal(t, partitions[0][0].Name, "cluster0rack0")

	assert.Equal(t, len(partitions[1]), 1)
	assert.Equal(t, partitions[1][0].Name, "cluster0rack1")

	assert.Equal(t, len(partitions[2]), 2)
	assert.Equal(t, partitions[2][0].Name, "cluster1rack0-1")
	assert.Equal(t, partitions[2][1].Name, "cluster1rack0-2")

	assert.Equal(t, len(partitions[3]), 2)
	assert.Equal(t, partitions[3][0].Name, "cluster1rack1-1")
	assert.Equal(t, partitions[3][1].Name, "cluster1rack1-2")
}

func TestVictimInvariantPrePredicateFailure(t *testing.T) {
	task := &pod_info.PodInfo{Name: "task-1"}
	expectedErr := errors.New("missing pvc")

	t.Run("returns nil when no functions are registered", func(t *testing.T) {
		ssn := &Session{}
		assert.Nil(t, ssn.VictimInvariantPrePredicateFailure(task))
	})

	t.Run("returns the first non-nil failure", func(t *testing.T) {
		ssn := &Session{}
		secondCalled := false
		ssn.AddVictimInvariantPrePredicateFn(func(_ *pod_info.PodInfo) *api.VictimInvariantPrePredicateFailure {
			return nil
		})
		ssn.AddVictimInvariantPrePredicateFn(func(gotTask *pod_info.PodInfo) *api.VictimInvariantPrePredicateFailure {
			assert.Same(t, task, gotTask)
			return &api.VictimInvariantPrePredicateFailure{
				Err: expectedErr,
			}
		})
		ssn.AddVictimInvariantPrePredicateFn(func(_ *pod_info.PodInfo) *api.VictimInvariantPrePredicateFailure {
			secondCalled = true
			return &api.VictimInvariantPrePredicateFailure{
				Err: errors.New("should not be returned"),
			}
		})

		failure := ssn.VictimInvariantPrePredicateFailure(task)
		if assert.NotNil(t, failure) {
			assert.Same(t, expectedErr, failure.Err)
		}
		assert.False(t, secondCalled)
	})
}

func TestAddScenarioGeneratorPreservesOrder(t *testing.T) {
	ssn := &Session{}

	ssn.AddScenarioGenerator("first", newTestScenarioGenerator("first"))
	ssn.AddScenarioGenerator("second", newTestScenarioGenerator("second"))
	ssn.AddScenarioGenerator("third", newTestScenarioGenerator("third"))

	require.Len(t, ssn.ScenarioGeneratorRegistrations, 3)
	require.Equal(t, "first", ssn.ScenarioGeneratorRegistrations[0].Name)
	require.Equal(t, "second", ssn.ScenarioGeneratorRegistrations[1].Name)
	require.Equal(t, "third", ssn.ScenarioGeneratorRegistrations[2].Name)

	generator := ssn.ScenarioGeneratorRegistrations[0].Factory(testScenarioGeneratorContext{action: Reclaim})
	require.Equal(t, "first", generator.Name())
}

func TestValidateScenarioGeneratorBudgetKeys(t *testing.T) {
	ssn := &Session{
		Config: &conf.SchedulerConfiguration{
			ScenarioSearchBudgets: &kaiv1.ScenarioSearchBudgets{
				MaxGeneratorSearchDuration: map[string]metav1.Duration{
					constants.ActionDefault: scenarioSearchDurationForTest("1s"),
					"first":                 scenarioSearchDurationForTest("2s"),
				},
			},
		},
	}
	ssn.AddScenarioGenerator("first", newTestScenarioGenerator("first"))

	require.NoError(t, ssn.ValidateScenarioGeneratorBudgetKeys())

	ssn.Config.ScenarioSearchBudgets.MaxGeneratorSearchDuration["missing"] = scenarioSearchDurationForTest("3s")
	require.EqualError(t, ssn.ValidateScenarioGeneratorBudgetKeys(),
		`unknown scenario generator budget key "missing"`)
}

func TestValidateScenarioGeneratorBudgetKeysAcceptsBuiltInGeneratorsWithoutPlugins(t *testing.T) {
	ssn := &Session{
		Config: &conf.SchedulerConfiguration{
			ScenarioSearchBudgets: &kaiv1.ScenarioSearchBudgets{
				MaxGeneratorSearchDuration: map[string]metav1.Duration{
					constants.GeneratorNodeLocalGreedy: scenarioSearchDurationForTest("30s"),
					constants.GeneratorMultiNodeGang:   scenarioSearchDurationForTest("2m"),
				},
			},
		},
	}

	require.NoError(t, ssn.ValidateScenarioGeneratorBudgetKeys())

	ssn.Config.ScenarioSearchBudgets.MaxGeneratorSearchDuration["missing"] = scenarioSearchDurationForTest("3s")
	require.EqualError(t, ssn.ValidateScenarioGeneratorBudgetKeys(),
		`unknown scenario generator budget key "missing"`)
}

func scenarioSearchDurationForTest(value string) metav1.Duration {
	duration, err := time.ParseDuration(value)
	if err != nil {
		panic(err)
	}
	return metav1.Duration{Duration: duration}
}

func TestVictimOrderFn_PreservesExistingJobOrderFnBehavior(t *testing.T) {
	priorityLikeJobOrderFn := func(l, r interface{}) int {
		lv := l.(*podgroup_info.PodGroupInfo)
		rv := r.(*podgroup_info.PodGroupInfo)
		if lv.Priority > rv.Priority {
			return -1
		}
		if lv.Priority < rv.Priority {
			return 1
		}
		return 0
	}

	vm := resource_info.NewResourceVectorMap()
	priorities := []int32{1, 1, 5, 5, 10, 20, 20, 100}

	for i, lp := range priorities {
		for j, rp := range priorities {
			if i == j {
				continue
			}
			ssn := &Session{JobOrderFns: []common_info.CompareFn{priorityLikeJobOrderFn}}
			lJob := podgroup_info.NewPodGroupInfoWithVectorMap(
				common_info.PodGroupID("l"), vm)
			lJob.Priority = lp
			rJob := podgroup_info.NewPodGroupInfoWithVectorMap(
				common_info.PodGroupID("r"), vm)
			rJob.Priority = rp

			oldBehavior := !ssn.JobOrderFn(lJob, rJob)
			newBehavior := ssn.VictimOrderFn(lJob, rJob)

			if oldBehavior != newBehavior {
				t.Errorf("mismatch for priorities l=%d r=%d: old !JobOrderFn=%v, new VictimOrderFn=%v",
					lp, rp, oldBehavior, newBehavior)
			}
		}
	}
}