def _zip(self, local_dir, tarball):
"""
Zip the model information and all files in local_dir into a tarball.
Args:
local_dir (str): the local directory.
tarball (str): the tarball path.
Returns:
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="./", filter=filter)
def test_tar(self):
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)
def save(self, datasource, table, cwd="./"):
"""This save function would archive all the files on work director
into a tarball, and saved it into DBMS with the specified table name.
Args:
datasource: string
the connection string to DBMS.
table: string
the saved table name.
"""
_dump_pkl(self, model_obj_file)
zip_dir(cwd, tarball)
def _bytes_reader(filename, buf_size=8 * 32):
def _gen():
with open(filename, "rb") as f:
while True:
data = f.read(buf_size)
if data:
yield data
else:
break
return _gen
write_with_generator(datasource, table, _bytes_reader(tarball))
def _zip(self, local_dir, tarball):
"""
Zip the model information and all files in local_dir into a tarball.
Args:
local_dir (str): the local directory.
tarball (str): the tarball path.
Returns:
None.
"""
model_obj_file = os.path.join(local_dir, MODEL_OBJ_FILE_NAME)
_dump_pkl(self, model_obj_file)
zip_dir(local_dir, tarball, arcname="./")
os.remove(model_obj_file)
def _zip(self, local_dir, tarball, save_to_db=False):
"""
Zip the model information and all files in local_dir into a tarball.
Args:
local_dir (str): the local directory.
tarball (str): the tarball path.
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
zip_dir(local_dir, tarball, arcname="./")
if model_obj_file:
os.remove(model_obj_file)