Make pai_local executor can run using experimental codegen (#3034)

* make pai_local run successfully

* fix circular reference

Author: sneaxiy

python/runtime/model/model.py

@@ -99,7 +99,7 @@ def _from_dict(d):
         typ = d.pop("model_type")
         return Model(typ, d)
 
-    def _zip(self, local_dir, tarball, save_to_db=False):
+    def _zip(self, local_dir, tarball):
         """
         Zip the model information and all files in local_dir into a tarball.
 
@@ -110,20 +110,23 @@ def _zip(self, local_dir, tarball, save_to_db=False):
         Returns:
             None.
         """
-        if not save_to_db:
-            model_obj_file = os.path.join(local_dir, MODEL_OBJ_FILE_NAME)
-            with open(model_obj_file, "w") as f:
-                d = self._to_dict()
-                f.write(json.dumps(d, cls=JSONEncoderWithFeatureColumn))
-        else:
-            model_obj_file = None
+        # NOTE: the unzip files of the job tarball should be skipped
+        from runtime.pai.prepare_archive import ALL_TAR_FILES
+
+        def filter(tarinfo):
+            name = tarinfo.name
+            if name.startswith("./"):
+                name = name[2:]
+
+            if name in ALL_TAR_FILES:
+                return None
+
+            return tarinfo
 
-        zip_dir(local_dir, tarball, arcname="./")
-        if model_obj_file:
-            os.remove(model_obj_file)
+        zip_dir(local_dir, tarball, arcname="./", filter=filter)
 
     @staticmethod
-    def _unzip(local_dir, tarball, load_from_db=False):
+    def _unzip(local_dir, tarball):
         """
         Unzip the tarball into local_dir and deserialize the model
         information.
@@ -137,13 +140,6 @@ def _unzip(local_dir, tarball, load_from_db=False):
             information.
         """
         unzip_dir(tarball, local_dir)
-        if not load_from_db:
-            model_obj_file = os.path.join(local_dir, MODEL_OBJ_FILE_NAME)
-            with open(model_obj_file, "r") as f:
-                d = json.loads(f.read(), cls=JSONDecoderWithFeatureColumn)
-                model = Model._from_dict(d)
-            os.remove(model_obj_file)
-            return model
 
     def save_to_db(self, datasource, table, local_dir=None):
         """
@@ -164,7 +160,7 @@ def save_to_db(self, datasource, table, local_dir=None):
 
         with temp_file.TemporaryDirectory() as tmp_dir:
             tarball = os.path.join(tmp_dir, TARBALL_NAME)
-            self._zip(local_dir, tarball, save_to_db=True)
+            self._zip(local_dir, tarball)
 
             def _bytes_reader(filename, buf_size=8 * 32):
                 def _gen():
@@ -212,7 +208,7 @@ def load_from_db(datasource, table, local_dir=None):
                 for data in gen():
                     f.write(bytes(data))
 
-            Model._unzip(local_dir, tarball, load_from_db=True)
+            Model._unzip(local_dir, tarball)
 
         return Model._from_dict(metadata)
 
@@ -237,6 +233,14 @@ def save_to_oss(self, oss_model_dir, local_dir=None):
             self._zip(local_dir, tarball)
             oss.save_file(oss_model_dir, tarball, TARBALL_NAME)
 
+        with temp_file.TemporaryDirectory() as tmp_dir:
+            model_obj_file = os.path.join(tmp_dir, MODEL_OBJ_FILE_NAME)
+            with open(model_obj_file, "w") as f:
+                f.write(
+                    json.dumps(self._to_dict(),
+                               cls=JSONEncoderWithFeatureColumn))
+            oss.save_file(oss_model_dir, model_obj_file, MODEL_OBJ_FILE_NAME)
+
     @staticmethod
     def load_from_oss(oss_model_dir, local_dir=None):
         """
@@ -257,7 +261,14 @@ def load_from_oss(oss_model_dir, local_dir=None):
         with temp_file.TemporaryDirectory() as tmp_dir:
             tarball = os.path.join(tmp_dir, TARBALL_NAME)
             oss.load_file(oss_model_dir, tarball, TARBALL_NAME)
-            return Model._unzip(local_dir, tarball)
+            Model._unzip(local_dir, tarball)
+
+            model_obj_file = os.path.join(tmp_dir, MODEL_OBJ_FILE_NAME)
+            oss.load_file(oss_model_dir, model_obj_file, MODEL_OBJ_FILE_NAME)
+            with open(model_obj_file, "r") as f:
+                d = json.loads(f.read(), cls=JSONDecoderWithFeatureColumn)
+                model = Model._from_dict(d)
+            return model
 
 
 def _decompose_model_name(name):

python/runtime/model/tar.py

@@ -17,7 +17,7 @@
 import tarfile
 
 
-def zip_dir(src_dir, tarball, arcname=None):
+def zip_dir(src_dir, tarball, arcname=None, filter=None):
     """To compress a directory into tarball.
 
     Args:
@@ -31,7 +31,7 @@ def zip_dir(src_dir, tarball, arcname=None):
             The output name of src_dir in the tarball.
     """
     with tarfile.open(tarball, "w:gz") as tar:
-        tar.add(src_dir, arcname=arcname, recursive=True)
+        tar.add(src_dir, arcname=arcname, recursive=True, filter=filter)
 
 
 def unzip_dir(tarball, dest_dir=None):

python/runtime/pai/entry.py

@@ -72,9 +72,7 @@ def call_fun(func, params):
     return func(**dict_args)
 
 
-def entrypoint():
-    with open("train_params.pkl", "rb") as file:
-        params = pickle.load(file)
+def entrypoint(params):
     if params["entry_type"] == "train_tf":
         call_fun(train_tf, params)
     elif params["entry_type"] == "train_xgb":
@@ -96,4 +94,6 @@ def entrypoint():
 if __name__ == "__main__":
     FLAGS = define_tf_flags()
     set_oss_environs(FLAGS)
-    entrypoint()
+    with open("train_params.pkl", "rb") as file:
+        params = pickle.load(file)
+    entrypoint(params)

python/runtime/pai/prepare_archive.py

@@ -41,6 +41,15 @@
 sklearn_pandas==1.6.0
 """
 
+ALL_TAR_FILES = [
+    JOB_ARCHIVE_FILE,
+    ENTRY_FILE,
+    "runtime",
+    "sqlflow_models",
+    "requirements.txt",
+    TRAIN_PARAMS_FILE,
+]
+
 
 def prepare_archive(cwd, estimator, model_save_path, train_params):
     """package needed resource into a tarball"""
@@ -60,10 +69,7 @@ def prepare_archive(cwd, estimator, model_save_path, train_params):
     _copy_python_package("sqlflow_models", cwd)
     _copy_custom_package(estimator, cwd)
 
-    args = [
-        "tar", "czf", JOB_ARCHIVE_FILE, ENTRY_FILE, "runtime",
-        "sqlflow_models", "requirements.txt", TRAIN_PARAMS_FILE
-    ]
+    args = ["tar", "czf"] + ALL_TAR_FILES
     if subprocess.call(args, cwd=cwd) != 0:
         raise SQLFlowDiagnostic("Can't zip resource")
 

python/runtime/pai/submitter_train.py

@@ -21,6 +21,7 @@
 from runtime.pai.pai_ml.random_forest import get_train_random_forest_pai_cmd
 from runtime.pai.prepare_archive import prepare_archive
 from runtime.pai.submit_pai_task import submit_pai_task
+from runtime.pai_local.try_run import try_pai_local_run
 
 
 def get_pai_train_cmd(datasource, estimator_string, model_name, train_table,
@@ -137,6 +138,9 @@ def submit_pai_train(datasource,
             pai_model.clean_oss_model_path(oss_path_to_save + "/")
         train_params["oss_path_to_load"] = oss_path_to_load
 
+        if try_pai_local_run(params, oss_path_to_save):
+            return
+
         with temp_file.TemporaryDirectory(prefix="sqlflow", dir="/tmp") as cwd:
             # zip all required resource to a tarball
             prepare_archive(cwd, estimator_string, oss_path_to_save, params)

python/runtime/pai_local/__init__.py

@@ -0,0 +1,30 @@
+# 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 os
+
+
+def _gen_pai_local_method(name):
+    def impl(*args, **kwargs):
+        import runtime.pai as pai
+        method = getattr(pai, name)
+        os.environ["SQLFLOW_submitter"] = "pai_local"
+        return method(*args, **kwargs)
+
+    return impl
+
+
+train = _gen_pai_local_method('train')
+predict = _gen_pai_local_method('predict')
+evaluate = _gen_pai_local_method('evaluate')
+explain = _gen_pai_local_method('explain')

python/runtime/pai_local/try_run.py

@@ -0,0 +1,40 @@
+# 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 os
+
+from runtime.pai import pai_model
+from runtime.pai.pai_distributed import define_tf_flags, set_oss_environs
+
+
+def init_pai_local_tf_flags_and_envs(oss_model_dir):
+    FLAGS = define_tf_flags()
+    FLAGS.sqlflow_oss_ak = os.getenv("SQLFLOW_OSS_AK")
+    FLAGS.sqlflow_oss_sk = os.getenv("SQLFLOW_OSS_SK")
+    FLAGS.sqlflow_oss_ep = os.getenv("SQLFLOW_OSS_MODEL_ENDPOINT")
+    if not oss_model_dir.startswith("oss://"):
+        oss_model_dir = pai_model.get_oss_model_url(oss_model_dir)
+    FLAGS.sqlflow_oss_modeldir = oss_model_dir
+    FLAGS.checkpointDir = os.getcwd()
+    set_oss_environs(FLAGS)
+
+
+def try_pai_local_run(params, oss_model_dir):
+    if os.getenv("SQLFLOW_submitter") == "pai_local":
+        from runtime.pai.entry import entrypoint
+        init_pai_local_tf_flags_and_envs(oss_model_dir)
+        print('start to run using pai_local submitter ...')
+        entrypoint(params)
+        return True
+    else:
+        return False

python/runtime/step/tensorflow/train.py

@@ -19,7 +19,7 @@
 from runtime.feature.compile import compile_ir_feature_columns
 from runtime.feature.derivation import (get_ordered_field_descs,
                                         infer_feature_columns)
-from runtime.model import EstimatorType, Model, collect_metadata, oss
+from runtime.model import EstimatorType, Model, collect_metadata
 from runtime.pai.pai_distributed import define_tf_flags, set_oss_environs
 from runtime.step.tensorflow.train_estimator import estimator_train_and_save
 from runtime.step.tensorflow.train_keras import keras_train_and_save
@@ -69,9 +69,15 @@ def train_step(original_sql,
     if validation_params is None:
         validation_params = {}
 
+    is_pai = True if pai_table else False
+    if is_pai:
+        actual_select = "SELECT * FROM %s" % pai_table
+    else:
+        actual_select = select
+
     conn = db.connect_with_data_source(datasource)
     fc_map_ir, fc_label_ir = infer_feature_columns(conn,
-                                                   select,
+                                                   actual_select,
                                                    feature_column_map,
                                                    label_column,
                                                    n=1000)
@@ -124,7 +130,6 @@ def train_step(original_sql,
     estimator = import_model(estimator_string)
     is_estimator = is_tf_estimator(estimator)
 
-    is_pai = True if pai_table else False
     # always use verbose == 1 when using PAI to get more logs
     if verbose < 1:
         verbose = 1
@@ -150,7 +155,7 @@ def train_step(original_sql,
                                       num_workers=num_workers,
                                       worker_id=worker_id)
     val_dataset_fn = None
-    if validation_select:
+    if validation_select or pai_val_table:
         val_dataset_fn = get_dataset_fn(validation_select, datasource,
                                         feature_column_names, feature_metas,
                                         label_meta, is_pai, pai_val_table,
@@ -185,19 +190,16 @@ def train_step(original_sql,
                                  save_checkpoints_steps, validation_metrics,
                                  load, model_meta)
 
-    # save model to DB
+    # save model to DB/OSS
+    model = Model(EstimatorType.TENSORFLOW, model_meta)
     if num_workers == 1 or worker_id == 0:
         if is_pai:
             oss_model_dir = FLAGS.sqlflow_oss_modeldir
-            oss.save_oss_model(oss_model_dir, estimator_string, is_estimator,
-                               feature_column_names, feature_column_names_map,
-                               feature_metas, label_meta, model_params,
-                               fc_map_ir, num_workers)
+            model.save_to_oss(oss_model_dir)
             print("Model saved to OSS: %s" % oss_model_dir)
         else:
-            model = Model(EstimatorType.TENSORFLOW, model_meta)
             model.save_to_db(datasource, save)
-            print("Model saved to db: %s" % save)
+            print("Model saved to DB: %s" % save)
 
     print("Done training")
     conn.close()

python/runtime/xgboost/train.py

@@ -13,13 +13,13 @@
 
 import sys
 
-import runtime.pai.pai_distributed as pai_dist
 import six
 import xgboost as xgb
 from runtime.local.xgboost_submitter.save import save_model_to_local_file
 from runtime.model import collect_metadata
 from runtime.model import oss as pai_model_store
 from runtime.model import save_metadata
+from runtime.pai.pai_distributed import make_distributed_info_without_evaluator
 from runtime.xgboost.dataset import xgb_dataset
 from runtime.xgboost.pai_rabit import PaiXGBoostTracker, PaiXGBoostWorker
 
@@ -50,8 +50,7 @@ def dist_train(flags,
             "XGBoost distributed training is only supported on PAI")
 
     num_workers = len(flags.worker_hosts.split(","))
-    cluster, node, task_id = pai_dist.make_distributed_info_without_evaluator(
-        flags)
+    cluster, node, task_id = make_distributed_info_without_evaluator(flags)
     master_addr = cluster["ps"][0].split(":")
     master_host = master_addr[0]
     master_port = int(master_addr[1]) + 1