fix: change async to threading

Author: justinmerrell

runpod/serverless/modules/rp_fastapi.py

@@ -3,12 +3,13 @@
 
 import os
 import uuid
-import asyncio
+import time
+import threading
 from dataclasses import dataclass
 from typing import Union, Optional, Dict, Any
 
-import aiohttp
 import uvicorn
+import requests
 from fastapi import FastAPI, APIRouter
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import RedirectResponse
@@ -85,31 +86,31 @@ class StreamOutput:
 
 
 # ------------------------------ Webhook Sender ------------------------------ #
-async def _send_webhook_async(url: str, payload: Dict[str, Any]) -> None:
+def _send_webhook(url: str, payload: Dict[str, Any]) -> None:
     """
-    Sends a webhook to the provided URL asynchronously. Retries once if the first attempt fails.
+    Sends a webhook to the provided URL. Retries once if the first attempt fails.
 
     Args:
         url (str): The URL to send the webhook to.
         payload (Dict[str, Any]): The JSON payload to send.
     """
-    async def attempt_send(session, url, payload):
+    def attempt_send(session, url, payload):
         try:
-            async with session.post(url, json=payload, timeout=10) as response:
-                response.raise_for_status()  # Raises exception for 4xx/5xx responses
-                return True
-        except (aiohttp.ClientError, aiohttp.http_exceptions.HttpProcessingError) as err:
+            response = session.post(url, json=payload, timeout=10)
+            response.raise_for_status()  # Raises exception for 4xx/5xx responses
+            return True
+        except requests.RequestException as err:
             print(f"Request to {url} failed: {err}")
             return False
 
-    async with aiohttp.ClientSession() as session:
-        if await attempt_send(session, url, payload):
+    with requests.Session() as session:
+        if attempt_send(session, url, payload):
             return True
 
         print("Retrying...")
-        await asyncio.sleep(1)  # Wait for 1 second before retrying
+        time.sleep(1)  # Wait for 1 second before retrying
 
-        if await attempt_send(session, url, payload):
+        if attempt_send(session, url, payload):
             return True
 
         print("Failed to send webhook after retry.")
@@ -240,7 +241,10 @@ async def _sim_runsync(self, job_request: DefaultRequest) -> JobOutput:
             })
 
         if job_request.webhook:
-            asyncio.create_task(_send_webhook_async(job_request.webhook, job_output))
+            thread = threading.Thread(
+                target=_send_webhook,
+                args=(job_request.webhook, job_output), daemon=True)
+            thread.start()
 
         return jsonable_encoder({
             "id": job.id,
@@ -276,7 +280,10 @@ async def _sim_stream(self, job_id: str) -> StreamOutput:
         job_list.remove_job(job.id)
 
         if stashed_job.webhook:
-            asyncio.create_task(_send_webhook_async(stashed_job.webhook, stream_accumulator))
+            thread = threading.Thread(
+                target=_send_webhook,
+                args=(stashed_job.webhook, stream_accumulator), daemon=True)
+            thread.start()
 
         return jsonable_encoder({
             "id": job_id,
@@ -315,7 +322,10 @@ async def _sim_status(self, job_id: str) -> JobOutput:
             })
 
         if stashed_job.webhook:
-            asyncio.create_task(_send_webhook_async(stashed_job.webhook, job_output))
+            thread = threading.Thread(
+                target=_send_webhook,
+                args=(stashed_job.webhook, job_output), daemon=True)
+            thread.start()
 
         return jsonable_encoder({
             "id": job_id,

tests/test_serverless/test_modules/test_fastapi.py

@@ -5,10 +5,10 @@
 import asyncio
 
 import unittest
-from unittest.mock import patch, Mock, AsyncMock
+from unittest.mock import patch, Mock, MagicMock
 import pytest
 
-import aiohttp
+import requests
 import runpod
 from runpod.serverless.modules import rp_fastapi
 
@@ -54,43 +54,31 @@ def test_start_serverless_with_realtime(self):
             os.environ.pop("RUNPOD_REALTIME_PORT")
             os.environ.pop("RUNPOD_ENDPOINT_ID")
 
-    @pytest.mark.asyncio
     def test_webhook_sender_success(self):
         """Test the webhook sender when the request is successful."""
-        loop = asyncio.get_event_loop()
-
-        module_location = "runpod.serverless.modules.rp_fastapi.aiohttp.ClientSession"
+        module_location = "runpod.serverless.modules.rp_fastapi.requests.Session.post"
 
-        with patch(f"{module_location}.post", new_callable=AsyncMock) as mock_post:
+        with patch(module_location, new_callable=MagicMock) as mock_post:
             # Simulate a successful response
-            mock_post.return_value.__aenter__.return_value.status = 200
+            mock_post.return_value.status_code = 200
 
-            # Directly await the function
-            success = asyncio.run(rp_fastapi._send_webhook_async(
-                "test_webhook", {"test": "output"}))
+            # Call the function
+            success = rp_fastapi._send_webhook("test_webhook", {"test": "output"})
             assert success is True
 
-        loop.close()
-
-    @pytest.mark.asyncio
     def test_webhook_sender_failure(self):
         """Test the webhook sender when the request fails."""
-        loop = asyncio.get_event_loop()
-
-        module_location = "runpod.serverless.modules.rp_fastapi.aiohttp.ClientSession"
+        module_location = "runpod.serverless.modules.rp_fastapi.requests.Session.post"
 
-        with patch(f"{module_location}.post", new_callable=AsyncMock) as mock_post:
-            # Configure the mock to raise an exception to simulate a 500 error
-            mock_post.return_value.__aenter__.return_value.raise_for_status.side_effect = aiohttp.ClientResponseError(  # pylint: disable=line-too-long
-                request_info=None, history=None, status=500)
+        with patch(module_location, new_callable=MagicMock) as mock_post:
+            # Configure the mock to simulate a failure (e.g., a 500 status code)
+            mock_post.return_value.raise_for_status.side_effect = requests.HTTPError()
+            mock_post.return_value.status_code = 500
 
-            # Directly await the function
-            success = asyncio.run(rp_fastapi._send_webhook_async(
-                "test_webhook", {"test": "output"}))
+            # Call the function
+            success = rp_fastapi._send_webhook("test_webhook", {"test": "output"})
             assert success is False
 
-        loop.close()
-
     @pytest.mark.asyncio
     def test_run(self):
         '''
@@ -195,9 +183,9 @@ def generator_handler(job):
             assert "error" in error_runsync_return
 
             # Test webhook caller sent
-            with patch(f"{module_location}._send_webhook_async", return_value=True) as mock_send_webhook:  # pylint: disable=line-too-long
+            with patch(f"{module_location}.threading") as mock_thread:
                 asyncio.run(worker_api._sim_runsync(input_object_with_webhook))
-                assert mock_send_webhook.called
+                assert mock_thread.Thread.called
 
         loop.close()
 
@@ -212,7 +200,8 @@ def test_stream(self):
         with patch(f"{module_location}.FastAPI", Mock()), \
                 patch(f"{module_location}.APIRouter", return_value=Mock()), \
                 patch(f"{module_location}.uvicorn", Mock()), \
-                patch(f"{module_location}.uuid.uuid4", return_value="123"):
+                patch(f"{module_location}.uuid.uuid4", return_value="123"), \
+                patch(f"{module_location}.threading"):
 
             default_input_object = rp_fastapi.DefaultRequest(
                 input={"test_input": "test_input"}
@@ -243,10 +232,10 @@ def test_stream(self):
             }
 
             # Test webhook caller sent
-            with patch(f"{module_location}._send_webhook_async", return_value=True) as mock_send_webhook:  # pylint: disable=line-too-long
+            with patch(f"{module_location}.threading", return_value=True) as mock_threading:
                 asyncio.run(worker_api._sim_run(input_object_with_webhook))
                 asyncio.run(worker_api._sim_stream("test-123"))
-                assert mock_send_webhook.called
+                assert mock_threading.Thread.called
 
             # Test with generator handler
             def generator_handler(job):
@@ -275,7 +264,8 @@ def test_status(self):
         with patch(f"{module_location}.FastAPI", Mock()), \
                 patch(f"{module_location}.APIRouter", return_value=Mock()), \
                 patch(f"{module_location}.uvicorn", Mock()), \
-                patch(f"{module_location}.uuid.uuid4", return_value="123"):
+                patch(f"{module_location}.uuid.uuid4", return_value="123"), \
+                patch(f"{module_location}.threading"):
 
             worker_api = rp_fastapi.WorkerAPI({"handler": self.handler})
 
@@ -306,10 +296,10 @@ def test_status(self):
             }
 
             # Test webhook caller sent
-            with patch(f"{module_location}._send_webhook_async", return_value=True) as mock_send_webhook:  # pylint: disable=line-too-long
+            with patch(f"{module_location}.threading", return_value=True) as mock_threading:
                 asyncio.run(worker_api._sim_run(input_object_with_webhook))
                 asyncio.run(worker_api._sim_status("test-123"))
-                assert mock_send_webhook.called
+                assert mock_threading.Thread.called
 
             # Test with generator handler
             def generator_handler(job):