[Programming]Add optimize attribute checker (#2482)

sneaxiy · web-flow · commit c8b4432cb891 · 2020-06-18T11:00:03.000+08:00
* add optimize attr checker

* add optimize attr checker

* remove is_pai

* follow parts of comments
diff --git a/pkg/codegen/optimize/codegen.go b/pkg/codegen/optimize/codegen.go
@@ -15,21 +15,85 @@ package optimize
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
+	"sqlflow.org/sqlflow/pkg/attribute"
 	"sqlflow.org/sqlflow/pkg/ir"
 	pb "sqlflow.org/sqlflow/pkg/proto"
 	"strings"
 	"text/template"
 )
 
+func checkIsPositiveInteger(i interface{}, name string) error {
+	if v, ok := i.(int); !ok || v <= 0 {
+		return fmt.Errorf("%s should be positive integer", name)
+	}
+	return nil
+}
+
+// TODO(sneaxiy): polish attribute codes
+var attributeDictionary = attribute.Dictionary{
+	"data.enable_slice": {attribute.Bool, false, "Whether to enable data slicing", nil},
+	"data.batch_size":   {attribute.Int, -1, "Batch size when training", nil},
+	"worker.num": {attribute.Int, 1, "Worker number", func(i interface{}) error {
+		return checkIsPositiveInteger(i, "worker.num")
+	}},
+	"worker.core": {attribute.Int, 8, "Worker core number", func(i interface{}) error {
+		return checkIsPositiveInteger(i, "worker.core")
+	}},
+	"worker.memory": {attribute.Int, 4096, "Worker memory", func(i interface{}) error {
+		return checkIsPositiveInteger(i, "worker.memory")
+	}},
+	"solver.*": {attribute.Unknown, nil, "Solver options", nil},
+}
+
+// InitializeAttributes initialize attributes in optimize clause IR
+func InitializeAttributes(stmt *ir.OptimizeStmt) error {
+	attributeDictionary.FillDefaults(stmt.Attributes)
+	err := attributeDictionary.Validate(stmt.Attributes)
+	return err
+}
+
 // GenerateOptFlowOptimizeCode generates optimize codes for execution
 // The returned value is (runnerProgramCode, submitProgramCode, error)
-func GenerateOptFlowOptimizeCode(optimStmt *ir.OptimizeStmt, session *pb.Session, dbName, tableName, runnerModuleName string, isPai bool) (string, string, error) {
+func GenerateOptFlowOptimizeCode(optimStmt *ir.OptimizeStmt, session *pb.Session, dbName, tableName, runnerModuleName string) (string, string, error) {
+	const (
+		dataAttrPrefix   = "data."
+		solverAttrPrefix = "solver."
+		workerAttrPrefix = "worker."
+	)
+
 	resultTable := optimStmt.ResultTable
 	if !strings.Contains(resultTable, ".") {
 		resultTable = fmt.Sprintf("%s.%s", dbName, resultTable)
 	}
 
+	attrs := make(map[string]map[string]interface{})
+	for k, v := range optimStmt.Attributes {
+		prefix := ""
+		if strings.HasPrefix(k, dataAttrPrefix) {
+			prefix = dataAttrPrefix
+		} else if strings.HasPrefix(k, solverAttrPrefix) {
+			prefix = solverAttrPrefix
+		} else if strings.HasPrefix(k, workerAttrPrefix) {
+			prefix = workerAttrPrefix
+		} else {
+			return "", "", fmt.Errorf("unrecognized attribute %s", k)
+		}
+
+		k = k[len(prefix):]
+		prefixKey := prefix[0 : len(prefix)-1]
+		if _, ok := attrs[prefixKey]; !ok {
+			attrs[prefixKey] = make(map[string]interface{})
+		}
+		attrs[prefixKey][k] = v
+	}
+
+	attrJSON, err := json.Marshal(attrs)
+	if err != nil {
+		return "", "", err
+	}
+
 	filler := optimizeFiller{
 		UserID:          session.UserId,
 		Variables:       optimStmt.Variables,
@@ -39,9 +103,9 @@ func GenerateOptFlowOptimizeCode(optimStmt *ir.OptimizeStmt, session *pb.Session
 		Direction:       optimStmt.Direction,
 		Constraints:     optimStmt.Constraints,
 		Solver:          optimStmt.Solver,
+		AttributeJSON:   string(attrJSON),
 		TrainTable:      fmt.Sprintf("%s.%s", dbName, tableName),
 		ResultTable:     resultTable,
-		IsPAI:           isPai,
 		RunnerModule:    runnerModuleName,
 	}
 
diff --git a/pkg/codegen/optimize/template_optimize.go b/pkg/codegen/optimize/template_optimize.go
@@ -24,9 +24,9 @@ type optimizeFiller struct {
 	Direction       string
 	Constraints     []*ir.OptimizeExpr
 	Solver          string
+	AttributeJSON   string
 	TrainTable      string
 	ResultTable     string
-	IsPAI           bool
 	RunnerModule    string
 }
 
@@ -56,63 +56,23 @@ class CustomOptFlowRunner(BaseOptFlowRunner):
 `
 
 const optFlowSubmitText = `
-import os
-from optflow.core.api.config import (InputConf, OdpsItemConf, OdpsConf, OutputConf, RunnerConf, SolverConf,
-                                     SolverExperiment, OptflowLocalEngine, OptflowKubemakerEngine, OptionConf)
-from optflow.core.submit import submit_experiment
-
-if "{{.IsPAI}}":
-    from alps.framework.engine import ResourceConf
-
-def submit():
-    options = {}  # solver options
-    solver_conf = SolverConf(name="{{.Solver}}", options=OptionConf(options))
-
-    pai_project = "{{.TrainTable}}".split('.')[0]
-    odps_conf = OdpsConf(project=pai_project,
-                         accessid=os.environ.get("SQLFLOW_TEST_DB_MAXCOMPUTE_AK"),
-                         accesskey=os.environ.get("SQLFLOW_TEST_DB_MAXCOMPUTE_SK"),
-                         partitions=None)
-    
-    runner = RunnerConf(cls="{{.RunnerModule}}.CustomOptFlowRunner")
-    
-    output_table = OdpsItemConf(path="odps://{{.ResultTable}}", odps=odps_conf)
-    output = OutputConf(df1=output_table)
-
-    df1 = OdpsItemConf(path="odps://{{.TrainTable}}",
-                       odps=odps_conf,
-                       enable_slice=False)
-    
-    input_conf = InputConf(df1=df1)
-
-    optflow_version = os.environ.get("SQLFLOW_OPTFLOW_VERSION")
-    if not optflow_version:
-        raise ValueError("Environment variable SQLFLOW_OPTFLOW_VERSION must be set")
-    	
-    if "{{.IsPAI}}":
-        cluster = os.environ.get("SQLFLOW_OPTFLOW_KUBEMAKER_CLUSTER")
-        if not cluster:
-            raise ValueError("Environment variable SQLFLOW_OPTFLOW_KUBEMAKER_CLUSTER must be set")
-    
-        # TODO(sneaxiy): move ResourceConf setting to WITH statements
-        engine = OptflowKubemakerEngine(worker=ResourceConf(core=8, memory=20000, num=1), cluster=cluster)
-    else:
-        # TODO(sneaxiy): support local engine
-        engine = OptflowLocalEngine()
-    
-    user_id = "{{.UserID}}"
-    if not user_id:
-        user_id = "jinle.zjl"
-
-    experiment = SolverExperiment(user=user_id,
-                                  engine=engine,
-                                  runner=runner,
-                                  solver=solver_conf,
-                                  input_conf=input_conf,
-                                  output_conf=output)
-
-    submit_experiment(experiment, optflow_version=optflow_version)
-
-if __name__ == '__main__':
-    submit()
+import json
+from sqlflow_submitter.optimize import submit
+
+runner = "{{.RunnerModule}}.CustomOptFlowRunner"
+solver = "{{.Solver}}"
+attributes = json.loads('''{{.AttributeJSON}}''')
+train_table = "{{.TrainTable}}"
+result_table = "{{.ResultTable}}"
+
+user_id = "{{.UserID}}"
+if not user_id:
+    user_id = "jinle.zjl"
+
+submit(runner=runner, 
+       solver=solver, 
+       attributes=attributes, 
+       train_table=train_table,
+       result_table=result_table,
+       user_id=user_id)
 `
diff --git a/pkg/parser/extended_syntax_parser.y b/pkg/parser/extended_syntax_parser.y
@@ -498,7 +498,9 @@ func (e *Expr) ToTokens() []string {
     case '-':
         switch len(e.Sexp) {
         case 2:
-            return []string{fmt.Sprintf("-%s", e.Sexp[1])}
+            result = append(result, "-")
+            result = append(result, e.Sexp[1].ToTokens()...)
+            return result
         case 3:
             result = append(result, e.Sexp[1].ToTokens()...)
             result = append(result, e.Sexp[0].Value)
diff --git a/pkg/sql/ir_generator.go b/pkg/sql/ir_generator.go
@@ -18,6 +18,7 @@ import (
 	"strconv"
 	"strings"
 
+	"sqlflow.org/sqlflow/pkg/codegen/optimize"
 	"sqlflow.org/sqlflow/pkg/codegen/pai"
 	"sqlflow.org/sqlflow/pkg/codegen/tensorflow"
 	"sqlflow.org/sqlflow/pkg/codegen/xgboost"
@@ -1126,6 +1127,11 @@ func generateOptimizeStmt(optimizeStmt *parser.SQLFlowSelectStmt) (*ir.OptimizeS
 		}
 	}
 
+	solver := optimizeStmt.Solver
+	if solver == "" {
+		solver = "glpk" // find a better way to set default value
+	}
+
 	stmt := &ir.OptimizeStmt{
 		Select:          optimizeStmt.StandardSelect.String(),
 		Variables:       vars,
@@ -1135,8 +1141,14 @@ func generateOptimizeStmt(optimizeStmt *parser.SQLFlowSelectStmt) (*ir.OptimizeS
 		Objective:       objective,
 		Direction:       strings.ToLower(optimizeStmt.Direction),
 		Constraints:     constraints,
-		Solver:          optimizeStmt.Solver,
+		Solver:          solver,
 		ResultTable:     optimizeStmt.OptimizeInto,
 	}
+
+	err = optimize.InitializeAttributes(stmt)
+	if err != nil {
+		return nil, err
+	}
+
 	return stmt, nil
 }
diff --git a/pkg/sql/ir_generator_test.go b/pkg/sql/ir_generator_test.go
@@ -507,8 +507,14 @@ SELECT * FROM alifin_jtest_dev.zjl_shipment_test
 TO MINIMIZE SUM(distance * shipment * 90 / 1000)
 CONSTRAINT SUM(shipment) <= capacity GROUP BY plants,
            SUM(shipment) >= demand GROUP BY markets
-WITH variables="shipment(plants,markets)",
-     var_type="NonNegativeReals"
+WITH variables = "shipment(plants,markets)",
+     var_type = "NonNegativeReals",
+     data.enable_slice = True,
+     data.batch_size = 1,
+     worker.core = 16,
+     worker.num = 4,
+     worker.memory = 8192,
+     solver.max_iter = 10
 USING glpk
 INTO shipment_result_table;
 `
@@ -534,4 +540,11 @@ INTO shipment_result_table;
 	a.Equal("shipment", stmt.ResultValueName)
 	a.Equal("NonNegativeReals", stmt.VariableType)
 	a.Equal("shipment_result_table", stmt.ResultTable)
+
+	a.Equal(true, stmt.Attributes["data.enable_slice"])
+	a.Equal(1, stmt.Attributes["data.batch_size"])
+	a.Equal(16, stmt.Attributes["worker.core"])
+	a.Equal(4, stmt.Attributes["worker.num"])
+	a.Equal(8192, stmt.Attributes["worker.memory"])
+	a.Equal(10, stmt.Attributes["solver.max_iter"])
 }
diff --git a/pkg/sql/submitter.go b/pkg/sql/submitter.go
@@ -284,7 +284,7 @@ func generateOptFlowOptimizeCodeAndExecute(cl *ir.OptimizeStmt, submitter *defau
 	// Generate optimization code
 	runnerFileName := "custom_optimize_runner"
 	runnerCode, submitCode, err := optimize.GenerateOptFlowOptimizeCode(cl, session, dbName, tableName,
-		runnerFileName, isPai)
+		runnerFileName)
 
 	if err != nil {
 		return err
diff --git a/python/sqlflow_submitter/optimize/__init__.py b/python/sqlflow_submitter/optimize/__init__.py
@@ -14,6 +14,6 @@
 from sqlflow_submitter.optimize.optimize import generate_model_with_data_frame
 
 try:
-    from sqlflow_submitter.optimize.runner import BaseOptFlowRunner
+    from sqlflow_submitter.optimize.runner import BaseOptFlowRunner, submit
 except:
     pass
diff --git a/python/sqlflow_submitter/optimize/optimize.py b/python/sqlflow_submitter/optimize/optimize.py
@@ -11,8 +11,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import re
-
 import numpy as np
 import pyomo.environ as pyomo_env
 import six
diff --git a/python/sqlflow_submitter/optimize/runner.py b/python/sqlflow_submitter/optimize/runner.py
@@ -11,11 +11,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import sys
+
 import numpy as np
 import pandas as pd
+from alps.framework.engine import ResourceConf
+from optflow.core.api.config import (InputConf, OdpsConf, OdpsItemConf,
+                                     OptflowKubemakerEngine,
+                                     OptflowLocalEngine, OptionConf,
+                                     OutputConf, RunnerConf, SolverConf,
+                                     SolverExperiment)
+from optflow.core.submit import submit_experiment
 from optflow.workflow.runner.custom_solver_runner import CustomSolverRunner
 from sqlflow_submitter.optimize.optimize import generate_model_with_data_frame
 
+__all__ = [
+    'BaseOptFlowRunner',
+    'submit',
+]
+
 
 class BaseOptFlowRunner(CustomSolverRunner):
     def init_parameters(self):
@@ -101,3 +116,62 @@ def solver_run(self):
         print('Output data is \n', output)
         output_dfs = {'df1': output}
         self.dump_outputs(output_dfs)
+
+
+def submit(runner, solver, attributes, train_table, result_table, user_id):
+    solver_options = attributes.get("solver", {})
+    sys.stderr.write('solver options: {}\n'.format(solver_options))
+    solver_conf = SolverConf(name=solver, options=OptionConf(solver_options))
+
+    pai_project = train_table.split('.')[0]
+    odps_conf = OdpsConf(
+        project=pai_project,
+        accessid=os.environ.get("SQLFLOW_TEST_DB_MAXCOMPUTE_AK"),
+        accesskey=os.environ.get("SQLFLOW_TEST_DB_MAXCOMPUTE_SK"),
+        partitions=None)
+
+    runner = RunnerConf(cls=runner)
+
+    data_options = attributes.get("data", {})
+    sys.stderr.write('data options: {}\n'.format(data_options))
+    enable_slice = data_options.get("enable_slice", False)
+    batch_size = data_options.get("batch_size", None)
+    if batch_size is not None and batch_size <= 0:
+        batch_size = None
+
+    output_table = OdpsItemConf(path="odps://{}".format(result_table),
+                                odps=odps_conf)
+    output_conf = OutputConf(df1=output_table)
+
+    df1 = OdpsItemConf(path="odps://{}".format(train_table),
+                       odps=odps_conf,
+                       enable_slice=enable_slice,
+                       batch_size=batch_size)
+
+    input_conf = InputConf(df1=df1)
+
+    optflow_version = os.environ.get("SQLFLOW_OPTFLOW_VERSION")
+    if not optflow_version:
+        raise ValueError(
+            "Environment variable SQLFLOW_OPTFLOW_VERSION must be set")
+
+    cluster = os.environ.get("SQLFLOW_OPTFLOW_KUBEMAKER_CLUSTER")
+    if not cluster:
+        raise ValueError(
+            "Environment variable SQLFLOW_OPTFLOW_KUBEMAKER_CLUSTER must be set"
+        )
+
+    worker_options = attributes.get("worker", {})
+    sys.stderr.write('worker options: {}\n'.format(worker_options))
+    # TODO(sneaxiy): support local engine
+    engine = OptflowKubemakerEngine(worker=ResourceConf(**worker_options),
+                                    cluster=cluster)
+
+    experiment = SolverExperiment(user=user_id,
+                                  engine=engine,
+                                  runner=runner,
+                                  solver=solver_conf,
+                                  input_conf=input_conf,
+                                  output_conf=output_conf)
+
+    submit_experiment(experiment, optflow_version=optflow_version)
