add early stoppping service (#41)

* add early stoppping service Signed-off-by: YujiOshima <yuji.oshima0x3fd@gmail.com> * delete debug message Signed-off-by: YujiOshima <yuji.oshima0x3fd@gmail.com> * add medianstopping to test script Signed-off-by: YujiOshima <yuji.oshima0x3fd@gmail.com>
kubeflow · Apr 13, 2018 · 5733bd8 · 5733bd8
1 parent af2c2fa
commit 5733bd8
Show file tree

Hide file tree

Showing 21 changed files with 936 additions and 337 deletions.
diff --git a/api/api.pb.go b/api/api.pb.go
diff --git a/api/api.proto b/api/api.proto
@@ -8,7 +8,7 @@ service Manager {
 	rpc GetStudies(GetStudiesRequest) returns (GetStudiesReply);
 	rpc SuggestTrials(SuggestTrialsRequest) returns (SuggestTrialsReply);
 	rpc CompleteTrial(CompleteTrialRequest) returns (CompleteTrialReply);
-	rpc ShouldTrialStop(ShouldTrialStopRequest) returns (ShouldTrialStopReply);
+	rpc EarlyStopping(EarlyStoppingRequest) returns (EarlyStoppingReply);
 	rpc GetObjectValue(GetObjectValueRequest) returns (GetObjectValueReply);
 	rpc AddMeasurementToTrials(AddMeasurementToTrialsRequest) returns (AddMeasurementToTrialsReply);
     rpc InitializeSuggestService(InitializeSuggestServiceRequest) returns(InitializeSuggestServiceReply);
@@ -20,7 +20,9 @@ service Suggestion {
     rpc StopSuggestion(StopSuggestionRequest) returns (StopSuggestionReply);
 }
 
-service AutoStopping {
+service EarlyStopping {
+    rpc ShouldTrialStop(ShouldTrialStopRequest) returns (ShouldTrialStopReply);
+    rpc SetEarlyStoppingParameter(SetEarlyStoppingParameterRequest) returns (SetEarlyStoppingParameterReply);
 }
 
 enum ParameterType {
@@ -84,6 +86,11 @@ message SuggestionParameter {
     string value = 2;
 }
 
+message EarlyStoppingParameter {
+    string name = 1;
+    string value = 2;
+}
+
 message Tag {
     string name = 1;
     string value = 2;
@@ -115,18 +122,19 @@ message StudyConfig {
 	ParameterConfigs parameter_configs = 5;
 	repeated string access_permissions = 6;
 	string suggest_algorithm = 7;
-	string autostop_algorithm = 8;
+	string early_stopping_algorithm = 8;
 	string study_task_name = 9;
-    repeated SuggestionParameter suggestion_parameters =10;
-    repeated Tag tags = 11;
-    string objective_value_name = 12;
-    repeated string metrics = 13;
-    string image = 14;
-    repeated string command = 15;
-    int32 gpu = 16;
-    string scheduler = 17;
-    MountConf mount = 18;
-    string pull_secret = 19;
+    repeated SuggestionParameter suggestion_parameters = 10;
+    repeated EarlyStoppingParameter early_stopping_parameters= 11;
+    repeated Tag tags = 12;
+    string objective_value_name = 13;
+    repeated string metrics = 14;
+    string image = 15;
+    repeated string command = 16;
+    int32 gpu = 17;
+    string scheduler = 18;
+    MountConf mount = 19;
+    string pull_secret = 20;
 	//string log_collector = 10; // XXX
 }
 
@@ -172,21 +180,21 @@ message SuggestTrialsReply {
 }
 
 message CompleteTrialRequest {
-	string worker_id = 1;
-	bool is_complete = 2;
+	string study_id = 1;
+	string trial_id = 2;
+	bool is_complete = 3;
 }
 
 message CompleteTrialReply {
 }
 
-message ShouldTrialStopRequest {
+message EarlyStoppingRequest {
 	string study_id = 1;
-	string autostop_algorithm = 2;
+	string early_stopping_algorithm = 2;
 }
 
-message ShouldTrialStopReply {
+message EarlyStoppingReply {
 	repeated Trial trials = 1;
-	repeated string worker_ids = 2;
 }
 
 message GetObjectValueRequest {
@@ -243,3 +251,19 @@ message StopSuggestionRequest {
 
 message StopSuggestionReply {
 }
+
+message ShouldTrialStopRequest {
+	string study_id = 1;
+}
+
+message ShouldTrialStopReply {
+	repeated Trial trials = 1;
+}
+
+message SetEarlyStoppingParameterRequest {
+	string study_id = 1;
+    repeated EarlyStoppingParameter early_stopping_parameters =2;
+}
+
+message SetEarlyStoppingParameterReply {
+}
diff --git a/build.sh b/build.sh
@@ -6,6 +6,7 @@ docker build -t ${PREFIX}vizier-core -f manager/Dockerfile .
 docker build -t ${PREFIX}suggestion-random -f suggestion/random/Dockerfile .
 docker build -t ${PREFIX}suggestion-grid -f suggestion/grid/Dockerfile .
 docker build -t ${PREFIX}suggestion-hyperband -f suggestion/hyperband/Dockerfile .
+docker build -t ${PREFIX}earlystopping-medianstopping -f earlystopping/medianstopping/Dockerfile .
 docker build -t ${PREFIX}dlk-manager -f dlk/Dockerfile .
 docker build -t ${PREFIX}katib-frontend -f manager/modeldb/Dockerfile .
 docker build -t ${PREFIX}katib-cli -f cli/Dockerfile .

diff --git a/db/db_init.go b/db/db_init.go
@@ -14,9 +14,10 @@ func (d *db_conn) DB_Init() {
 		"optimization_goal DOUBLE, " +
 		"parameter_configs TEXT, " +
 		"suggest_algo VARCHAR(255), " +
-		"autostop_algo VARCHAR(255), " +
+		"early_stop_algo VARCHAR(255), " +
 		"study_task_name VARCHAR(255), " +
 		"suggestion_parameters TEXT, " +
+		"early_stopping_parameters TEXT, " +
 		"tags TEXT, " +
 		"objective_value_name VARCHAR(255), " +
 		"metrics TEXT, " +

diff --git a/db/interface.go b/db/interface.go
@@ -98,17 +98,18 @@ func (d *db_conn) GetStudyConfig(id string) (*api.StudyConfig, error) {
 	row := d.db.QueryRow("SELECT * FROM studies WHERE id = ?", id)
 
 	study := new(api.StudyConfig)
-	var dummy_id, configs, suggestion_parameters, tags, metrics, command, mconf string
+	var dummy_id, configs, suggestion_parameters, early_stopping_parameters, tags, metrics, command, mconf string
 	err := row.Scan(&dummy_id,
 		&study.Name,
 		&study.Owner,
 		&study.OptimizationType,
 		&study.OptimizationGoal,
 		&configs,
 		&study.SuggestAlgorithm,
-		&study.AutostopAlgorithm,
+		&study.EarlyStoppingAlgorithm,
 		&study.StudyTaskName,
 		&suggestion_parameters,
+		&early_stopping_parameters,
 		&tags,
 		&study.ObjectiveValueName,
 		&metrics,
@@ -204,6 +205,13 @@ func (d *db_conn) CreateStudy(in *api.StudyConfig) (string, error) {
 			log.Printf("Error marshalling %v: %v", elem, err)
 		}
 	}
+	earlystopping_parameters := make([]string, len(in.EarlyStoppingParameters))
+	for i, elem := range in.EarlyStoppingParameters {
+		earlystopping_parameters[i], err = (&jsonpb.Marshaler{}).MarshalToString(elem)
+		if err != nil {
+			log.Printf("Error marshalling %v: %v", elem, err)
+		}
+	}
 	var mconf string = ""
 	if in.Mount != nil {
 		mconf, err = (&jsonpb.Marshaler{}).MarshalToString(in.Mount)
@@ -236,17 +244,18 @@ func (d *db_conn) CreateStudy(in *api.StudyConfig) (string, error) {
 	for true {
 		study_id = generate_randid()
 		_, err := d.db.Exec(
-			"INSERT INTO studies VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+			"INSERT INTO studies VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
 			study_id,
 			in.Name,
 			in.Owner,
 			in.OptimizationType,
 			in.OptimizationGoal,
 			configs,
 			in.SuggestAlgorithm,
-			in.AutostopAlgorithm,
+			in.EarlyStoppingAlgorithm,
 			in.StudyTaskName,
 			strings.Join(suggestion_parameters, ",\n"),
+			strings.Join(earlystopping_parameters, ",\n"),
 			strings.Join(tags, ",\n"),
 			in.ObjectiveValueName,
 			strings.Join(in.Metrics, ",\n"),

diff --git a/db/interface_test.go b/db/interface_test.go
@@ -56,8 +56,28 @@ func TestGetStudyConfig(t *testing.T) {
 	}
 	//	mock.ExpectExec("SELECT * FROM studies WHERE id").WithArgs(id).WillReturnRows(sqlmock.NewRows())
 	mock.ExpectQuery("SELECT").WillReturnRows(
-		sqlmock.NewRows([]string{"id", "name", "owner", "optimization_type", "optimization_goal", "parameter_configs", "suggest_algo", "autostop_algo", "study_task_name", "suggestion_parameters", "tags", "objective_value_name", "metrics", "image", "command", "gpu", "scheduler", "mount", "pull_secret"}).
-			AddRow("abc", "test", "admin", 1, 0.99, "{}", "random", "test", "", "", "", "", "", "", "", 1, "", "", ""))
+		sqlmock.NewRows([]string{"id",
+			"name",
+			"owner",
+			"optimization_type",
+			"optimization_goal",
+			"parameter_configs",
+			"suggest_algo",
+			"early_stop_algo",
+			"study_task_name",
+			"suggestion_parameters",
+			"early_stopping_parameters",
+			"tags",
+			"objective_value_name",
+			"metrics",
+			"image",
+			"command",
+			"gpu",
+			"scheduler",
+			"mount",
+			"pull_secret",
+		}).
+			AddRow("abc", "test", "admin", 1, 0.99, "{}", "random", "test", "", "", "", "", "", "", "", "", 1, "", "", ""))
 	study, err := db_interface.GetStudyConfig(id)
 	if err != nil {
 		t.Errorf("GetStudyConfig failed: %v", err)

diff --git a/earlystopping/earlyStoppingService.go b/earlystopping/earlyStoppingService.go