improvement: Introduce the idea of GangScaling

cmd/server.go

@@ -207,6 +207,7 @@ func initServer(ctx context.Context, info witchcraft.InitInfo) (func(), error) {
 			install.ExecutorPrioritizedNodeLabel,
 		),
 		wasteMetricsReporter,
+		install.GangScaling,
 	)
 
 	resourceReporter := metrics.NewResourceReporter(

config/config.go

@@ -26,7 +26,8 @@ type Install struct {
 	config.Runtime                                      `yaml:",inline"`
 	Kubeconfig                                          string            `yaml:"kube-config,omitempty"`
 	FIFO                                                bool              `yaml:"fifo,omitempty"`
-	FifoConfig                                          FifoConfig        `yaml:"fifo-config,omitempty"`
+	FifoConfig                                          FifoConfig        `yaml:"t,omitempty"`
+	GangScaling                                         bool              `yaml:"gang-scaling,omitempty"`
 	QPS                                                 float32           `yaml:"qps,omitempty"`
 	Burst                                               int               `yaml:"burst,omitempty"`
 	BinpackAlgo                                         string            `yaml:"binpack,omitempty"`

internal/extender/demand.go

@@ -59,6 +59,28 @@ func (s *SparkSchedulerExtender) createDemandForExecutorInAnyZone(ctx context.Co
 	s.createDemandForExecutorInSpecificZone(ctx, executorPod, executorResources, nil)
 }
 
+func (s *SparkSchedulerExtender) createPartialDemand(
+	ctx context.Context,
+	driverPod *v1.Pod,
+	executorResources *resources.Resources,
+	executorCount int,
+	zone *demandapi.Zone) {
+	if !s.demands.CRDExists() {
+		return
+	}
+	units := []demandapi.DemandUnit{
+		{
+			Count: executorCount,
+			Resources: demandapi.ResourceList{
+				demandapi.ResourceCPU:       executorResources.CPU,
+				demandapi.ResourceMemory:    executorResources.Memory,
+				demandapi.ResourceNvidiaGPU: executorResources.NvidiaGPU,
+			},
+		},
+	}
+	s.createDemand(ctx, driverPod, units, zone)
+}
+
 func (s *SparkSchedulerExtender) createDemandForExecutorInSpecificZone(ctx context.Context, executorPod *v1.Pod, executorResources *resources.Resources, zone *demandapi.Zone) {
 	if !s.demands.CRDExists() {
 		return

internal/extender/extendertest/extender_test_utils.go

@@ -151,6 +151,7 @@ func NewTestExtender(binpackAlgo string, objects ...runtime.Object) (*Harness, e
 		instanceGroupLabel,
 		sort.NewNodeSorter(nil, nil),
 		wasteMetricsReporter,
+		false,
 	)
 
 	unschedulablePodMarker := extender.NewUnschedulablePodMarker(

internal/extender/resource.go

@@ -79,6 +79,7 @@ type SparkSchedulerExtender struct {
 	overheadComputer                                    *OverheadComputer
 	lastRequest                                         time.Time
 	instanceGroupLabel                                  string
+	isGangScaling                                       bool
 
 	wasteMetricsReporter *metrics.WasteMetricsReporter
 }
@@ -100,7 +101,8 @@ func NewExtender(
 	overheadComputer *OverheadComputer,
 	instanceGroupLabel string,
 	nodeSorter *ns.NodeSorter,
-	wasteMetricsReporter *metrics.WasteMetricsReporter) *SparkSchedulerExtender {
+	wasteMetricsReporter *metrics.WasteMetricsReporter,
+	isGangScaling bool) *SparkSchedulerExtender {
 	return &SparkSchedulerExtender{
 		nodeLister:                 nodeLister,
 		podLister:                  podLister,
@@ -118,6 +120,7 @@ func NewExtender(
 		instanceGroupLabel:   instanceGroupLabel,
 		nodeSorter:           nodeSorter,
 		wasteMetricsReporter: wasteMetricsReporter,
+		isGangScaling:        isGangScaling,
 	}
 }
 
@@ -316,14 +319,7 @@ func (s *SparkSchedulerExtender) selectDriverNode(
 		}
 	}
 
-	packingResult := s.binpacker.BinpackFunc(
-		ctx,
-		applicationResources.driverResources,
-		applicationResources.executorResources,
-		applicationResources.minExecutorCount,
-		driverNodeNames,
-		executorNodeNames,
-		availableNodesSchedulingMetadata)
+	packingResult := s.getPackingResult(ctx, applicationResources, driverNodeNames, executorNodeNames, availableNodesSchedulingMetadata)
 	efficiency := computeAvgPackingEfficiencyForResult(availableNodesSchedulingMetadata, packingResult)
 
 	svc1log.FromContext(ctx).Debug("binpacking result",
@@ -364,6 +360,14 @@ func (s *SparkSchedulerExtender) selectDriverNode(
 	if err != nil {
 		return "", failureInternal, err
 	}
+	// Then we consider create partial demands
+	if s.isGangScaling {
+		differenceBetweenAllocationAndDesired := applicationResources.minExecutorCount - len(packingResult.ExecutorNodes)
+		if differenceBetweenAllocationAndDesired > 0 {
+			zone := s.getZone(packingResult.DriverNode, availableNodesSchedulingMetadata)
+			s.createPartialDemand(ctx, driver, applicationResources.executorResources, differenceBetweenAllocationAndDesired, zone)
+		}
+	}
 	return packingResult.DriverNode, success, nil
 }
 
@@ -703,3 +707,60 @@ func (s *SparkSchedulerExtender) rescheduleExecutorWithMinimalFragmentation(
 func (s *SparkSchedulerExtender) isSuccessOutcome(outcome string) bool {
 	return outcome == success || outcome == successAlreadyBound || outcome == successRescheduled || outcome == successScheduledExtraExecutor
 }
+
+func (s *SparkSchedulerExtender) getPackingResult(
+	ctx context.Context,
+	applicationResources *sparkApplicationResources,
+	driverNodeNames []string,
+	executorNodeNames []string,
+	availableNodesSchedulingMetadata resources.NodeGroupSchedulingMetadata) *binpack.PackingResult {
+	if !s.isGangScaling {
+		return s.binpacker.BinpackFunc(
+			ctx,
+			applicationResources.driverResources,
+			applicationResources.executorResources,
+			applicationResources.minExecutorCount,
+			driverNodeNames,
+			executorNodeNames,
+			availableNodesSchedulingMetadata)
+	}
+	i := applicationResources.minExecutorCount
+	if i < 0 {
+		i = 0
+	}
+	var packingResult *binpack.PackingResult
+	for {
+		if i == -1 {
+			break
+		}
+		packingResult = s.binpacker.BinpackFunc(
+			ctx,
+			applicationResources.driverResources,
+			applicationResources.executorResources,
+			i,
+			driverNodeNames,
+			executorNodeNames,
+			availableNodesSchedulingMetadata)
+
+		if packingResult.HasCapacity {
+			break
+		}
+		i = i - 1
+	}
+	return packingResult
+}
+
+func (s *SparkSchedulerExtender) getZone(driverNodeName string, metadata resources.NodeGroupSchedulingMetadata) *demandapi.Zone {
+	if !s.binpacker.IsSingleAz {
+		return nil
+	}
+	nodeSchedulingMetadata, ok := metadata[driverNodeName]
+	if !ok || nodeSchedulingMetadata == nil {
+		return nil
+	}
+	if nodeSchedulingMetadata.ZoneLabel == "" {
+		return nil
+	}
+	demandZone := demandapi.Zone(nodeSchedulingMetadata.ZoneLabel)
+	return &demandZone
+}