From 301071fc87271fbee5d363e90b4c41dcadba1a8a Mon Sep 17 00:00:00 2001
From: Newnius <newnius.cn@gmail.com>
Date: Mon, 13 Apr 2020 18:26:40 +0800
Subject: [PATCH] add switcher for schduler

---
 README.md             | 10 ++++++++++
 src/history_logger.go |  4 ++--
 src/main.go           | 14 ++++++++++++++
 src/scheduler.go      |  4 ++++
 src/scheduler_fair.go | 14 +++++++++++++-
 5 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 11deb68..f4c19f2 100644
--- a/README.md
+++ b/README.md
@@ -21,4 +21,14 @@ GPU is occupied by which job(s)
 
 ```
 ?action=get_bindings
+```
+
+**EnableSchedule**
+```
+?action=debug_enable
+```
+
+**DisableSchedule**
+```
+?action=debug_disable
 ```
\ No newline at end of file
diff --git a/src/history_logger.go b/src/history_logger.go
index 53ac340..928390f 100644
--- a/src/history_logger.go
+++ b/src/history_logger.go
@@ -41,14 +41,14 @@ func (jhl *JobHistoryLogger) submitJob(job Job) {
 }
 
 func (jhl *JobHistoryLogger) updateJobStatus(jobName string, state State) {
-	log.Info("update job status", jobName)
+	log.Debug("update job status", jobName)
 	if job, ok := jhl.jobs[jobName]; ok {
 		job.Status = state
 	}
 }
 
 func (jhl *JobHistoryLogger) submitTaskStatus(jobName string, task TaskStatus) {
-	log.Info("submit job task status", jobName)
+	log.Debug("submit job task status", jobName)
 	if tasks, ok := jhl.tasks[jobName]; ok {
 		jhl.tasks[jobName] = append(tasks, task)
 	}
diff --git a/src/main.go b/src/main.go
index a14c4d1..ec17852 100644
--- a/src/main.go
+++ b/src/main.go
@@ -166,6 +166,20 @@ func serverAPI(w http.ResponseWriter, r *http.Request) {
 		w.Write(js)
 		break
 
+	case "debug_enable":
+		log.Debug("enable schedule")
+		js, _ := json.Marshal(scheduler.Enable)
+		w.Header().Set("Content-Type", "application/json")
+		w.Write(js)
+		break
+
+	case "debug_disable":
+		log.Debug("disable schedule")
+		js, _ := json.Marshal(scheduler.Disable)
+		w.Header().Set("Content-Type", "application/json")
+		w.Write(js)
+		break
+
 	default:
 		http.Error(w, "Not Found", http.StatusNotFound)
 		break
diff --git a/src/scheduler.go b/src/scheduler.go
index afa149c..c4b40b1 100644
--- a/src/scheduler.go
+++ b/src/scheduler.go
@@ -28,4 +28,8 @@ type Scheduler interface {
 	Attach(GPU string, job string)
 
 	Detach(GPU string, job string)
+
+	Enable()
+
+	Disable()
 }
diff --git a/src/scheduler_fair.go b/src/scheduler_fair.go
index e0db6cf..c2b263e 100644
--- a/src/scheduler_fair.go
+++ b/src/scheduler_fair.go
@@ -22,6 +22,7 @@ type SchedulerFair struct {
 	jobs                map[string]*JobManager
 	nextQueue           string
 	resourceAllocations map[string]*ResourceCount
+	enabled             bool
 }
 
 type FairJobSorter []Job
@@ -48,6 +49,9 @@ func (scheduler *SchedulerFair) Start() {
 		for {
 			log.Debug("Scheduling")
 			time.Sleep(time.Millisecond * 100)
+			if !scheduler.enabled {
+				continue
+			}
 			scheduler.scheduling.Lock()
 			scheduler.mu.Lock()
 			queue := scheduler.nextQueue
@@ -207,7 +211,7 @@ func (scheduler *SchedulerFair) ReleaseResource(job Job, agent NodeStatus) {
 			if gpu.UUID == nodes.Status[j].UUID {
 				nodes.Status[j].MemoryAllocated -= gpu.MemoryTotal
 				if nodes.Status[j].MemoryAllocated < 0 {
-					// in case error
+					// in case of error
 					nodes.Status[j].MemoryAllocated = 0
 				}
 			}
@@ -381,3 +385,11 @@ func (scheduler *SchedulerFair) Attach(GPU string, job string) {
 func (scheduler *SchedulerFair) Detach(GPU string, job string) {
 	pool.detach(GPU, job)
 }
+
+func (scheduler *SchedulerFair) Enable() {
+	scheduler.enabled = true
+}
+
+func (scheduler *SchedulerFair) Disable() {
+	scheduler.enabled = false
+}