Expose private APIs to enable/disable pickling ScriptModules without RPC (#39631)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/39631
Background:
Currently, we cannot send ScriptModule over RPC as an argument.
Otherwise, it would hit the following error:
> _pickle.PickleError: ScriptModules cannot be deepcopied using
> copy.deepcopy or saved using torch.save. Mixed serialization of
> script and non-script modules is not supported. For purely
> script modules use my_script_module.save(<filename>) instead.
Failed attempt:
tried to install `torch.jit.ScriptModule` to RPC's
dispatch table, but it does not work as the dispatch table only
matches exact types and using base type `torch.jit.ScriptModule`
does not work for derived typed.
Current solution:
The current solution exposes `_enable_jit_rref_pickle` and
`_disable_jit_rref_pickle` APIs to toggle the `allowJitRRefPickle`
flag. See `test_pickle_script_module_with_rref` as an example.
Test Plan: Imported from OSS
Differential Revision: D21920870
Pulled By: mrshenli
fbshipit-source-id: 4d58afce5d0b4b81249b383c173488820b1a47d6
