PR: Expand pickle importing to support non-package C-modules
I am looking for reviewers for a patch I desire to how pickle loads non-package modules, like those created via PyModule_Create. Erroring while trying to serialize PyModule_Create attributes has been a recurring issue in PyTorch for at least 4 years:
- Functions in torch._C._nn and torch._C._onnx are not pickleable · Issue #38137 · pytorch/pytorch · GitHub
- Optimizer classes not `dill` picklable after using `torch.compile` · Issue #126154 · pytorch/pytorch · GitHub (my issue most recently)
The workaround for this issue has been to manually insert these modules into sys.modules:
- add submodules to sys.modules so their attributes can be pickled by mattip · Pull Request #53107 · pytorch/pytorch · GitHub
- Bind VariableFunctions as a module, not a class with static methods. by ezyang · Pull Request #38136 · pytorch/pytorch · GitHub
I am proposing that instead of always doing what equates to import package.c_module, we can try from package import c_module, and if we notice that "package.c_module" not in sys.modules after __import__("package", fromlist=["c_module"]), we can return getattr(package, "c_module").