Merge pull request #295 from runpod/280-enable-webhooks-for-local-testing-of-serverless-workers

280 enable webhooks for local testing of serverless workers

Author: justinmerrell

runpod/serverless/modules/rp_fastapi.py

@@ -3,10 +3,12 @@
 
 import os
 import uuid
+import threading
 from dataclasses import dataclass
 from typing import Union, Optional, Dict, Any
 
 import uvicorn
+import requests
 from fastapi import FastAPI, APIRouter
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import RedirectResponse
@@ -53,12 +55,14 @@ class TestJob:
     '''
     id: Optional[str] = None
     input: Optional[Union[dict, list, str, int, float, bool]] = None
+    webhook: Optional[str] = None
 
 
 @dataclass
-class DefaultInput:
+class DefaultRequest:
     """ Represents a test input. """
     input: Dict[str, Any]
+    webhook: Optional[str] = None
 
 
 # ------------------------------ Output Objects ------------------------------ #
@@ -80,6 +84,28 @@ class StreamOutput:
     error: Optional[str] = None
 
 
+# ------------------------------ Webhook Sender ------------------------------ #
+def _send_webhook(url: str, payload: Dict[str, Any]) -> bool:
+    """
+    Sends a webhook to the provided URL.
+
+    Args:
+        url (str): The URL to send the webhook to.
+        payload (Dict[str, Any]): The JSON payload to send.
+
+    Returns:
+        bool: True if the request was successful, False otherwise.
+    """
+    with requests.Session() as session:
+        try:
+            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
+
+
 # ---------------------------------------------------------------------------- #
 #                                  API Worker                                  #
 # ---------------------------------------------------------------------------- #
@@ -176,17 +202,17 @@ async def _realtime(self, job: Job):
     # ---------------------------------------------------------------------------- #
 
     # ------------------------------------ run ----------------------------------- #
-    async def _sim_run(self, job_input: DefaultInput) -> JobOutput:
+    async def _sim_run(self, job_request: DefaultRequest) -> JobOutput:
         """ Development endpoint to simulate run behavior. """
         assigned_job_id = f"test-{uuid.uuid4()}"
-        job_list.add_job(assigned_job_id, job_input.input)
+        job_list.add_job(assigned_job_id, job_request.input, job_request.webhook)
         return jsonable_encoder({"id": assigned_job_id, "status": "IN_PROGRESS"})
 
     # ---------------------------------- runsync --------------------------------- #
-    async def _sim_runsync(self, job_input: DefaultInput) -> JobOutput:
+    async def _sim_runsync(self, job_request: DefaultRequest) -> JobOutput:
         """ Development endpoint to simulate runsync behavior. """
         assigned_job_id = f"test-{uuid.uuid4()}"
-        job = TestJob(id=assigned_job_id, input=job_input.input)
+        job = TestJob(id=assigned_job_id, input=job_request.input)
 
         if is_generator(self.config["handler"]):
             generator_output = run_job_generator(self.config["handler"], job.__dict__)
@@ -203,6 +229,12 @@ async def _sim_runsync(self, job_input: DefaultInput) -> JobOutput:
                 "error": job_output['error']
             })
 
+        if job_request.webhook:
+            thread = threading.Thread(
+                target=_send_webhook,
+                args=(job_request.webhook, job_output), daemon=True)
+            thread.start()
+
         return jsonable_encoder({
             "id": job.id,
             "status": "COMPLETED",
@@ -212,15 +244,15 @@ async def _sim_runsync(self, job_input: DefaultInput) -> JobOutput:
     # ---------------------------------- stream ---------------------------------- #
     async def _sim_stream(self, job_id: str) -> StreamOutput:
         """ Development endpoint to simulate stream behavior. """
-        job_input = job_list.get_job_input(job_id)
-        if job_input is None:
+        stashed_job = job_list.get_job(job_id)
+        if stashed_job is None:
             return jsonable_encoder({
                 "id": job_id,
                 "status": "FAILED",
                 "error": "Job ID not found"
             })
 
-        job = TestJob(id=job_id, input=job_input)
+        job = TestJob(id=job_id, input=stashed_job.input)
 
         if is_generator(self.config["handler"]):
             generator_output = run_job_generator(self.config["handler"], job.__dict__)
@@ -236,6 +268,12 @@ async def _sim_stream(self, job_id: str) -> StreamOutput:
 
         job_list.remove_job(job.id)
 
+        if stashed_job.webhook:
+            thread = threading.Thread(
+                target=_send_webhook,
+                args=(stashed_job.webhook, stream_accumulator), daemon=True)
+            thread.start()
+
         return jsonable_encoder({
             "id": job_id,
             "status": "COMPLETED",
@@ -245,15 +283,15 @@ async def _sim_stream(self, job_id: str) -> StreamOutput:
     # ---------------------------------- status ---------------------------------- #
     async def _sim_status(self, job_id: str) -> JobOutput:
         """ Development endpoint to simulate status behavior. """
-        job_input = job_list.get_job_input(job_id)
-        if job_input is None:
+        stashed_job = job_list.get_job(job_id)
+        if stashed_job is None:
             return jsonable_encoder({
                 "id": job_id,
                 "status": "FAILED",
                 "error": "Job ID not found"
             })
 
-        job = TestJob(id=job_id, input=job_input)
+        job = TestJob(id=stashed_job.id, input=stashed_job.input)
 
         if is_generator(self.config["handler"]):
             generator_output = run_job_generator(self.config["handler"], job.__dict__)
@@ -272,6 +310,12 @@ async def _sim_status(self, job_id: str) -> JobOutput:
                 "error": job_output['error']
             })
 
+        if stashed_job.webhook:
+            thread = threading.Thread(
+                target=_send_webhook,
+                args=(stashed_job.webhook, job_output), daemon=True)
+            thread.start()
+
         return jsonable_encoder({
             "id": job_id,
             "status": "COMPLETED",

runpod/serverless/modules/worker_state.py

@@ -25,24 +25,40 @@ def get_auth_header():
 
 # ------------------------------- Job Tracking ------------------------------- #
 class Job:
-    """ Represents a job. """
-
-    def __init__(self, job_id: str, job_input: Optional[Dict[str, Any]] = None) -> None:
-        self.job_id = job_id
-        self.job_input = job_input
+    """
+    Represents a job object.
+
+    Args:
+        job_id: The id of the job, a unique string.
+        job_input: The input to the job.
+        webhook: The webhook to send the job output to.
+    """
+
+    def __init__(
+        self,
+        job_id: str,
+        job_input: Optional[Dict[str, Any]] = None,
+        webhook: Optional[str] = None,
+    ) -> None:
+        self.id = job_id
+        self.input = job_input
+        self.webhook = webhook
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, Job):
-            return self.job_id == other.job_id
+            return self.id == other.id
         return False
 
     def __hash__(self) -> int:
-        return hash(self.job_id)
+        return hash(self.id)
 
     def __str__(self) -> str:
-        return self.job_id
+        return self.id
 
 
+# ---------------------------------------------------------------------------- #
+#                                    Tracker                                   #
+# ---------------------------------------------------------------------------- #
 class Jobs:
     ''' Track the state of current jobs.'''
 
@@ -55,26 +71,26 @@ def __new__(cls):
             Jobs._instance.jobs = set()
         return Jobs._instance
 
-    def add_job(self, job_id, job_input=None):
+    def add_job(self, job_id, job_input=None, webhook=None):
         '''
         Adds a job to the list of jobs.
         '''
-        self.jobs.add(Job(job_id, job_input))
+        self.jobs.add(Job(job_id, job_input, webhook))
 
     def remove_job(self, job_id):
         '''
         Removes a job from the list of jobs.
         '''
         self.jobs.remove(Job(job_id))
 
-    def get_job_input(self, job_id) -> Optional[Union[dict, list, str, int, float, bool]]:
+    def get_job(self, job_id) -> Optional[Union[dict, list, str, int, float, bool]]:
         '''
         Returns the job with the given id.
         Used within rp_fastapi.py for local testing.
         '''
         for job in self.jobs:
-            if job.job_id == job_id:
-                return job.job_input
+            if job.id == job_id:
+                return job
 
         return None
 

tests/test_serverless/test_modules/test_fastapi.py

@@ -5,9 +5,10 @@
 import asyncio
 
 import unittest
-from unittest.mock import patch, Mock
+from unittest.mock import patch, Mock, MagicMock
 import pytest
 
+import requests
 import runpod
 from runpod.serverless.modules import rp_fastapi
 
@@ -53,6 +54,31 @@ def test_start_serverless_with_realtime(self):
             os.environ.pop("RUNPOD_REALTIME_PORT")
             os.environ.pop("RUNPOD_ENDPOINT_ID")
 
+    def test_webhook_sender_success(self):
+        """Test the webhook sender when the request is successful."""
+        module_location = "runpod.serverless.modules.rp_fastapi.requests.Session.post"
+
+        with patch(module_location, new_callable=MagicMock) as mock_post:
+            # Simulate a successful response
+            mock_post.return_value.status_code = 200
+
+            # Call the function
+            success = rp_fastapi._send_webhook("test_webhook", {"test": "output"})
+            assert success is True
+
+    def test_webhook_sender_failure(self):
+        """Test the webhook sender when the request fails."""
+        module_location = "runpod.serverless.modules.rp_fastapi.requests.Session.post"
+
+        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
+
+            # Call the function
+            success = rp_fastapi._send_webhook("test_webhook", {"test": "output"})
+            assert success is False
+
     @pytest.mark.asyncio
     def test_run(self):
         '''
@@ -72,7 +98,7 @@ def test_run(self):
                 input={"test_input": "test_input"}
             )
 
-            default_input_object = rp_fastapi.DefaultInput(
+            default_input_object = rp_fastapi.DefaultRequest(
                 input={"test_input": "test_input"}
             )
 
@@ -115,12 +141,18 @@ def test_runsync(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") as mock_threading:
 
-            default_input_object = rp_fastapi.DefaultInput(
+            default_input_object = rp_fastapi.DefaultRequest(
                 input={"test_input": "test_input"}
             )
 
+            input_object_with_webhook = rp_fastapi.DefaultRequest(
+                input={"test_input": "test_input"},
+                webhook="test_webhook"
+            )
+
             # Test with handler
             worker_api = rp_fastapi.WorkerAPI({"handler": self.handler})
 
@@ -151,6 +183,10 @@ def generator_handler(job):
                 error_worker_api._sim_runsync(default_input_object))
             assert "error" in error_runsync_return
 
+            # Test webhook caller sent
+            asyncio.run(worker_api._sim_runsync(input_object_with_webhook))
+            assert mock_threading.Thread.called
+
         loop.close()
 
     @pytest.mark.asyncio
@@ -164,12 +200,18 @@ 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") as mock_threading:
 
-            default_input_object = rp_fastapi.DefaultInput(
+            default_input_object = rp_fastapi.DefaultRequest(
                 input={"test_input": "test_input"}
             )
 
+            input_object_with_webhook = rp_fastapi.DefaultRequest(
+                input={"test_input": "test_input"},
+                webhook="test_webhook"
+            )
+
             worker_api = rp_fastapi.WorkerAPI({"handler": self.handler})
 
             # Add job to job_list
@@ -203,6 +245,11 @@ def generator_handler(job):
                 "stream": [{"output": {"result": "success"}}]
             }
 
+            # Test webhook caller sent
+            asyncio.run(generator_worker_api._sim_run(input_object_with_webhook))
+            asyncio.run(generator_worker_api._sim_stream("test-123"))
+            assert mock_threading.Thread.called
+
         loop.close()
 
     @pytest.mark.asyncio
@@ -216,14 +263,20 @@ 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") as mock_threading:
 
             worker_api = rp_fastapi.WorkerAPI({"handler": self.handler})
 
-            default_input_object = rp_fastapi.DefaultInput(
+            default_input_object = rp_fastapi.DefaultRequest(
                 input={"test_input": "test_input"}
             )
 
+            input_object_with_webhook = rp_fastapi.DefaultRequest(
+                input={"test_input": "test_input"},
+                webhook="test_webhook"
+            )
+
             # Add job to job_list
             asyncio.run(worker_api._sim_run(default_input_object))
 
@@ -241,6 +294,11 @@ def test_status(self):
                 "output": {"result": "success"}
             }
 
+            # Test webhook caller sent
+            asyncio.run(worker_api._sim_run(input_object_with_webhook))
+            asyncio.run(worker_api._sim_status("test-123"))
+            assert mock_threading.Thread.called
+
             # Test with generator handler
             def generator_handler(job):
                 del job