detect mocked module on saving pass (#70641)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/70641

Raises a not implemented error if we attempt to pickle an object which uses a mocked module. Now we no longer have to load the object to get this check, and instead happens right on the saving path.

Review History is on https://github.com/pytorch/pytorch/pull/69793 PR was moved to a different branch due to original branch getting corrupted.

Test Plan: Imported from OSS

Reviewed By: suo

Differential Revision: D33414365

Pulled By: PaliC

fbshipit-source-id: 6d72ddb05c47a3d060e9622ec0b6e5cd6c6c71c8

torch/package/package_exporter.py

@@ -592,6 +592,17 @@ class PackageExporter:
                     module, field = arg.split(" ")
                     if module not in all_dependencies:
                         all_dependencies.append(module)
+                    for pattern, pattern_info in self.patterns.items():
+                        if pattern.matches(module):
+                            if pattern_info.action == _ModuleProviderAction.MOCK:
+                                raise NotImplementedError(
+                                    f"Object '{field}' from module {module} was mocked out during packaging "
+                                    f"but is being used in resource - {resource} in package {package}. "
+                                    "If this error is happening during 'save_pickle', please ensure that your "
+                                    "pickled object doesn't contain any mocked objects."
+                                )
+                            else:
+                                break
 
             for module_name in all_dependencies:
                 self.dependency_graph.add_edge(name_in_dependency_graph, module_name)