Fix pai training with optimizer config

go/cmd/sqlflowserver/e2e_pai_maxcompute_test.go

@@ -282,7 +282,7 @@ func CasePAIMaxComputeDNNTrainPredictExplain(t *testing.T) {
 	a := assert.New(t)
 	trainSQL := fmt.Sprintf(`SELECT * FROM %s
 TO TRAIN DNNClassifier
-WITH model.n_classes = 3, model.hidden_units = [10, 20]
+WITH model.n_classes = 3, model.hidden_units = [10, 20], optimizer.learning_rate=0.01
 LABEL class
 INTO e2etest_pai_dnn;`, caseTrainTable)
 	_, _, _, err := connectAndRunSQL(trainSQL)

go/codegen/pai/template_tf.go

@@ -92,40 +92,6 @@ else:
     oss.load_file("{{.OSSModelDir}}", "model_save")
 `
 
-const tfSaveModelTmplText = tfImportsText + `
-import types
-
-estimator = import_model('''{{.Estimator}}''')
-is_estimator = is_tf_estimator(estimator)
-
-# Keras single node is using h5 format to save the model, no need to deal with export model format.
-# Keras distributed mode will use estimator, so this is also needed.
-FLAGS = tf.app.flags.FLAGS
-if is_estimator:
-    if FLAGS.task_index == 0:
-        with open("exported_path", "r") as fn:
-            saved_model_path = fn.read()
-        oss.save_dir("{{.OSSModelDir}}", saved_model_path)
-        oss.save_file("{{.OSSModelDir}}", "exported_path")
-else:
-    if len(FLAGS.worker_hosts.split(",")) > 1:
-        if FLAGS.task_index == 0:
-            oss.save_file("{{.OSSModelDir}}", "exported_path")
-    else:
-        oss.save_file("{{.OSSModelDir}}", "model_save")
-
-oss.save_metas("{{.OSSModelDir}}",
-           {{.NumWorkers}},
-           "tensorflow_model_desc",
-           "{{.Estimator}}",
-           feature_column_names,
-           feature_column_names_map,
-           feature_metas,
-           label_meta,
-           model_params,
-           feature_columns_code)
-`
-
 // install sklearn-pandas==1.8.0 to fix deps for sklearn2pmml with Python2 on PAI.
 const paiRequirementsTmplText = `
 adanet==0.8.0

go/codegen/pai/tensorflow.go

@@ -70,20 +70,7 @@ func TFTrainWithLoadAndSave(ir *ir.TrainStmt, session *pb.Session, modelPathToSa
 		return "", err
 	}
 
-	// append code snippet to save model
-	checkpointDir := OSSModelURL(modelPathToSave)
-	var tpl = template.Must(template.New("SaveModel").Parse(tfSaveModelTmplText))
-	filler := saveModelFiller{
-		OSSModelDir: checkpointDir,
-		Estimator:   ir.Estimator,
-		NumWorkers:  cc.Worker.Count,
-	}
-	var saveCode bytes.Buffer
-	if err = tpl.Execute(&saveCode, filler); err != nil {
-		return "", err
-	}
-
-	fullCode := fmt.Sprintf("%s\n%s\n%s", loadCode, trainCode, saveCode.String())
+	fullCode := fmt.Sprintf("%s\n%s", loadCode, trainCode)
 	return fullCode, nil
 }
 

go/codegen/tensorflow/template_train.go

@@ -149,6 +149,7 @@ train(datasource="{{.DataSource}}",
       pai_table="{{.PAITrainTable}}",
       pai_val_table="{{.PAIValidateTable}}",
       feature_columns_code=feature_columns_code,
+      model_params_code_map=model_params,
       model_repo_image="{{.ModelRepoImage}}",
       original_sql='''{{.OriginalSQL}}''')
 `

python/runtime/model/oss.py

@@ -283,7 +283,9 @@ def save_oss_model(oss_model_dir, model_name, is_estimator,
         save_file(oss_model_dir, "exported_path")
     else:
         if num_workers > 1:
-            save_file(oss_model_dir, "exported_path")
+            FLAGS = tf.app.flags.FLAGS
+            if FLAGS.task_index == 0:
+                save_file(oss_model_dir, "exported_path")
         else:
             save_file(oss_model_dir, "model_save")
 

python/runtime/pai/tensorflow/train.py

@@ -52,6 +52,7 @@ def train(datasource,
           pai_table="",
           pai_val_table="",
           feature_columns_code="",
+          model_params_code_map={},
           model_repo_image="",
           original_sql="",
           feature_column_names_map=None):
@@ -119,7 +120,7 @@ def train(datasource,
         oss_model_dir = FLAGS.sqlflow_oss_modeldir
         oss.save_oss_model(oss_model_dir, estimator_string, is_estimator,
                            feature_column_names, feature_column_names_map,
-                           feature_metas, label_meta, model_params,
+                           feature_metas, label_meta, model_params_code_map,
                            feature_columns_code, num_workers)
         print("Model saved to oss: %s" % oss_model_dir)
     print("Done training")

python/runtime/tensorflow/train.py

@@ -52,6 +52,7 @@ def train(datasource,
           pai_table="",
           pai_val_table="",
           feature_columns_code="",
+          model_params_code_map={},
           model_repo_image="",
           original_sql=""):
     # TODO(sneaxiy): collect features and label