add custom logger to support mute certain modules

src/allocator.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"sync"
-	log "github.com/sirupsen/logrus"
 	"math"
 	"time"
 	"github.com/MaxHalford/eaopt"

src/collector.go

@@ -4,7 +4,6 @@ import (
 	"sync"
 	"github.com/Shopify/sarama"
 	"encoding/json"
-	log "github.com/sirupsen/logrus"
 	"time"
 )
 

src/configuration.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"sync"
-	log "github.com/sirupsen/logrus"
 	"os"
 	"strings"
 	"strconv"
@@ -126,24 +125,32 @@ func (config *Configuration) SetMockEnabled(enabled bool) bool {
 }
 
 func (config *Configuration) SetShareRatio(ratio float64) bool {
+	config.mu.Lock()
+	defer config.mu.Unlock()
 	config.EnableShareRatio = ratio
 	log.Info("enableShareRatio is updated to ", ratio)
 	return true
 }
 
 func (config *Configuration) SetPreScheduleRatio(ratio float64) bool {
+	config.mu.Lock()
+	defer config.mu.Unlock()
 	config.EnablePreScheduleRatio = ratio
 	log.Info("enablePreScheduleRatio is updated to ", ratio)
 	return true
 }
 
 func (config *Configuration) SetShareMaxUtilization(value float64) bool {
+	config.mu.Lock()
+	defer config.mu.Unlock()
 	config.ShareMaxUtilization = value
 	log.Info("ShareMaxUtilization is set to ", value)
 	return true
 }
 
 func (config *Configuration) Dump() map[string]interface{} {
+	config.mu.Lock()
+	defer config.mu.Unlock()
 	res := map[string]interface{}{}
 	res["KafkaBrokers"] = config.KafkaBrokers
 	res["KafkaTopic"] = config.KafkaTopic
@@ -158,5 +165,7 @@ func (config *Configuration) Dump() map[string]interface{} {
 	res["EnablePreScheduleRatio"] = config.EnablePreScheduleRatio
 	res["PreScheduleExtraTime"] = config.PreScheduleExtraTime
 	res["PreScheduleTimeout"] = config.PreScheduleTimeout
+	res["logger.level"] = log.LoggerLevel
+	res["logger.modules_disabled"] = log.LoggerModuleDisabled
 	return res
 }

src/ga.go

@@ -4,7 +4,6 @@ import (
 	"math/rand"
 	"github.com/MaxHalford/eaopt"
 	"math"
-	log "github.com/sirupsen/logrus"
 )
 
 // A resource allocation

src/ga_test.go

@@ -3,7 +3,6 @@ package main
 import (
 	"strconv"
 	"time"
-	log "github.com/sirupsen/logrus"
 	"testing"
 )
 
@@ -50,8 +49,8 @@ func TestBestFit(t *testing.T) {
 	}
 	s := time.Now()
 	allocation := InstanceOfAllocator().fastBestFit(nodes, tasks)
-	log.Println(time.Since(s))
-	log.Println(allocation)
+	log.Info(time.Since(s))
+	log.Info(allocation)
 }
 
 func TestGA(t *testing.T) {

src/history_logger.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	log "github.com/sirupsen/logrus"
 	"sync"
 )
 

src/job_manager.go

@@ -6,7 +6,6 @@ import (
 	"strings"
 	"io/ioutil"
 	"encoding/json"
-	log "github.com/sirupsen/logrus"
 	"sync"
 	"strconv"
 	"math/rand"
@@ -346,7 +345,7 @@ func (jm *JobManager) logs(taskName string) MsgLog {
 	var res MsgLog
 	err = json.Unmarshal([]byte(string(body)), &res)
 	if err != nil {
-		log.Println(err)
+		log.Warn(err)
 		return MsgLog{Code: 3, Error: "Unknown"}
 	}
 	return res

src/main.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"net/http"
-	log "github.com/sirupsen/logrus"
 	"encoding/json"
 	"time"
 	"strconv"
@@ -10,6 +9,8 @@ import (
 	"os"
 )
 
+var log Logger
+
 var scheduler Scheduler
 
 func serverAPI(w http.ResponseWriter, r *http.Request) {
@@ -332,6 +333,18 @@ func serverAPI(w http.ResponseWriter, r *http.Request) {
 			ok = InstanceOfAllocator().updateStrategy(value)
 			break
 
+		case "logger.level":
+			ok = log.SetLoggerLevel(value)
+			break
+
+		case "logger.enable_module":
+			ok = log.LoggerEnableModule(value)
+			break
+
+		case "logger.disable_module":
+			ok = log.LoggerDisableModule(value)
+			break
+
 		}
 		var msg MsgConfUpdate
 		msg.Code = 0
@@ -365,7 +378,6 @@ func main() {
 		}
 		log.SetOutput(f)
 	}
-	//log.SetLevel(log.InfoLevel)
 
 	/* parse configuration */
 	config := *InstanceOfConfiguration()

src/optimizer.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	log "github.com/sirupsen/logrus"
 	"sync"
 	"strings"
 	"io/ioutil"

src/pool_test.go

@@ -3,7 +3,6 @@ package main
 import (
 	"testing"
 	"strconv"
-	log "github.com/sirupsen/logrus"
 	"time"
 )
 

src/resource_pool.go

@@ -5,7 +5,6 @@ import (
 	"time"
 	"net/url"
 	"strings"
-	log "github.com/sirupsen/logrus"
 	"math/rand"
 	"strconv"
 	"sort"
@@ -30,7 +29,8 @@ type ResourcePool struct {
 	pools      []PoolSeg
 	poolsMu    sync.Mutex
 
-	history []PoolStatus
+	history   []PoolStatus
+	historyMu sync.Mutex
 
 	heartBeat    map[string]time.Time
 	heartBeatMu  sync.Mutex
@@ -67,6 +67,8 @@ type ResourcePool struct {
 func (pool *ResourcePool) init(conf Configuration) {
 	log.Info("RM started ")
 
+	InstanceOfConfiguration().LoggerEnableModule("RM")
+
 	pool.networks = map[string]bool{}
 	pool.networksFree = map[string]bool{}
 
@@ -289,11 +291,12 @@ func (pool *ResourcePool) saveStatusHistory() {
 		summary.TotalMemGPU = TotalMemGPU
 		summary.AvailableMemGPU = AvailableMemGPU
 
+		pool.historyMu.Lock()
 		pool.history = append(pool.history, summary)
-
 		if len(pool.history) > 60 {
 			pool.history = pool.history[len(pool.history)-60:]
 		}
+		pool.historyMu.Unlock()
 
 		pool.TotalGPUMu.Lock()
 		pool.TotalGPU = TotalGPU
@@ -490,7 +493,10 @@ func (pool *ResourcePool) list() MsgResource {
 }
 
 func (pool *ResourcePool) statusHistory() MsgPoolStatusHistory {
-	return MsgPoolStatusHistory{Code: 0, Data: pool.history}
+	pool.historyMu.Lock()
+	defer pool.historyMu.Unlock()
+	history := pool.history
+	return MsgPoolStatusHistory{Code: 0, Data: history}
 }
 
 func (pool *ResourcePool) getCounter() map[string]int {
@@ -514,7 +520,7 @@ func (pool *ResourcePool) acquireNetwork() string {
 			v.Set("name", network)
 			resp, err := doRequest("POST", "http://yao-agent-master:8000/create", strings.NewReader(v.Encode()), "application/x-www-form-urlencoded", "")
 			if err != nil {
-				log.Println(err.Error())
+				log.Warn(err.Error())
 				continue
 			}
 			resp.Body.Close()

src/scheduler_FCFS.go

@@ -3,7 +3,6 @@ package main
 import (
 	"sync"
 	"time"
-	log "github.com/sirupsen/logrus"
 )
 
 type SchedulerFCFS struct {

src/scheduler_capacity.go

@@ -3,7 +3,6 @@ package main
 import (
 	"sync"
 	"time"
-	log "github.com/sirupsen/logrus"
 	"sort"
 	"math"
 )

src/scheduler_fair.go

@@ -3,7 +3,6 @@ package main
 import (
 	"sync"
 	"time"
-	log "github.com/sirupsen/logrus"
 	"math"
 	"sort"
 )

src/scheduler_priority.go

@@ -3,7 +3,6 @@ package main
 import (
 	"sync"
 	"time"
-	log "github.com/sirupsen/logrus"
 	"sort"
 	"math"
 )