From 0feccd0a514b97871d2af0407f431cc7ef353fab Mon Sep 17 00:00:00 2001
From: Newnius <newnius.cn@gmail.com>
Date: Wed, 24 Jun 2020 23:27:04 +0800
Subject: [PATCH] update

---
 src/optimizer.go | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/optimizer.go b/src/optimizer.go
index 56dc35c..3ef7da6 100644
--- a/src/optimizer.go
+++ b/src/optimizer.go
@@ -391,7 +391,12 @@ func (optimizer *Optimizer) predict(job string, seq int) (OptimizerJobExecutionT
 }
 
 func (optimizer *Optimizer) PredictReq(job Job, role string) MsgJobReq {
-	cmd := job.Tasks[0].Cmd
+	var jobName string
+	str := strings.Split(job.Name, "-")
+	if len(str) == 2 {
+		jobName = str[0]
+	}
+	cmd := ""
 	params := map[string]int{}
 
 	psNumber := 0
@@ -443,7 +448,7 @@ func (optimizer *Optimizer) PredictReq(job Job, role string) MsgJobReq {
 
 	spider := Spider{}
 	spider.Method = "GET"
-	spider.URL = "http://yao-optimizer:8080/predict?job=" + job.Name + "&features=" + string(features)
+	spider.URL = "http://yao-optimizer:8080/predict?job=" + jobName + "&features=" + string(features)
 
 	err := spider.do()
 	if err != nil {
@@ -474,6 +479,9 @@ func (optimizer *Optimizer) PredictReq(job Job, role string) MsgJobReq {
 		if v, ok := tmp["gpu_mem"]; ok {
 			req.MemGPU = int(math.Ceil(v/1024)) * 1024
 		}
+		if v, ok := tmp["bw"]; ok {
+			req.BW = int(math.Ceil(v/10)) * 10
+		}
 	}
 	return req
 }