Skip to content

Commit 75a387f

Browse files
committed
Embellish the exception message on model error
* Tell us the prediction ID and status in the error message
1 parent e5569a2 commit 75a387f

File tree

2 files changed

+4
-2
lines changed

2 files changed

+4
-2
lines changed

replicate/exceptions.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,9 @@ class ModelError(ReplicateException):
1717

1818
def __init__(self, prediction: "Prediction") -> None:
1919
self.prediction = prediction
20-
super().__init__(prediction.error)
20+
super().__init__(
21+
f"Prediction {prediction.id} {prediction.status}: {prediction.error}"
22+
)
2123

2224

2325
class ReplicateError(ReplicateException):

tests/test_run.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -581,7 +581,7 @@ async def test_run_with_model_error(mock_replicate_api_token):
581581
},
582582
)
583583

584-
assert str(excinfo.value) == "OOM"
584+
assert str(excinfo.value) == "Prediction p1 failed: OOM"
585585
assert excinfo.value.prediction.error == "OOM"
586586
assert excinfo.value.prediction.status == "failed"
587587

0 commit comments

Comments
 (0)