debounce at HTTP 429 response (#367)

Author: deanq

runpod/http_client.py

@@ -5,13 +5,17 @@
 import os
 
 import requests
-from aiohttp import ClientSession, ClientTimeout, TCPConnector
+from aiohttp import ClientSession, ClientTimeout, TCPConnector, ClientResponseError
 
 from .cli.groups.config.functions import get_credentials
 from .tracer import create_aiohttp_tracer, create_request_tracer
 from .user_agent import USER_AGENT
 
 
+class TooManyRequests(ClientResponseError):
+    pass
+
+
 def get_auth_header():
     """
     Produce a header dict with the `Authorization` key derived from

runpod/serverless/modules/rp_job.py

@@ -8,7 +8,7 @@
 import traceback
 from typing import Any, AsyncGenerator, Callable, Dict, Optional, Union, List
 
-from runpod.http_client import ClientSession
+from runpod.http_client import ClientSession, TooManyRequests
 from runpod.serverless.modules.rp_logger import RunPodLogger
 
 from ...version import __version__ as runpod_version
@@ -70,15 +70,20 @@ async def get_job(
             log.debug("- Received 400 status, expected when FlashBoot is enabled.")
             return
 
-        try:
-            response.raise_for_status()
-        except Exception:
-            log.error(f"- Failed to get job, status code: {response.status}")
-            return
+        if response.status == 429:
+            raise TooManyRequests(
+                response.request_info,
+                response.history,
+                status=response.status,
+                message=response.reason
+            )
+
+        # All other errors should raise an exception
+        response.raise_for_status()
 
         # Verify if the content type is JSON
         if response.content_type != "application/json":
-            log.error(f"- Unexpected content type: {response.content_type}")
+            log.debug(f"- Unexpected content type: {response.content_type}")
             return
 
         # Check if there is a non-empty content to parse

runpod/serverless/modules/rp_scale.py

@@ -7,7 +7,7 @@
 import signal
 from typing import Any, Dict
 
-from ...http_client import AsyncClientSession, ClientSession
+from ...http_client import AsyncClientSession, ClientSession, TooManyRequests
 from .rp_job import get_job, handle_job
 from .rp_logger import RunPodLogger
 from .worker_state import JobsQueue, JobsProgress
@@ -150,6 +150,10 @@ async def get_jobs(self, session: ClientSession):
                 acquired_jobs = await asyncio.wait_for(
                     get_job(session, jobs_needed), timeout=30
                 )
+            except TooManyRequests:
+                log.debug(f"JobScaler.get_jobs | Too many requests. Debounce for 5 seconds.")
+                await asyncio.sleep(5)  # debounce for 5 seconds
+                continue
             except asyncio.CancelledError:
                 log.debug("JobScaler.get_jobs | Request was cancelled.")
                 continue

tests/test_serverless/test_modules/test_job.py

@@ -2,13 +2,13 @@
 Test Serverless Job Module
 """
 
-import asyncio
 from unittest.mock import Mock, patch
 
 from unittest import IsolatedAsyncioTestCase
-from aiohttp import ClientResponse
+from aiohttp import ClientResponse, ClientResponseError
 from aiohttp.test_utils import make_mocked_coro
 
+from runpod.http_client import TooManyRequests
 from runpod.serverless.modules import rp_job
 
 
@@ -63,6 +63,38 @@ async def test_get_job_400(self):
             job = await rp_job.get_job(mock_session)
             self.assertIsNone(job)
 
+    async def test_get_job_429(self):
+        """Tests the get_job function with a 429 response."""
+        response = Mock(ClientResponse)
+        response.raise_for_status.side_effect = TooManyRequests(
+            request_info=None,
+            history=(),
+            status=429,
+        )
+
+        with patch("aiohttp.ClientSession") as mock_session, patch(
+            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
+        ):
+            mock_session.get.return_value.__aenter__.return_value = response
+            with self.assertRaises(ClientResponseError):
+                await rp_job.get_job(mock_session)
+
+    async def test_get_job_500(self):
+        """Tests the get_job function with a 500 response."""
+        # Mock 500 response
+        response = Mock(ClientResponse)
+        response.raise_for_status.side_effect = TooManyRequests(
+            request_info=None,  # Not needed for the test
+            history=(),  # Not needed for the test
+            status=500, 
+        )
+        with patch("aiohttp.ClientSession") as mock_session, patch(
+            "runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"
+        ):
+            mock_session.get.return_value.__aenter__.return_value = response
+            with self.assertRaises(Exception):
+                await rp_job.get_job(mock_session)
+
     async def test_get_job_no_id(self):
         """Tests the get_job function with a 200 response but no 'id' field."""
         response = Mock(ClientResponse)