| 28 | |
| 29 | # from https://github.com/pytorch/pytorch/blob/master/torch/utils/data/_utils/__init__.py |
| 30 | class _ExceptionWrapper: |
| 31 | MAGIC = b"EXC_MAGIC" |
| 32 | """Wraps an exception plus traceback to communicate across threads""" |
| 33 | def __init__(self, exc_info): |
| 34 | # It is important that we don't store exc_info, see |
| 35 | # NOTE [ Python Traceback Reference Cycle Problem ] |
| 36 | self.exc_type = exc_info[0] |
| 37 | self.exc_msg = "".join(traceback.format_exception(*exc_info)) |
| 38 | |
| 39 | def pack(self): |
| 40 | return self.MAGIC + pickle.dumps(self) |
| 41 | |
| 42 | @staticmethod |
| 43 | def unpack(dp): |
| 44 | if isinstance(dp, bytes) and dp.startswith(_ExceptionWrapper.MAGIC): |
| 45 | return pickle.loads(dp[len(_ExceptionWrapper.MAGIC):]) |
| 46 | |
| 47 | |
| 48 | def _repeat_iter(get_itr): |