add temp_file.py (#2838)

Author: sneaxiy

go/codegen/experimental/xgboost.go

@@ -141,7 +141,7 @@ const xgbTrainTemplate = `
 def step_entry_{{.StepIndex}}():
     import json
     import os
-    import tempfile
+    import runtime.temp_file as temp_file
     import runtime.feature.column as fc
     import runtime.feature.field_desc as fd
     import runtime.{{.Submitter}}.xgboost as xgboost_submitter
@@ -156,8 +156,7 @@ def step_entry_{{.StepIndex}}():
     model_params = json.loads('''{{.ModelParamsJSON}}''')
     train_params = json.loads('''{{.TrainParamsJSON}}''')
 
-    with tempfile.TemporaryDirectory() as temp_dir:
-        os.chdir(temp_dir)
+    with temp_file.TemporaryDirectory(as_cwd=True) as temp_dir:
         xgboost_submitter.train(original_sql='''{{.OriginalSQL}}''',
                                 model_image='''{{.ModelImage}}''',
                                 estimator='''{{.Estimator}}''',

python/runtime/local/xgboost/train.py

@@ -14,10 +14,10 @@
 This module launches a XGBoost training task on host.
 """
 import os
-import tempfile
 import types
 
 import runtime.db as db
+import runtime.temp_file as temp_file
 import runtime.xgboost as xgboost_extended
 import xgboost as xgb
 from runtime.feature.compile import compile_ir_feature_columns
@@ -106,7 +106,7 @@ def build_dataset(fn, slct):
     else:
         bst = None
 
-    with tempfile.TemporaryDirectory() as tmp_dir_name:
+    with temp_file.TemporaryDirectory() as tmp_dir_name:
         train_fn = os.path.join(tmp_dir_name, 'train.txt')
         val_fn = os.path.join(tmp_dir_name, 'val.txt')
         train_dataset = build_dataset(train_fn, select)

python/runtime/model/model.py

@@ -13,9 +13,9 @@
 """This module saves or loads the SQLFlow model.
 """
 import os
-import tempfile
 from enum import Enum
 
+import runtime.temp_file as temp_file
 from runtime.model import oss
 from runtime.model.db import read_with_generator, write_with_generator
 from runtime.model.tar import unzip_dir, zip_dir
@@ -143,7 +143,7 @@ def save_to_db(self, datasource, table, local_dir=None):
         if local_dir is None:
             local_dir = os.getcwd()
 
-        with tempfile.TemporaryDirectory() as tmp_dir:
+        with temp_file.TemporaryDirectory() as tmp_dir:
             tarball = os.path.join(tmp_dir, TARBALL_NAME)
             self._zip(local_dir, tarball)
 
@@ -178,7 +178,7 @@ def load_from_db(datasource, table, local_dir=None):
         if local_dir is None:
             local_dir = os.getcwd()
 
-        with tempfile.TemporaryDirectory() as tmp_dir:
+        with temp_file.TemporaryDirectory() as tmp_dir:
             tarball = os.path.join(tmp_dir, TARBALL_NAME)
             gen = read_with_generator(datasource, table)
             with open(tarball, "wb") as f:
@@ -203,7 +203,7 @@ def save_to_oss(self, oss_model_dir, local_dir=None):
         if local_dir is None:
             local_dir = os.getcwd()
 
-        with tempfile.TemporaryDirectory() as tmp_dir:
+        with temp_file.TemporaryDirectory() as tmp_dir:
             tarball = os.path.join(tmp_dir, TARBALL_NAME)
             self._zip(local_dir, tarball)
             oss.save_file(oss_model_dir, tarball, TARBALL_NAME)
@@ -225,7 +225,7 @@ def load_from_oss(oss_model_dir, local_dir=None):
         if local_dir is None:
             local_dir = os.getcwd()
 
-        with tempfile.TemporaryDirectory() as tmp_dir:
+        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)

python/runtime/model/model_test.py

@@ -12,33 +12,27 @@
 # limitations under the License.
 
 import os
-import tempfile
 import unittest
 
 import runtime.model.oss as oss
+import runtime.temp_file as temp_file
 from runtime.model import EstimatorType, Model
 from runtime.testing import get_datasource
 
 
 class TestModel(unittest.TestCase):
-    def setUp(self):
-        self.cur_dir = os.getcwd()
-
-    def tearDown(self):
-        os.chdir(self.cur_dir)
-
     def test_save_load_db(self):
         table = "sqlflow_models.test_model"
         meta = {"model_params": {"n_classes": 3}}
         m = Model(EstimatorType.XGBOOST, meta)
         datasource = get_datasource()
 
         # save mode
-        with tempfile.TemporaryDirectory() as d:
+        with temp_file.TemporaryDirectory() as d:
             m.save_to_db(datasource, table, d)
 
         # load model
-        with tempfile.TemporaryDirectory() as d:
+        with temp_file.TemporaryDirectory() as d:
             m = Model.load_from_db(datasource, table, d)
             self.assertEqual(m._meta, meta)
 
@@ -57,12 +51,12 @@ def test_save_load_oss(self):
 
         # save model
         def save_to_oss():
-            with tempfile.TemporaryDirectory() as d:
+            with temp_file.TemporaryDirectory() as d:
                 m.save_to_oss(oss_model_path, d)
 
         # load model
         def load_from_oss():
-            with tempfile.TemporaryDirectory() as d:
+            with temp_file.TemporaryDirectory() as d:
                 return Model.load_from_oss(oss_model_path, d)
 
         with self.assertRaises(Exception):

python/runtime/model/tar_test.py

@@ -12,45 +12,38 @@
 # limitations under the License.
 
 import os
-import tempfile
 import unittest
 
+import runtime.temp_file as temp_file
 from runtime.model.tar import unzip_dir, zip_dir
 
 
 class TestTarOperator(unittest.TestCase):
-    def setUp(self):
-        self.test_dir = tempfile.TemporaryDirectory()
-        self.cur_dir = os.getcwd()
-        os.chdir(self.test_dir.name)
-
-    def tearDown(self):
-        self.test_dir.cleanup()
-        os.chdir(self.cur_dir)
-
     def test_tar(self):
-        # create the test file tree:
-        #
-        # |-sqlflow_tar
-        #   |-sqlflow_sub_dir
-        #     |-hello.py
-        test_dir = "sqlflow_tar"
-        test_sub_dir = "sqlflow_sub_dir"
-        test_py_file = "hello.py"
-        test_py_content = "print('hello SQLFlow!')"
-
-        fullpath = os.path.join(test_dir, test_sub_dir)
-        os.makedirs(fullpath)
-        with open(os.path.join(fullpath, test_py_file), "w") as f:
-            f.write(test_py_content)
-
-        zip_dir(fullpath, "sqlflow.tar.gz")
-        unzip_dir("sqlflow.tar.gz", "output")
-        self.assertTrue(os.path.isdir("output/sqlflow_tar/sqlflow_sub_dir"))
-        self.assertTrue(
-            os.path.isfile("output/sqlflow_tar/sqlflow_sub_dir/hello.py"))
-        with open(os.path.join(fullpath, test_py_file), "r") as f:
-            self.assertEqual(f.read(), test_py_content)
+        with temp_file.TemporaryDirectory(as_cwd=True):
+            # create the test file tree:
+            #
+            # |-sqlflow_tar
+            #   |-sqlflow_sub_dir
+            #     |-hello.py
+            test_dir = "sqlflow_tar"
+            test_sub_dir = "sqlflow_sub_dir"
+            test_py_file = "hello.py"
+            test_py_content = "print('hello SQLFlow!')"
+
+            fullpath = os.path.join(test_dir, test_sub_dir)
+            os.makedirs(fullpath)
+            with open(os.path.join(fullpath, test_py_file), "w") as f:
+                f.write(test_py_content)
+
+            zip_dir(fullpath, "sqlflow.tar.gz")
+            unzip_dir("sqlflow.tar.gz", "output")
+            self.assertTrue(
+                os.path.isdir("output/sqlflow_tar/sqlflow_sub_dir"))
+            self.assertTrue(
+                os.path.isfile("output/sqlflow_tar/sqlflow_sub_dir/hello.py"))
+            with open(os.path.join(fullpath, test_py_file), "r") as f:
+                self.assertEqual(f.read(), test_py_content)
 
 
 if __name__ == '__main__':

python/runtime/temp_file.py

@@ -0,0 +1,46 @@
+# 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
+import shutil
+import tempfile
+
+
+# NOTE: Python 2 does not have tempfile.TemporaryDirectory. To unify the code
+# of Python 2 and 3, we make the following class.
+class TemporaryDirectory(object):
+    def __init__(self, as_cwd=False, suffix=None, prefix=None, dir=None):
+        """
+        Create a temporary directory.
+
+        Args:
+            as_cwd (bool): whether to change the current working directory
+                as the created temporary directory.
+            suffix (str): the suffix of the created temporary directory.
+            prefix (str): the prefix of the created temporary directory.
+            dir (str): where to create the temporary directory.
+        """
+        self.tmp_dir = tempfile.mkdtemp(suffix=suffix, prefix=prefix, dir=dir)
+        self.as_cwd = as_cwd
+        if self.as_cwd:
+            self.old_dir = os.getcwd()
+
+    def __enter__(self, *args, **kwargs):
+        if self.as_cwd:
+            os.chdir(self.tmp_dir)
+        return self.tmp_dir
+
+    def __exit__(self, *args, **kwargs):
+        if self.as_cwd:
+            os.chdir(self.old_dir)
+        shutil.rmtree(self.tmp_dir, ignore_errors=True)