SQLFlow Code Generator

SQLFlow is a compiler that compiles a SQL program to an Argo workflow as the following pipeline:

  1. parser -> AST -> sematic -> IR -> optimizer -> code generator
  2.    ↑                                                ↓
  3. sql program                                       .YAML

The Argo controller running on Kubernetes is the executor that executes the workflow. This is a design doc about how to implement the code generator.

The High-level Design of the Code Generator

As mentioned above, SQLFlow compiler generates the .YAML file as the following, you can check more detail about SQLFlow workflow from here.

  1. steps:
  2.     name: step-1
  3.     command: ["python", "-c"]
  4.     args: |
  5.         from runtime import tensorflow
  6.         tensorflow.train(....)
  7.     env:
  8.       name: SQLFLOW_OSS_AK
  9.       value: "xxxxxx"

From the above workflow .YAML file, each workflow step contains three parts:

  1. The execution command as the command spec to execute the program.
  2. The execution program, which can be written in Python, R, or Bash. The program submits an AI task on an AI platform .e.g, ElasticDL, Alibaba PAI or just runs on a host by involving the SQLFlow runtime library.
  3. The runtime environment variables with the env spec.

SQLFlow compiler provides the code generator component to generate the step program, the code generation is divided into the following stages:

  1. Target Submitter Registry, register a Code Generator in SQLFlow compiler.
  2. CodeGenerator Interface is a Go interface that all code generators should implement.
  3. Code Generation provides an assembler API to generate a step program.

Target Submitter Register

For a new code generator, develops should register it in SQLFlow compiler as the following pseudo-code:

  2. cgMapping = map[string]CodeGenerator {
  3.   "paiTensorFlow": PAITensorFlow{},
  4.   "paiXGBoost", PAIXGBoost{},
  5.   ...
  6. }

Code Generator Interface

For each code generator implementation, you should care about all IR types, different IR types have different behaviors and generate different submitter program. Each code generator owns an ExecutionCtx instance to tell Argo workflow on how to execute the target code.

  1. type ExecutionCtx struct {
  2.   ExecCommand []string      // How to execute the target code, .e.g ["python" "-c"]
  3.   Env map[string]string     // The environment variables for execution
  4. }
  6. type CodeGenerator interface {
  7.   GenerateExecCtx(*ir.SQLStmt) ExecutionCtx
  8.   EmitNormal(*ir.NormalStmt) (string, error)
  9.   EmitTrain(*ir.TrainStmt) (string, error)
  10.   EmitPredict(*ir.PredictStmt) (string, error)
  11.   EmitExplain(*ir.ExplainStmt) (string, error)
  12.   EmitEvaluate(*ir.EvaluateStmt) (string, error)
  13.   EmitShowTrain(*ir.ShowTrainStmt) (string, error)
  14.   EmitOptimize(*ir.OptimizeStmt) (string, error)
  15.   EmitRun(*ir.RunStmt) (string, error)
  16. }

Code Generation

The code generation phase is responsible for generating target code from a SQL statement IR, this is an assembler API that routes to a specified code generator, the pseudo-code is as the following:

  1. func Generate(session *pb.Session, stmt *ir.SQLStatement) (string, error) {
  2.   // routing to a specified code generator from session.submitter
  3.   cf := cgMapping[session.submitter]
  5.   switch v := stmt.(type) {
  6.   case *ir.TrainStmt:
  7.     return cg.EmitTrain(stmt.(*ir.TrainStmt)), cg.GenerateExecCtx(), nil
  8.   case *ir.PredictStmt:
  9.     return cg.EmitPredict(stmt.(*ir.TrainStmt)), cg.GenerateExecCtx(), nil
  10.   ...
  11.   }
  12. }