From 71ab7bbb75042057414fc79bd08a0f6680fd4ef3 Mon Sep 17 00:00:00 2001
From: Newnius <newnius.cn@gmail.com>
Date: Sat, 11 Jul 2020 11:17:38 +0800
Subject: [PATCH] robust, make sure allocatinngGPU is decreased even if job is
 killed before requesting resource

---
 src/scheduler_fair.go | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/src/scheduler_fair.go b/src/scheduler_fair.go
index 7fe7c99..87caddf 100644
--- a/src/scheduler_fair.go
+++ b/src/scheduler_fair.go
@@ -367,7 +367,11 @@ func (scheduler *SchedulerFair) UpdateProgress(job Job, state State) {
 	defer scheduler.historyMu.Unlock()
 
 	scheduler.schedulingMu.Lock()
-	delete(scheduler.schedulingJobs, job.Name)
+	if _, ok := scheduler.schedulingJobs[job.Name]; ok {
+		delete(scheduler.schedulingJobs, job.Name)
+		scheduler.allocatingGPU -= job.NumberGPU
+		log.Info("allocatingGPU is ", scheduler.allocatingGPU)
+	}
 	scheduler.schedulingMu.Unlock()
 
 	switch state {
@@ -456,15 +460,6 @@ func (scheduler *SchedulerFair) Schedule(job Job) {
 
 func (scheduler *SchedulerFair) AcquireResource(job Job) []NodeStatus {
 	res := InstanceOfResourcePool().acquireResource(job)
-	if len(res) != 0 {
-		for _, task := range job.Tasks {
-
-			scheduler.allocatingGPUMu.Lock()
-			scheduler.allocatingGPU -= task.NumberGPU
-			scheduler.allocatingGPUMu.Unlock()
-		}
-		log.Info("allocatingGPU is ", scheduler.allocatingGPU)
-	}
 	go func() {
 		scheduler.UpdateQuota()
 	}()