pytorch

958d9a83 - [fx/package] make GraphModules packageable (#51976)

[fx/package] make GraphModules packageable (#51976) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/51976 FX serializes things by serializing Python code as a string and exec'ing it on load. This accomplishes one goal (we don't have to pickle the graph object directly) but breaks the pickle abstraction in ways that are not composable with `torch.package`. In particular: 1. `forward` is serialized by saving Python code. On load, it's installed by `exec`ing that code. This `exec` call needs to have the right importer installed, otherwise it will not import modules from the `torch.package` but instead import from the Python environment. 2. Any types/functions used are emitted as `import` statement in the generated Python code. These are effectively dynamic dependencies of the `GraphModule` being saved, and need to be registered as such so that the `PackageImporter` will package them. To address these, this PR introduces a new protocol for the importer/exporter: `__reduce_package__`. A class can implement `__reduce_package__` to customize how it is placed in the importer/exproter. It functions very similarly to `__reduce__`, except: - `__reduce_package__` takes one argument, which is the `PackageExporter` instance. Users can use this instance to save stuff to the package to implement their serialization. `__reduce__` takes no args. - Only the 2-element tuple version of the return value for `__reduce__` is supported (this could be extended if necessary). - When the reduction function is called on load, an additional argument is added to the beginning of the args tuple. This is the `PackageImporter` instance doing the loading. The `__reduce_package__` protocol is defined using `persistent_id` and `persistent_load`, which ensures that we can still use the cpickle implementation of the pickler by default. Pull Request resolved: #51971 Test Plan: Imported from OSS Reviewed By: zdevito Differential Revision: D26340591 Pulled By: suo fbshipit-source-id: 5872a7d22e832056399a7372bae8a57807717882