1import builtins
2import sys
3import types
4import warnings
5
6import dill
7
8from omegaml.backends.basedata import BaseDataBackend
9from omegaml.util import tryOr
10
11
[docs]
12class VirtualObjectBackend(BaseDataBackend):
13 """
14 Support arbitrary functions as object handlers
15
16 Virtual object functions can be any callable that provides a __omega_virtual__
17 attribute. The callable must support the following signature::
18
19 @virtualobj
20 def virtualobjfn(data=None, method='get|put|drop',
21 meta=None, store=None, **kwargs):
22 ...
23 return data
24
25 Note that there is a distinction between storing the function as a virtual object,
26 and passing data in or getting data out of the store. It is the responsibility
27 of the function to implement the appropriate code for get, put, drop, as well as
28 to keep track of the data it actually stores.
29
30 As a convenience virtual object handlers can be implemented as a subclass of
31 VirtualObjectHandler
32
33 Usage::
34
35 1) as a virtual data handler
36
37 # create the 'foo' virtual object
38 om.datasets.put(virtualobjfn, 'foo')
39
40 # get data from the virtualobj
41 om.datasets.get('foo')
42 => will call virtualobjfn(method='get')
43
44 # put data into the virtualobj
45 om.datasets.put(data, 'foo')
46 => will call virtualobjfn(data=data, method='put')
47
48 # drop the virtualfn
49 om.datasets.drop('name')
50 => will call virtualobjfn(method='drop') and then
51 drop the virtual object completely from the storage
52
53 2) as a virtual model
54
55 # create the mymodel model as a virtualobj
56 om.models.put(virtualobjfn, 'mymodel')
57
58 # run the model's predict() function
59 om.runtime.model('mymodel').predict(X)
60 => will call virtualobjfn(method='predict')
61
62 3) as a virtual script
63
64 # create the myscript script as a virtualobj
65 om.models.put(virtualobjfn, 'myscript')
66
67 # run the script
68 om.runtime.script('myscript').run()
69 => will call virtualobjfn(method='run')
70
71 WARNING:
72
73 Virtual objects are executed in the address space of the client or
74 runtime context. Make sure that the source of the code is trustworthy.
75 Note that this is different from Backends and Mixins as these are
76 pro-actively enabled by the administrator of the client or runtime
77 context, respectively - virtual objects can be injected by anyone
78 who are authorized to write data.
79 """
80
81 # TODO split VirtualObjectBackend into VirtualModelBackend and VirtualDataBackend
82 # to avoid confusion between the two (currently the same class is used for both)
83 KIND = 'virtualobj.dill'
84 PROMOTE = 'export'
85 TOOLS_PATHS = 'tools/', 'policy/'
86
[docs]
87 @classmethod
88 def supports(self, obj, name, **kwargs):
89 is_virtual = callable(obj) and getattr(obj, '_omega_virtual', False)
90 is_tool = self._is_tool_virtualobj(obj, name, **kwargs)
91 return is_virtual or is_tool
92
93 @classmethod
94 def _is_tool_virtualobj(self, obj, name, **kwargs):
95 return isinstance(obj, types.FunctionType) and any(name.startswith(v) for v in self.TOOLS_PATHS)
96
97 @property
98 def _call_handler(self):
99 # the model store handles _pre and _post methods in self.perform()
100 return self.model_store
101
[docs]
102 def put(self, obj, name, attributes=None, dill_kwargs=None, as_source=False, **kwargs):
103 # TODO add obj signing so that only trustworthy sources can put functions
104 # since 0.15.6: only __main__ objects are stored as bytecodes,
105 # all module code is stored as source code. This
106 # removes the dependency on opcode parity between client
107 # and server. Source objects are compiled into __main__
108 # within the runtime. This is a tradeoff compatibility
109 # v.s. execution time. Use as_source=False to force
110 # storing bytecodes.
111 data = dilldip.dumps(obj, as_source=as_source, **(dill_kwargs or {}))
112 filename = self.model_store.object_store_key(name, '.dill', hashed=True)
113 gridfile = self._store_to_file(self.model_store, data, filename)
114 return self.model_store._make_metadata(
115 name=name,
116 prefix=self.model_store.prefix,
117 bucket=self.model_store.bucket,
118 kind=self.KIND,
119 attributes=attributes,
120 gridfile=gridfile,
121 ).save()
122
[docs]
123 def get(self, name, version=-1, force_python=False, lazy=False, **kwargs):
124 meta = self.model_store.metadata(name)
125 outf = meta.gridfile
126 data = outf.read()
127 obj = dilldip.loads(data)
128 outf.close()
129 return self._ensure_handler_instance(obj)
130
131 def _ensure_handler_instance(self, obj):
132 # ensure VirtualObjectHandler classes are transformed to a virtualobj
133 return obj() if isinstance(obj, type) and issubclass(obj, VirtualObjectHandler) else obj
134
135 def predict(self, modelname, xName, rName=None, **kwargs):
136 # make this work as a model backend too
137 meta = self.model_store.metadata(modelname)
138 handler = self._ensure_handler_instance(self.get(modelname))
139 X = self.data_store.get(xName)
140 return handler(
141 method='predict', data=X, meta=meta, store=self.model_store, rName=rName, tracking=self.tracking, **kwargs
142 )
143
144 def fit(self, modelname, xName, yName=None, rName=None, **kwargs):
145 # make this work as a model backend too
146 meta = self.model_store.metadata(modelname)
147 handler = self._ensure_handler_instance(self.get(modelname))
148 X = self.data_store.get(xName)
149 y = self.data_store.get(yName) if yName else None
150 return handler(
151 method='fit', data=(X, y), meta=meta, store=self.model_store, rName=rName, tracking=self.tracking, **kwargs
152 )
153
154 def score(self, modelname, xName, yName=None, rName=None, **kwargs):
155 # make this work as a model backend too
156 meta = self.model_store.metadata(modelname)
157 handler = self._ensure_handler_instance(self.get(modelname))
158 X = self.data_store.get(xName)
159 y = self.data_store.get(yName) if yName else None
160 return handler(
161 method='score',
162 data=(X, y),
163 meta=meta,
164 store=self.model_store,
165 rName=rName,
166 tracking=self.tracking,
167 **kwargs,
168 )
169
170 def run(self, scriptname, *args, **kwargs):
171 # run as a script
172 meta = self.model_store.metadata(scriptname)
173 handler = self._ensure_handler_instance(self.get(scriptname))
174 data = args[0] if args else None
175 kwargs['args'] = args
176 return handler(method='run', data=data, meta=meta, store=self.data_store, tracking=self.tracking, **kwargs)
177
[docs]
178 def reduce(self, modelname, results, rName=None, **kwargs):
179 """
180 reduce a list of results to a single result
181
182 Use this as the last step in a task canvas
183
184 Args:
185 modelname (str): the name of the virtualobj
186 results (list): the list of results forwarded by task canvas
187 rName (result): the name of the result object
188 **kwargs:
189
190 Returns:
191 result of the virtualobj handler
192
193 See Also
194 om.runtime.mapreduce
195 """
196 meta = self.model_store.metadata(modelname)
197 handler = self._ensure_handler_instance(self.get(modelname))
198 return handler(
199 method='reduce',
200 data=results,
201 meta=meta,
202 store=self.model_store,
203 rName=rName,
204 tracking=self.tracking,
205 **kwargs,
206 )
207
208
209def virtualobj(fn):
210 """
211 function decorator to create a virtual object handler from any
212 callable
213
214 Args:
215 fn: the virtual handler function
216
217 Usage:
218
219 .. code::
220
221 @virtualobj
222 def myvirtualobj(data=None, method=None, meta=None, store=None, **kwargs):
223 ...
224
225 See:
226 VirtualObjectBackend for details on virtual object handlers
227
228 Returns:
229 fn
230 """
231 fn._omega_virtual = True
232 return fn
233
234
235class VirtualObjectHandler:
236 """
237 Object-oriented API for virtual object functions
238 """
239
240 _omega_virtual = True
241
242 def _vobj_call_map(self):
243 return {
244 'drop': self.drop,
245 'get': self.get,
246 'put': self.put,
247 'predict': self.predict,
248 'run': self.run,
249 }
250
251 def load(self):
252 pass
253
254 def get(self, data=None, meta=None, store=None, **kwargs):
255 raise NotImplementedError
256
257 def put(self, data=None, meta=None, store=None, **kwargs):
258 raise NotImplementedError
259
260 def drop(self, data=None, meta=None, store=None, **kwargs):
261 raise NotImplementedError
262
263 def predict(self, data=None, meta=None, store=None, **kwargs):
264 raise NotImplementedError
265
266 def run(self, data=None, meta=None, store=None, **kwargs):
267 raise NotImplementedError
268
269 def __call__(self, data=None, method=None, meta=None, store=None, tracking=None, **kwargs):
270 MAP = self._vobj_call_map()
271 methodfn = MAP[method]
272 self.load() # ensure loading of imports and other dependencies
273 return methodfn(data=data, meta=meta, store=store, tracking=tracking, **kwargs)
274
275
276class _DillDip:
277 # enhanced dill for functions and classes
278 # - warns about bound variables that could cause issues when undilling
279 # - checks types and functions are defined in __main__ before dilling (so we will code, not references),
280 # dynamically recompiles in __main__ before dilling if source is available
281 # - saves source code for improved python cross-version compatibility
282 # - falls back on standard dill for any other object type
283 # Rationale:
284 # - https://stackoverflow.com/questions/26389981/serialize-a-python-function-with-dependencies
285 # - https://github.com/uqfoundation/dill/issues/424#issuecomment-1598887257
286
287 # magic id for dipped dill objects
288 # -- why 0x1565? 15g of dill dip have 65 kcal
289 # -- see https://www.nutritionvalue.org/Dill_dip%2C_regular_12350210_nutritional_value.html
290 # -- also 42 is overused
291 __calories = 0x1565
292
293 def dumps(self, obj, as_source=False, **dill_kwargs):
294 # enhanced flavor of dill that stores source code for cross-version compatibility
295 # ensure we have a dill'able object
296 # if isinstance(obj, type):
297 # obj = obj()
298 self._check(obj)
299 # FIXME respect as_source before _dill_main
300 data = (
301 self._dill_main(obj, **dill_kwargs)
302 or self._dill_types_or_function(obj, as_source=as_source, **dill_kwargs)
303 or self._dill_dill(obj, **dill_kwargs)
304 )
305 return data
306
307 def loads(self, data):
308 # compat: Python 3.8.x < 3.8.2
309 # https://github.com/python/cpython/commit/b19f7ecfa3adc6ba1544225317b9473649815b38
310 # https://docs.python.org/3.8/whatsnew/changelog.html#python-3-8-2-final
311 try:
312 obj = self._dynamic_compile(dill.loads(data), module='__main__')
313 except ModuleNotFoundError as e:
314 # if the functions original module is not known, simulate it
315 # this is to deal with functions created outside of __main__
316 # see https://stackoverflow.com/q/26193102/890242
317 # https://stackoverflow.com/a/70513630/890242
318 mod = types.ModuleType(e.name, '__dynamic__')
319 sys.modules[e.name] = mod # sys.modules['__main__']
320 obj = dill.loads(data)
321 return obj
322
323 def _check(self, obj):
324 # check for freevars
325 freevars = dill.detect.nestedglobals(obj)
326 freevars += list(dill.detect.freevars(obj).keys())
327 freevars += list(dill.detect.referredglobals(obj))
328 freevars = [n for n in set(freevars) if n not in dir(builtins)]
329 if len(freevars):
330 warnings.warn(
331 f'The {obj!r} module references {freevars}, this may lead to errors at runtime; import/declare all variables within method/function scope'
332 )
333
334 def _dill_dill(self, obj, **dill_kwargs):
335 # fallback to standard dill
336 # e.g. class instances cannot be dumped unless they come from __main__
337 return dill.dumps(obj, **dill_kwargs)
338
339 def _dill_main(self, obj, **dill_kwargs):
340 # dynamic __main__ objects can be dilled directly, there is no source code
341 if dill.source.isfrommain(obj) or dill.source.isdynamic(obj):
342 return dill.dumps(obj, **dill_kwargs)
343 return None
344
345 def _dill_types_or_function(self, obj, as_source=False, **dill_kwargs):
346 # classes or functions should be dilled as source, unless they come from __main__
347 if isinstance(obj, type) or isinstance(obj, types.FunctionType):
348 return self._dill_source(obj, as_source=as_source, **dill_kwargs)
349 return None
350
351 def _dill_source(self, obj, as_source=False, **dill_kwargs):
352 # include source code along dill
353 try:
354 # FIXME: include python version
355 source = dill.source.getsource(obj, lstrip=True)
356 source_obj = {
357 '__dipped__': self.__calories,
358 'source': ''.join(source),
359 'name': obj.__name__,
360 '__dict__': getattr(obj, '__dict__', {}),
361 }
362 except:
363 source_obj = {}
364 else:
365 # check obvious references in source
366 if '__main__' in source_obj.get('source', []):
367 warnings.warn(f'The {obj!r} references __main__, this may lead to unexpected results')
368 if as_source and source_obj:
369 # if source code was requested, transport as source code
370 data = dill.dumps(source_obj, **dill_kwargs)
371 elif source_obj and dill.detect.getmodule(obj) != '__main__':
372 # we have a source obj, make sure we can dill it and have source to revert from
373 # compile to __main__ module to enable full serialization
374 warnings.warn(f'The {obj!r} is defined outside of __main__, recompiling in __main__.')
375 obj = self._dynamic_compile(source_obj, module='__main__')
376 source_obj['dill'] = dill.dumps(obj, **dill_kwargs)
377 data = dill.dumps(source_obj, **dill_kwargs)
378 else:
379 # we have no source object, revert to standard dill
380 if as_source:
381 warnings.warn(f'Cannot save {obj!r} as source code, reverting to dill')
382 # could not get source code, revert to dill
383 data = dill.dumps(obj, **dill_kwargs)
384 return data
385
386 def _dynamic_compile(self, obj, module='__main__'):
387 from omegaml.backends.genai import GenAIModelHandler, virtual_genai
388
389 # re-compile source obj in __main__
390 if self.isdipped(obj):
391 # FIXME and check python version; if not the same don't load the dill object (it may load but still not be executable)
392 if 'dill' in obj:
393 try:
394 obj = dill.loads(obj['dill'])
395 except:
396 warnings.warn('could not undill, reverting to dynamic compile source code')
397 else:
398 return obj
399 source, data = obj.get('source'), obj.get('__dict__', {})
400 mod = types.ModuleType(module)
401 mod.__dict__.update({
402 '__compiling__': True,
403 'virtualobj': virtualobj,
404 'virtual_genai': virtual_genai,
405 'VirtualObjectHandler': VirtualObjectHandler,
406 'GenAIModelHandler': GenAIModelHandler,
407 })
408 sys.modules[module] = mod
409 code = compile(source, '<string>', 'exec')
410 exec(code, mod.__dict__)
411 obj = getattr(mod, obj['name'])
412 # restore instance data, if any
413 try:
414 getattr(obj, '__dict__', {}).update(data)
415 except AttributeError:
416 # we ignore attribute errors on class types
417 if not isinstance(obj, type):
418 warnings.warn(f'could not restore instance data for {obj}')
419 return obj
420
421 def isdipped(self, data_or_obj):
422 obj = tryOr(lambda: dill.loads(data_or_obj), None) if not isinstance(data_or_obj, dict) else data_or_obj
423 return isinstance(obj, dict) and obj.get('__dipped__') == self.__calories
424
425 def dump(self, filelike, obj, **kwargs):
426 filelike.write(self.dumps(obj, **kwargs))
427
428 def load(self, filelike):
429 return self.loads(filelike.read())
430
431
432def tool(fn):
433 fn._omega_virtual = True
434 return fn
435
436
437dilldip = _DillDip()
438# enable recursive tracing of globals
439# -- fixes https://github.com/uqfoundation/dill/issues/255
440# -- see also https://github.com/uqfoundation/dill/issues/537
441# -- the issue only happens when the openai module is loaded (root cause unknown)
442dill.settings['recurse'] = True