Fix: failed requests due to race conditions in the job queue vs job progress (#376)

* fix: JobsProgress is now asyncio-safe. This prevents any race conditions when job_progress.get_job_count() was checked before getting more jobs.
* fix: strict jobs count for evaluating if new jobs can be taken `jobs_needed = concurrency - queue - in progress`
* debug: better debug logs
* improved unit tests

Author: deanq

.github/workflows/CD-publish_to_pypi.yml

@@ -13,10 +13,10 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python 3.11.0
+      - name: Set up Python 3.11.10
         uses: actions/setup-python@v5
         with:
-          python-version: 3.11.0
+          python-version: 3.11.10
 
       - name: Install pypa/build
         run: >-

.github/workflows/CD-test_publish_to_pypi.yml

@@ -14,10 +14,10 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python 3.11.0
+      - name: Set up Python 3.11.10
         uses: actions/setup-python@v5
         with:
-          python-version: 3.11.0
+          python-version: 3.11.10
 
       - name: Install pypa/build
         run: >-

.github/workflows/CI-pytests.yml

@@ -15,7 +15,7 @@ jobs:
   run_tests:
     strategy:
       matrix:
-        python-version: [3.8, 3.9, 3.10.12, 3.11.0]
+        python-version: [3.8, 3.9, 3.10.15, 3.11.10]
     runs-on: ubuntu-latest
 
     steps:

pyproject.toml

@@ -54,7 +54,6 @@ runpod = "runpod.cli.entry:runpod_cli"
 test = [
     "asynctest",
     "nest_asyncio",
-    "pylint==3.3.1",
     "pytest-asyncio",
     "pytest-cov",
     "pytest-timeout",

runpod/serverless/modules/rp_fastapi.py

@@ -286,12 +286,12 @@ async def _realtime(self, job: Job):
         Performs model inference on the input data using the provided handler.
         If handler is not provided, returns an error message.
         """
-        job_list.add(job.id)
+        await job_list.add(job.id)
 
         # Process the job using the provided handler, passing in the job input.
         job_results = await run_job(self.config["handler"], job.__dict__)
 
-        job_list.remove(job.id)
+        await job_list.remove(job.id)
 
         # Return the results of the job processing.
         return jsonable_encoder(job_results)
@@ -304,7 +304,7 @@ async def _realtime(self, job: Job):
     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({
+        await job_list.add({
             "id": assigned_job_id,
             "input": job_request.input,
             "webhook": job_request.webhook
@@ -345,7 +345,7 @@ async def _sim_runsync(self, job_request: DefaultRequest) -> JobOutput:
     # ---------------------------------- stream ---------------------------------- #
     async def _sim_stream(self, job_id: str) -> StreamOutput:
         """Development endpoint to simulate stream behavior."""
-        stashed_job = job_list.get(job_id)
+        stashed_job = await job_list.get(job_id)
         if stashed_job is None:
             return jsonable_encoder(
                 {"id": job_id, "status": "FAILED", "error": "Job ID not found"}
@@ -367,7 +367,7 @@ async def _sim_stream(self, job_id: str) -> StreamOutput:
                 }
             )
 
-        job_list.remove(job.id)
+        await job_list.remove(job.id)
 
         if stashed_job.webhook:
             thread = threading.Thread(
@@ -384,7 +384,7 @@ 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."""
-        stashed_job = job_list.get(job_id)
+        stashed_job = await job_list.get(job_id)
         if stashed_job is None:
             return jsonable_encoder(
                 {"id": job_id, "status": "FAILED", "error": "Job ID not found"}
@@ -400,7 +400,7 @@ async def _sim_status(self, job_id: str) -> JobOutput:
         else:
             job_output = await run_job(self.config["handler"], job.__dict__)
 
-        job_list.remove(job.id)
+        await job_list.remove(job.id)
 
         if job_output.get("error", None):
             return jsonable_encoder(

runpod/serverless/modules/rp_job.py

@@ -8,6 +8,8 @@
 import traceback
 from typing import Any, AsyncGenerator, Callable, Dict, Optional, Union, List
 
+import aiohttp
+
 from runpod.http_client import ClientSession, TooManyRequests
 from runpod.serverless.modules.rp_logger import RunPodLogger
 
@@ -34,15 +36,18 @@ def _job_get_url(batch_size: int = 1):
     Returns:
         str: The prepared URL for the 'get' request to the serverless API.
     """
-    job_in_progress = "1" if job_progress.get_job_count() else "0"
 
     if batch_size > 1:
         job_take_url = JOB_GET_URL.replace("/job-take/", "/job-take-batch/")
         job_take_url += f"&batch_size={batch_size}"
     else:
         job_take_url = JOB_GET_URL
 
-    return job_take_url + f"&job_in_progress={job_in_progress}"
+    job_in_progress = "1" if job_progress.get_job_list() else "0"
+    job_take_url += f"&job_in_progress={job_in_progress}"
+
+    log.debug(f"rp_job | get_job: {job_take_url}")
+    return job_take_url
 
 
 async def get_job(
@@ -60,14 +65,14 @@ async def get_job(
         num_jobs (int): The number of jobs to get.
     """
     async with session.get(_job_get_url(num_jobs)) as response:
-        log.debug(f"- Response: {type(response).__name__} {response.status}")
+        log.debug(f"rp_job | Response: {type(response).__name__} {response.status}")
 
         if response.status == 204:
-            log.debug("- No content, no job to process.")
+            log.debug("rp_job | Received 204 status, no jobs.")
             return
 
         if response.status == 400:
-            log.debug("- Received 400 status, expected when FlashBoot is enabled.")
+            log.debug("rp_job | Received 400 status, expected when FlashBoot is enabled.")
             return
 
         if response.status == 429:
@@ -83,16 +88,23 @@ async def get_job(
 
         # Verify if the content type is JSON
         if response.content_type != "application/json":
-            log.debug(f"- Unexpected content type: {response.content_type}")
+            log.debug(f"rp_job | Unexpected content type: {response.content_type}")
             return
 
         # Check if there is a non-empty content to parse
         if response.content_length == 0:
-            log.debug("- No content to parse.")
+            log.debug("rp_job | No content to parse.")
             return
 
-        jobs = await response.json()
-        log.debug(f"- Received Job(s)")
+        try:
+            jobs = await response.json()
+            log.debug("rp_job | Received Job(s)")
+        except aiohttp.ContentTypeError:
+            log.debug(f"rp_job | Response content is not valid JSON. {response.content}")
+            return
+        except ValueError as json_error:
+            log.debug(f"rp_job | Failed to parse JSON response: {json_error}")
+            return
 
         # legacy job-take API
         if isinstance(jobs, dict):

runpod/serverless/modules/rp_scale.py

@@ -114,14 +114,20 @@ async def get_jobs(self, session: ClientSession):
         Adds jobs to the JobsQueue
         """
         while self.is_alive():
-            log.debug(f"JobScaler.get_jobs | Jobs in progress: {job_progress.get_job_count()}")
+            log.debug("JobScaler.get_jobs | Starting job acquisition.")
 
             self.current_concurrency = self.concurrency_modifier(
                 self.current_concurrency
             )
-            log.debug(f"JobScaler.get_jobs | Concurrency set to: {self.current_concurrency}")
+            log.debug(f"JobScaler.get_jobs | current Concurrency set to: {self.current_concurrency}")
 
-            jobs_needed = self.current_concurrency - job_progress.get_job_count()
+            current_progress_count = await job_progress.get_job_count()
+            log.debug(f"JobScaler.get_jobs | current Jobs in progress: {current_progress_count}")
+
+            current_queue_count = job_list.get_job_count()
+            log.debug(f"JobScaler.get_jobs | current Jobs in queue: {current_queue_count}")
+
+            jobs_needed = self.current_concurrency - current_progress_count - current_queue_count
             if jobs_needed <= 0:
                 log.debug("JobScaler.get_jobs | Queue is full. Retrying soon.")
                 await asyncio.sleep(1)  # don't go rapidly
@@ -197,10 +203,9 @@ async def handle_job(self, session: ClientSession, job: dict):
         """
         Process an individual job. This function is run concurrently for multiple jobs.
         """
-        log.debug(f"JobScaler.handle_job | {job}")
-        job_progress.add(job)
-
         try:
+            await job_progress.add(job)
+
             await handle_job(session, self.config, job)
 
             if self.config.get("refresh_worker", False):
@@ -215,4 +220,4 @@ async def handle_job(self, session: ClientSession, job: dict):
             job_list.task_done()
 
             # Job is no longer in progress
-            job_progress.remove(job["id"])
+            await job_progress.remove(job["id"])

runpod/serverless/modules/worker_state.py

@@ -6,7 +6,7 @@
 import time
 import uuid
 from typing import Any, Dict, Optional
-from asyncio import Queue
+from asyncio import Queue, Lock
 
 from .rp_logger import RunPodLogger
 
@@ -72,10 +72,19 @@ def __new__(cls):
             JobsProgress._instance = set.__new__(cls)
         return JobsProgress._instance
 
+    def __init__(self):
+        if not hasattr(self, "_lock"):
+            # Initialize the lock once
+            self._lock = Lock()
+
     def __repr__(self) -> str:
         return f"<{self.__class__.__name__}>: {self.get_job_list()}"
 
-    def add(self, element: Any):
+    async def clear(self) -> None:
+        async with self._lock:
+            return super().clear()
+
+    async def add(self, element: Any):
         """
         Adds a Job object to the set.
 
@@ -92,16 +101,17 @@ def add(self, element: Any):
         if not isinstance(element, Job):
             raise TypeError("Only Job objects can be added to JobsProgress.")
 
-        log.debug(f"JobsProgress.add | {element}")
-        return super().add(element)
+        async with self._lock:
+            log.debug(f"JobsProgress.add", element.id)
+            super().add(element)
 
-    def remove(self, element: Any):
+    async def remove(self, element: Any):
         """
-        Adds a Job object to the set.
+        Removes a Job object from the set.
 
-        If the added element is a string, then `Job(id=element)` is added
+        If the element is a string, then `Job(id=element)` is removed
         
-        If the added element is a dict, that `Job(**element)` is added
+        If the element is a dict, then `Job(**element)` is removed
         """
         if isinstance(element, str):
             element = Job(id=element)
@@ -112,34 +122,37 @@ def remove(self, element: Any):
         if not isinstance(element, Job):
             raise TypeError("Only Job objects can be removed from JobsProgress.")
 
-        log.debug(f"JobsProgress.remove | {element}")
-        return super().remove(element)
+        async with self._lock:
+            log.debug(f"JobsProgress.remove", element.id)
+            return super().discard(element)
 
-    def get(self, element: Any) -> Job:
+    async def get(self, element: Any) -> Job:
         if isinstance(element, str):
             element = Job(id=element)
 
         if not isinstance(element, Job):
             raise TypeError("Only Job objects can be retrieved from JobsProgress.")
 
-        for job in self:
-            if job == element:
-                return job
+        async with self._lock:
+            for job in self:
+                if job == element:
+                    return job
 
     def get_job_list(self) -> str:
         """
         Returns the list of job IDs as comma-separated string.
         """
-        if not self.get_job_count():
+        if not len(self):
             return None
 
         return ",".join(str(job) for job in self)
 
-    def get_job_count(self) -> int:
+    async def get_job_count(self) -> int:
         """
-        Returns the number of jobs.
+        Returns the number of jobs in a thread-safe manner.
         """
-        return len(self)
+        async with self._lock:
+            return len(self)
 
 
 class JobsQueue(Queue):
@@ -162,7 +175,7 @@ async def add_job(self, job: dict):
         If the queue is full, wait until a free
         slot is available before adding item.
         """
-        log.debug(f"JobsQueue.add_job | {job}")
+        log.debug(f"JobsQueue.add_job", job["id"])
         return await self.put(job)
 
     async def get_job(self) -> dict:

setup.py

@@ -17,7 +17,6 @@
     "test": [
         "asynctest",
         "nest_asyncio",
-        "pylint==3.3.1",
         "pytest",
         "pytest-cov",
         "pytest-timeout",

tests/test_serverless/test_modules/test_http.py

@@ -30,11 +30,11 @@ def __init__(self, *args, **kwargs):
 class TestHTTP(unittest.IsolatedAsyncioTestCase):
     """Test HTTP module."""
 
-    def setUp(self) -> None:
+    async def asyncSetUp(self) -> None:
         self.job = {"id": "test_id"}
         self.job_data = {"output": "test_output"}
 
-    def tearDown(self) -> None:
+    async def asyncTearDown(self) -> None:
         gc.collect()
 
     async def test_send_result(self):