Error with typing.AbstractContextManager in Python 3.8 and newer dill versions (>=0.3.6)

[cnellington]

Encountered only in Python 3.8, occurs in dill>=0.3.6. No issue with Python 3.9+ or dill<0.3.6. Normal python pickle succeeds while dill fails. Came across this while working on the contextualized-ml package. Quickest steps to reproduce:

Setup

# Using Python 3.8.19
pip install git+https://github.com/cnellington/Contextualized@12843bbab6bd2f63db1dfdc5247ab865ca47cfa6
pip install dill==0.3.8

Run

import dill
import torch
from contextualized.easy import ContextualizedRegressor

# Succeeds
model = ContextualizedRegressor()
torch.save(model, open('test_torch_model.pkl', 'wb'))
torch.load(open('test_torch_model.pkl', 'rb'))

# Fails for Python 3.8 and dill>=0.3.6. Succeeds for dill==0.3.5
torch.save(model, open('test_dill_model.pkl', 'wb'), pickle_module=dill)
torch.load(open('test_dill_model.pkl', 'rb'), pickle_module=dill)

Error

Traceback (most recent call last):
  File "test_dill.py", line 10, in <module>
    torch.save(model, open('test_dill_model.pkl', 'wb'), pickle_module=dill)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/torch/serialization.py", line 619, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol, _disable_byteorder_record)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/torch/serialization.py", line 831, in _save
    pickler.dump(obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 420, in dump
    StockPickler.dump(self, obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 487, in dump
    self.save(obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 414, in save
    StockPickler.save(self, obj, save_persistent_id)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 603, in save
    self.save_reduce(obj=obj, *rv)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 687, in save_reduce
    save(cls)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 414, in save
    StockPickler.save(self, obj, save_persistent_id)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 1832, in save_type
    _save_with_postproc(pickler, (_create_type, (
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 1098, in _save_with_postproc
    pickler.save_reduce(*reduction, obj=obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 692, in save_reduce
    save(args)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 414, in save
    StockPickler.save(self, obj, save_persistent_id)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 901, in save_tuple
    save(element)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 414, in save
    StockPickler.save(self, obj, save_persistent_id)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 886, in save_tuple
    save(element)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 414, in save
    StockPickler.save(self, obj, save_persistent_id)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 1832, in save_type
    _save_with_postproc(pickler, (_create_type, (
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 1112, in _save_with_postproc
    pickler._batch_setitems(iter(source.items()))
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 414, in save
    StockPickler.save(self, obj, save_persistent_id)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/site-packages/dill/_dill.py", line 1297, in save_generic_alias
    StockPickler.save_global(pickler, obj, name=obj.__reduce__())
  File "/opt/homebrew/Caskroom/miniforge/base/envs/testdill/lib/python3.8/pickle.py", line 1070, in save_global
    raise PicklingError(
_pickle.PicklingError: Can't pickle typing.AbstractContextManager: it's not found as typing.AbstractContextManager

Happy to help more with digging into any underlying reasons in contextualized-ml since I manage this package. Two recent changes I made seem related to the error: replaced lambdas with partials and began adding type hints in the code.

[mmckerns]

I'm happy to work with you to figure it out, and determine if it's something that should be patched in your code or in dill (i.e. if there's an issue with pickling something from the standard library).

The first thing I noticed is that typing.AbstractContextManager doesn't exist in any current of the python releases, as far as I can tell (after a cursory look). There's ContextManager...

Python 3.12.3 (main, Apr 12 2024, 17:29:56) [Clang 13.1.6 (clang-1316.0.21.2.5)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import typing
>>> typing.ContextManager.mro()
[<class 'contextlib.AbstractContextManager'>, <class 'abc.ABC'>, <class 'object'>]
>>>

Same for the most recent 3.8 - 3.12.

Can you run your code again, but with dill.detect on?

>>> import dill
>>> dill.detect.trace(True)

The above should show the recursive pickling traceback, and potentially clarify where the difference is for 3.8 versus newer. Same for dill 0.3.5 versus 0.3.6.