diff --git a/runpod/endpoint/asyncio/asyncio_runner.py b/runpod/endpoint/asyncio/asyncio_runner.py index cb75144d1..e76ad1971 100644 --- a/runpod/endpoint/asyncio/asyncio_runner.py +++ b/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,15 @@ 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: + if self.job_status == "COMPLETED" and self.job_output is None: + # stream() caches a terminal status without the aggregate + # output; fetch it from /status before returning the result + await self._fetch_job() + raise_for_job_state(self.job_id, self.job_status, self.job_error) return self.job_output try: @@ -86,6 +94,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: diff --git a/runpod/endpoint/helpers.py b/runpod/endpoint/helpers.py index cbdb83f9b..d15dfeff0 100644 --- a/runpod/endpoint/helpers.py +++ b/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) diff --git a/runpod/endpoint/runner.py b/runpod/endpoint/runner.py index d8c487599..4bf6650b7 100644 --- a/runpod/endpoint/runner.py +++ b/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,17 @@ 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) + elif self.job_status == "COMPLETED" and self.job_output is None: + # stream() caches a terminal status without the aggregate + # output; fetch it from /status before returning the result + self._fetch_job() - 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 +246,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 +260,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( diff --git a/runpod/error.py b/runpod/error.py index 2410edf61..36bca3e3c 100644 --- a/runpod/error.py +++ b/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 diff --git a/tests/test_endpoint/test_asyncio_runner.py b/tests/test_endpoint/test_asyncio_runner.py index 77237956a..0312df89f 100644 --- a/tests/test_endpoint/test_asyncio_runner.py +++ b/tests/test_endpoint/test_asyncio_runner.py @@ -142,6 +142,42 @@ async def json_side_effect(): assert outputs == ["OUTPUT1", "OUTPUT2"] + async def test_output_after_stream_fetches_aggregate_result(self): + """output() after a completed stream fetches the aggregate result + from /status instead of returning the stream's empty output.""" + with patch( + "aiohttp.ClientSession", new_callable=AsyncMock + ) as mock_session_class: + mock_session = mock_session_class.return_value + mock_get = mock_session.get + mock_resp = AsyncMock() + + responses = [ + {"stream": [{"output": "OUTPUT1"}], "status": "IN_PROGRESS"}, + {"stream": [], "status": "COMPLETED"}, + {"output": "FINAL", "status": "COMPLETED"}, + ] + + async def json_side_effect(): + return responses.pop(0) if responses else { + "output": "FINAL", + "status": "COMPLETED", + } + + mock_resp.json.side_effect = json_side_effect + mock_get.return_value = mock_resp + + headers = { + "Content-Type": "application/json", + "Authorization": "Bearer MOCK_API_KEY", + "X-Request-ID": "job_id", + } + job = Job("endpoint_id", "job_id", mock_session, headers) + + outputs = [chunk async for chunk in job.stream()] + assert outputs == ["OUTPUT1"] + assert await job.output() == "FINAL" + async def test_cancel(self): """ Tests Job.cancel diff --git a/tests/test_endpoint/test_failed_jobs.py b/tests/test_endpoint/test_failed_jobs.py new file mode 100644 index 000000000..b87f288bc --- /dev/null +++ b/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" diff --git a/tests/test_endpoint/test_runner.py b/tests/test_endpoint/test_runner.py index 4fd199e44..bfca40012 100644 --- a/tests/test_endpoint/test_runner.py +++ b/tests/test_endpoint/test_runner.py @@ -373,3 +373,17 @@ def test_job_stream(self, mock_client): job = runner.Job("endpoint_id", "job_id", mock_client) output = list(job.stream()) self.assertEqual(output, ["Job output 1", "Job output 2"]) + + @patch("runpod.endpoint.runner.RunPodClient") + def test_output_after_stream_fetches_aggregate_result(self, mock_client): + """output() after a completed stream must fetch the aggregate result + from /status, not return the stream response's empty output.""" + mock_client.get.side_effect = [ + {"status": "IN_PROGRESS", "stream": [{"output": "chunk 1"}]}, + {"status": "COMPLETED", "stream": []}, + {"status": "COMPLETED", "output": "final result"}, + ] + + job = runner.Job("endpoint_id", "job_id", mock_client) + self.assertEqual(list(job.stream()), ["chunk 1"]) + self.assertEqual(job.output(), "final result")