From 75a387fe9bc5816a3e6865d15babfa2f8f6d8db7 Mon Sep 17 00:00:00 2001 From: Will Sackfield Date: Wed, 13 Aug 2025 13:09:55 -0400 Subject: [PATCH] Embellish the exception message on model error * Tell us the prediction ID and status in the error message --- replicate/exceptions.py | 4 +++- tests/test_run.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/replicate/exceptions.py b/replicate/exceptions.py index f52f9fb4..dbf563f6 100644 --- a/replicate/exceptions.py +++ b/replicate/exceptions.py @@ -17,7 +17,9 @@ class ModelError(ReplicateException): def __init__(self, prediction: "Prediction") -> None: self.prediction = prediction - super().__init__(prediction.error) + super().__init__( + f"Prediction {prediction.id} {prediction.status}: {prediction.error}" + ) class ReplicateError(ReplicateException): diff --git a/tests/test_run.py b/tests/test_run.py index 93f7248b..fd666450 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -581,7 +581,7 @@ async def test_run_with_model_error(mock_replicate_api_token): }, ) - assert str(excinfo.value) == "OOM" + assert str(excinfo.value) == "Prediction p1 failed: OOM" assert excinfo.value.prediction.error == "OOM" assert excinfo.value.prediction.status == "failed"