perf(scheduler): avoid duplicate full job solver probe

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

Author: enoodle

CHANGELOG.md

@@ -22,6 +22,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 - Removed redundant `PodDisruptionBudgetImplemented` guard from operator PDB creation helper [#1613](https://github.com/kai-scheduler/KAI-Scheduler/pull/1613) [dttung2905](https://github.com/dttung2905)
 - Updated Go toolchain and base build images to v1.26.3.
 - **Breaking:** The podgroup produced for JobSet is now produces as a single PodGroup per JobSet with a two-level SubGroup hierarchy (one parent SubGroup per `replicatedJob`, one leaf SubGroup per replica) regardless of `startupPolicyOrder`. The `kai.scheduler/batch-min-member` annotation on the JobSet now overrides the root `minSubGroup`; the same annotation on `replicatedJobs[].template.metadata.annotations` overrides the leaf `minMember` (defaulting to `template.spec.parallelism`). [#1617](https://github.com/kai-scheduler/KAI-Scheduler/pull/1617) [davidLif](https://github.com/davidLif)
+- Optimized the job solver to run the full allocation probe only once after partial search finds at least one solvable pending task.
 
 ### Fixed
 - Reduced scheduler heap retention after scheduling cycles by clearing completed session snapshots and callback references, and by releasing the node scoring pool without waiting for finalizers.

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

@@ -178,16 +178,18 @@ func (s *JobSolver) solvePendingJobWithGenerator(
 ) *SearchResult {
 	n := len(tasksToAllocate)
 	enteredSearch := false
-	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, jobBudget.ReducedBudget(), false)
+	if n > 1 {
+		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, jobBudget.ReducedBudget(), false)
+			}
+			preserveEnteredSearch(searchResult, enteredSearch)
+			return searchResult
 		}
-		preserveEnteredSearch(searchResult, enteredSearch)
-		return searchResult
 	}
 
 	result := s.probeAtK(ssn, state, pendingJob, tasksToAllocate, n, jobBudget, registration, generatorBudget)
@@ -197,7 +199,7 @@ func (s *JobSolver) solvePendingJobWithGenerator(
 	return result
 }
 
-// searchMaxSolvableK returns the largest k in [0, n] for which a probe at k succeeds.
+// searchMaxSolvableK returns the largest k in [0, n) for which a probe at k succeeds.
 // Each probe is discarded before returning, so session state is clean on return.
 // Successful probes update hints in state for use by subsequent probes.
 // Complexity: O(log n) probes — exponential doubling to locate a failing k (or reach n),
@@ -212,7 +214,7 @@ func (s *JobSolver) searchMaxSolvableK(
 	generatorBudget *generatorSearchBudget,
 ) (int, *SearchResult) {
 	n := len(tasksToAllocate)
-	if n == 0 {
+	if n <= 1 {
 		return 0, nil
 	}
 
@@ -224,16 +226,16 @@ func (s *JobSolver) searchMaxSolvableK(
 }
 
 func searchMaxSolvableK(n int, probe func(k int) *SearchResult) (int, *SearchResult) {
-	if n == 0 {
+	if n <= 1 {
 		return 0, nil
 	}
 
 	lo := 0
-	var hi int
+	hi := n
 	var lastUnsolvedResult *SearchResult
 	enteredSearch := false
 	k := 1
-	for {
+	for k < n {
 		result := probe(k)
 		enteredSearch = enteredSearch || searchResultEntered(result) || resultSolved(result)
 		if shouldStopSearch(result) {
@@ -246,12 +248,10 @@ func searchMaxSolvableK(n int, probe func(k int) *SearchResult) (int, *SearchRes
 			break
 		}
 		lo = k
-		if k == n {
-			return n, lastUnsolvedResult
-		}
 		k *= 2
-		if k > n {
-			k = n
+		if k >= n {
+			hi = n
+			break
 		}
 	}
 

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

@@ -243,7 +243,7 @@ func TestSolveWithResultRunsCompletePartialSearchForOneGeneratorBeforeNext(t *te
 
 	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)
+	require.Equal(t, []string{"first:1", "second:1", "second:2", "second:3"}, factoryCalls)
 }
 
 func TestSolveWithResultReportsDeadlineBeforeScenarioSimulation(t *testing.T) {

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

@@ -0,0 +1,36 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package solvers
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSearchMaxSolvableKSkipsFullProbe(t *testing.T) {
+	probes := []int{}
+
+	maxSolvedK, result := searchMaxSolvableK(4, func(k int) *SearchResult {
+		probes = append(probes, k)
+		return solvedSearchResult(&solutionResult{solved: true}, false)
+	})
+
+	require.Equal(t, 3, maxSolvedK)
+	require.Nil(t, result)
+	require.Equal(t, []int{1, 2, 3}, probes)
+}
+
+func TestSearchMaxSolvableKSkipsSingleTaskFullProbe(t *testing.T) {
+	probes := []int{}
+
+	maxSolvedK, result := searchMaxSolvableK(1, func(k int) *SearchResult {
+		probes = append(probes, k)
+		return solvedSearchResult(&solutionResult{solved: true}, false)
+	})
+
+	require.Equal(t, 0, maxSolvedK)
+	require.Nil(t, result)
+	require.Empty(t, probes)
+}