[TorchScript] clearer debug for ConcreteModuleType::findSubmoduleConcreteType (#166192)

Summary:
right now the log is just
```
RuntimeError: it != data_.modules_.end() INTERNAL ASSERT FAILED at "fbcode/caffe2/torch/csrc/jit/frontend/concrete_module_type.cpp":207, please report a bug to PyTorch.
```
we have no clue where the error happens
https://fb.workplace.com/groups/gpuinference/posts/789257990578348/?comment_id=789284783909002&reply_comment_id=789415260562621

Test Plan: UT

Reviewed By: jcwchen

Differential Revision: D80020093

Pull Request resolved: https://github.com/pytorch/pytorch/pull/166192
Approved by: https://github.com/gmagogsfm
This commit is contained in:
Chang Pan 2025-10-25 14:07:54 +00:00 committed by PyTorch MergeBot
parent 798a6d2be1
commit 74e53d0761
2 changed files with 55 additions and 1 deletions

View File

@ -0,0 +1,53 @@
# Owner(s): ["oncall: jit"]
import unittest
import torch
from torch.testing._internal.common_utils import raise_on_run_directly
class TestConcreteModuleTypeFindSubmodule(unittest.TestCase):
def test_error_message_includes_submodule_name(self):
class ChildModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 3)
def forward(self, x):
return self.linear(x)
class ParentModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.existing_child = ChildModule()
def forward(self, x):
return self.existing_child(x)
module = ParentModule()
scripted_module = torch.jit.script(module)
self.assertIsNotNone(scripted_module.existing_child)
# Now try to trigger the error by accessing a non-existent submodule
# through the internal ConcreteModuleType mechanism. This happens
# when the TorchScript compiler tries to resolve submodule references.
class BrokenModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.parent = ParentModule()
def forward(self, x):
return self.parent.missing_submodule(x)
broken_module = BrokenModule()
with self.assertRaises(RuntimeError) as context:
torch.jit.script(broken_module)
error_msg = str(context.exception)
self.assertIn("missing_submodule", error_msg.lower())
if __name__ == "__main__":
raise_on_run_directly("test/test_jit.py")

View File

@ -204,7 +204,8 @@ std::shared_ptr<ConcreteModuleType> ConcreteModuleType::
[&](const ConcreteModuleTypeBuilder::ModuleInfo& info) { [&](const ConcreteModuleTypeBuilder::ModuleInfo& info) {
return info.name_ == name; return info.name_ == name;
}); });
TORCH_INTERNAL_ASSERT(it != data_.modules_.end()); TORCH_INTERNAL_ASSERT(
it != data_.modules_.end(), "Cannot find submodule with name/key ", name);
return it->meta_; return it->meta_;
} }