__init__.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  1. from __future__ import absolute_import
  2. import os
  3. from pickle import PicklingError
  4. from ray.cloudpickle.cloudpickle import * # noqa
  5. from ray.cloudpickle.cloudpickle_fast import CloudPickler, dumps, dump # noqa
  6. # Conform to the convention used by python serialization libraries, which
  7. # expose their Pickler subclass at top-level under the "Pickler" name.
  8. Pickler = CloudPickler
  9. __version__ = '3.0.0'
  10. def _warn_msg(obj, method, exc):
  11. return (
  12. f"{method}({str(obj)}) failed."
  13. "\nTo check which non-serializable variables are captured "
  14. "in scope, re-run the ray script with 'RAY_PICKLE_VERBOSE_DEBUG=1'.")
  15. def dump_debug(obj, *args, **kwargs):
  16. try:
  17. return dump(obj, *args, **kwargs)
  18. except (TypeError, PicklingError) as exc:
  19. if os.environ.get("RAY_PICKLE_VERBOSE_DEBUG"):
  20. from ray.util.check_serialize import inspect_serializability
  21. inspect_serializability(obj)
  22. raise
  23. else:
  24. msg = _warn_msg(obj, "ray.cloudpickle.dump", exc)
  25. raise type(exc)(msg)
  26. def dumps_debug(obj, *args, **kwargs):
  27. try:
  28. return dumps(obj, *args, **kwargs)
  29. except (TypeError, PicklingError) as exc:
  30. if os.environ.get("RAY_PICKLE_VERBOSE_DEBUG"):
  31. from ray.util.check_serialize import inspect_serializability
  32. inspect_serializability(obj)
  33. raise
  34. else:
  35. msg = _warn_msg(obj, "ray.cloudpickle.dumps", exc)
  36. raise type(exc)(msg)