Allow torch._C to be recognized a module in torch.package in pytorch/multipy (#91)

Summary:
Pull Request resolved: #91

This pr addresses #82 and #44.

A C extension module behaves a bit differently than a normal python package as it does not contain a `__path__` attribute. However, these modules still have information about their submodules. This PR also checks if a module is a C extension module  and checks if the module we are looking for is in it's children.

For example, if we are importing `torch._C._nn` we check if the parent `torch._C` is a C extension module if necessary, and then check if `torch._C._nn` is a proper child of `torch._C`.

The corresponding PR in torch.package is pytorch/pytorch#80917

Reviewed By: d4l3k

Differential Revision: D37630606

fbshipit-source-id: 59b5b7c291ccecef9f598a93c33719d348a339f5

Author: PaliC

multipy/package/package_importer_no_torch.py

@@ -6,6 +6,7 @@
 
 import builtins
 import importlib
+import importlib.machinery
 import inspect
 import io
 import linecache
@@ -378,18 +379,45 @@ def _install_on_parent(self, parent: str, name: str, module: types.ModuleType):
     def _do_find_and_load(self, name):
         path = None
         parent = name.rpartition(".")[0]
+        module_name_no_parent = name.rpartition(".")[-1]
         if parent:
             if parent not in self.modules:
                 self._gcd_import(parent)
             # Crazy side-effects!
             if name in self.modules:
                 return self.modules[name]
             parent_module = self.modules[parent]
+
             try:
                 path = parent_module.__path__  # type: ignore[attr-defined] # noqa
+
             except AttributeError:
-                msg = (_ERR_MSG + "; {!r} is not a package").format(name, parent)
-                raise ModuleNotFoundError(msg, name=name) from None
+                # when we attempt to import a package only containing pybinded files,
+                # the parent directory isn't always a package as defined by python,
+                # so we search if the package is actually there or not before calling the error.
+                if isinstance(
+                    parent_module.__loader__,
+                    importlib.machinery.ExtensionFileLoader,
+                ):
+                    if name not in self.extern_modules:
+                        msg = (
+                            _ERR_MSG
+                            + "; {!r} is a c extension module which was not externed. C extension modules \
+                            need to be externed by the PackageExporter in order to be used as we do not support interning them.}."
+                        ).format(name, name)
+                        raise ModuleNotFoundError(msg, name=name) from None
+                    if not isinstance(
+                        parent_module.__dict__.get(module_name_no_parent),
+                        types.ModuleType,
+                    ):
+                        msg = (
+                            _ERR_MSG
+                            + "; {!r} is a c extension package which does not contain {!r}."
+                        ).format(name, parent, name)
+                        raise ModuleNotFoundError(msg, name=name) from None
+                else:
+                    msg = (_ERR_MSG + "; {!r} is not a package").format(name, parent)
+                    raise ModuleNotFoundError(msg, name=name) from None
 
         module = self._load_module(name, parent)
 

multipy/test/package/test_dependency_api.py

@@ -7,11 +7,14 @@
 # Owner(s): ["oncall: package/deploy"]
 
 import importlib
+
 from io import BytesIO
 from sys import version_info
 from textwrap import dedent
 from unittest import skipIf
 
+import torch.nn
+
 from multipy.package import (
     EmptyMatchError,
     Importer,
@@ -369,6 +372,29 @@ def test_repackage_mocked_module(self):
         with self.assertRaises(NotImplementedError):
             foo2.package_a.get_something()
 
+    def test_externing_c_extension(self):
+        """Externing c extensions modules should allow us to still access them especially those found in torch._C."""
+
+        buffer = BytesIO()
+        # The C extension module in question is F.gelu which comes from torch._C._nn
+        model = torch.nn.TransformerEncoderLayer(
+            d_model=64,
+            nhead=2,
+            dim_feedforward=64,
+            dropout=1.0,
+            batch_first=True,
+            activation="gelu",
+            norm_first=True,
+        )
+        with PackageExporter(buffer) as e:
+            e.extern("torch.**")
+            e.intern("**")
+
+            e.save_pickle("model", "model.pkl", model)
+        buffer.seek(0)
+        imp = PackageImporter(buffer)
+        imp.load_pickle("model", "model.pkl")
+
 
 class TestDependencyAPINoTorch(TestDependencyAPI):
     def __init__(self, *args, **kwargs):