diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 48fc85be..fb9491d1 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -549,7 +549,21 @@ def _load_checkpoint(torch_module: Any, checkpoint_path: Path, *, map_location: kwargs: dict[str, Any] = {"weights_only": False} if map_location is not None: kwargs["map_location"] = map_location - return torch_module.load(checkpoint_path, **kwargs) + try: + return torch_module.load(checkpoint_path, **kwargs) + except ModuleNotFoundError as error: + # yolov5/v7/v9 checkpoints pickle their model classes by reference (models.yolo, + # utils.*), so unpickling only resolves when the training repository is + # importable. Raised bare, this surfaces as "No module named 'models'" with no + # indication that the checkpoint, not roboflow, is what needs the extra module. + raise ModelPackagingError( + f"Could not load {checkpoint_path}: the checkpoint references the module " + f"'{error.name}', which is not importable here. Checkpoints produced by the " + "yolov5, yolov7 and yolov9 training repositories store their model classes by " + "reference, so they can only be unpickled from an environment where that " + "repository is importable. Run the upload from the training repository's " + "directory, or add it to PYTHONPATH." + ) from error def _legacy_yolo_args(opts: dict[str, Any], opt_path: Path) -> dict[str, Any]: diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index da39b96e..75cce399 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -1031,3 +1031,38 @@ def test_keypoint_checkpoint_rejected_before_export(self): if __name__ == "__main__": unittest.main() + + +class TestLoadCheckpointErrors(unittest.TestCase): + """`_load_checkpoint` translates unpicklable-module failures into a usable error.""" + + def test_module_not_found_becomes_actionable_model_packaging_error(self): + # yolov5/v7/v9 checkpoints pickle `models.yolo` by reference, so torch raises + # ModuleNotFoundError when the training repo is not importable. + torch_module = mock.Mock() + torch_module.load.side_effect = ModuleNotFoundError("No module named 'models'", name="models") + + with self.assertRaises(ModelPackagingError) as context: + model_processor._load_checkpoint(torch_module, Path("weights/best.pt")) + + message = str(context.exception) + self.assertIn("models", message) + self.assertIn("yolov5", message) + self.assertIn("PYTHONPATH", message) + self.assertIsInstance(context.exception.__cause__, ModuleNotFoundError) + + def test_successful_load_is_passed_through_unchanged(self): + torch_module = mock.Mock() + torch_module.load.return_value = {"model": "sentinel"} + + result = model_processor._load_checkpoint(torch_module, Path("weights/best.pt"), map_location="cpu") + + self.assertEqual(result, {"model": "sentinel"}) + torch_module.load.assert_called_once_with(Path("weights/best.pt"), weights_only=False, map_location="cpu") + + def test_other_errors_are_not_swallowed(self): + torch_module = mock.Mock() + torch_module.load.side_effect = RuntimeError("corrupt archive") + + with self.assertRaises(RuntimeError): + model_processor._load_checkpoint(torch_module, Path("weights/best.pt"))