feat(scheduler): drive solver from generator portfolio

Signed-off-by: Erez Freiberger <enoodle@gmail.com>

Author: enoodle

pkg/scheduler/actions/common/solvers/job_solver.go

@@ -8,7 +8,6 @@ import (
 	"strings"
 	"time"
 
-	solverscenario "github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/common/solvers/scenario"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/utils"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/node_info"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
@@ -95,52 +94,97 @@ func (s *JobSolver) Solve(
 // describing why the scenario search stopped.
 func (s *JobSolver) SolveWithResult(
 	ssn *framework.Session, pendingJob *podgroup_info.PodGroupInfo,
-) (bool, *framework.Statement, []string, *SearchResult) {
-	state := solvingState{}
+) (solved bool, statement *framework.Statement, victimTaskNames []string, searchResult *SearchResult) {
 	originalNumActiveTasks := pendingJob.GetNumActiveUsedTasks()
 
 	tasksToAllocate := podgroup_info.GetTasksToAllocate(pendingJob, ssn.SubGroupOrderFn, ssn.TaskOrderFn, false)
 	n := len(tasksToAllocate)
 	if n == 0 {
-		return false, nil, calcVictimNames(state.recordedVictimsTasks),
-			terminalSearchResult(SearchResultGeneratorsExhausted, false, false)
+		return false, nil, nil, terminalSearchResult(SearchResultGeneratorsExhausted, false, false)
 	}
 
 	actionBudget := s.ensureActionBudget()
 	jobBudget := actionBudget.BeginJob()
 	if actionBudget.Exhausted() {
-		return false, nil, calcVictimNames(state.recordedVictimsTasks),
-			terminalSearchResult(SearchResultNotAttempted, false, false)
+		return false, nil, nil, terminalSearchResult(SearchResultNotAttempted, false, false)
 	}
 
+	if s.generateVictimsQueue == nil {
+		return false, nil, nil, terminalSearchResult(SearchResultNoGenerator, jobBudget.ReducedBudget(), false)
+	}
+	registrations := applicableScenarioGeneratorRegistrations(ssn, s.actionType)
+	if len(registrations) == 0 {
+		return false, nil, nil, terminalSearchResult(SearchResultNoGenerator, jobBudget.ReducedBudget(), false)
+	}
+
+	enteredSearch := false
+	var lastVictimTasks []*pod_info.PodInfo
+	var lastResult *SearchResult
+	for _, registration := range registrations {
+		state := solvingState{}
+		generatorBudget := jobBudget.BeginGenerator(registration.Name)
+		result := s.solvePendingJobWithGenerator(
+			ssn, &state, pendingJob, tasksToAllocate, jobBudget, registration, generatorBudget,
+		)
+		enteredSearch = enteredSearch || searchResultEntered(result) || resultSolved(result)
+		lastVictimTasks = state.recordedVictimsTasks
+		lastResult = result
+
+		if resultSolved(result) {
+			solution := result.solution
+			numActiveTasks := pendingJob.GetNumActiveUsedTasks()
+			jobSolved := pendingJob.IsGangSatisfied()
+			if originalNumActiveTasks >= numActiveTasks {
+				jobSolved = false
+			}
+
+			log.InfraLogger.V(4).Infof(
+				"Scenario solved for %d tasks to allocate for %s. Victims: %s",
+				n, pendingJob.Name, victimPrintingStruct{solution.victimsTasks})
+			return jobSolved, solution.statement, calcVictimNames(solution.victimsTasks), result
+		}
+
+		if shouldStopSearch(result) {
+			preserveEnteredSearch(result, enteredSearch)
+			return false, nil, calcVictimNames(lastVictimTasks), result
+		}
+	}
+
+	if lastResult == nil {
+		lastResult = terminalSearchResult(SearchResultGeneratorsExhausted, jobBudget.ReducedBudget(), enteredSearch)
+	}
+	preserveEnteredSearch(lastResult, enteredSearch)
+	return false, nil, calcVictimNames(lastVictimTasks), lastResult
+}
+
+func (s *JobSolver) solvePendingJobWithGenerator(
+	ssn *framework.Session,
+	state *solvingState,
+	pendingJob *podgroup_info.PodGroupInfo,
+	tasksToAllocate []*pod_info.PodInfo,
+	jobBudget *jobSearchBudget,
+	registration framework.ScenarioGeneratorRegistration,
+	generatorBudget *generatorSearchBudget,
+) *SearchResult {
+	n := len(tasksToAllocate)
 	enteredSearch := false
-	maxSolvedK, searchResult := s.searchMaxSolvableK(ssn, &state, pendingJob, tasksToAllocate, jobBudget)
+	maxSolvedK, searchResult := s.searchMaxSolvableK(
+		ssn, state, pendingJob, tasksToAllocate, jobBudget, registration, generatorBudget,
+	)
 	enteredSearch = searchResultEntered(searchResult) || maxSolvedK > 0
 	if maxSolvedK == 0 {
 		if searchResult == nil {
-			searchResult = terminalSearchResult(SearchResultGeneratorsExhausted, false, false)
+			searchResult = terminalSearchResult(SearchResultGeneratorsExhausted, jobBudget.ReducedBudget(), false)
 		}
 		preserveEnteredSearch(searchResult, enteredSearch)
-		return false, nil, calcVictimNames(state.recordedVictimsTasks), searchResult
+		return searchResult
 	}
 
-	result := s.probeAtK(ssn, &state, pendingJob, tasksToAllocate, n, jobBudget)
+	result := s.probeAtK(ssn, state, pendingJob, tasksToAllocate, n, jobBudget, registration, generatorBudget)
 	if !resultSolved(result) {
 		preserveEnteredSearch(result, enteredSearch)
-		return false, nil, calcVictimNames(state.recordedVictimsTasks), result
 	}
-
-	solution := result.solution
-	numActiveTasks := pendingJob.GetNumActiveUsedTasks()
-	jobSolved := pendingJob.IsGangSatisfied()
-	if originalNumActiveTasks >= numActiveTasks {
-		jobSolved = false
-	}
-
-	log.InfraLogger.V(4).Infof(
-		"Scenario solved for %d tasks to allocate for %s. Victims: %s",
-		n, pendingJob.Name, victimPrintingStruct{solution.victimsTasks})
-	return jobSolved, solution.statement, calcVictimNames(solution.victimsTasks), result
+	return result
 }
 
 // searchMaxSolvableK returns the largest k in [0, n] for which a probe at k succeeds.
@@ -154,14 +198,18 @@ func (s *JobSolver) searchMaxSolvableK(
 	pendingJob *podgroup_info.PodGroupInfo,
 	tasksToAllocate []*pod_info.PodInfo,
 	jobBudget *jobSearchBudget,
+	registration framework.ScenarioGeneratorRegistration,
+	generatorBudget *generatorSearchBudget,
 ) (int, *SearchResult) {
 	n := len(tasksToAllocate)
 	if n == 0 {
 		return 0, nil
 	}
 
 	return searchMaxSolvableK(n, func(k int) *SearchResult {
-		return s.tryProbeAndDiscard(ssn, state, pendingJob, tasksToAllocate, k, jobBudget)
+		return s.tryProbeAndDiscard(
+			ssn, state, pendingJob, tasksToAllocate, k, jobBudget, registration, generatorBudget,
+		)
 	})
 }
 
@@ -224,8 +272,10 @@ func (s *JobSolver) tryProbeAndDiscard(
 	tasksToAllocate []*pod_info.PodInfo,
 	k int,
 	jobBudget *jobSearchBudget,
+	registration framework.ScenarioGeneratorRegistration,
+	generatorBudget *generatorSearchBudget,
 ) *SearchResult {
-	result := s.probeAtK(ssn, state, pendingJob, tasksToAllocate, k, jobBudget)
+	result := s.probeAtK(ssn, state, pendingJob, tasksToAllocate, k, jobBudget, registration, generatorBudget)
 	if !resultSolved(result) {
 		log.InfraLogger.V(5).Infof("No solution found for %d tasks out of %d tasks to allocate for %s",
 			k, len(tasksToAllocate), pendingJob.Name)
@@ -250,15 +300,18 @@ func (s *JobSolver) probeAtK(
 	tasksToAllocate []*pod_info.PodInfo,
 	k int,
 	jobBudget *jobSearchBudget,
+	registration framework.ScenarioGeneratorRegistration,
+	generatorBudget *generatorSearchBudget,
 ) *SearchResult {
 	pendingTasks := tasksToAllocate[:k]
 	partialPendingJob := getPartialJobRepresentative(pendingJob, pendingTasks)
-	return s.solvePartialJob(ssn, state, partialPendingJob, jobBudget)
+	return s.solvePartialJob(ssn, state, partialPendingJob, jobBudget, registration, generatorBudget, k)
 }
 
 func (s *JobSolver) solvePartialJob(
 	ssn *framework.Session, state *solvingState, partialPendingJob *podgroup_info.PodGroupInfo,
-	jobBudget *jobSearchBudget,
+	jobBudget *jobSearchBudget, registration framework.ScenarioGeneratorRegistration,
+	generatorBudget *generatorSearchBudget, probeK int,
 ) *SearchResult {
 	actionBudget := s.ensureActionBudget()
 	if jobBudget == nil {
@@ -274,40 +327,33 @@ func (s *JobSolver) solvePartialJob(
 		feasibleNodeMap[task.NodeName] = node
 	}
 
-	if s.generateVictimsQueue == nil {
-		return terminalSearchResult(SearchResultNoGenerator, jobBudget.ReducedBudget(), false)
-	}
-	victimsQueue := s.generateVictimsQueue()
-	if victimsQueue == nil {
-		return terminalSearchResult(SearchResultNoGenerator, jobBudget.ReducedBudget(), false)
+	solveCtx := &SolveContext{
+		Session:              ssn,
+		ActionType:           s.actionType,
+		PartialPendingJob:    partialPendingJob,
+		RecordedVictimsJobs:  state.recordedVictimsJobs,
+		RecordedVictimsTasks: state.recordedVictimsTasks,
+		GenerateVictimsQueue: s.generateVictimsQueue,
+		FeasibleNodes:        feasibleNodeMap,
+		ProbeK:               probeK,
 	}
+	portfolio := newSingleGeneratorScenarioPortfolio(solveCtx, jobBudget, registration, generatorBudget)
 
-	scenarioBuilder := NewPodAccumulatedScenarioBuilder(
-		ssn, partialPendingJob, state.recordedVictimsJobs, victimsQueue, feasibleNodeMap)
-
-	enteredSearch := false
-	firstScenario := true
 	for {
 		if actionBudget.Exhausted() || jobBudget.Remaining() <= 0 {
-			return terminalSearchResult(SearchResultDeadlineExhausted, jobBudget.ReducedBudget(), enteredSearch)
-		}
-		var scenarioToSolve *solverscenario.ByNodeScenario
-		if firstScenario {
-			scenarioToSolve = scenarioBuilder.GetValidScenario()
-			firstScenario = false
-		} else {
-			scenarioToSolve = scenarioBuilder.GetNextScenario()
+			return terminalSearchResult(
+				SearchResultDeadlineExhausted, jobBudget.ReducedBudget(), portfolio.enteredSearch,
+			)
 		}
+		scenarioToSolve := portfolio.Next()
 		if actionBudget.Exhausted() || jobBudget.Remaining() <= 0 {
-			return terminalSearchResult(SearchResultDeadlineExhausted, jobBudget.ReducedBudget(), enteredSearch)
+			return terminalSearchResult(
+				SearchResultDeadlineExhausted, jobBudget.ReducedBudget(), portfolio.enteredSearch,
+			)
 		}
 		if scenarioToSolve == nil {
-			if actionBudget.Exhausted() || jobBudget.Remaining() <= 0 {
-				return terminalSearchResult(SearchResultDeadlineExhausted, jobBudget.ReducedBudget(), enteredSearch)
-			}
 			break
 		}
-		enteredSearch = true
 		scenarioSolver := newByPodSolver(feasibleNodeMap, s.solutionValidator, ssn.AllowConsolidatingReclaim(),
 			s.actionType)
 
@@ -320,7 +366,7 @@ func (s *JobSolver) solvePartialJob(
 		}
 	}
 
-	return terminalSearchResult(SearchResultGeneratorsExhausted, jobBudget.ReducedBudget(), enteredSearch)
+	return terminalSearchResult(portfolio.StopReason(), jobBudget.ReducedBudget(), portfolio.enteredSearch)
 }
 
 func searchResultEntered(result *SearchResult) bool {

pkg/scheduler/actions/common/solvers/job_solver_result_test.go

@@ -4,6 +4,7 @@
 package solvers
 
 import (
+	"fmt"
 	"testing"
 	"time"
 
@@ -13,10 +14,12 @@ import (
 
 	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/actions/common/solvers/scenario"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/utils"
 	"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/queue_info"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/framework"
@@ -88,6 +91,54 @@ func TestSolveWithResultReturnsNoGeneratorWhenGeneratorReturnsNil(t *testing.T)
 	require.False(t, result.EnteredSearch())
 }
 
+func TestSolveWithResultRunsCompletePartialSearchForOneGeneratorBeforeNext(t *testing.T) {
+	ssn := newGeneratorTestSession(t, map[string]int{
+		"node-1": 1,
+		"node-2": 1,
+		"node-3": 1,
+	})
+	require.NoError(t, ssn.InitNodeScoringPool())
+	pendingJob := addGeneratorTestPendingJob(t, ssn, 3, 10, "team-pending")
+	setGeneratorTestMinAvailable(pendingJob, 3)
+	victimJob, victimTasks := addGeneratorTestJob(t, ssn, 3, 20, "team-victim", "node-1", "node-2", "node-3")
+	factoryCalls := []string{}
+
+	ssn.AddScenarioGenerator("first", func(ctx framework.ScenarioGeneratorContext) framework.ScenarioGenerator {
+		solveCtx := ctx.(*SolveContext)
+		factoryCalls = append(factoryCalls, fmt.Sprintf("first:%d", solveCtx.ProbeK))
+		return &portfolioTestGenerator{name: "first"}
+	}, framework.Reclaim)
+	ssn.AddScenarioGenerator("second", func(ctx framework.ScenarioGeneratorContext) framework.ScenarioGenerator {
+		solveCtx := ctx.(*SolveContext)
+		factoryCalls = append(factoryCalls, fmt.Sprintf("second:%d", solveCtx.ProbeK))
+		pendingTasks := podgroup_info.GetTasksToAllocate(
+			solveCtx.PartialPendingJob, ssn.SubGroupOrderFn, ssn.TaskOrderFn, false,
+		)
+		sn := scenario.NewByNodeScenario(
+			ssn, solveCtx.PartialPendingJob, pendingTasks,
+			unrecordedVictimsForProbe(victimTasks, solveCtx.RecordedVictimsTasks, solveCtx.ProbeK),
+			solveCtx.RecordedVictimsJobs,
+		)
+		return &portfolioTestGenerator{name: "second", scenarios: []api.ScenarioInfo{sn}}
+	}, framework.Reclaim)
+	solver := NewJobsSolver(
+		jobSolverResultTestFeasibleNodes(ssn),
+		nil,
+		generatorTestVictimsQueueFactory(ssn, victimJob),
+		framework.Reclaim,
+		nil,
+	)
+
+	solved, statement, _, result := solver.SolveWithResult(ssn, pendingJob)
+	if statement != nil {
+		defer statement.Discard()
+	}
+
+	require.True(t, solved)
+	require.Equal(t, SearchResultSolved, result.Reason())
+	require.Equal(t, []string{"first:1", "second:1", "second:2", "second:3", "second:3"}, factoryCalls)
+}
+
 func TestSolveWithResultReportsDeadlineBeforeScenarioSimulation(t *testing.T) {
 	clock := &fakeClock{now: time.Unix(0, 0)}
 	actionBudget, err := newActionSearchBudgetWithClock(
@@ -102,6 +153,7 @@ func TestSolveWithResultReportsDeadlineBeforeScenarioSimulation(t *testing.T) {
 	)
 	require.NoError(t, err)
 	ssn, pendingJob := newJobSolverResultTestSession(t, 1)
+	ssn.AddScenarioGenerator("deadline-test", NewMultiNodeGangGenerator, framework.Reclaim)
 	solver := NewJobsSolver(
 		nil,
 		nil,
@@ -137,6 +189,40 @@ func TestSearchMaxSolvableKPreservesEnteredSearchAfterTerminalPartialProbe(t *te
 	require.True(t, result.EnteredSearch())
 }
 
+func jobSolverResultTestFeasibleNodes(ssn *framework.Session) []*node_info.NodeInfo {
+	nodes := make([]*node_info.NodeInfo, 0, len(ssn.ClusterInfo.Nodes))
+	for _, node := range ssn.ClusterInfo.Nodes {
+		nodes = append(nodes, node)
+	}
+	return nodes
+}
+
+func unrecordedVictimsForProbe(
+	victimTasks []*pod_info.PodInfo, recordedVictims []*pod_info.PodInfo, probeK int,
+) []*pod_info.PodInfo {
+	recordedByUID := map[common_info.PodID]struct{}{}
+	for _, task := range recordedVictims {
+		recordedByUID[task.UID] = struct{}{}
+	}
+
+	neededVictims := probeK - len(recordedVictims)
+	if neededVictims <= 0 {
+		return nil
+	}
+
+	selectedVictims := make([]*pod_info.PodInfo, 0, neededVictims)
+	for _, task := range victimTasks {
+		if _, alreadyRecorded := recordedByUID[task.UID]; alreadyRecorded {
+			continue
+		}
+		selectedVictims = append(selectedVictims, task)
+		if len(selectedVictims) == neededVictims {
+			return selectedVictims
+		}
+	}
+	return selectedVictims
+}
+
 func TestPreserveEnteredSearchMarksTerminalResult(t *testing.T) {
 	result := terminalSearchResult(SearchResultDeadlineExhausted, false, false)