From 2d1133a8f96237fd3490955d974d125175c5c1b8 Mon Sep 17 00:00:00 2001 From: Roshan Sharma Date: Thu, 27 Aug 2026 14:02:09 -0400 Subject: [PATCH] fix: explain why a legacy YOLO checkpoint fails to unpickle yolov5/v7/v9 checkpoints pickle their model classes by reference, so torch.load only resolves them when the training repository is importable. Outside it, deploy() failed with a bare ModuleNotFoundError: No module named 'models' which names neither the checkpoint nor anything the user can act on. Translate it into ModelPackagingError, naming the missing module and what to do about it, and chain the original with `from error`. Other load failures are left to propagate unchanged. --- roboflow/util/model_processor.py | 16 +++++++++++++- tests/util/test_model_processor.py | 35 ++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) 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"))