Test cleanup

.github/workflows/CI-pylint.yml

@@ -5,9 +5,9 @@ on:
     branches-ignore:
       - "main-ci"
       - "release"
-  pull_request:
-    branches:
-      - main
+  #   pull_request:
+  #     branches:
+  #       - main
 
   workflow_dispatch:
 

.github/workflows/CI-tests.yml

@@ -6,9 +6,9 @@ on:
       - main-ci
       - release
 
-  pull_request:
-    branches:
-      - main
+  #   pull_request:
+  #     branches:
+  #       - main
 
   workflow_dispatch:
 
@@ -38,4 +38,4 @@ jobs:
 
       - name: Run Tests
         run: |
-          pytest --timeout=120 --timeout_method=thread --cov=runpod --cov-report=xml --cov-report=term-missing --cov-fail-under=100
+          pytest --timeout=120 --timeout_method=thread --cov=runpod --cov-report=xml --cov-report=term-missing --cov-fail-under=100 -W error -p no:cacheprovider -p no:unraisableexception

requirements.txt

@@ -1,4 +1,5 @@
 aiohttp >= 3.8.4
+aiohttp-retry >= 2.8.3
 
 backoff == 2.2.1
 boto3 >= 1.26.165

runpod/serverless/modules/rp_http.py

@@ -4,67 +4,71 @@
 
 import os
 import json
+import aiohttp
+from aiohttp_retry import RetryClient, ExponentialRetry
 
 from runpod.serverless.modules.rp_logger import RunPodLogger
-from .worker_state import Jobs
-from .retry import retry
-from .worker_state import WORKER_ID
+from .worker_state import Jobs, WORKER_ID
 
-JOB_DONE_URL_TEMPLATE = str(os.environ.get('RUNPOD_WEBHOOK_POST_OUTPUT'))
-JOB_DONE_URL_TEMPLATE = JOB_DONE_URL_TEMPLATE.replace('$RUNPOD_POD_ID', WORKER_ID)
+JOB_DONE_URL_TEMPLATE = str(os.environ.get('RUNPOD_WEBHOOK_POST_OUTPUT', 'JOB_DONE_URL'))
+JOB_DONE_URL = JOB_DONE_URL_TEMPLATE.replace('$RUNPOD_POD_ID', WORKER_ID)
 
-JOB_STREAM_URL_TEMPLATE = str(os.environ.get('RUNPOD_WEBHOOK_POST_STREAM'))
-JOB_STREAM_URL_TEMPLATE = JOB_STREAM_URL_TEMPLATE.replace('$RUNPOD_POD_ID', WORKER_ID)
+JOB_STREAM_URL_TEMPLATE = str(os.environ.get('RUNPOD_WEBHOOK_POST_STREAM', 'JOB_STREAM_URL'))
+JOB_STREAM_URL = JOB_STREAM_URL_TEMPLATE.replace('$RUNPOD_POD_ID', WORKER_ID)
 
 log = RunPodLogger()
 job_list = Jobs()
 
 
-@retry(max_attempts=3, base_delay=1, max_delay=3)
-async def transmit(session, job_data, url):
+async def _transmit(client_session, url, job_data ):
     """
-    Wrapper for sending results.
+    Wrapper for transmitting results via POST.
     """
-    headers = {
-        "charset": "utf-8",
-        "Content-Type": "application/x-www-form-urlencoded"
-    }
+    retry_options = ExponentialRetry(attempts=3)
+    retry_client = RetryClient(client_session=client_session, retry_options=retry_options)
 
-    async with session.post(url,
-                            data=job_data,
-                            headers=headers,
-                            raise_for_status=True) as resp:
-        await resp.text()
+    kwargs = {
+        "data": job_data,
+        "headers": {"charset": "utf-8", "Content-Type": "application/x-www-form-urlencoded"},
+        "raise_for_status": True
+        }
 
+    async with retry_client.post(url, **kwargs) as client_response:
+        await client_response.text()
 
-async def send_result(session, job_data, job):
-    '''
-    Return the job results.
-    '''
+
+async def _handle_result(session, job_data, job, url_template, log_message):
+    """
+    A helper function to handle the result, either for sending or streaming.
+    """
     try:
-        job_data = json.dumps(job_data, ensure_ascii=False)
-        job_done_url = JOB_DONE_URL_TEMPLATE.replace('$ID', job['id'])
+        serialized_job_data = json.dumps(job_data, ensure_ascii=False)
+        url = url_template.replace('$ID', job['id'])
 
-        await transmit(session, job_data, job_done_url)
-        log.debug(f"{job['id']} | Results sent.")
+        await _transmit(session, url, serialized_job_data)
+        log.debug(f"{job['id']} | {log_message}")
 
-        job_list.remove_job(job["id"])
-        log.info(f'{job["id"]} | Finished')
+    except aiohttp.ClientResponseError as err:
+        log.error(f"{job['id']} | Client response error while transmitting job. | {err}")
 
-    except Exception as err:  # pylint: disable=broad-except
+    except (TypeError, RuntimeError) as err:
         log.error(f"Error while returning job result {job['id']}: {err}")
 
+    finally:
+        if url_template == JOB_DONE_URL:
+            job_list.remove_job(job["id"])
+            log.info(f'{job["id"]} | Finished')
 
-async def stream_result(session, job_data, job):
-    '''
-    Return the stream job results.
-    '''
-    try:
-        job_data = json.dumps(job_data, ensure_ascii=False)
-        job_done_url = JOB_STREAM_URL_TEMPLATE.replace('$ID', job['id'])
 
-        await transmit(session, job_data, job_done_url)
-        log.debug(f"{job['id']} | Intermediate Results sent.")
+async def send_result(session, job_data, job):
+    """
+    Return the job results.
+    """
+    await _handle_result(session, job_data, job, JOB_DONE_URL, "Results sent.")
+
 
-    except Exception as err:  # pylint: disable=broad-except
-        log.error(f"Error while returning job result {job['id']}: {err}")
+async def stream_result(session, job_data, job):
+    """
+    Return the stream job results.
+    """
+    await _handle_result(session, job_data, job, JOB_STREAM_URL, "Intermediate results sent.")

setup.cfg

@@ -27,6 +27,7 @@ packages = find:
 python_requires = >= 3.10
 install_requires =
     aiohttp >= 3.8.4
+    aiohttp-retry >= 2.8.3
     backoff >= 2.2.1
     boto3 >= 1.26.165
     click >= 8.1.7

tests/test_endpoint/test_asyncio_runner.py

@@ -1,13 +1,11 @@
 ''' Unit tests for the asyncio_runner module. '''
 # pylint: disable=too-few-public-methods
 
-
 import tracemalloc
 import asyncio
 import unittest
 from unittest.mock import patch, MagicMock
 from unittest import IsolatedAsyncioTestCase
-import pytest
 
 from runpod.endpoint.asyncio.asyncio_runner import Job, Endpoint
 
@@ -17,7 +15,6 @@
 class TestJob(IsolatedAsyncioTestCase):
     ''' Tests the Job class. '''
 
-    @pytest.mark.asyncio
     async def test_status(self):
         '''
         Tests Job.status
@@ -32,7 +29,6 @@ async def test_status(self):
             status = await job.status()
             assert status == "COMPLETED"
 
-    @pytest.mark.asyncio
     async def test_output(self):
         '''
         Tests Job.output
@@ -57,7 +53,6 @@ async def json_side_effect():
             output = await output_task
             assert output == "OUTPUT"
 
-    @pytest.mark.asyncio
     async def test_cancel(self):
         '''
         Tests Job.cancel
@@ -72,10 +67,35 @@ async def test_cancel(self):
             cancel_result = await job.cancel()
             assert cancel_result == {"result": "CANCELLED"}
 
-class TestEndpoint:
+    async def test_output_in_progress_then_completed(self):
+        '''Tests Job.output when status is initially IN_PROGRESS and then changes to COMPLETED'''
+        with (
+            patch("runpod.endpoint.asyncio.asyncio_runner.asyncio.sleep") as mock_sleep,
+            patch("aiohttp.ClientSession") as mock_session
+        ):
+            mock_resp = MagicMock()
+            responses = [
+                {"status": "IN_PROGRESS"},
+                {"status": "COMPLETED"},
+                {"output": "OUTPUT"}
+            ]
+
+            async def json_side_effect():
+                if responses:
+                    return responses.pop(0)
+                return {"status": "IN_PROGRESS"}
+
+            mock_resp.json = json_side_effect
+            mock_session.get.return_value.__aenter__.return_value = mock_resp
+
+            job = Job("endpoint_id", "job_id", mock_session)
+            output = await job.output()
+            assert output == "OUTPUT"
+            mock_sleep.assert_called_once_with(1)
+
+class TestEndpoint(IsolatedAsyncioTestCase):
     ''' Unit tests for the Endpoint class. '''
 
-    @pytest.mark.asyncio
     async def test_run(self):
         '''
         Tests Endpoint.run
@@ -90,6 +110,17 @@ async def test_run(self):
             job = await endpoint.run({"input": "INPUT"})
             assert job.job_id == "job_id"
 
+class TestEndpointInitialization(unittest.TestCase):
+    '''Tests for the Endpoint class initialization.'''
+
+    def test_endpoint_initialization(self):
+        '''Tests initialization of Endpoint class.'''
+        with patch("aiohttp.ClientSession"):
+            endpoint = Endpoint("endpoint_id", MagicMock())
+            self.assertEqual(endpoint.endpoint_url, "https://api.runpod.ai/v2/endpoint_id/run")
+            self.assertIn("Content-Type", endpoint.headers)
+            self.assertIn("Authorization", endpoint.headers)
+
 if __name__ == "__main__":
     unittest.main()