detect mocked module on saving pass (pytorch#70641)

Summary:
Pull Request resolved: pytorch#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 pytorch#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

Author: PaliC

test/package/test_dependency_api.py

@@ -182,16 +182,24 @@ def test_pickle_mocked(self):
         obj2 = package_a.PackageAObject(obj)
 
         buffer = BytesIO()
-        with PackageExporter(buffer) as he:
-            he.mock(include="package_a.subpackage")
-            he.intern("**")
-            he.save_pickle("obj", "obj.pkl", obj2)
+        with self.assertRaises(NotImplementedError):
+            with PackageExporter(buffer) as he:
+                he.mock(include="package_a.subpackage")
+                he.intern("**")
+                he.save_pickle("obj", "obj.pkl", obj2)
 
-        buffer.seek(0)
+    @skipIf(version_info < (3, 7), "mock uses __getattr__ a 3.7 feature")
+    def test_pickle_mocked_all(self):
+        import package_a.subpackage
 
-        hi = PackageImporter(buffer)
-        with self.assertRaises(NotImplementedError):
-            hi.load_pickle("obj", "obj.pkl")
+        obj = package_a.subpackage.PackageASubpackageObject()
+        obj2 = package_a.PackageAObject(obj)
+
+        buffer = BytesIO()
+        with PackageExporter(buffer) as he:
+            he.intern(include="package_a.**")
+            he.mock("**")
+            he.save_pickle("obj", "obj.pkl", obj2)
 
     def test_allow_empty_with_error(self):
         """If an error occurs during packaging, it should not be shadowed by the allow_empty error."""

torch/package/package_exporter.py

@@ -592,6 +592,17 @@ def save_pickle(
                     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)