diff --git a/roboflow/core/version.py b/roboflow/core/version.py index a4633a3b..f1db47b0 100644 --- a/roboflow/core/version.py +++ b/roboflow/core/version.py @@ -662,6 +662,9 @@ def deploy(self, model_type: str, model_path: str, filename: str = "weights/best model_type (str): The type of the model to be deployed. model_path (str): File path to the model weights to be uploaded. filename (str, optional): The name of the weights file. Defaults to "weights/best.pt". + + Raises: + RoboflowError: If Roboflow refuses the upload or the upload of the weights fails. """ bundle = package_custom_weights_interactive(model_type, model_path, filename, build_dir=model_path) @@ -676,41 +679,36 @@ def _upload_zip(self, model_type: str, model_path: str, model_file_name: str): f"{API_URL}/{self.workspace}/{self.project}/{self.version}" f"/uploadModel?api_key={self.__api_key}&modelType={model_type}&nocache=true" ) + if res.status_code == 429: + raise rfapi.RoboflowError( + "This version already has a trained model. Please generate and" + f" train a new version in order to upload model to Roboflow. Response: {res.text}", + status_code=res.status_code, + ) try: - if res.status_code == 429: - raise RuntimeError( - "This version already has a trained model. Please generate and" - " train a new version in order to upload model to Roboflow." - ) - else: - res.raise_for_status() - except Exception as e: - print(f"An error occured when getting the model upload URL: {e}") - return - - res = requests.put( - res.json()["url"], - data=open(os.path.join(model_path, model_file_name), "rb"), - ) + res.raise_for_status() + except requests.HTTPError as e: + raise rfapi.RoboflowError( + f"An error occurred when getting the model upload URL: {e}. Response: {res.text}", + status_code=res.status_code, + ) from e + + with open(os.path.join(model_path, model_file_name), "rb") as model_file: + res = requests.put(res.json()["url"], data=model_file) try: res.raise_for_status() - - if self.public: - print( - f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/{self.version}" - ) - print( - "Share your model with the world at:" - f" {UNIVERSE_URL}/{self.workspace}/{self.project}/" - f"model/{self.version}" - ) - else: - print( - f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/{self.version}" - ) - - except Exception as e: - print(f"An error occured when uploading the model: {e}") + except requests.HTTPError as e: + raise rfapi.RoboflowError( + f"An error occurred when uploading the model: {e}", status_code=res.status_code + ) from e + + print(f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/{self.version}") + if self.public: + print( + "Share your model with the world at:" + f" {UNIVERSE_URL}/{self.workspace}/{self.project}/" + f"model/{self.version}" + ) def __download_zip(self, link, location, format): """ diff --git a/roboflow/core/workspace.py b/roboflow/core/workspace.py index 081e9ac6..e9c6d308 100644 --- a/roboflow/core/workspace.py +++ b/roboflow/core/workspace.py @@ -907,6 +907,9 @@ def deploy_model( model_path (str): File path to the model weights to be uploaded. project_ids (list[str]): List of project IDs to deploy the model to. filename (str, optional): The name of the weights file. Defaults to "weights/best.pt". + + Raises: + RoboflowError: If Roboflow refuses the upload or the upload of the weights fails. """ from roboflow.util.model_processor import ( @@ -945,32 +948,24 @@ def _upload_zip( ) try: res.raise_for_status() - except Exception as e: - error_message = str(e) - status_code = str(res.status_code) - - print("\n\033[91m❌ ERROR\033[0m: Failed to get model deployment URL") - print("\033[93mDetails\033[0m:", error_message) - print("\033[93mStatus\033[0m:", status_code) - print(f"\033[93mResponse\033[0m:\n{res.text}\n") - return + except HTTPError as e: + raise RoboflowError( + f"Failed to get model deployment URL: {e}. Response: {res.text}", status_code=res.status_code + ) from e # Upload the model to the signed URL - res = requests.put( - res.json()["url"], - data=open(os.path.join(model_path, model_file_name), "rb"), - ) + with open(os.path.join(model_path, model_file_name), "rb") as model_file: + res = requests.put(res.json()["url"], data=model_file) try: res.raise_for_status() + except HTTPError as e: + raise RoboflowError(f"An error occurred when uploading the model: {e}", status_code=res.status_code) from e - for project_id in project_ids: - print( - f"View the status of your deployment for project {project_id} at:" - f" {APP_URL}/{self.url}/{project_id}/models" - ) - - except Exception as e: - print(f"An error occured when uploading the model: {e}") + for project_id in project_ids: + print( + f"View the status of your deployment for project {project_id} at:" + f" {APP_URL}/{self.url}/{project_id}/models" + ) def search( self, diff --git a/tests/test_model_upload.py b/tests/test_model_upload.py new file mode 100644 index 00000000..bf1c46b1 --- /dev/null +++ b/tests/test_model_upload.py @@ -0,0 +1,155 @@ +"""Failed model uploads raise instead of printing an error and returning normally.""" + +import contextlib +import io +import json +import os +import tempfile +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import responses + +from roboflow.adapters.rfapi import RoboflowError +from roboflow.config import API_URL +from roboflow.core.workspace import Workspace +from tests.helpers import get_version + +SIGNED_URL = "https://storage.example.com/upload" + + +class _WeightsDirMixin(unittest.TestCase): + def setUp(self): + tmpdir = tempfile.TemporaryDirectory() + self.addCleanup(tmpdir.cleanup) + self.model_path = tmpdir.name + with open(os.path.join(self.model_path, "roboflow_deploy.zip"), "wb") as f: + f.write(b"weights") + + +class TestVersionDeploy(_WeightsDirMixin): + UPLOAD_MODEL_URL = f"{API_URL}/Test Workspace Name/Test Project Name/1/uploadModel" + + def deploy(self): + bundle = SimpleNamespace(model_type="yolov8n", archive_path=SimpleNamespace(name="roboflow_deploy.zip")) + with patch("roboflow.core.version.package_custom_weights_interactive", return_value=bundle): + get_version().deploy("yolov8n", self.model_path, "weights/best.pt") + + @responses.activate + def test_upload_url_error_raises(self): + responses.add(responses.GET, self.UPLOAD_MODEL_URL, json={"error": "nope"}, status=500) + + with self.assertRaises(RoboflowError) as ctx: + self.deploy() + + self.assertEqual(ctx.exception.status_code, 500) + self.assertIn("model upload URL", str(ctx.exception)) + + @responses.activate + def test_429_raises(self): + responses.add(responses.GET, self.UPLOAD_MODEL_URL, status=429) + + with self.assertRaises(RoboflowError) as ctx: + self.deploy() + + self.assertEqual(ctx.exception.status_code, 429) + self.assertIn("already has a trained model", str(ctx.exception)) + + @responses.activate + def test_weights_upload_error_raises(self): + responses.add(responses.GET, self.UPLOAD_MODEL_URL, json={"url": SIGNED_URL}) + responses.add(responses.PUT, SIGNED_URL, status=403) + + with self.assertRaises(RoboflowError) as ctx: + self.deploy() + + self.assertEqual(ctx.exception.status_code, 403) + + @responses.activate + def test_success_prints_status_url(self): + responses.add(responses.GET, self.UPLOAD_MODEL_URL, json={"url": SIGNED_URL}) + responses.add(responses.PUT, SIGNED_URL, status=200) + + with patch("builtins.print") as mock_print: + self.deploy() + + printed = " ".join(str(call.args[0]) for call in mock_print.call_args_list) + self.assertIn("View the status of your deployment at:", printed) + + +class TestWorkspaceUploadZip(_WeightsDirMixin): + PREPARE_URL = f"{API_URL}/test-ws/models/prepareUpload" + + def upload(self): + workspace = Workspace( + {"workspace": {"name": "Test", "projects": [], "url": "test-ws"}}, + api_key="test-key", + default_workspace="test-ws", + model_format="yolov8", + ) + workspace._upload_zip("yolov8n", self.model_path, ["proj"], "my-model", "roboflow_deploy.zip") + + @responses.activate + def test_deployment_url_error_raises(self): + responses.add(responses.POST, self.PREPARE_URL, json={"error": "nope"}, status=400) + + with self.assertRaises(RoboflowError) as ctx: + self.upload() + + self.assertEqual(ctx.exception.status_code, 400) + self.assertIn("nope", str(ctx.exception)) + + @responses.activate + def test_weights_upload_error_raises(self): + responses.add(responses.POST, self.PREPARE_URL, json={"url": SIGNED_URL}) + responses.add(responses.PUT, SIGNED_URL, status=500) + + with self.assertRaises(RoboflowError) as ctx: + self.upload() + + self.assertEqual(ctx.exception.status_code, 500) + + @responses.activate + def test_success_prints_status_url(self): + responses.add(responses.POST, self.PREPARE_URL, json={"url": SIGNED_URL}) + responses.add(responses.PUT, SIGNED_URL, status=200) + + with patch("builtins.print") as mock_print: + self.upload() + + printed = " ".join(str(call.args[0]) for call in mock_print.call_args_list) + self.assertIn("View the status of your deployment for project proj", printed) + + +class TestCliModelUpload(_WeightsDirMixin): + @patch("roboflow.Roboflow") + def test_failed_upload_exits_non_zero(self, mock_rf_cls): + from roboflow.cli.handlers.model import _upload_model + + version = mock_rf_cls.return_value.workspace.return_value.project.return_value.version.return_value + version.deploy.side_effect = RoboflowError("An error occurred when uploading the model: 500", status_code=500) + args = SimpleNamespace( + json=True, + api_key="test-key", + workspace="test-ws", + project=["proj"], + version_number=1, + model_type="yolov8n", + model_path=self.model_path, + filename="weights/best.pt", + model_name=None, + ) + stdout, stderr = io.StringIO(), io.StringIO() + + with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + with self.assertRaises(SystemExit) as ctx: + _upload_model(args) + + self.assertEqual(ctx.exception.code, 1) + self.assertNotIn("uploaded", stdout.getvalue()) + self.assertIn("uploading the model", json.loads(stderr.getvalue())["error"]["message"]) + + +if __name__ == "__main__": + unittest.main()