-
Notifications
You must be signed in to change notification settings - Fork 704
Expand file tree
/
Copy pathmodel.py
More file actions
341 lines (279 loc) · 11 KB
/
Copy pathmodel.py
File metadata and controls
341 lines (279 loc) · 11 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
# 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.
"""This module saves or loads the SQLFlow model.
"""
import json
import os
import runtime.temp_file as temp_file
from runtime.dbapi import connect as connect_with_data_source
from runtime.feature.column import (JSONDecoderWithFeatureColumn,
JSONEncoderWithFeatureColumn)
from runtime.model import oss
from runtime.model.db import (read_with_generator_and_metadata,
write_with_generator_and_metadata)
from runtime.model.modelzoo import load_model_from_model_zoo
from runtime.model.tar import unzip_dir, zip_dir
# archive the current work director into a tarball
TARBALL_NAME = "model.tar.gz"
# serialize the Model object into file
MODEL_OBJ_FILE_NAME = "metadata.json"
class EstimatorType(object):
"""The enum type for various SQLFlow estimator.
"""
# To stay compitable with old models, we start at 0
TENSORFLOW = 0
XGBOOST = 1
# PAIML is the model type that trained by PAI machine learning algorithm
# toolkit
PAIML = 2
class Model(object):
"""Model module represents a SQLFlow trained model, which includes
three parts:
1. the estimator type indicates which SQLFlow estimator comes from.
2. the model meta indicates the meta data of training .e.g attributions,
feature column types.
3. the model data indicated the trained model, which generated by the AI
engine, .e.g TensorFlow, XGBoost.
Usage:
meta = runtime.model.collect_metadata(attributes={...}, ...)
m = runtime.model.Model(ModelType.XGBOOST, meta)
m.save(datasource="mysql://", "sqlflow_models.my_model")
"""
def __init__(self, typ, meta):
"""
Args:
typ: EstimatorType
the enum value of EstimatorType.
meta: JSON
the training meta with JSON format.
"""
self._typ = typ
self._meta = meta
def get_type(self):
"""
Returns the model type.
"""
return self._typ
def get_meta(self, name, default=None):
"""
Get the metadata by name.
Args:
name (str): the metadata name.
default: the default value if the name does not exist.
Returns:
Return the metadata with the given name if exists.
Otherwise, return the default value.
"""
return self._meta.get(name, default)
def _to_dict(self):
meta = dict(self._meta)
meta["model_type"] = self._typ
return meta
@staticmethod
def _from_dict(d):
typ = d.pop("model_type")
return Model(typ, d)
@staticmethod
def estimator_type(estimator):
estimator = estimator.lower()
if estimator in ["kmeans", "randomforests"]:
return EstimatorType.PAIML
elif estimator.startswith("xgboost"):
return EstimatorType.XGBOOST
else:
return EstimatorType.TENSORFLOW
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)
@staticmethod
def _unzip(local_dir, tarball):
"""
Unzip the tarball into local_dir and deserialize the model
information.
Args:
local_dir (str): the local directory.
tarball (str): the tarball path.
Returns:
Model: a Model object represent the model type and meta
information.
"""
unzip_dir(tarball, local_dir)
def save_to_db(self,
datasource,
table,
local_dir=None,
oss_model_dir=None):
"""
This save function would archive all the files on local_dir
into a tarball, and save it into DBMS with the specified table
name.
Args:
datasource (str): the connection string to DBMS.
table (str): the saved table name.
local_dir (str): the local directory to save.
Returns:
None.
"""
if local_dir is None:
local_dir = os.getcwd()
conn = connect_with_data_source(datasource)
if oss_model_dir:
cur_dir = os.getcwd()
os.chdir(local_dir)
oss.load_dir(oss_model_dir)
os.chdir(cur_dir)
if "." not in table:
project_name = conn.param("database")
table = project_name + "." + table
with temp_file.TemporaryDirectory() as tmp_dir:
tarball = os.path.join(tmp_dir, TARBALL_NAME)
self._zip(local_dir, 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_and_metadata(datasource, table,
_bytes_reader(tarball),
self._to_dict())
conn.persist_table(table)
conn.close()
return table
@staticmethod
def load_from_db(datasource, table, local_dir=None):
"""
Load the saved model from DBMS and unzip it on local_dir.
Args:
datasource (str): the connection string to DBMS
table (str): the table name which saved in DBMS
local_dir (str): the local directory to load.
Returns:
Model: a Model object represent the model type and meta
information.
"""
if local_dir is None:
local_dir = os.getcwd()
model_zoo_addr, table, tag = _decompose_model_name(table)
if model_zoo_addr:
gen, metadata = load_model_from_model_zoo(model_zoo_addr, table,
tag)
else:
gen, metadata = read_with_generator_and_metadata(datasource, table)
with temp_file.TemporaryDirectory() as tmp_dir:
tarball = os.path.join(tmp_dir, TARBALL_NAME)
with open(tarball, "wb") as f:
for data in gen():
f.write(bytes(data))
Model._unzip(local_dir, tarball)
return Model._from_dict(metadata)
@staticmethod
def load_metadata_from_db(datasource, table):
try:
return Model._load_metadata_from_db_impl(datasource, table)
except: # noqa: E722
return Model._load_metadata_from_db_impl(
datasource, table + "_sqlflow_pai_model")
@staticmethod
def _load_metadata_from_db_impl(datasource, table):
model_zoo_addr, table, tag = _decompose_model_name(table)
if model_zoo_addr:
_, metadata = load_model_from_model_zoo(model_zoo_addr,
table,
tag,
meta_only=True)
else:
_, metadata = read_with_generator_and_metadata(datasource,
table,
meta_only=True)
return Model._from_dict(metadata)
def save_to_oss(self, oss_model_dir, local_dir=None):
"""
This save function would archive all the files on local_dir
into a tarball, and save it into OSS model directory.
Args:
oss_model_dir (str): the OSS model directory to save.
It is in the format of oss://bucket/path/to/dir/.
local_dir (str): the local directory to save.
Returns:
None.
"""
if local_dir is None:
local_dir = os.getcwd()
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)
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):
"""
Load the saved model from OSS and unzip it on local_dir.
Args:
oss_model_dir (str): the OSS model directory to load.
It is in the format of oss://bucket/path/to/dir/.
local_dir (str): the local directory to load.
Returns:
Model: a Model object represent the model type and meta
information.
"""
if local_dir is None:
local_dir = os.getcwd()
with temp_file.TemporaryDirectory() as tmp_dir:
tarball = os.path.join(tmp_dir, TARBALL_NAME)
oss.load_file(oss_model_dir, tarball, TARBALL_NAME)
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):
idx = name.rfind("/")
if idx < 0:
return "", name, ""
model_zoo_addr = name[0:idx]
name = name[idx + 1:]
tag = ""
idx = name.rfind(":")
if idx >= 0:
tag = name[idx + 1:]
name = name[0:idx]
return model_zoo_addr, name, tag