fix: explain why a legacy YOLO checkpoint fails to unpickle

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]:

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"))