mirror of
https://github.com/wassname/ray.git
synced 2026-06-29 22:51:19 +08:00
Philipp Moritz
a708e36225
* switch to CMake completely ... * cleanup * Run C tests, update installation instructions.
68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
# Note that a little bit of code here is taken and slightly modified from the pickler because it was not possible to change its behavior otherwise.
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import sys
|
|
from ctypes import c_void_p
|
|
from cloudpickle import pickle, cloudpickle, CloudPickler, load, loads
|
|
|
|
try:
|
|
from ctypes import pythonapi
|
|
pythonapi.PyCell_Set # Make sure this exists
|
|
except:
|
|
pythonapi = None
|
|
|
|
def dump(obj, file, protocol=2):
|
|
return BetterPickler(file, protocol).dump(obj)
|
|
|
|
def dumps(obj):
|
|
stringio = cloudpickle.StringIO()
|
|
dump(obj, stringio)
|
|
return stringio.getvalue()
|
|
|
|
def _make_skel_func(code, closure, base_globals = None):
|
|
""" Creates a skeleton function object that contains just the provided
|
|
code and the correct number of cells in func_closure. All other
|
|
func attributes (e.g. func_globals) are empty.
|
|
"""
|
|
if base_globals is None: base_globals = {}
|
|
base_globals['__builtins__'] = __builtins__
|
|
return _make_skel_func.__class__(code, base_globals, None, None, tuple(closure))
|
|
|
|
def _fill_function(func, globals, defaults, closure, dict):
|
|
""" Fills in the rest of function data into the skeleton function object
|
|
that were created via _make_skel_func(), including closures.
|
|
"""
|
|
result = cloudpickle._fill_function(func, globals, defaults, dict)
|
|
if pythonapi is not None:
|
|
for i, v in enumerate(closure):
|
|
pythonapi.PyCell_Set(c_void_p(id(result.__closure__[i])), c_void_p(id(v)))
|
|
return result
|
|
|
|
class BetterPickler(CloudPickler):
|
|
def save_function_tuple(self, func):
|
|
code, f_globals, defaults, closure, dct, base_globals = self.extract_func_data(func)
|
|
|
|
self.save(_fill_function)
|
|
self.write(pickle.MARK)
|
|
|
|
self.save(_make_skel_func if pythonapi else cloudpickle._make_skel_func)
|
|
self.save((code, map(lambda _: cloudpickle._make_cell(None), closure) if closure and pythonapi is not None else closure, base_globals))
|
|
self.write(pickle.REDUCE)
|
|
self.memoize(func)
|
|
|
|
self.save(f_globals)
|
|
self.save(defaults)
|
|
self.save(closure)
|
|
self.save(dct)
|
|
self.write(pickle.TUPLE)
|
|
self.write(pickle.REDUCE)
|
|
def save_cell(self, obj):
|
|
self.save(cloudpickle._make_cell)
|
|
self.save((obj.cell_contents,))
|
|
self.write(pickle.REDUCE)
|
|
dispatch = CloudPickler.dispatch.copy()
|
|
dispatch[(lambda _: lambda: _)(0).__closure__[0].__class__] = save_cell
|