[jit] fix segfault in attribute lookup on loaded ScriptModules

suo · suo · commit a0a74f0a0589 · 2020-08-19T12:44:02.000-07:00
The IR emitter looks for attributes on modules like: 1. Check the JIT type for the attribute 2. Check the originating Python class, in order to fulfill requests for, e.g. static methods or ignored methods. In the case where you do: ``` inner_module = torch.jit.load("inner.pt") wrapped = Wrapper(inner_module) # wrap the loaded ScriptModule in an nn.Module torch.jit.script(wrapped) ``` The IR emitter may check for attributes on `inner_module`. There is no originating Python class for `inner_module`, since it was directly compiled from the serialized format. Due to a bug in the code, we don't guard for this case an a segfault results if the wrapper asks for an undefined attribute. The lookup in this case looks like: 1. Check the JIT type for the attribute (not there!) 2. Check the originating Python class (this is a nullptr! segfault!) This PR guards this case and properly just raises an attribute missing compiler error instead of segfaulting. ghstack-source-id: a27e4a4 Pull Request resolved: #43284
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -15158,6 +15158,30 @@ def parameter_script(x: torch.nn.Parameter):
         input = torch.ones(2, 2)
         self.assertEqual(input, parameter_script(input))
 
+    def test_save_load_attr_error(self):
+        class Inner(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return x
+
+        class Wrapper(nn.Module):
+            def __init__(self, inner):
+                super().__init__()
+                self.inner = inner
+
+            def forward(self, x):
+                # this attribute doesn't exist on `Inner`
+                return self.inner.b(x)
+
+        inner_module = torch.jit.script(Inner())
+        inner_module = self.getExportImportCopy(inner_module)
+        wrapped = Wrapper(inner_module)
+        # This should properly complain that `self.inner` doesn't have the attribute `b`
+        with self.assertRaisesRegex(RuntimeError, 'has no attribute'):
+            torch.jit.script(wrapped)
+
 
 # known to be failing in tracer
 EXCLUDE_TRACED = {
diff --git a/torch/csrc/jit/frontend/concrete_module_type.cpp b/torch/csrc/jit/frontend/concrete_module_type.cpp
@@ -128,7 +128,10 @@ TypePtr ConcreteModuleType::getJitType() const {
   return jitType_;
 }
 
-py::object ConcreteModuleType::getPyClass() const {
+c10::optional<py::object> ConcreteModuleType::getPyClass() const {
+  if (!data_.pyClass_) {
+    return c10::nullopt;
+  }
   return data_.pyClass_;
 }
 
diff --git a/torch/csrc/jit/frontend/concrete_module_type.h b/torch/csrc/jit/frontend/concrete_module_type.h
@@ -58,6 +58,7 @@ class ConcreteModuleType;
 class VISIBILITY_HIDDEN ConcreteModuleTypeBuilder {
  public:
   explicit ConcreteModuleTypeBuilder(py::object pyClass) {
+    TORCH_INTERNAL_ASSERT(pyClass);
     pyClass_ = std::move(pyClass);
   }
   void addConstant(std::string name, py::object value);
@@ -192,7 +193,7 @@ class VISIBILITY_HIDDEN ConcreteModuleType {
   static std::shared_ptr<ConcreteModuleType> fromJitType(TypePtr type);
 
   TypePtr getJitType() const;
-  py::object getPyClass() const;
+  c10::optional<py::object> getPyClass() const;
   IterableModuleKind getIterableModuleKind() const;
   c10::optional<std::vector<std::string>> findOverloads(
       const std::string& name) const;
diff --git a/torch/csrc/jit/python/python_sugared_value.cpp b/torch/csrc/jit/python/python_sugared_value.cpp
@@ -535,21 +535,26 @@ std::shared_ptr<SugaredValue> ModuleValue::tryGetAttr(
   // 5. Check if it's an attribute of the original Python class that this
   // ScriptModule was derived from. The only class attributes we handle are
   // methods.
+  const auto maybePyClass = concreteType_->getPyClass();
+  if (!maybePyClass) {
+    // ConcreteType doesn't always have an originating Python class, e.g. if it
+    // was derived from a serialized ScriptModule. In this case, we've exhausted
+    // our options for attr lookup.
+    return nullptr;
+  }
   py::object unboundMethod = py::getattr(
-      concreteType_->getPyClass(),
-      field.c_str(),
-      pybind11::cast<pybind11::none>(Py_None));
+      *maybePyClass, field.c_str(), pybind11::cast<pybind11::none>(Py_None));
 
   if (py::isinstance<py::function>(unboundMethod)) {
-    bool isStaticFn = py::cast<bool>(
-        py::module::import("torch._jit_internal")
-            .attr("is_static_fn")(concreteType_->getPyClass(), field.c_str()));
+    bool isStaticFn =
+        py::cast<bool>(py::module::import("torch._jit_internal")
+                           .attr("is_static_fn")(*maybePyClass, field.c_str()));
     if (isStaticFn) {
       // Functions within the module annotated with @staticmethod do not need
       // binding.
-      py::object staticFn = py::module::import("torch._jit_internal")
-                                .attr("get_static_fn")(
-                                    concreteType_->getPyClass(), field.c_str());
+      py::object staticFn =
+          py::module::import("torch._jit_internal")
+              .attr("get_static_fn")(*maybePyClass, field.c_str());
       return toSugaredValue(staticFn, m, loc);
     }
     // For Python methods that we're trying to call directly, we need to bind

Original file line number	Diff line number	Diff line change
`@@ -128,7 +128,10 @@ TypePtr ConcreteModuleType::getJitType() const {`
`128`	`128`	`return jitType_;`
`129`	`129`	`}`
`130`	`130`
`131`		`-py::object ConcreteModuleType::getPyClass() const {`
	`131`	`+c10::optional<py::object> ConcreteModuleType::getPyClass() const {`
	`132`	`+ if (!data_.pyClass_) {`
	`133`	`+ return c10::nullopt;`
	`134`	`+ }`
`132`	`135`	`return data_.pyClass_;`
`133`	`136`	`}`
`134`	`137`