support pai random forests predict and explain (#1706)

Author: typhoonzero

cmd/sqlflowserver/main_test.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io"
@@ -1280,16 +1281,62 @@ USING %s;`, caseTestTable, casePredictTable, caseInto)
 
 }
 
+func dropPAIModel(dataSource, modelName string) error {
+	code := fmt.Sprintf(`import subprocess
+import sqlflow_submitter.db
+
+driver, dsn = "%s".split("://")
+assert driver == "maxcompute"
+user, passwd, address, database = sqlflow_submitter.db.parseMaxComputeDSN(dsn)
+
+cmd = "drop offlinemodel if exists %s"
+subprocess.run(["odpscmd", "-u", user,
+                           "-p", passwd,
+                           "--project", database,
+                           "--endpoint", address,
+                           "-e", cmd],
+               check=True)	
+	`, dataSource, modelName)
+	cmd := exec.Command("python", "-u")
+	cmd.Stdin = bytes.NewBufferString(code)
+	if e := cmd.Run(); e != nil {
+		return e
+	}
+	return nil
+}
+
 func CaseTrainPAIRandomForests(t *testing.T) {
 	a := assert.New(t)
-	trainSQL := fmt.Sprintf(`
-	SELECT * FROM %s.%s
+	err := dropPAIModel(dbConnStr, "my_rf_model")
+	a.NoError(err)
+
+	trainSQL := fmt.Sprintf(`SELECT * FROM %s
 	TO TRAIN randomforests
 	WITH tree_num = 3
 	LABEL class
 	INTO my_rf_model;
-	`, caseDB, caseTrainTable)
-	_, _, err := connectAndRunSQL(trainSQL)
+	`, caseTrainTable)
+	_, _, err = connectAndRunSQL(trainSQL)
+	if err != nil {
+		a.Fail("Run trainSQL error: %v", err)
+	}
+
+	predSQL := fmt.Sprintf(`SELECT * FROM %s
+	TO PREDICT %s.class
+	USING my_rf_model;
+	`, caseTestTable, casePredictTable)
+	_, _, err = connectAndRunSQL(predSQL)
+	if err != nil {
+		a.Fail("Run trainSQL error: %v", err)
+	}
+
+	explainSQL := fmt.Sprintf(`SELECT * FROM %s
+	TO EXPLAIN my_rf_model
+	WITH label_column = class
+	USING TreeExplainer
+	INTO %s.rf_model_explain;
+	`, caseTestTable, caseDB)
+	_, _, err = connectAndRunSQL(explainSQL)
 	if err != nil {
 		a.Fail("Run trainSQL error: %v", err)
 	}

pkg/sql/codegen/pai/codegen.go

@@ -17,15 +17,19 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"log"
 	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
 	"text/template"
 
+	"github.com/aliyun/aliyun-oss-go-sdk/oss"
+	"sqlflow.org/sqlflow/pkg/database"
 	pb "sqlflow.org/sqlflow/pkg/proto"
 	"sqlflow.org/sqlflow/pkg/sql/codegen/tensorflow"
 	"sqlflow.org/sqlflow/pkg/sql/ir"
+	"sqlflow.org/sqlflow/pkg/verifier"
 )
 
 const entryFile = "entry.py"
@@ -160,7 +164,6 @@ func trainRandomForests(ir *ir.TrainStmt, session *pb.Session) (string, error) {
 	}
 	filler := &randomForestsTrainFiller{
 		DataSource:     session.DbConnStr,
-		Select:         ir.Select,
 		TmpTrainTable:  ir.TmpTrainTable,
 		FeatureColumns: featureCols,
 		LabelColumn:    ir.Label.GetFieldDesc()[0].Name,
@@ -172,10 +175,44 @@ func trainRandomForests(ir *ir.TrainStmt, session *pb.Session) (string, error) {
 	if err := tpl.Execute(&rfCode, filler); err != nil {
 		return "", err
 	}
-	fmt.Println(rfCode.String())
 	return rfCode.String(), nil
 }
 
+// getColumnTypes is quiet like verify but accept a SQL string as input, and returns
+// an ordered list of the field types.
+// FIXME(typhoonzero): copied from executor_ir.go
+func getColumnTypes(slct string, db *database.DB) ([]string, []string, error) {
+	rows, err := db.Query(slct)
+	if err != nil {
+		return nil, nil, err
+	}
+	defer rows.Close()
+
+	if !rows.Next() {
+		return nil, nil, fmt.Errorf("query %s gives 0 row", slct)
+	}
+
+	if rows.Err() != nil {
+		return nil, nil, err
+	}
+
+	columnTypes, err := rows.ColumnTypes()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	ft := []string{}
+	flds := []string{}
+	for _, ct := range columnTypes {
+		_, fld := verifier.Decomp(ct.Name())
+		typeName := ct.DatabaseTypeName()
+		flds = append(flds, fld)
+		ft = append(ft, typeName)
+	}
+
+	return flds, ft, nil
+}
+
 // Train generates a Python program for train a TensorFlow model.
 func Train(ir *ir.TrainStmt, session *pb.Session, modelName, cwd string) (string, error) {
 	if strings.ToLower(ir.Estimator) == "randomforests" {
@@ -214,8 +251,74 @@ func TFTrainAndSave(ir *ir.TrainStmt, session *pb.Session, modelName string) (st
 	return code + saveCode.String(), nil
 }
 
+func ossFileExists(modelName string) (bool, error) {
+	endpoint := os.Getenv("SQLFLOW_OSS_ENDPOINT")
+	ak := os.Getenv("SQLFLOW_OSS_AK")
+	sk := os.Getenv("SQLFLOW_OSS_SK")
+	// NOTE(typhoonzero): PAI Tensorflow need SQLFLOW_OSS_CHECKPOINT_DIR, get bucket name from it
+	ossCheckpointDir := os.Getenv("SQLFLOW_OSS_CHECKPOINT_DIR")
+	ckptParts := strings.Split(ossCheckpointDir, "?")
+	if len(ckptParts) != 2 {
+		return false, fmt.Errorf("SQLFLOW_OSS_CHECKPOINT_DIR got wrong format")
+	}
+	urlParts := strings.Split(ckptParts[0], "://")
+	if len(urlParts) != 2 {
+		return false, fmt.Errorf("SQLFLOW_OSS_CHECKPOINT_DIR got wrong format")
+	}
+	bucketName := strings.Split(urlParts[1], "/")[0]
+
+	cli, err := oss.New(endpoint, ak, sk)
+	if err != nil {
+		return false, err
+	}
+	bucket, err := cli.Bucket(bucketName)
+	if err != nil {
+		return false, err
+	}
+	return bucket.IsObjectExist(modelName + "/sqlflow_model_desc")
+}
+
+func predictRandomForests(ir *ir.PredictStmt, session *pb.Session) (string, error) {
+	// NOTE(typhoonzero): for PAI random forests predicting, we can not load the TrainStmt
+	// since the model saving is fully done by PAI. We directly use the columns in SELECT
+	// statement for prediction, error will be reported by PAI job if the columns not match.
+	db, err := database.OpenAndConnectDB(session.DbConnStr)
+	if err != nil {
+		return "", err
+	}
+	flds, _, err := getColumnTypes(ir.Select, db)
+	if err != nil {
+		return "", err
+	}
+	// drop result table if exists
+	db.Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s;", ir.ResultTable))
+	filler := &randomForestsPredictFiller{
+		DataSource:      session.DbConnStr,
+		TmpPredictTable: ir.TmpPredictTable,
+		FeatureColumns:  flds,
+		Save:            ir.Using,
+		ResultTable:     ir.ResultTable,
+	}
+	var tpl = template.Must(template.New("RandomForestsPredict").Parse(randomForestsPredictTemplate))
+	var rfCode bytes.Buffer
+	if err := tpl.Execute(&rfCode, filler); err != nil {
+		return "", err
+	}
+	return rfCode.String(), nil
+}
+
 // Predict generates a Python program for train a TensorFlow model.
 func Predict(ir *ir.PredictStmt, session *pb.Session, modelName, cwd string) (string, error) {
+	// FIXME(typhoonzero): if the model not exist on OSS, assume it's a random forest model
+	// should use a general method to fetch the model and see the model type.
+	exists, err := ossFileExists(modelName)
+	if err != nil {
+		return "", err
+	}
+	if !exists {
+		log.Printf("predicting using pai random forests")
+		return predictRandomForests(ir, session)
+	}
 	cc, err := GetClusterConfig(ir.Attributes)
 	if err != nil {
 		return "", err
@@ -253,3 +356,53 @@ func TFLoadAndPredict(ir *ir.PredictStmt, session *pb.Session, modelName string)
 	}
 	return code.String(), nil
 }
+
+func explainRandomForests(ir *ir.ExplainStmt, session *pb.Session) (string, error) {
+	// NOTE(typhoonzero): for PAI random forests predicting, we can not load the TrainStmt
+	// since the model saving is fully done by PAI. We directly use the columns in SELECT
+	// statement for prediction, error will be reported by PAI job if the columns not match.
+	db, err := database.OpenAndConnectDB(session.DbConnStr)
+	if err != nil {
+		return "", err
+	}
+	flds, _, err := getColumnTypes(ir.Select, db)
+	if err != nil {
+		return "", err
+	}
+	// drop result table if exists
+	db.Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s;", ir.Into))
+	labelCol, ok := ir.Attributes["label_column"]
+	if !ok {
+		return "", fmt.Errorf("must specify WITH label_column when using pai random forest to explain models")
+	}
+	featureFileds := []string{}
+	for _, f := range flds {
+		if f != labelCol {
+			featureFileds = append(featureFileds, f)
+		}
+	}
+
+	filler := &randomForestsExplainFiller{
+		DataSource:      session.DbConnStr,
+		TmpExplainTable: ir.TmpExplainTable,
+		FeatureColumns:  featureFileds,
+		LabelColumn:     labelCol.(string),
+		Save:            ir.ModelName,
+		ResultTable:     ir.Into,
+	}
+	var tpl = template.Must(template.New("RandomForestsExplain").Parse(randomForestsExplainTemplate))
+	var rfCode bytes.Buffer
+	if err := tpl.Execute(&rfCode, filler); err != nil {
+		return "", err
+	}
+	return rfCode.String(), nil
+}
+
+// Explain generates a Python program for train a TensorFlow model.
+func Explain(ir *ir.ExplainStmt, session *pb.Session, modelName, cwd string) (string, error) {
+	// NOTE(typhoonzero): only support random forests explain.
+	if ir.Into == "" {
+		return "", fmt.Errorf("explain PAI random forests model need INTO clause to output the explain result to a table")
+	}
+	return explainRandomForests(ir, session)
+}

pkg/sql/codegen/pai/template_random_forests.go

@@ -15,14 +15,30 @@ package pai
 
 type randomForestsTrainFiller struct {
 	DataSource     string
-	Select         string
 	TmpTrainTable  string
 	FeatureColumns []string
 	LabelColumn    string
 	Save           string
 	TreeNum        int
 }
 
+type randomForestsPredictFiller struct {
+	DataSource      string
+	TmpPredictTable string
+	FeatureColumns  []string
+	Save            string
+	ResultTable     string
+}
+
+type randomForestsExplainFiller struct {
+	DataSource      string
+	TmpExplainTable string
+	FeatureColumns  []string
+	LabelColumn     string
+	Save            string
+	ResultTable     string
+}
+
 const randomForestsTrainTemplate = `
 import os
 import subprocess
@@ -48,3 +64,55 @@ subprocess.run(["odpscmd", "-u", user,
                            "-e", pai_cmd],
                check=True)
 `
+
+const randomForestsPredictTemplate = `
+import os
+import subprocess
+import sqlflow_submitter.db
+
+driver, dsn = "{{.DataSource}}".split("://")
+assert driver == "maxcompute"
+user, passwd, address, database = sqlflow_submitter.db.parseMaxComputeDSN(dsn)
+
+column_names = []
+{{ range $colname := .FeatureColumns }}
+column_names.append("{{$colname}}")
+{{ end }}
+pai_cmd = 'pai -name prediction -project algo_public -DmodelName="{{.Save}}" -DinputTableName="{{.TmpPredictTable}}"  -DoutputTableName="{{.ResultTable}}" -DfeatureColNames="%s" ' % (
+    ",".join(column_names)
+)
+
+# Submit the tarball to PAI
+subprocess.run(["odpscmd", "-u", user,
+                           "-p", passwd,
+                           "--project", database,
+                           "--endpoint", address,
+                           "-e", pai_cmd],
+               check=True)
+`
+
+const randomForestsExplainTemplate = `
+import os
+import subprocess
+import sqlflow_submitter.db
+
+driver, dsn = "{{.DataSource}}".split("://")
+assert driver == "maxcompute"
+user, passwd, address, database = sqlflow_submitter.db.parseMaxComputeDSN(dsn)
+
+column_names = []
+{{ range $colname := .FeatureColumns }}
+column_names.append("{{$colname}}")
+{{ end }}
+pai_cmd = 'pai -name feature_importance -project algo_public -DmodelName="{{.Save}}" -DinputTableName="{{.TmpExplainTable}}"  -DoutputTableName="{{.ResultTable}}" -DlabelColName="{{.LabelColumn}}" -DfeatureColNames="%s" ' % (
+    ",".join(column_names)
+)
+
+# Submit the tarball to PAI
+subprocess.run(["odpscmd", "-u", user,
+                           "-p", passwd,
+                           "--project", database,
+                           "--endpoint", address,
+                           "-e", pai_cmd],
+               check=True)
+`

pkg/sql/ir/ir.go

@@ -182,8 +182,13 @@ type ExplainStmt struct {
 	Attributes map[string]interface{}
 	// Explainer types. For example TreeExplainer.
 	Explainer string
+	// ModelName is the model to be explained, e.g. TO EXPLAIN model_name
+	ModelName string
 	// Into stores the model explain result. Note that this field is optional.
 	Into string
+	// When SQLFLOW_submitter == "pai", tmp tables will be created for predicting task
+	// see: pai_submitter.go
+	TmpExplainTable string
 	// TrainStmt is the TrainStmt used for generating the training job of the corresponding model
 	TrainStmt *TrainStmt
 }

pkg/sql/ir_generator.go

@@ -305,6 +305,7 @@ func generateExplainStmt(slct *parser.SQLFlowSelectStmt, connStr, modelDir strin
 		Attributes: attrs,
 		Explainer:  slct.Explainer,
 		TrainStmt:  trainStmt,
+		ModelName:  slct.TrainedModel,
 		Into:       slct.ExplainInto,
 	}