feat(scheduler): add built-in scenario generators

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

Author: enoodle

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 ## [Unreleased]
 
 ### Added
+- Added built-in `NodeLocalGreedy` and `MultiNodeGang` scenario generator implementations for bounded reclaim, preempt, and consolidation search.
 - Added an opt-in `deviceaccess` admission plugin (`--block-nvidia-visible-devices`, config field `admission.blockNvidiaVisibleDevices`, default disabled) that (1) rejects pods overriding the `NVIDIA_VISIBLE_DEVICES` environment variable with values other than `void`/`none` (or via a `valueFrom` reference), and (2) injects `NVIDIA_VISIBLE_DEVICES=void` into containers that do not request a GPU, blocking their access to GPUs on the node.
 - Added support for configuring admission Pod Disruption Budget via Helm values (`admission.podDisruptionBudget`) [#1490](https://github.com/kai-scheduler/KAI-Scheduler/pull/1490) [dttung2905](https://github.com/dttung2905)
 - Added an opt-in `hamicore` binder plugin (depends on `gpusharing`) to write the HAMI-core GPU memory limit (`CUDA_DEVICE_MEMORY_LIMIT`) for fractional GPU pods.

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

@@ -0,0 +1,49 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package solvers
+
+import (
+	"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/framework"
+)
+
+type multiNodeGangGenerator struct {
+	builder *PodAccumulatedScenarioBuilder
+	first   bool
+}
+
+func NewMultiNodeGangGenerator(ctx framework.ScenarioGeneratorContext) framework.ScenarioGenerator {
+	solveCtx, generateVictimsQueue, ok := validateScenarioGeneratorContext(ctx)
+	if !ok {
+		return nil
+	}
+	victimsQueue := generateVictimsQueue()
+	if victimsQueue == nil {
+		return nil
+	}
+
+	return &multiNodeGangGenerator{
+		builder: NewPodAccumulatedScenarioBuilder(
+			solveCtx.Session,
+			solveCtx.PartialPendingJob,
+			solveCtx.RecordedVictimsJobs,
+			victimsQueue,
+			solveCtx.FeasibleNodes,
+		),
+		first: true,
+	}
+}
+
+func (g *multiNodeGangGenerator) Name() string {
+	return constants.GeneratorMultiNodeGang
+}
+
+func (g *multiNodeGangGenerator) Next() api.ScenarioInfo {
+	if g.first {
+		g.first = false
+		return g.builder.GetValidScenario()
+	}
+	return g.builder.GetNextScenario()
+}

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

@@ -0,0 +1,183 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package solvers
+
+import (
+	"sort"
+	"strings"
+
+	"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/api"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_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/framework"
+)
+
+type nodeLocalGreedyGenerator struct {
+	solveCtx             *SolveContext
+	generateVictimsQueue GenerateVictimsQueue
+	builder              *PodAccumulatedScenarioBuilder
+	scenarios            []*scenario.ByNodeScenario
+	advanceNext          bool
+}
+
+func NewNodeLocalGreedyGenerator(ctx framework.ScenarioGeneratorContext) framework.ScenarioGenerator {
+	solveCtx, generateVictimsQueue, ok := validateScenarioGeneratorContext(ctx)
+	if !ok {
+		return nil
+	}
+	return &nodeLocalGreedyGenerator{
+		solveCtx:             solveCtx,
+		generateVictimsQueue: generateVictimsQueue,
+	}
+}
+
+func (g *nodeLocalGreedyGenerator) Name() string {
+	return constants.GeneratorNodeLocalGreedy
+}
+
+func (g *nodeLocalGreedyGenerator) Next() api.ScenarioInfo {
+	if !g.ensureBuilder() {
+		return nil
+	}
+	for {
+		if sn := g.popScenario(); sn != nil {
+			return sn
+		}
+		accumulated := g.nextAccumulatedScenario()
+		if accumulated == nil {
+			return nil
+		}
+		g.scenarios = nodeLocalScenarios(g.builder.session, accumulated)
+	}
+}
+
+func (g *nodeLocalGreedyGenerator) ensureBuilder() bool {
+	if g.builder != nil {
+		return true
+	}
+	victimsQueue := g.generateVictimsQueue()
+	if victimsQueue == nil {
+		return false
+	}
+	g.builder = NewPodAccumulatedScenarioBuilder(
+		g.solveCtx.Session,
+		g.solveCtx.PartialPendingJob,
+		g.solveCtx.RecordedVictimsJobs,
+		victimsQueue,
+		g.solveCtx.FeasibleNodes,
+	)
+	return true
+}
+
+func addPotentialVictimsGroupedByJob(sn *scenario.ByNodeScenario, tasks []*pod_info.PodInfo) {
+	groupedTasks := map[common_info.PodGroupID][]*pod_info.PodInfo{}
+	var jobOrder []common_info.PodGroupID
+	for _, task := range tasks {
+		if _, found := groupedTasks[task.Job]; !found {
+			jobOrder = append(jobOrder, task.Job)
+		}
+		groupedTasks[task.Job] = append(groupedTasks[task.Job], task)
+	}
+	for _, jobID := range jobOrder {
+		sn.AddPotentialVictimsTasks(groupedTasks[jobID])
+	}
+}
+
+func (g *nodeLocalGreedyGenerator) popScenario() *scenario.ByNodeScenario {
+	if len(g.scenarios) == 0 {
+		return nil
+	}
+	sn := g.scenarios[0]
+	g.scenarios = g.scenarios[1:]
+	return sn
+}
+
+func (g *nodeLocalGreedyGenerator) nextAccumulatedScenario() *scenario.ByNodeScenario {
+	for {
+		if g.advanceNext {
+			if g.builder.victimsJobsQueue.IsEmpty() {
+				return nil
+			}
+			if !g.builder.addNextPotentialVictims() {
+				continue
+			}
+		}
+		g.advanceNext = true
+
+		if g.builder.lastScenario == nil {
+			return nil
+		}
+		if !g.builder.outerScenarioValid() {
+			continue
+		}
+		return g.builder.lastScenario
+	}
+}
+
+func nodeLocalScenarios(session *framework.Session, base *scenario.ByNodeScenario) []*scenario.ByNodeScenario {
+	if base == nil {
+		return nil
+	}
+	if len(base.PotentialVictimsTasks()) == 0 {
+		if len(base.RecordedVictimsTasks()) == 0 {
+			return nil
+		}
+		return []*scenario.ByNodeScenario{base}
+	}
+
+	var scenarios []*scenario.ByNodeScenario
+	seen := map[string]struct{}{}
+	for _, nodeName := range nodeNamesOfJob(base.LatestPotentialVictim()) {
+		victimTasks := base.VictimsTasksFromNodes([]string{nodeName})
+		if len(victimTasks) == 0 {
+			continue
+		}
+		key := victimUIDSetKey(victimTasks)
+		if _, found := seen[key]; found {
+			continue
+		}
+		seen[key] = struct{}{}
+		sn := scenario.NewByNodeScenario(
+			session,
+			base.GetPreemptor(),
+			base.PendingTasks(),
+			nil,
+			base.RecordedVictimsJobs(),
+		)
+		addPotentialVictimsGroupedByJob(sn, victimTasks)
+		scenarios = append(scenarios, sn)
+	}
+	return scenarios
+}
+
+func nodeNamesOfJob(job *podgroup_info.PodGroupInfo) []string {
+	if job == nil {
+		return nil
+	}
+	seen := map[string]struct{}{}
+	for _, task := range job.GetAllPodsMap() {
+		if task.NodeName == "" {
+			continue
+		}
+		seen[task.NodeName] = struct{}{}
+	}
+	nodeNames := make([]string, 0, len(seen))
+	for nodeName := range seen {
+		nodeNames = append(nodeNames, nodeName)
+	}
+	sort.Strings(nodeNames)
+	return nodeNames
+}
+
+func victimUIDSetKey(tasks []*pod_info.PodInfo) string {
+	uids := make([]string, 0, len(tasks))
+	for _, task := range tasks {
+		uids = append(uids, string(task.UID))
+	}
+	sort.Strings(uids)
+	return strings.Join(uids, "\x00")
+}