1import shutil
2import tarfile
3from pathlib import Path
4
5import joblib
6import smart_open
7
8from omegaml.backends.basecommon import BackendBaseCommon
9from omegaml.util import reshaped
10
11
12class BaseModelBackend(BackendBaseCommon):
13 """
14 OmegaML BaseModelBackend to be subclassed by other arbitrary backends
15
16 This provides the abstract interface for any model backend to be implemented
17 Subclass to implement custom backends.
18
19 Essentially a model backend:
20
21 * provides methods to serialize and deserialize a machine learning model for a given ML framework
22 * offers fit() and predict() methods to be called by the runtime
23 * offers additional methods such as score(), partial_fit(), transform()
24
25 Model backends are the middleware that connects the om.models API to specific frameworks. This class
26 makes it simple to implement a model backend by offering a common syntax as well as a default implementation
27 for get() and put().
28
29 Methods to implement:
30 # for model serialization (mandatory)
31 @classmethod supports() - determine if backend supports given model instance
32 _package_model() - serialize a model instance into a temporary file
33 _extract_model() - deserialize the model from a file-like
34
35 By default BaseModelBackend uses joblib.dumps/loads to store the model as serialized
36 Python objects. If this is not sufficient or applicable to your type models, override these
37 methods.
38
39 Both methods provide readily set up temporary file names so that all you have to do is actually
40 save the model to the given output file and restore the model from the given input file, respectively.
41 All other logic has already been implemented (see get_model and put_model methods).
42
43 # for fitting and predicting (mandatory)
44 fit()
45 predict()
46
47 # other methods (optional)
48 fit_transform() - fit and return a transformed dataset
49 partial_fit() - fit incrementally
50 predict_proba() - predict probabilities
51 score() - score fitted classifier vv test dataset
52
53 """
54
55 _backend_version_tag = '_om_backend_version'
56 _backend_version = '1'
57
58 #: model serializer - save a model instance to a filename (or directory)
59 serializer = lambda store, model, filename, **kwargs: joblib.dump(model, filename)[0]
60 #: model deserializer - load a model instance from a filename (or directory)
61 loader = lambda store, infile, filename=None, **kwargs: joblib.load(infile or filename)
62 #: return the equivalent of the predict() function, defaults to obj.predict()
63 infer = lambda obj, **kwargs: obj.predict
64 #: transform the input data into model predict() format, defaults to np.array(data).reshape(-1, -1)
65 reshape = lambda data, **kwargs: reshaped(data)
66 #: the types this backend can store, defaults to (object,),
67 #: used as isinstance(obj, self.types). override .supports() for more complex logic
68 types = (object,)
69 #: tgz magic file, used to signal a serialized model was packaged by this class
70 _magictgz = '.omegaml.tgz'
71
72 def __init__(self, model_store=None, data_store=None, tracking=None, **kwargs):
73 assert model_store, "Need a model store"
74 assert data_store, "Need a data store"
75 self.model_store = model_store
76 self.data_store = data_store
77 self.tracking = tracking
78
79 @classmethod
80 def supports(self, obj, name, **kwargs):
81 """
82 test if this backend supports this obj
83 """
84 return isinstance(obj, self.types) if self.types else False
85
86 @property
87 def _call_handler(self):
88 # the model store handles _pre and _post methods in self.perform()
89 return self.model_store
90
91 def get(self, name, uri=None, **kwargs):
92 """
93 retrieve a model
94
95 :param name: the name of the object
96 :param uri: optional, /path/to/file, defaults to meta.gridfile, may use /path/{key} as placeholder
97 for the file's name
98 :param version: the version of the object (not supported)
99
100 .. versionadded: 0.18.0
101 uri specifies a target filename to store the serialized model
102 """
103 # support new backend architecture while keeping back compatibility
104 return self.get_model(name, uri=uri, **kwargs)
105
106 def put(self, obj, name, uri=None, **kwargs):
107 """
108 store a model
109
110 :param obj: the model object to be stored
111 :param name: the name of the object
112 :param uri: optional, /path/to/file, defaults to meta.gridfile, may use /path/{key} as placeholder
113 for the file's name
114 :param attributes: attributes for meta data
115
116 .. versionadded: 0.18.0
117 local specifies a local filename to store the serialized model
118 """
119 # support new backend architecture while keeping back compatibility
120 return self.put_model(obj, name, uri=uri, **kwargs)
121
122 def drop(self, name, force=False, version=-1, **kwargs):
123 return self.model_store._drop(name, force=force, version=version)
124
125 def _package_model(self, model, key, tmpfn, serializer=None, **kwargs):
126 """
127 implement this method to serialize a model to the given tmpfn
128
129 Args:
130 model (object): the model object to serialize to a file
131 key (str): the object store's key for this object
132 tmpfn (str): the filename to store the serialized object to
133 serializer (callable): optional, a callable as serializer(store, model, filename, **kwargs),
134 defaults to self.serializer, using joblib.dump()
135 **kwargs (dict): optional, keyword arguments passed to the serializer
136
137 Returns:
138 tmpfn or absolute path of serialized file
139
140 .. versionchanged:: 0.18.0
141 enable custom serializer
142
143 .. versionchanged: 0.18.2
144 enable multi-file serialized models
145 """
146 serializer = serializer or self.serializer.__func__ # __func__ is the unbound method
147 kwargs.setdefault('key', key)
148 tmpfn = serializer(self, model, tmpfn, **kwargs) or tmpfn
149 modelfn = Path(tmpfn)
150 if modelfn.is_dir():
151 fn_tgz = Path(tmpfn).parent / 's' / self._magictgz
152 shutil.rmtree(fn_tgz, ignore_errors=True)
153 fn_tgz.parent.mkdir(parents=True, exist_ok=True)
154 with open(modelfn / self._magictgz, 'w') as fout:
155 fout.write('# tgz packaged by omegaml')
156 with tarfile.open(fn_tgz, 'w:') as fout:
157 fout.add(modelfn, arcname=modelfn.relative_to(modelfn.parent))
158 modelfn = fn_tgz
159 return str(modelfn)
160
161 def _extract_model(self, infile, key, tmpfn, loader=None, **kwargs):
162 """
163 implement this method to deserialize a model from the given infile
164
165 Args:
166 infile (filelike): this is a file-like object supporting read() and seek(). if
167 deserializing from this does not work directly, use tmpfn
168 key (str): the object store's key for this object
169 tmpfn (str): the filename from which to extract the object
170 loader (callable): optional, a callable as loader(store, filename, **kwargs),
171 defaults to self.loader, using joblib.load()
172 **kwargs (dict): optional, keyword arguments passed to the loader
173
174 Returns:
175 model instance
176
177 .. versionchanged:: 0.18.0
178 enable custom loader
179
180 .. versionchanged: 0.18.2
181 enable multi-file serialized models
182 """
183 loader = loader or self.loader.__func__ # __func__ is the unbound method
184
185 def is_omegaml_tgz(infile):
186 try:
187 with tarfile.open(fileobj=infile) as tf:
188 if self._magictgz not in (Path(fn).name for fn in tf.getnames()):
189 raise ValueError(f'no {self._magictgz} file found')
190 except (tarfile.TarError, ValueError):
191 return False
192 infile.open() if hasattr(infile, 'open') else None
193 infile.seek(0) if hasattr(infile, 'seek') else None
194 return True
195
196 if is_omegaml_tgz(infile):
197 fn_tgz = Path(tmpfn).parent / 's' / 'modelfile.tgz'
198 fn_dir = Path(tmpfn).parent / 'x'
199 shutil.rmtree(fn_tgz.parent, ignore_errors=True)
200 shutil.rmtree(fn_dir, ignore_errors=True)
201 fn_tgz.parent.mkdir(parents=True, exist_ok=True)
202 fn_dir.mkdir(parents=True, exist_ok=True)
203 # -- write the tgz created by _package_model to a temp tgz file
204 with open(fn_tgz, mode='wb') as fout:
205 fout.write(infile.read())
206 # -- extract contents, we get back what the serializer wrote
207 with tarfile.open(fn_tgz, mode='r') as fin:
208 fin.extractall(fn_dir)
209 # -- remove the omegaml magic to avoid manifest issues
210 (fn_dir / self._magictgz).unlink(missing_ok=True)
211 # call the loader
212 files = list(fn_dir.glob('**/*'))
213 infile = open(files[-1], 'rb') if len(files) == 1 else None
214 infile.close() if len(files) == 1 else None
215 kwargs.setdefault('key', key)
216 kwargs.setdefault('filename', files[0])
217 else:
218 kwargs.setdefault('key', key)
219 kwargs.setdefault('filename', tmpfn)
220 # finally, load the model
221 obj = loader(self, infile, **kwargs)
222 shutil.rmtree(tmpfn, ignore_errors=True)
223 return obj
224
225 def _remove_path(self, path):
226 """
227 Remove a path, either a file or a directory
228
229 Args:
230 path (str): filename or path to remove. If path is a directory,
231 it will be removed recursively. If path is a file, it will be
232 removed.
233
234 Returns:
235 None
236 """
237 if Path(path).is_dir():
238 shutil.rmtree(path, ignore_errors=True)
239 else:
240 Path(path).unlink(missing_ok=True)
241
242 def get_model(self, name, version=-1, uri=None, loader=None, **kwargs):
243 """
244 Retrieves a pre-stored model
245 """
246 meta = self.model_store.metadata(name)
247 storekey = self.model_store.object_store_key(name, 'omm', hashed=True)
248 uri = uri or meta.uri
249 if uri:
250 uri = str(uri).format(key=storekey)
251 infile = smart_open.open(uri, 'rb')
252 else:
253 infile = meta.gridfile
254 model = self._extract_model(
255 infile, storekey, self._tmp_packagefn(self.model_store, storekey), loader=loader, **kwargs
256 )
257 infile.close()
258 return model
259
260 def put_model(self, obj, name, attributes=None, _kind_version=None, uri=None, **kwargs):
261 """
262 Packages a model using joblib and stores in GridFS
263 """
264 storekey = self.model_store.object_store_key(name, 'omm', hashed=True)
265 tmpfn = self._tmp_packagefn(self.model_store, storekey)
266 packagefname = self._package_model(obj, storekey, tmpfn, **kwargs) or tmpfn
267 gridfile = self._store_to_file(self.model_store, packagefname, storekey, uri=uri)
268 self._remove_path(packagefname)
269 kind_meta = {self._backend_version_tag: self._backend_version}
270 return self.model_store._make_metadata(
271 name=name,
272 prefix=self.model_store.prefix,
273 bucket=self.model_store.bucket,
274 kind=self.KIND,
275 kind_meta=kind_meta,
276 attributes=attributes,
277 uri=str(uri or ''),
278 gridfile=gridfile,
279 ).save()
280
[docs]
281 def predict(self, modelname, Xname, rName=None, pure_python=True, **kwargs):
282 """
283 predict using data stored in Xname
284
285 :param modelname: the name of the model object
286 :param Xname: the name of the X data set
287 :param rName: the name of the result data object or None
288 :param pure_python: if True return a python object. If False return
289 a dataframe. Defaults to True to support any client.
290 :param kwargs: kwargs passed to the model's predict method
291 :return: return the predicted outcome
292 """
293 model = self.model_store.get(modelname)
294 data = self._resolve_input_data('predict', Xname, 'X', **kwargs)
295 infer = self.infer.__func__ # __func__ is the unbound method
296 reshape = self.reshape.__func__
297 result = infer(model)(reshape(data))
298 return self._prepare_result('predict', result, rName=rName, pure_python=pure_python, **kwargs)
299
300 def _resolve_input_data(self, method, Xname, key, **kwargs):
301 data = self.data_store.get(Xname)
302 meta = self.data_store.metadata(Xname)
303 if self.tracking and getattr(self.tracking, 'autotrack', False):
304 self.tracking.log_data(key, data, dataset=Xname, kind=meta.kind, event=method)
305 return data
306
307 def _prepare_result(self, method, result, rName=None, pure_python=False, **kwargs):
308 if pure_python:
309 result = result.tolist()
310 if rName:
311 meta = self.data_store.put(result, rName)
312 result = meta
313 if self.tracking and getattr(self.tracking, 'autotrack', False):
314 self.tracking.log_data(
315 'Y', result, dataset=rName, kind=str(type(result)) if rName is None else meta.kind, event=method
316 )
317 return result
318
319 def predict_proba(self, modelname, Xname, rName=None, pure_python=True, **kwargs):
320 """
321 predict the probability using data stored in Xname
322
323 :param modelname: the name of the model object
324 :param Xname: the name of the X data set
325 :param rName: the name of the result data object or None
326 :param pure_python: if True return a python object. If False return
327 a dataframe. Defaults to True to support any client.
328 :param kwargs: kwargs passed to the model's predict method
329 :return: return the predicted outcome
330 """
331 raise NotImplementedError
332
[docs]
333 def fit(self, modelname, Xname, Yname=None, pure_python=True, **kwargs):
334 """
335 fit the model with data
336
337 :param modelname: the name of the model object
338 :param Xname: the name of the X data set
339 :param Yname: the name of the Y data set
340 :param pure_python: if True return a python object. If False return
341 a dataframe. Defaults to True to support any client.
342 :param kwargs: kwargs passed to the model's predict method
343 :return: return the meta data object of the model
344 """
345 raise NotImplementedError
346
347 def partial_fit(self, modelname, Xname, Yname=None, pure_python=True, **kwargs):
348 """
349 partially fit the model with data (online)
350
351 :param modelname: the name of the model object
352 :param Xname: the name of the X data set
353 :param Yname: the name of the Y data set
354 :param pure_python: if True return a python object. If False return
355 a dataframe. Defaults to True to support any client.
356 :param kwargs: kwargs passed to the model's predict method
357 :return: return the meta data object of the model
358 """
359
360 raise NotImplementedError
361
362 def fit_transform(self, modelname, Xname, Yname=None, rName=None, pure_python=True, **kwargs):
363 """
364 fit and transform using data
365
366 :param modelname: the name of the model object
367 :param Xname: the name of the X data set
368 :param Yname: the name of the Y data set
369 :param rName: the name of the transforms's result data object or None
370 :param pure_python: if True return a python object. If False return
371 a dataframe. Defaults to True to support any client.
372 :param kwargs: kwargs passed to the model's transform method
373 :return: return the meta data object of the model
374 """
375 raise NotImplementedError
376
388
389 def score(self, modelname, Xname, Yname=None, rName=True, pure_python=True, **kwargs):
390 """
391 score using data
392
393 :param modelname: the name of the model object
394 :param Xname: the name of the X data set
395 :param Yname: the name of the Y data set
396 :param rName: the name of the transforms's result data object or None
397 :param pure_python: if True return a python object. If False return
398 a dataframe. Defaults to True to support any client.
399 :param kwargs: kwargs passed to the model's predict method
400 :return: return the score result
401 """
402 raise NotImplementedError