Save model meta into 'model_meta.json' (#2476)

* merge files

* Design doc for model metadata storage

* refine wording

* add field description to metadata

* add metadata definition

* Save model metadata into model_meta.json

* sort import

Author: lhw362950217

pkg/codegen/tensorflow/codegen.go

@@ -18,10 +18,11 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
-	"sqlflow.org/sqlflow/pkg/codegen"
 	"strings"
 	"text/template"
 
+	"sqlflow.org/sqlflow/pkg/codegen"
+
 	"sqlflow.org/sqlflow/pkg/attribute"
 	"sqlflow.org/sqlflow/pkg/ir"
 	pb "sqlflow.org/sqlflow/pkg/proto"
@@ -333,6 +334,7 @@ func Train(trainStmt *ir.TrainStmt, session *pb.Session) (string, error) {
 		IsPAI:               IsPAI(),
 		PAITrainTable:       paiTrainTable,
 		PAIValidateTable:    paiValidateTable,
+		ModelRepoImage:      trainStmt.ModelImage,
 	}
 	var program bytes.Buffer
 	var trainTemplate = template.Must(template.New("Train").Funcs(template.FuncMap{

pkg/codegen/tensorflow/codegen_test.go

@@ -53,6 +53,16 @@ func TestTrainCodegen(t *testing.T) {
 	a.Equal(r.FindStringSubmatch(code)[1], "sqlflow_pass")
 }
 
+func TestTrainWithModelRepoImage(t *testing.T) {
+	a := assert.New(t)
+	tir := ir.MockTrainStmt(false)
+	tir.ModelImage = "myRepo/MyDNNClassifier:v1.0"
+	code, err := Train(tir, mockSession())
+	a.NoError(err)
+	r, _ := regexp.Compile(`model_repo_image="(.*)"`)
+	a.Equal(r.FindStringSubmatch(code)[1], tir.ModelImage)
+}
+
 func TestTrainWithOptimizer(t *testing.T) {
 	a := assert.New(t)
 	tir := ir.MockTrainStmt(false)

pkg/codegen/tensorflow/template_train.go

@@ -32,6 +32,7 @@ type trainFiller struct {
 	IsPAI               bool
 	PAITrainTable       string
 	PAIValidateTable    string
+	ModelRepoImage      string
 }
 
 const tfTrainTemplateText = `
@@ -138,5 +139,7 @@ train(datasource="{{.DataSource}}",
       load_pretrained_model="{{.LoadPreTrainedModel}}" == "true",
       is_pai="{{.IsPAI}}" == "true",
       pai_table="{{.PAITrainTable}}",
-      pai_val_table="{{.PAIValidateTable}}")
+      pai_val_table="{{.PAIValidateTable}}",
+      feature_columns_code=feature_columns_code,
+      model_repo_image="{{.ModelRepoImage}}")
 `

pkg/codegen/xgboost/codegen.go

@@ -18,9 +18,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"regexp"
-	"sqlflow.org/sqlflow/pkg/codegen"
 	"strings"
 
+	"sqlflow.org/sqlflow/pkg/codegen"
+
 	"sqlflow.org/sqlflow/pkg/attribute"
 	tf "sqlflow.org/sqlflow/pkg/codegen/tensorflow"
 	"sqlflow.org/sqlflow/pkg/ir"
@@ -336,7 +337,9 @@ func newTrainFiller(trainStmt *ir.TrainStmt, session *pb.Session, ossURIToSave,
 		LoadPreTrainedModel: trainStmt.PreTrainedModel != "",
 		IsPAI:               tf.IsPAI(),
 		PAITrainTable:       paiTrainTable,
-		PAIValidateTable:    paiValidateTable}, nil
+		PAIValidateTable:    paiValidateTable,
+		ModelRepoImage:      trainStmt.ModelImage,
+	}, nil
 }
 
 // Pred generates a Python program for predict a xgboost model.

pkg/codegen/xgboost/codegen_test.go

@@ -82,3 +82,21 @@ func TestResolveModelParams(t *testing.T) {
 		a.Equal(objectiveName[i], tir.Attributes["objective"])
 	}
 }
+
+func TestTrainWithModelRepoImage(t *testing.T) {
+	a := assert.New(t)
+	tir := ir.MockTrainStmt(true)
+	a.NoError(InitializeAttributes(tir))
+	tir.ModelImage = "myRepo/MyXGBClassifier:v1.0"
+	code, err := Train(tir, mockSession())
+	a.NoError(err)
+	r, _ := regexp.Compile(`model_repo_image="(.*)"`)
+	a.Equal(r.FindStringSubmatch(code)[1], tir.ModelImage)
+
+	// dist train
+	code, err = DistTrain(tir, mockSession(), 2, "", "")
+	a.NoError(err)
+	r, _ = regexp.Compile(`model_repo_image="(.*)"`)
+	a.Equal(r.FindStringSubmatch(code)[1], tir.ModelImage)
+
+}

pkg/codegen/xgboost/template_train.go

@@ -34,6 +34,7 @@ type trainFiller struct {
 	IsPAI               bool
 	PAITrainTable       string
 	PAIValidateTable    string
+	ModelRepoImage      string
 }
 
 const trainTemplateText = `
@@ -81,7 +82,8 @@ train(datasource='''{{.DataSource}}''',
       pai_validate_table="{{.PAIValidateTable}}",
       oss_model_dir="{{.OSSModelDirToSave}}",
       transform_fn=transform_fn,
-      feature_column_code='''{{.FeatureColumnCode}}''')
+      feature_column_code='''{{.FeatureColumnCode}}''',
+      model_repo_image="{{.ModelRepoImage}}")
 `
 
 const distTrainTemplateText = `
@@ -128,7 +130,8 @@ dist_train(flags=FLAGS,
       pai_validate_table="{{.PAIValidateTable}}",
       oss_model_dir="{{.OSSModelDirToSave}}",
       transform_fn=transform_fn,
-      feature_column_code='''{{.FeatureColumnCode}}''')
+      feature_column_code='''{{.FeatureColumnCode}}''',
+      model_repo_image="{{.ModelRepoImage}}")
 `
 
 var trainTemplate = template.Must(template.New("Train").Parse(trainTemplateText))

python/sqlflow_submitter/model_metadata.py

@@ -0,0 +1,43 @@
+# 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
+
+import copy
+import json
+
+
+def collect_model_metadata(select, validate_select, estimator, attributes,
+                           feature_columns, field_descs, label, evaluation,
+                           model_repo_image):
+    """ collect kinds of model metadata and put them in a dict """
+    metadata = dict(locals())
+    attr_copy = copy.deepcopy(attributes)
+    for (k, v) in attr_copy.items():
+        try:
+            json.dumps(v)
+        except:
+            attr_copy[k] = str(v)
+    metadata['attributes'] = attr_copy
+    return metadata
+
+
+def save_model_metadata(path, metadata):
+    """save_model_metdata saves given params into 'path'"""
+    with open(path, mode="w") as meta_file:
+        meta_file.write(json.dumps(metadata, indent=2))
+
+
+def load_model_metadata(path):
+    """load_model_metadata load metadata from given 'path'"""
+    with open(path, mode="r") as meta_file:
+        lines = meta_file.readlines()
+        return json.loads(lines)

python/sqlflow_submitter/tensorflow/train.py

@@ -31,6 +31,7 @@
                                   DNNLinearCombinedRegressor, DNNRegressor,
                                   LinearClassifier, LinearRegressor)
 
+from ..model_metadata import collect_model_metadata
 from .get_tf_version import tf_is_version2
 from .input_fn import get_dataset_fn
 from .pai_distributed import define_tf_flags, set_oss_environs
@@ -70,7 +71,13 @@ def train(datasource,
           load_pretrained_model=False,
           is_pai=False,
           pai_table="",
-          pai_val_table=""):
+          pai_val_table="",
+          feature_columns_code="",
+          model_repo_image=""):
+    model_meta = collect_model_metadata(select, validation_select,
+                                        estimator_string, model_params,
+                                        feature_columns_code, feature_metas,
+                                        label_meta, None, model_repo_image)
     # import custom model package
     sqlflow_submitter.import_model_def(estimator_string, globals())
     estimator = eval(estimator_string)
@@ -123,13 +130,16 @@ def train(datasource,
         keras_train_and_save(estimator, model_params, save, is_pai, FLAGS,
                              train_dataset_fn, val_dataset_fn, label_meta,
                              epoch, verbose, validation_metrics,
-                             validation_steps, load_pretrained_model)
+                             validation_steps, load_pretrained_model,
+                             model_meta)
     else:
-        estimator_train_and_save(
-            estimator, model_params, save, is_pai, FLAGS, train_dataset_fn,
-            val_dataset_fn, log_every_n_iter, max_steps,
-            validation_start_delay_secs, validation_throttle_secs,
-            save_checkpoints_steps, validation_metrics, load_pretrained_model)
+        estimator_train_and_save(estimator, model_params, save, is_pai, FLAGS,
+                                 train_dataset_fn, val_dataset_fn,
+                                 log_every_n_iter, max_steps,
+                                 validation_start_delay_secs,
+                                 validation_throttle_secs,
+                                 save_checkpoints_steps, validation_metrics,
+                                 load_pretrained_model, model_meta)
 
     # remove cache files
     any(map(os.remove, glob.glob('cache_train.*')))

python/sqlflow_submitter/tensorflow/train_estimator.py

@@ -11,8 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from os import path
+
 import tensorflow as tf
 
+from ..model_metadata import save_model_metadata
 from . import metrics
 from .diag import check_and_load_estimator
 from .get_tf_version import tf_is_version2
@@ -25,7 +28,7 @@ def estimator_train_and_save(estimator, model_params, save, is_pai, FLAGS,
                              log_every_n_iter, train_max_steps,
                              eval_start_delay_secs, eval_throttle_secs,
                              save_checkpoints_steps, metric_names,
-                             load_pretrained_model):
+                             load_pretrained_model, model_meta):
     print("Start training using estimator model...")
 
     is_distributed = False
@@ -73,9 +76,13 @@ def estimator_train_and_save(estimator, model_params, save, is_pai, FLAGS,
         tf.feature_column.make_parse_example_spec(all_feature_columns))
     export_path = classifier.export_saved_model(save, serving_input_fn)
     # write the path under current directory
+    export_path_str = str(export_path.decode("utf-8"))
     with open("exported_path", "w") as fn:
-        fn.write(str(export_path.decode("utf-8")))
-    print("Done training, model exported to: %s" % export_path)
+        fn.write(export_path_str)
+    # write model metadata to model_meta.json
+    save_model_metadata(path.join(export_path_str, "model_meta.json"),
+                        model_meta)
+    print("Done training, model exported to: %s" % export_path_str)
 
 
 def estimator_train_compiled(estimator, is_pai, FLAGS, train_dataset_fn,

python/sqlflow_submitter/tensorflow/train_keras.py

@@ -14,12 +14,14 @@
 import inspect
 import sys
 import warnings
+from os import path
 
 import six
 import tensorflow as tf
 from sqlflow_submitter.pai import model
 from sqlflow_submitter.seeding import get_tf_random_seed
 
+from ..model_metadata import save_model_metadata
 from . import metrics
 from .diag import check_and_load_estimator
 from .get_tf_version import tf_is_version2
@@ -33,7 +35,7 @@
 def keras_train_and_save(estimator, model_params, save, is_pai, FLAGS,
                          train_dataset_fn, val_dataset_fn, label_meta, epochs,
                          verbose, metric_names, validation_steps,
-                         load_pretrained_model):
+                         load_pretrained_model, model_meta):
     print("Start training using keras model...")
     # remove optimizer param from model_params and use it when call "compile()"
     optimizer = None
@@ -171,10 +173,15 @@ def keras_train_and_save(estimator, model_params, save, is_pai, FLAGS,
             tf.feature_column.make_parse_example_spec(all_feature_columns))
         export_path = keras_estimator.export_saved_model(
             save, serving_input_fn)
+
         # write the path under current directory
+        export_path_str = str(export_path.decode("utf-8"))
         with open("exported_path", "w") as fn:
-            fn.write(str(export_path.decode("utf-8")))
-        print("Done training, model exported to: %s" % export_path)
+            fn.write(export_path_str)
+        # write model metadata to model_meta.json
+        save_model_metadata(path.join(export_path_str, "model_meta.json"),
+                            model_meta)
+        print("Done training, model exported to: %s" % export_path_str)
         return
 
     if hasattr(classifier, 'sqlflow_train_loop'):
@@ -200,25 +207,29 @@ def keras_train_and_save(estimator, model_params, save, is_pai, FLAGS,
                                      epochs=epochs if epochs else
                                      classifier.default_training_epochs(),
                                      verbose=verbose)
-        train_keys = []
-        val_keys = []
+        train_metrics = dict()
+        val_metrics = dict()
         for k in history.history.keys():
             if k.startswith("val_"):
-                val_keys.append(k)
+                val_metrics[k] = float(history.history[k][-1])
             else:
-                train_keys.append(k)
+                train_metrics[k] = float(history.history[k][-1])
         print("====== Result for training set: ======")
-        for k in train_keys:
-            print("%s: %s" % (k, history.history[k][-1]))
+        for k, v in train_metrics.items():
+            print("%s: %s" % (k, v))
         print("====== Result for validation set: ======")
-        for k in val_keys:
-            print("%s: %s" % (k, history.history[k][-1]))
+        for k, v in val_metrics.items():
+            print("%s: %s" % (k, v))
+        model_meta["evaluation"] = val_metrics
 
     try:
         classifier.save_weights(save, save_format="h5")
+        # write model metadata to model_meta.json
+        save_model_metadata("model_meta.json", model_meta)
         if is_pai:
             print("saving keras model to: %s" % FLAGS.sqlflow_oss_modeldir)
             model.save_file(FLAGS.sqlflow_oss_modeldir, save)
+            model.save_file(FLAGS.sqlflow_oss_modeldir, "model_meta.json")
     except:
         if has_none_optimizer:
             warnings.warn("Saving model with None optimizer fails")