| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You 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. |
| # |
| |
| """Pickler for values, functions, and classes. |
| |
| For internal use only. No backwards compatibility guarantees. |
| |
| Pickles created by the pickling library contain non-ASCII characters, so |
| we base64-encode the results so that we can put them in a JSON objects. |
| The pickler is used to embed FlatMap callable objects into the workflow JSON |
| description. |
| |
| The pickler module should be used to pickle functions and modules; for values, |
| the coders.*PickleCoder classes should be used instead. |
| """ |
| |
| # pytype: skip-file |
| |
| import base64 |
| import bz2 |
| import logging |
| import sys |
| import threading |
| import traceback |
| import types |
| import zlib |
| from typing import Any |
| from typing import Dict |
| from typing import Tuple |
| |
| import dill |
| |
| from apache_beam.internal.code_object_pickler import get_normalized_path |
| from apache_beam.internal.set_pickler import save_frozenset |
| from apache_beam.internal.set_pickler import save_set |
| |
| settings = {'dill_byref': None} |
| |
| patch_save_code = sys.version_info >= (3, 10) and dill.__version__ == "0.3.1.1" |
| |
| if patch_save_code: |
| # The following function is based on 'save_code' from 'dill' |
| # Author: Mike McKerns (mmckerns @caltech and @uqfoundation) |
| # Copyright (c) 2008-2015 California Institute of Technology. |
| # Copyright (c) 2016-2023 The Uncertainty Quantification Foundation. |
| # License: 3-clause BSD. The full license text is available at: |
| # - https://github.com/uqfoundation/dill/blob/master/LICENSE |
| |
| # The following function is also based on 'save_codeobject' from 'cloudpickle' |
| # Copyright (c) 2012, Regents of the University of California. |
| # Copyright (c) 2009 `PiCloud, Inc. <http://www.picloud.com>`_. |
| # License: 3-clause BSD. The full license text is available at: |
| # - https://github.com/cloudpipe/cloudpickle/blob/master/LICENSE |
| |
| from types import CodeType |
| |
| @dill.register(CodeType) |
| def save_code(pickler, obj): |
| co_filename = get_normalized_path(obj.co_filename) |
| if hasattr(obj, "co_endlinetable"): # python 3.11a (20 args) |
| args = ( |
| obj.co_argcount, |
| obj.co_posonlyargcount, |
| obj.co_kwonlyargcount, |
| obj.co_nlocals, |
| obj.co_stacksize, |
| obj.co_flags, |
| obj.co_code, |
| obj.co_consts, |
| obj.co_names, |
| obj.co_varnames, |
| co_filename, |
| obj.co_name, |
| obj.co_qualname, |
| obj.co_firstlineno, |
| obj.co_linetable, |
| obj.co_endlinetable, |
| obj.co_columntable, |
| obj.co_exceptiontable, |
| obj.co_freevars, |
| obj.co_cellvars) |
| elif hasattr(obj, "co_exceptiontable"): # python 3.11 (18 args) |
| args = ( |
| obj.co_argcount, |
| obj.co_posonlyargcount, |
| obj.co_kwonlyargcount, |
| obj.co_nlocals, |
| obj.co_stacksize, |
| obj.co_flags, |
| obj.co_code, |
| obj.co_consts, |
| obj.co_names, |
| obj.co_varnames, |
| co_filename, |
| obj.co_name, |
| obj.co_qualname, |
| obj.co_firstlineno, |
| obj.co_linetable, |
| obj.co_exceptiontable, |
| obj.co_freevars, |
| obj.co_cellvars) |
| elif hasattr(obj, "co_linetable"): # python 3.10 (16 args) |
| args = ( |
| obj.co_argcount, |
| obj.co_posonlyargcount, |
| obj.co_kwonlyargcount, |
| obj.co_nlocals, |
| obj.co_stacksize, |
| obj.co_flags, |
| obj.co_code, |
| obj.co_consts, |
| obj.co_names, |
| obj.co_varnames, |
| co_filename, |
| obj.co_name, |
| obj.co_firstlineno, |
| obj.co_linetable, |
| obj.co_freevars, |
| obj.co_cellvars) |
| elif hasattr(obj, "co_posonlyargcount"): # python 3.8 (16 args) |
| args = ( |
| obj.co_argcount, |
| obj.co_posonlyargcount, |
| obj.co_kwonlyargcount, |
| obj.co_nlocals, |
| obj.co_stacksize, |
| obj.co_flags, |
| obj.co_code, |
| obj.co_consts, |
| obj.co_names, |
| obj.co_varnames, |
| co_filename, |
| obj.co_name, |
| obj.co_firstlineno, |
| obj.co_lnotab, |
| obj.co_freevars, |
| obj.co_cellvars) |
| else: # python 3.7 (15 args) |
| args = ( |
| obj.co_argcount, |
| obj.co_kwonlyargcount, |
| obj.co_nlocals, |
| obj.co_stacksize, |
| obj.co_flags, |
| obj.co_code, |
| obj.co_consts, |
| obj.co_names, |
| obj.co_varnames, |
| obj.co_filename, |
| obj.co_name, |
| obj.co_firstlineno, |
| obj.co_lnotab, |
| obj.co_freevars, |
| obj.co_cellvars) |
| pickler.save_reduce(CodeType, args, obj=obj) |
| |
| dill._dill.save_code = save_code |
| |
| |
| class _NoOpContextManager(object): |
| def __enter__(self): |
| pass |
| |
| def __exit__(self, *unused_exc_info): |
| pass |
| |
| |
| # Pickling, especially unpickling, causes broken module imports on Python 3 |
| # if executed concurrently, see: BEAM-8651, http://bugs.python.org/issue38884. |
| _pickle_lock = threading.RLock() |
| # Dill 0.28.0 renamed dill.dill to dill._dill: |
| # https://github.com/uqfoundation/dill/commit/f0972ecc7a41d0b8acada6042d557068cac69baa |
| # TODO: Remove this once Beam depends on dill >= 0.2.8 |
| if not getattr(dill, 'dill', None): |
| dill.dill = dill._dill |
| sys.modules['dill.dill'] = dill._dill |
| |
| # TODO: Remove once Dataflow has containers with a preinstalled dill >= 0.2.8 |
| if not getattr(dill, '_dill', None): |
| dill._dill = dill.dill |
| sys.modules['dill._dill'] = dill.dill |
| |
| dill_log = getattr(dill.dill, 'log', None) |
| |
| # dill v0.3.6 changed the attribute name from 'log' to 'logger' |
| if not dill_log: |
| dill_log = getattr(dill.dill, 'logger') |
| |
| |
| def _is_nested_class(cls): |
| """Returns true if argument is a class object that appears to be nested.""" |
| return ( |
| isinstance(cls, type) and cls.__module__ is not None and |
| cls.__module__ != 'builtins' and |
| cls.__name__ not in sys.modules[cls.__module__].__dict__) |
| |
| |
| def _find_containing_class(nested_class): |
| """Finds containing class of a nested class passed as argument.""" |
| |
| seen = set() |
| |
| def _find_containing_class_inner(outer): |
| if outer in seen: |
| return None |
| seen.add(outer) |
| for k, v in outer.__dict__.items(): |
| if v is nested_class: |
| return outer, k |
| elif isinstance(v, type) and hasattr(v, '__dict__'): |
| res = _find_containing_class_inner(v) |
| if res: return res |
| |
| return _find_containing_class_inner(sys.modules[nested_class.__module__]) |
| |
| |
| def _dict_from_mappingproxy(mp): |
| d = mp.copy() |
| d.pop('__dict__', None) |
| d.pop('__prepare__', None) |
| d.pop('__weakref__', None) |
| return d |
| |
| |
| def _nested_type_wrapper(fun): |
| """A wrapper for the standard pickler handler for class objects. |
| |
| Args: |
| fun: Original pickler handler for type objects. |
| |
| Returns: |
| A wrapper for type objects that handles nested classes. |
| |
| The wrapper detects if an object being pickled is a nested class object. |
| For nested class object only it will save the containing class object so |
| the nested structure is recreated during unpickle. |
| """ |
| def wrapper(pickler, obj): |
| # When the nested class is defined in the __main__ module we do not have to |
| # do anything special because the pickler itself will save the constituent |
| # parts of the type (i.e., name, base classes, dictionary) and then |
| # recreate it during unpickling. |
| if _is_nested_class(obj) and obj.__module__ != '__main__': |
| containing_class_and_name = _find_containing_class(obj) |
| if containing_class_and_name is not None: |
| return pickler.save_reduce(getattr, containing_class_and_name, obj=obj) |
| try: |
| return fun(pickler, obj) |
| except dill.dill.PicklingError: |
| # pylint: disable=protected-access |
| return pickler.save_reduce( |
| dill.dill._create_type, |
| ( |
| type(obj), |
| obj.__name__, |
| obj.__bases__, |
| _dict_from_mappingproxy(obj.__dict__)), |
| obj=obj) |
| # pylint: enable=protected-access |
| |
| return wrapper |
| |
| |
| # Monkey patch the standard pickler dispatch table entry for type objects. |
| # Dill, for certain types, defers to the standard pickler (including type |
| # objects). We wrap the standard handler using type_wrapper() because |
| # for nested class we want to pickle the actual enclosing class object so we |
| # can recreate it during unpickling. |
| # TODO(silviuc): Make sure we submit the fix upstream to GitHub dill project. |
| dill.dill.Pickler.dispatch[type] = _nested_type_wrapper( |
| dill.dill.Pickler.dispatch[type]) |
| |
| |
| # Dill pickles generators objects without complaint, but unpickling produces |
| # TypeError: object.__new__(generator) is not safe, use generator.__new__() |
| # on some versions of Python. |
| def _reject_generators(unused_pickler, unused_obj): |
| raise TypeError("can't (safely) pickle generator objects") |
| |
| |
| dill.dill.Pickler.dispatch[types.GeneratorType] = _reject_generators |
| |
| # This if guards against dill not being full initialized when generating docs. |
| if 'save_module' in dir(dill.dill): |
| |
| # Always pickle non-main modules by name. |
| old_save_module = dill.dill.save_module |
| |
| @dill.dill.register(dill.dill.ModuleType) |
| def save_module(pickler, obj): |
| if dill.dill.is_dill(pickler) and obj is pickler._main: |
| return old_save_module(pickler, obj) |
| else: |
| dill_log.info('M2: %s' % obj) |
| # pylint: disable=protected-access |
| pickler.save_reduce( |
| dill.dill._import_module, (obj.__name__, ), obj=obj) |
| # pylint: enable=protected-access |
| dill_log.info('# M2') |
| |
| # Pickle module dictionaries (commonly found in lambda's globals) |
| # by referencing their module. |
| old_save_module_dict = dill.dill.save_module_dict |
| known_module_dicts: Dict[int, Tuple[types.ModuleType, Dict[str, Any]]] = {} |
| |
| @dill.dill.register(dict) |
| def new_save_module_dict(pickler, obj): |
| obj_id = id(obj) |
| if not known_module_dicts or '__file__' in obj or '__package__' in obj: |
| if obj_id not in known_module_dicts: |
| # Trigger loading of lazily loaded modules (such as pytest vendored |
| # modules). |
| # This pass over sys.modules needs to iterate on a copy of sys.modules |
| # since lazy loading modifies the dictionary, hence the use of list(). |
| for m in list(sys.modules.values()): |
| try: |
| _ = m.__dict__ |
| except AttributeError: |
| pass |
| |
| for m in list(sys.modules.values()): |
| try: |
| if (m and m.__name__ != '__main__' and |
| isinstance(m, dill.dill.ModuleType)): |
| d = m.__dict__ |
| known_module_dicts[id(d)] = m, d |
| except AttributeError: |
| # Skip modules that do not have the __name__ attribute. |
| pass |
| if obj_id in known_module_dicts and dill.dill.is_dill(pickler): |
| m = known_module_dicts[obj_id][0] |
| try: |
| # pylint: disable=protected-access |
| dill.dill._import_module(m.__name__) |
| return pickler.save_reduce( |
| getattr, (known_module_dicts[obj_id][0], '__dict__'), obj=obj) |
| except (ImportError, AttributeError): |
| return old_save_module_dict(pickler, obj) |
| else: |
| return old_save_module_dict(pickler, obj) |
| |
| dill.dill.save_module_dict = new_save_module_dict |
| |
| def _nest_dill_logging(): |
| """Prefix all dill logging with its depth in the callstack. |
| |
| Useful for debugging pickling of deeply nested structures. |
| """ |
| old_log_info = dill_log.info |
| |
| def new_log_info(msg, *args, **kwargs): |
| old_log_info( |
| ('1 2 3 4 5 6 7 8 9 0 ' * 10)[:len(traceback.extract_stack())] + msg, |
| *args, |
| **kwargs) |
| |
| dill_log.info = new_log_info |
| |
| |
| # Turn off verbose logging from the dill pickler. |
| logging.getLogger('dill').setLevel(logging.WARN) |
| |
| |
| def dumps( |
| o, |
| enable_trace=True, |
| use_zlib=False, |
| enable_best_effort_determinism=False) -> bytes: |
| """For internal use only; no backwards-compatibility guarantees.""" |
| with _pickle_lock: |
| if enable_best_effort_determinism: |
| old_save_set = dill.dill.Pickler.dispatch[set] |
| old_save_frozenset = dill.dill.Pickler.dispatch[frozenset] |
| dill.dill.pickle(set, save_set) |
| dill.dill.pickle(frozenset, save_frozenset) |
| try: |
| s = dill.dumps(o, byref=settings['dill_byref']) |
| except Exception: # pylint: disable=broad-except |
| if enable_trace: |
| dill.dill._trace(True) # pylint: disable=protected-access |
| s = dill.dumps(o, byref=settings['dill_byref']) |
| else: |
| raise |
| finally: |
| dill.dill._trace(False) # pylint: disable=protected-access |
| if enable_best_effort_determinism: |
| dill.dill.pickle(set, old_save_set) |
| dill.dill.pickle(frozenset, old_save_frozenset) |
| |
| # Compress as compactly as possible (compresslevel=9) to decrease peak memory |
| # usage (of multiple in-memory copies) and to avoid hitting protocol buffer |
| # limits. |
| # WARNING: Be cautious about compressor change since it can lead to pipeline |
| # representation change, and can break streaming job update compatibility on |
| # runners such as Dataflow. |
| if use_zlib: |
| c = zlib.compress(s, 9) |
| else: |
| c = bz2.compress(s, compresslevel=9) |
| del s # Free up some possibly large and no-longer-needed memory. |
| |
| return base64.b64encode(c) |
| |
| |
| def loads(encoded, enable_trace=True, use_zlib=False): |
| """For internal use only; no backwards-compatibility guarantees.""" |
| |
| c = base64.b64decode(encoded) |
| |
| if use_zlib: |
| s = zlib.decompress(c) |
| else: |
| s = bz2.decompress(c) |
| |
| del c # Free up some possibly large and no-longer-needed memory. |
| |
| with _pickle_lock: |
| try: |
| return dill.loads(s) |
| except Exception: # pylint: disable=broad-except |
| if enable_trace: |
| dill.dill._trace(True) # pylint: disable=protected-access |
| return dill.loads(s) |
| else: |
| raise |
| finally: |
| dill.dill._trace(False) # pylint: disable=protected-access |
| |
| |
| def dump_session(file_path): |
| """For internal use only; no backwards-compatibility guarantees. |
| |
| Pickle the current python session to be used in the worker. |
| |
| Note: Due to the inconsistency in the first dump of dill dump_session we |
| create and load the dump twice to have consistent results in the worker and |
| the running session. Check: https://github.com/uqfoundation/dill/issues/195 |
| """ |
| with _pickle_lock: |
| dill.dump_session(file_path) |
| dill.load_session(file_path) |
| return dill.dump_session(file_path) |
| |
| |
| def load_session(file_path): |
| with _pickle_lock: |
| return dill.load_session(file_path) |
| |
| |
| def override_pickler_hooks(extend=True): |
| """ Extends the dill library hooks into that of the standard pickler library. |
| |
| If false all hooks that dill overrides will be removed. |
| If true dill hooks will be injected into the pickler library dispatch_table. |
| """ |
| dill.extend(extend) |
