From 87211ef66b38036aeca18f006b7bd391046e681e Mon Sep 17 00:00:00 2001
From: Newnius <newnius.cn@gmail.com>
Date: Sat, 30 May 2020 19:39:19 +0800
Subject: [PATCH] update fair

---
 src/scheduler_fair.go | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/scheduler_fair.go b/src/scheduler_fair.go
index 5662557..391bea2 100644
--- a/src/scheduler_fair.go
+++ b/src/scheduler_fair.go
@@ -16,6 +16,7 @@ type SchedulerFair struct {
 	queues   map[string][]Job
 	queuesMu sync.Mutex
 
+	drfyarn       bool
 	enableBorrow  bool
 	IOUs          map[string]map[string]*ResourceCount
 	queuesQuota   map[string]*ResourceCount
@@ -41,7 +42,8 @@ func (scheduler *SchedulerFair) Start() {
 	scheduler.history = []*Job{}
 	scheduler.queues = map[string][]Job{}
 	scheduler.queues["default"] = []Job{}
-	scheduler.enableBorrow = true
+	scheduler.drfyarn = true
+	scheduler.enableBorrow = false
 	scheduler.IOUs = map[string]map[string]*ResourceCount{}
 	scheduler.queuesQuota = map[string]*ResourceCount{}
 	scheduler.resourceAllocations = map[string]*ResourceCount{}
@@ -212,6 +214,19 @@ func (scheduler *SchedulerFair) Start() {
 				}
 			}
 
+			/* drf of yarn/kube-batch */
+			if bestQueue == "" && scheduler.drfyarn {
+				least := math.MaxInt32
+				for queue, allocate := range scheduler.resourceAllocations {
+					if jobs, ok := scheduler.queues[queue]; ok && len(jobs) > 0 {
+						if bestQueue == "" || allocate.NumberGPU < least {
+							bestQueue = queue
+							least = allocate.NumberGPU
+						}
+					}
+				}
+			}
+
 			/* launch that job */
 			if bestQueue != "" {
 				numberGPUtmp := 0