fix(endpoint): raise JobError when a job fails

runpod/endpoint/asyncio/asyncio_runner.py

@@ -5,7 +5,7 @@
 import asyncio
 from typing import Any, Dict, Optional
 
-from runpod.endpoint.helpers import FINAL_STATES, is_completed
+from runpod.endpoint.helpers import FINAL_STATES, is_completed, raise_for_job_state
 from runpod.http_client import ClientSession
 
 
@@ -34,6 +34,7 @@ def __init__(self, endpoint_id: str, job_id: str, session: ClientSession, header
 
         self.job_status = None
         self.job_output = None
+        self.job_error = None
 
     async def _fetch_job(self, source: str = "status") -> Dict[str, Any]:
         """Returns the raw json of the status, raises an exception if invalid.
@@ -50,6 +51,7 @@ async def _fetch_job(self, source: str = "status") -> Dict[str, Any]:
         if is_completed(job_state["status"]):
             self.job_status = job_state["status"]
             self.job_output = job_state.get("output", None)
+            self.job_error = job_state.get("error", None)
 
         return job_state
 
@@ -75,9 +77,11 @@ async def output(self, timeout: int = 0) -> Any:
         Returns:
             Output of job
         Raises:
-            KeyError if job Failed
+            TimeoutError: The job did not finish within timeout.
+            JobError: The job ended FAILED, TIMED_OUT or CANCELLED.
         """
-        if self.job_output is not None:
+        if self.job_status is not None:
+            raise_for_job_state(self.job_id, self.job_status, self.job_error)
             return self.job_output
 
         try:
@@ -86,6 +90,7 @@ async def output(self, timeout: int = 0) -> Any:
             raise TimeoutError("Job timed out.") from exc
 
         job_data = await self._fetch_job()
+        raise_for_job_state(self.job_id, job_data["status"], job_data.get("error"))
         return job_data.get("output", None)
 
     async def stream(self) -> Any:

runpod/endpoint/helpers.py

@@ -1,5 +1,9 @@
 """ Helper functions for the Runpod Endpoint API. """
 
+from typing import Optional
+
+from runpod.error import JobError
+
 FINAL_STATES = ["COMPLETED", "FAILED", "TIMED_OUT"]
 
 # Exception Messages
@@ -15,3 +19,9 @@
 def is_completed(status: str) -> bool:
     """Returns true if status is one of the possible final states for a serverless request."""
     return status in ["COMPLETED", "FAILED", "TIMED_OUT", "CANCELLED"]
+
+
+def raise_for_job_state(job_id: str, status: str, error: Optional[str] = None) -> None:
+    """Raise JobError when a job ended in a final state other than COMPLETED."""
+    if status != "COMPLETED" and is_completed(status):
+        raise JobError(job_id, status, error)

runpod/endpoint/runner.py

@@ -13,6 +13,7 @@
     FINAL_STATES,
     UNAUTHORIZED_MSG,
     is_completed,
+    raise_for_job_state,
 )
 
 
@@ -115,6 +116,7 @@ def __init__(self, endpoint_id: str, job_id: str, client: RunPodClient):
 
         self.job_status = None
         self.job_output = None
+        self.job_error = None
 
     def _fetch_job(self, source: str = "status") -> Dict[str, Any]:
         """Returns the raw json of the status, raises an exception if invalid"""
@@ -124,6 +126,7 @@ def _fetch_job(self, source: str = "status") -> Dict[str, Any]:
         if is_completed(job_state["status"]):
             self.job_status = job_state["status"]
             self.job_output = job_state.get("output", None)
+            self.job_error = job_state.get("error", None)
 
         return job_state
 
@@ -140,6 +143,10 @@ def output(self, timeout: int = 0) -> Any:
 
         Args:
             timeout: The number of seconds to wait for the server to send data before giving up.
+
+        Raises:
+            TimeoutError: The job did not finish within timeout.
+            JobError: The job ended FAILED, TIMED_OUT or CANCELLED.
         """
         if timeout > 0:
             while not is_completed(self.status()):
@@ -148,10 +155,13 @@ def output(self, timeout: int = 0) -> Any:
                 if timeout <= 0:
                     raise TimeoutError("Job timed out.")
 
-        if self.job_output is not None:
-            return self.job_output
+        if self.job_status is None:
+            job_state = self._fetch_job()
+            if not is_completed(job_state["status"]):
+                return job_state.get("output", None)
 
-        return self._fetch_job().get("output", None)
+        raise_for_job_state(self.job_id, self.job_status, self.job_error)
+        return self.job_output
 
     def stream(self) -> Any:
         """Returns a generator that yields the output of the job request."""
@@ -232,6 +242,10 @@ def run_sync(
         Args:
             request_input: The input to pass into the endpoint.
             timeout: Maximum time to wait for the job to complete.
+
+        Raises:
+            TimeoutError: The job did not finish within timeout.
+            JobError: The job ended FAILED, TIMED_OUT or CANCELLED.
         """
         if not request_input.get("input"):
             request_input = {"input": request_input}
@@ -242,7 +256,10 @@ def run_sync(
             timeout=timeout
         )
 
-        if job_request["status"] in FINAL_STATES:
+        if is_completed(job_request["status"]):
+            raise_for_job_state(
+                job_request.get("id"), job_request["status"], job_request.get("error")
+            )
             return job_request.get("output", None)
 
         return Job(

runpod/error.py

@@ -58,3 +58,18 @@ def __init__(
     ):
         super().__init__(message, query, status_code=429, errors=errors)
         self.retry_after = retry_after
+
+
+class JobError(RunPodError):
+    """A serverless job ended FAILED, TIMED_OUT or CANCELLED.
+
+    Raised by Endpoint.run_sync and Job.output instead of returning None,
+    so a failed job can't be mistaken for a handler that returned None.
+    """
+
+    def __init__(self, job_id: str, status: str, error: Optional[str] = None):
+        detail = f": {error}" if error else ""
+        super().__init__(f"job {job_id} {status}{detail}")
+        self.job_id = job_id
+        self.status = status
+        self.error = error

tests/test_endpoint/test_failed_jobs.py

@@ -0,0 +1,89 @@
+"""a job that ends in a non-COMPLETED final state raises JobError."""
+
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+import runpod
+from runpod.endpoint.asyncio.asyncio_runner import Job as AsyncJob
+from runpod.endpoint.runner import Endpoint, Job
+from runpod.error import JobError
+
+UNSUCCESSFUL_FINAL_STATES = ["FAILED", "TIMED_OUT", "CANCELLED"]
+
+
+@pytest.fixture
+def endpoint():
+    runpod.api_key = "MOCK_API_KEY"
+    return Endpoint("ENDPOINT_ID")
+
+
+@pytest.mark.parametrize("status", UNSUCCESSFUL_FINAL_STATES)
+def test_run_sync_raises_with_status_and_error(endpoint, status):
+    response = {"id": "job-1", "status": status, "error": "handler blew up"}
+
+    with patch.object(runpod.endpoint.runner.RunPodClient, "_request") as request:
+        request.return_value = response
+        with pytest.raises(JobError) as caught:
+            endpoint.run_sync({"prompt": "hi"})
+
+    assert caught.value.job_id == "job-1"
+    assert caught.value.status == status
+    assert caught.value.error == "handler blew up"
+    assert "handler blew up" in str(caught.value)
+
+
+def test_run_sync_raises_when_polled_job_fails(endpoint):
+    responses = [
+        {"id": "job-1", "status": "IN_QUEUE"},
+        {"id": "job-1", "status": "FAILED", "error": "out of memory"},
+    ]
+
+    with (
+        patch.object(
+            runpod.endpoint.runner.RunPodClient, "_request", side_effect=responses
+        ),
+        patch("runpod.endpoint.runner.time.sleep"),
+    ):
+        with pytest.raises(JobError, match="out of memory"):
+            endpoint.run_sync({"prompt": "hi"}, timeout=5)
+
+
+def test_run_sync_still_returns_none_output_of_a_completed_job(endpoint):
+    with patch.object(runpod.endpoint.runner.RunPodClient, "_request") as request:
+        request.return_value = {"id": "job-1", "status": "COMPLETED"}
+        assert endpoint.run_sync({"prompt": "hi"}) is None
+
+
+def test_job_output_raises_on_every_call_after_failure():
+    client = Mock()
+    client.get.return_value = {"status": "FAILED", "error": "bad input"}
+    job = Job("ENDPOINT_ID", "job-1", client)
+
+    for _ in range(2):
+        with pytest.raises(JobError, match="bad input"):
+            job.output()
+
+
+def test_job_error_without_error_field_still_names_the_status():
+    client = Mock()
+    client.get.return_value = {"status": "TIMED_OUT"}
+    job = Job("ENDPOINT_ID", "job-1", client)
+
+    with pytest.raises(JobError, match="TIMED_OUT") as caught:
+        job.output()
+
+    assert caught.value.error is None
+
+
+async def test_async_job_output_raises_on_failure():
+    response = AsyncMock()
+    response.json.return_value = {"status": "FAILED", "error": "bad input"}
+    session = AsyncMock()
+    session.get.return_value = response
+    job = AsyncJob("ENDPOINT_ID", "job-1", session, headers={})
+
+    with pytest.raises(JobError, match="bad input") as caught:
+        await job.output(timeout=3)
+
+    assert caught.value.status == "FAILED"