diff --git a/runpod/serverless/modules/rp_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index 5451ae40e..997470bca 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -2,6 +2,7 @@ import os import threading +import time import uuid from dataclasses import dataclass from typing import Any, Dict, Optional, Union @@ -323,6 +324,8 @@ async def _sim_runsync(self, job_request: DefaultRequest) -> JobOutput: assigned_job_id = f"test-{uuid.uuid4()}" job = TestJob(id=assigned_job_id, input=job_request.input) + job_start = time.time() + if is_generator(self.config["handler"]): generator_output = run_job_generator(self.config["handler"], job.__dict__) job_output = {"output": []} @@ -331,19 +334,37 @@ async def _sim_runsync(self, job_request: DefaultRequest) -> JobOutput: else: job_output = await run_job(self.config["handler"], job.__dict__) - if job_output.get("error", None): - return jsonable_encoder( - {"id": job.id, "status": "FAILED", "error": job_output["error"]} - ) + execution_time = int((time.time() - job_start) * 1000) + + is_error = job_output.get("error", None) + status = "FAILED" if is_error else "COMPLETED" if job_request.webhook: + webhook_payload = { + "id": job.id, + "status": status, + "input": job_request.input, + "webhook": job_request.webhook, + "delayTime": 0, + "executionTime": execution_time, + } + if is_error: + webhook_payload["error"] = job_output["error"] + else: + webhook_payload["output"] = job_output["output"] + thread = threading.Thread( target=_send_webhook, - args=(job_request.webhook, job_output), + args=(job_request.webhook, webhook_payload), daemon=True, ) thread.start() + if is_error: + return jsonable_encoder( + {"id": job.id, "status": "FAILED", "error": job_output["error"]} + ) + return jsonable_encoder( {"id": job.id, "status": "COMPLETED", "output": job_output["output"]} ) @@ -359,6 +380,8 @@ async def _sim_stream(self, job_id: str) -> StreamOutput: job = TestJob(id=job_id, input=stashed_job.input) + job_start = time.time() + if is_generator(self.config["handler"]): generator_output = run_job_generator(self.config["handler"], job.__dict__) stream_accumulator = [] @@ -373,12 +396,22 @@ async def _sim_stream(self, job_id: str) -> StreamOutput: } ) + execution_time = int((time.time() - job_start) * 1000) job_list.remove(job.id) if stashed_job.webhook: + webhook_payload = { + "id": job_id, + "status": "COMPLETED", + "output": stream_accumulator, + "input": stashed_job.input, + "webhook": stashed_job.webhook, + "delayTime": 0, + "executionTime": execution_time, + } thread = threading.Thread( target=_send_webhook, - args=(stashed_job.webhook, stream_accumulator), + args=(stashed_job.webhook, webhook_payload), daemon=True, ) thread.start() @@ -398,6 +431,8 @@ async def _sim_status(self, job_id: str) -> JobOutput: job = TestJob(id=stashed_job.id, input=stashed_job.input) + job_start = time.time() + if is_generator(self.config["handler"]): generator_output = run_job_generator(self.config["handler"], job.__dict__) job_output = {"output": []} @@ -406,21 +441,38 @@ async def _sim_status(self, job_id: str) -> JobOutput: else: job_output = await run_job(self.config["handler"], job.__dict__) + execution_time = int((time.time() - job_start) * 1000) job_list.remove(job.id) - if job_output.get("error", None): - return jsonable_encoder( - {"id": job_id, "status": "FAILED", "error": job_output["error"]} - ) + is_error = job_output.get("error", None) + status = "FAILED" if is_error else "COMPLETED" if stashed_job.webhook: + webhook_payload = { + "id": job_id, + "status": status, + "input": stashed_job.input, + "webhook": stashed_job.webhook, + "delayTime": 0, + "executionTime": execution_time, + } + if is_error: + webhook_payload["error"] = job_output["error"] + else: + webhook_payload["output"] = job_output["output"] + thread = threading.Thread( target=_send_webhook, - args=(stashed_job.webhook, job_output), + args=(stashed_job.webhook, webhook_payload), daemon=True, ) thread.start() + if is_error: + return jsonable_encoder( + {"id": job_id, "status": "FAILED", "error": job_output["error"]} + ) + return jsonable_encoder( {"id": job_id, "status": "COMPLETED", "output": job_output["output"]} ) diff --git a/tests/test_serverless/test_modules/test_fastapi.py b/tests/test_serverless/test_modules/test_fastapi.py index 4830b9df0..a8234a0bf 100644 --- a/tests/test_serverless/test_modules/test_fastapi.py +++ b/tests/test_serverless/test_modules/test_fastapi.py @@ -220,9 +220,34 @@ def generator_handler(job): ) assert "error" in error_runsync_return - # Test webhook caller sent + # Test webhook caller sent with full cloud-parity payload asyncio.run(worker_api._sim_runsync(input_object_with_webhook)) assert mock_threading.Thread.called + webhook_kwargs = mock_threading.Thread.call_args + webhook_payload = webhook_kwargs[1]["args"][1] if "args" in webhook_kwargs[1] else webhook_kwargs[0][0][1] if webhook_kwargs[0] else webhook_kwargs[1]["args"][1] + # Extract payload from the positional args passed to Thread(target=..., args=(...)) + call_args = mock_threading.Thread.call_args + payload = call_args.kwargs.get("args", call_args[1].get("args", (None, None)))[1] + assert "id" in payload + assert payload["status"] == "COMPLETED" + assert "input" in payload + assert "webhook" in payload + assert "delayTime" in payload + assert "executionTime" in payload + assert "output" in payload + + # Test webhook fires on error too (fixes #410) + mock_threading.reset_mock() + error_input_with_webhook = rp_fastapi.DefaultRequest( + input={"test_input": "test_input"}, webhook="test_webhook" + ) + error_worker_api_wh = rp_fastapi.WorkerAPI({"handler": self.error_handler}) + asyncio.run(error_worker_api_wh._sim_runsync(error_input_with_webhook)) + assert mock_threading.Thread.called, "Webhook must fire on failed jobs" + call_args = mock_threading.Thread.call_args + payload = call_args.kwargs.get("args", call_args[1].get("args", (None, None)))[1] + assert payload["status"] == "FAILED" + assert "error" in payload @pytest.mark.asyncio @@ -330,10 +355,18 @@ def test_status(self): "output": {"result": "success"}, } - # Test webhook caller sent + # Test webhook caller sent with full cloud-parity payload asyncio.run(worker_api._sim_run(input_object_with_webhook)) asyncio.run(worker_api._sim_status("test-123")) assert mock_threading.Thread.called + call_args = mock_threading.Thread.call_args + payload = call_args.kwargs.get("args", call_args[1].get("args", (None, None)))[1] + assert "id" in payload + assert payload["status"] == "COMPLETED" + assert "input" in payload + assert "webhook" in payload + assert "delayTime" in payload + assert "executionTime" in payload # Test with generator handler def generator_handler(job):