Generate couler code of workflow steps (#2806)

* wip * fix yaml generate * fix tests * fix package deps * fix pip package deps * update

Author: typhoonzero

docker/ci/install-pips.bash

@@ -23,6 +23,7 @@ set -e
 # the saved PMML file is right.
 pip install --quiet \
  numpy==1.16.2 \
+ tensorflow-metadata==0.22.2 \
  tensorflow==2.0.1 \
  impyla==0.16.0 \
  pyodps==0.8.3 \

go/codegen/experimental/codegen.go

@@ -27,18 +27,32 @@ import (
 // 2. generate runtime code of each statement
 // 3. generate couler program to form a workflow
 func GenerateCodeCouler(sqlProgram string, session *pb.Session) (string, error) {
+var defaultDockerImage = "sqlflow/sqlflow:step"
 stmts, err := parseToIR(sqlProgram, session)
 if err != nil {
 return "", err
 }
-for _, stmt := range stmts {
-stepCode, err := generateStepCode(stmt, session)
+stepList := []*stepContext{}
+for idx, stmt := range stmts {
+stepCode, err := generateStepCode(stmt, idx, session)
 if err != nil {
 return "", err
 }
-fmt.Println(stepCode)
+image := defaultDockerImage
+if trainStmt, ok := stmt.(*ir.TrainStmt); ok {
+if trainStmt.ModelImage != "" {
+image = trainStmt.ModelImage
+}
+}
+// TODO(typhoonzero): find out the image that should be used by the predict statements.
+step := &stepContext{
+Code: stepCode,
+Image: image,
+StepIndex: idx,
+}
+stepList = append(stepList, step)
 }
-return "", nil
+return CodeGenCouler(stepList, session)
 }
 
 func parseToIR(sqlProgram string, session *pb.Session) ([]ir.SQLFlowStmt, error) {
@@ -95,12 +109,12 @@ func parseToIR(sqlProgram string, session *pb.Session) ([]ir.SQLFlowStmt, error)
 return result, nil
 }
 
-func generateStepCode(stmt ir.SQLFlowStmt, session *pb.Session) (string, error) {
+func generateStepCode(stmt ir.SQLFlowStmt, stepIndex int, session *pb.Session) (string, error) {
 switch stmt.(type) {
 case *ir.TrainStmt:
 trainStmt := stmt.(*ir.TrainStmt)
 if strings.HasPrefix(strings.ToUpper(trainStmt.Estimator), "XGBOOST.") {
-return XGBoostGenerateTrain(trainStmt, session)
+return XGBoostGenerateTrain(trainStmt, stepIndex, session)
 }
 return "", fmt.Errorf("not implemented estimator type %s", trainStmt.Estimator)
 default:

go/codegen/experimental/codegen_couler.go

@@ -0,0 +1,112 @@
+// Copyright 2020 The SQLFlow Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package experimental
+
+import (
+"bytes"
+"fmt"
+"os"
+"strconv"
+"text/template"
+
+pb "sqlflow.org/sqlflow/go/proto"
+"sqlflow.org/sqlflow/go/workflow/couler"
+)
+
+var workflowTTL = 24 * 3600
+
+type stepContext struct {
+Code string
+StepIndex int
+Image string
+}
+
+type coulerFiller struct {
+StepList []*stepContext
+DataSource string
+StepEnvs map[string]string
+WorkflowTTL int
+SecretName string
+SecretData string
+Resources string
+}
+
+// CodeGenCouler generate couler code to generate a workflow
+func CodeGenCouler(stepList []*stepContext, session *pb.Session) (string, error) {
+var workflowResourcesEnv = "SQLFLOW_WORKFLOW_RESOURCES"
+envs, err := couler.GetStepEnvs(session)
+if err != nil {
+return "", err
+}
+secretName, secretData, err := couler.GetSecret()
+if err != nil {
+return "", err
+}
+if err := couler.VerifyResources(os.Getenv(workflowResourcesEnv)); err != nil {
+return "", err
+}
+if os.Getenv("SQLFLOW_WORKFLOW_TTL") != "" {
+workflowTTL, err = strconv.Atoi(os.Getenv("SQLFLOW_WORKFLOW_TTL"))
+if err != nil {
+return "", fmt.Errorf("SQLFLOW_WORKFLOW_TTL: %s should be int", os.Getenv("SQLFLOW_WORKFLOW_TTL"))
+}
+}
+
+filler := &coulerFiller{
+StepList: stepList,
+DataSource: session.DbConnStr,
+StepEnvs: envs,
+WorkflowTTL: workflowTTL,
+SecretName: secretName,
+SecretData: secretData,
+Resources: os.Getenv(workflowResourcesEnv),
+}
+var program bytes.Buffer
+if err := coulerTemplate.Execute(&program, filler); err != nil {
+return "", err
+}
+return program.String(), nil
+}
+
+var coulerCodeTmpl = `
+import couler.argo as couler
+import json
+import re
+
+datasource = "{{ .DataSource }}"
+
+step_envs = dict()
+{{range $k, $v := .StepEnvs}}step_envs["{{$k}}"] = '''{{$v}}'''
+{{end}}
+
+sqlflow_secret = None
+if "{{.SecretName}}" != "":
+# note(yancey1989): set dry_run to true, just reference the secret meta to generate workflow YAML,
+# we should create the secret before launching sqlflowserver
+secret_data=json.loads('''{{.SecretData}}''')
+sqlflow_secret = couler.secret(secret_data, name="{{ .SecretName }}", dry_run=True)
+
+resources = None
+if '''{{.Resources}}''' != "":
+ resources=json.loads('''{{.Resources}}''')
+
+couler.clean_workflow_after_seconds_finished({{.WorkflowTTL}})
+
+{{ range $ss := .StepList }}
+{{.Code}}
+couler.run_script(image="{{.Image}}", source=step_entry_{{.StepIndex}}, env=step_envs, resources=resources)
+{{end}}
+`
+
+var coulerTemplate = template.Must(template.New("Couler").Parse(coulerCodeTmpl))

go/codegen/experimental/codegen_test.go

@@ -15,28 +15,34 @@ package experimental
 
 import (
 "os"
+"strings"
 "testing"
 
+"github.com/stretchr/testify/assert"
 "sqlflow.org/sqlflow/go/database"
 pb "sqlflow.org/sqlflow/go/proto"
 )
 
 func TestExperimentalXGBCodegen(t *testing.T) {
+a := assert.New(t)
 if os.Getenv("SQLFLOW_TEST_DB") != "mysql" {
 t.Skipf("skip TestExperimentalXGBCodegen of DB type %s", os.Getenv("SQLFLOW_TEST_DB"))
 }
 // test without COLUMN clause
 sql := "SELECT * FROM iris.train TO TRAIN xgboost.gbtree WITH objective=\"binary:logistic\",num_class=3 LABEL class INTO sqlflow_models.xgb_classification;"
 s := &pb.Session{DbConnStr: database.GetTestingMySQLURL()}
-_, err := GenerateCodeCouler(sql, s)
+coulerCode, err := GenerateCodeCouler(sql, s)
 if err != nil {
 t.Errorf("error %s", err)
 }
+a.True(strings.Contains(coulerCode, `couler.run_script(image="sqlflow/sqlflow:step", source=step_entry_0, env=step_envs, resources=resources)`))
 
 // test with COLUMN clause
 sql = "SELECT * FROM iris.train TO TRAIN xgboost.gbtree WITH objective=\"binary:logistic\",num_class=3 COLUMN petal_length LABEL class INTO sqlflow_models.xgb_classification;"
-_, err = GenerateCodeCouler(sql, s)
+coulerCode, err = GenerateCodeCouler(sql, s)
 if err != nil {
 t.Errorf("error %s", err)
 }
+expected := `feature_column_map = {"featuren_columns": [fc.NumericColumn(fd.FieldDesc(name="petal_length", dtype=fd.DataType.FLOAT32, delimiter="", format="", shape=[1], is_sparse=False, vocabulary=[]))]}`
+a.True(strings.Contains(coulerCode, expected))
 }

go/codegen/experimental/xgboost.go

@@ -28,6 +28,7 @@ import (
 )
 
 type xgbTrainFiller struct {
+StepIndex int
 DataSource string
 Select string
 ValidationSelect string
@@ -41,8 +42,8 @@ type xgbTrainFiller struct {
 Submitter string
 }
 
-// XGBoostGenerateTrain returns the step code
-func XGBoostGenerateTrain(trainStmt *ir.TrainStmt, session *pb.Session) (string, error) {
+// XGBoostGenerateTrain returns the step code.
+func XGBoostGenerateTrain(trainStmt *ir.TrainStmt, stepIndex int, session *pb.Session) (string, error) {
 var err error
 if err = resolveModelParams(trainStmt); err != nil {
 return "", err
@@ -93,6 +94,7 @@ func XGBoostGenerateTrain(trainStmt *ir.TrainStmt, session *pb.Session) (string,
 }
 
 filler := xgbTrainFiller{
+StepIndex: stepIndex,
 DataSource: session.DbConnStr,
 Select: trainStmt.Select,
 ValidationSelect: trainStmt.ValidationSelect,
@@ -115,7 +117,7 @@ func XGBoostGenerateTrain(trainStmt *ir.TrainStmt, session *pb.Session) (string,
 }
 
 var xgbTrainTemplate = `
-def step_entry():
+def step_entry_{{.StepIndex}}():
  import json
  import tempfile
  import os

go/workflow/couler/codegen.go

@@ -51,7 +51,8 @@ func fillEnvFromSession(envs *map[string]string, session *pb.Session) {
 fillMapIfValueNotEmpty(*envs, "SQLFLOW_submitter", session.Submitter)
 }
 
-func getStepEnvs(session *pb.Session) (map[string]string, error) {
+// GetStepEnvs returns a map of envs used for couler workflow.
+func GetStepEnvs(session *pb.Session) (map[string]string, error) {
 envs := make(map[string]string)
 // fill step envs from the environment variables on sqlflowserver
 for _, env := range os.Environ() {
@@ -67,7 +68,9 @@ func getStepEnvs(session *pb.Session) (map[string]string, error) {
 fillEnvFromSession(&envs, session)
 return envs, nil
 }
-func verifyResources(resources string) error {
+
+// VerifyResources verifies the SQLFLOW_WORKFLOW_RESOURCES env to be valid.
+func VerifyResources(resources string) error {
 if resources != "" {
 var r map[string]interface{}
 if e := json.Unmarshal([]byte(resources), &r); e != nil {
@@ -77,7 +80,8 @@ func verifyResources(resources string) error {
 return nil
 }
 
-func getSecret() (string, string, error) {
+// GetSecret returns the workflow secret name, value.
+func GetSecret() (string, string, error) {
 secretMap := make(map[string]map[string]string)
 secretCfg := os.Getenv("SQLFLOW_WORKFLOW_SECRET")
 if secretCfg == "" {
@@ -99,7 +103,7 @@ func getSecret() (string, string, error) {
 
 // GenFiller generates Filler to fill the template
 func GenFiller(programIR []ir.SQLFlowStmt, session *pb.Session) (*Filler, error) {
-stepEnvs, err := getStepEnvs(session)
+stepEnvs, err := GetStepEnvs(session)
 if err != nil {
 return nil, err
 }
@@ -109,11 +113,11 @@ func GenFiller(programIR []ir.SQLFlowStmt, session *pb.Session) (*Filler, error)
 return nil, fmt.Errorf("SQLFLOW_WORKFLOW_TTL: %s should be int", os.Getenv("SQLFLOW_WORKFLOW_TTL"))
 }
 }
-secretName, secretData, e := getSecret()
+secretName, secretData, e := GetSecret()
 if e != nil {
 return nil, e
 }
-if e := verifyResources(os.Getenv(envResource)); e != nil {
+if e := VerifyResources(os.Getenv(envResource)); e != nil {
 return nil, e
 }
 

go/workflow/couler/codegen_test.go

@@ -142,7 +142,8 @@ func TestStringInStringSQL(t *testing.T) {
 yaml, e := cg.GenYAML(code)
 a.NoError(e)
 println(yaml)
-expect := `validation.select=\"select * from iris.train where name like \\\"Versicolor\\\";\"`
+expect := `validation.select=\\\"select * from iris.train where\
+ \ name like \\\\\\\"Versicolor\\\\\\\";\\\"`
 a.True(strings.Contains(yaml, expect))
 }
 

python/couler/couler/argo.py

@@ -20,10 +20,10 @@
 import couler.pyfunc as pyfunc
 import pyaml
 
-_wf: dict = {}
-_secrets: dict = {}
-_steps: OrderedDict = OrderedDict()
-_templates: dict = {}
+_wf = dict()
+_secrets = dict()
+_steps = OrderedDict()
+_templates = dict()
 _update_steps_lock = True
 _run_concurrent_lock = False
 _concurrent_func_line = -1
@@ -36,7 +36,7 @@
 # '_condition_id' records the line number where the 'couler.when()' is invoked.
 _condition_id = None
 # '_while_steps' records the step of recursive logic
-_while_steps: OrderedDict = OrderedDict()
+_while_steps = OrderedDict()
 # '_while_lock' indicates the recursive call starts
 _while_lock = False
 # TTL_cleaned for the workflow
@@ -506,7 +506,7 @@ def concurrent(function_list):
  _run_concurrent_lock = False
 
 
-def yaml():
+def __dump_yaml():
  wf = copy.deepcopy(_wf)
  wf["apiVersion"] = "argoproj.io/v1alpha1"
  wf["kind"] = "Workflow"
@@ -529,10 +529,10 @@ def yaml():
 def _dump_yaml():
  yaml_str = ""
  if len(_secrets) > 0:
- yaml_str = pyaml.dump(_secrets, string_val_style="plain")
+ yaml_str = pyaml.dump(_secrets)
  yaml_str = "%s\n---\n" % yaml_str
  if len(_steps) > 0:
- yaml_str = yaml_str + pyaml.dump(yaml(), string_val_style="plain")
+ yaml_str = yaml_str + pyaml.dump(__dump_yaml())
  print(yaml_str)
 
 

scripts/test/prepare.sh

@@ -25,6 +25,7 @@ source build/env/bin/activate
 
 python -m pip install --quiet \
  numpy==1.16.2 \
+ tensorflow-metadata==0.22.2 \
  tensorflow==2.0.1 \
  impyla==0.16.0 \
  pyodps==0.8.3 \