Add workflow codegen for XGBoost predict (#2845)

* add codegen predict

* replace temp_file apis

* update

* update

Author: sneaxiy

go/cmd/sqlflowserver/e2e_workflow_test.go

@@ -373,7 +373,14 @@ SELECT * FROM iris.train
 TO TRAIN xgboost.gbtree
 WITH objective="multi:softmax",num_class=3
 LABEL class
-INTO sqlflow_models.xgb_classification;`
+INTO sqlflow_models.xgb_classification;
+
+SELECT * FROM iris.test
+TO PREDICT iris.test_result_table.class
+USING sqlflow_models.xgb_classification;
+
+SELECT * FROM iris.test_result_table;
+`
 
 	conn, err := createRPCConn()
 	if err != nil {

go/codegen/experimental/codegen_normal_stmt.go

@@ -20,7 +20,7 @@ import (
 	pb "sqlflow.org/sqlflow/go/proto"
 )
 
-var normalStmtStepTmpl = `
+const normalStmtStepTmpl = `
 def step_entry_{{.StepIndex}}():
     import runtime
     import runtime.dbapi

go/codegen/experimental/codegen_step.go

@@ -24,13 +24,19 @@ import (
 	pb "sqlflow.org/sqlflow/go/proto"
 )
 
-func generateStepCode(stmt ir.SQLFlowStmt, stepIndex int, session *pb.Session) (string, error) {
-	switch stmt.(type) {
+// TODO(sneaxiy): implement this method to distinguish whether
+// a model is a XGBoost model.
+func isTrainedXBoostModel(modelName string) bool {
+	return true
+}
+
+func generateStepCode(sqlStmt ir.SQLFlowStmt, stepIndex int, session *pb.Session) (string, error) {
+	switch stmt := sqlStmt.(type) {
 	case *ir.TrainStmt:
-		trainStmt := stmt.(*ir.TrainStmt)
-		return generateTrainCode(trainStmt, stepIndex, session)
+		return generateTrainCode(stmt, stepIndex, session)
+	case *ir.PredictStmt:
+		return generatePredictCode(stmt, stepIndex, session)
 	case *ir.NormalStmt:
-		stmt := stmt.(*ir.NormalStmt)
 		return GenerateNormalStmtStep(string(*stmt), session, stepIndex)
 	default:
 		return "", fmt.Errorf("not implemented stmt execution type %v", stmt)
@@ -44,6 +50,13 @@ func generateTrainCode(trainStmt *ir.TrainStmt, stepIndex int, session *pb.Sessi
 	return "", fmt.Errorf("not implemented estimator type %s", trainStmt.Estimator)
 }
 
+func generatePredictCode(predStmt *ir.PredictStmt, stepIndex int, session *pb.Session) (string, error) {
+	if isTrainedXBoostModel(predStmt.Using) {
+		return XGBoostGeneratePredict(predStmt, stepIndex, session)
+	}
+	return "", fmt.Errorf("not implemented model type")
+}
+
 func initializeAndCheckAttributes(stmt ir.SQLFlowStmt) error {
 	switch s := stmt.(type) {
 	case *ir.TrainStmt:

go/codegen/experimental/xgboost.go

@@ -99,10 +99,6 @@ func XGBoostGenerateTrain(trainStmt *ir.TrainStmt, stepIndex int, session *pb.Se
 	if err != nil {
 		return "", err
 	}
-	submitter := os.Getenv("SQLFLOW_submitter")
-	if submitter == "" {
-		submitter = "local"
-	}
 
 	dbConnStr, err := GeneratePyDbConnStr(session)
 	if err != nil {
@@ -126,7 +122,7 @@ func XGBoostGenerateTrain(trainStmt *ir.TrainStmt, stepIndex int, session *pb.Se
 		DiskCache:         diskCache,
 		BatchSize:         batchSize,
 		Epoch:             epoch,
-		Submitter:         submitter,
+		Submitter:         getSubmitter(session, "local"),
 	}
 	var program bytes.Buffer
 	var trainTemplate = template.Must(template.New("Train").Parse(xgbTrainTemplate))
@@ -140,7 +136,6 @@ func XGBoostGenerateTrain(trainStmt *ir.TrainStmt, stepIndex int, session *pb.Se
 const xgbTrainTemplate = `
 def step_entry_{{.StepIndex}}():
     import json
-    import os
     import runtime.temp_file as temp_file
     import runtime.feature.column as fc
     import runtime.feature.field_desc as fd
@@ -157,7 +152,6 @@ def step_entry_{{.StepIndex}}():
     train_params = json.loads('''{{.TrainParamsJSON}}''')
 
     with temp_file.TemporaryDirectory(as_cwd=True) as temp_dir:
-        os.chdir(temp_dir)
         train_params["original_sql"] = '''{{.OriginalSQL}}'''
         train_params["model_image"] = '''{{.ModelImage}}'''
         train_params["feature_column_map"] = feature_column_map
@@ -176,6 +170,67 @@ def step_entry_{{.StepIndex}}():
               train_params=train_params)
 `
 
+type xgbPredFiller struct {
+	StepIndex     int
+	DataSource    string
+	Select        string
+	PredLabelName string
+	ResultTable   string
+	Load          string
+	Submitter     string
+}
+
+// XGBoostGeneratePredict generates the XGBoost prediction code
+func XGBoostGeneratePredict(predStmt *ir.PredictStmt, stepIndex int, session *pb.Session) (string, error) {
+	dbConnStr, err := GeneratePyDbConnStr(session)
+	if err != nil {
+		return "", err
+	}
+
+	filler := &xgbPredFiller{
+		StepIndex:     stepIndex,
+		DataSource:    dbConnStr,
+		Select:        replaceNewLineRuneAndTrimSpace(predStmt.Select),
+		PredLabelName: predStmt.ResultColumn,
+		ResultTable:   predStmt.ResultTable,
+		Load:          predStmt.Using,
+		Submitter:     getSubmitter(session, "local"),
+	}
+
+	var program bytes.Buffer
+	predTmpl := template.Must(template.New("Train").Parse(xgbPredTemplate))
+	err = predTmpl.Execute(&program, filler)
+	if err != nil {
+		return "", err
+	}
+	return program.String(), nil
+}
+
+const xgbPredTemplate = `
+def step_entry_{{.StepIndex}}():
+    import runtime.temp_file as temp_file
+    from runtime.{{.Submitter}} import pred
+    
+    with temp_file.TemporaryDirectory(as_cwd=True):
+        pred(datasource='''{{.DataSource}}''', 
+             select='''{{.Select}}''', 
+             result_table='''{{.ResultTable}}''', 
+             pred_label_name='''{{.PredLabelName}}''', 
+             load='''{{.Load}}''')
+`
+
+func getSubmitter(session *pb.Session, defaultValue string) string {
+	if session.Submitter != "" {
+		return session.Submitter
+	}
+
+	submitter := os.Getenv("SQLFLOW_submitter")
+	if submitter != "" {
+		return submitter
+	}
+	return defaultValue
+}
+
 func generateFeatureColumnCode(fcList []ir.FeatureColumn) (string, error) {
 	fcCodes := make([]string, 0, len(fcList))
 	for _, fc := range fcList {

python/runtime/local/__init__.py

@@ -11,4 +11,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from runtime.local.submitter import submit_local_pred as pred  # noqa: F401
 from runtime.local.submitter import submit_local_train as train  # noqa: F401

python/runtime/local/submitter.py

@@ -11,7 +11,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from runtime.local.xgboost_submitter.predict import pred as xgboost_pred
 from runtime.local.xgboost_submitter.train import train as xgboost_train
+from runtime.model.model import EstimatorType, Model
 
 
 def submit_local_train(datasource, estimator_string, select, validation_select,
@@ -60,3 +62,12 @@ def submit_local_train(datasource, estimator_string, select, validation_select,
     else:
         raise NotImplementedError("not implemented model type: %s" %
                                   estimator_string)
+
+
+def submit_local_pred(datasource, select, result_table, pred_label_name, load):
+    model = Model.load_from_db(datasource, load)
+    if model.get_type() == EstimatorType.XGBOOST:
+        xgboost_pred(datasource, select, result_table, pred_label_name, model)
+    else:
+        raise NotImplementedError("not implemented model type: %s" %
+                                  model.get_type())

…/xgboost_submitter/train_predict_test.py …gboost_submitter/local_submitter_test.pypython/runtime/local/xgboost_submitter/train_predict_test.py renamed to python/runtime/local/xgboost_submitter/local_submitter_test.py python/runtime/local/xgboost_submitter/train_predict_test.py renamed to python/runtime/local/xgboost_submitter/local_submitter_test.py

@@ -12,10 +12,10 @@
 # limitations under the License.
 
 import os
-import tempfile
 import unittest
 
 import runtime.db as db
+import runtime.temp_file as temp_file
 import runtime.testing as testing
 from runtime.feature.column import NumericColumn
 from runtime.feature.field_desc import FieldDesc
@@ -57,8 +57,7 @@ def test_train_and_predict(self):
         class_name = "class"
 
         old_dir_name = os.getcwd()
-        with tempfile.TemporaryDirectory() as tmp_dir_name:
-            os.chdir(tmp_dir_name)
+        with temp_file.TemporaryDirectory(as_cwd=True):
             eval_result = train(original_sql=original_sql,
                                 model_image="sqlflow:step",
                                 estimator_string="xgboost.gbtree",

python/runtime/local/xgboost_submitter/predict.py

@@ -12,10 +12,11 @@
 # limitations under the License.
 
 import os
-import tempfile
 
 import numpy as np
+import runtime.temp_file as temp_file
 import runtime.xgboost as xgboost_extended
+import six
 import xgboost as xgb
 from runtime import db
 from runtime.feature.compile import compile_ir_feature_columns
@@ -25,7 +26,7 @@
 from runtime.xgboost.dataset import xgb_dataset
 
 
-def pred(datasource, select, result_table, pred_label_name, load):
+def pred(datasource, select, result_table, pred_label_name, model):
     """
     Do prediction using a trained model.
 
@@ -34,12 +35,17 @@ def pred(datasource, select, result_table, pred_label_name, load):
         select (str): the input data to predict.
         result_table (str): the output data table.
         pred_label_name (str): the output label name to predict.
-        load (str): where the trained model stores.
+        model (Model|str): the model object or where to load the model.
 
     Returns:
         None.
     """
-    model = Model.load_from_db(datasource, load)
+    if isinstance(model, six.string_types):
+        model = Model.load_from_db(datasource, model)
+    else:
+        assert isinstance(model,
+                          Model), "not supported model type %s" % type(model)
+
     model_params = model.get_meta("attributes")
     train_fc_map = model.get_meta("features")
     train_label_desc = model.get_meta("label").get_field_desc()[0]
@@ -62,7 +68,7 @@ def pred(datasource, select, result_table, pred_label_name, load):
     result_column_names, train_label_idx = _create_predict_table(
         conn, select, result_table, train_label_desc, pred_label_name)
 
-    with tempfile.TemporaryDirectory() as tmp_dir_name:
+    with temp_file.TemporaryDirectory() as tmp_dir_name:
         pred_fn = os.path.join(tmp_dir_name, "predict.txt")
         raw_data_dir = os.path.join(tmp_dir_name, "predict_raw_dir")
 

python/runtime/tensorflow/import_model.py

@@ -62,5 +62,5 @@ def import_model(model):
     Returns:
         An imported model class or function.
     """
-    import_model_package(model, globals())
+    import_model_package(model, locals())
     return eval(model)