Merge pull request #270 from runpod/concurrency-fix

Concurrency fix

Author: justinmerrell

runpod/serverless/modules/rp_http.py

@@ -57,6 +57,7 @@ async def _handle_result(session, job_data, job, url_template, log_message, is_s
         log.error(f"Error while returning job result. | {err}", job['id'])
 
     finally:
+        #job_data status is used for local development with FastAPI
         if url_template == JOB_DONE_URL and job_data.get('status', None) != 'IN_PROGRESS':
             job_list.remove_job(job["id"])
             log.info("Finished.", job['id'])

runpod/serverless/modules/rp_job.py

@@ -198,4 +198,4 @@ async def run_job_generator(
         log.error(err, job["id"])
         yield {"error": f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"}
     finally:
-        log.info('Finished', job["id"])
+        log.info('Finished running generator.', job["id"])

runpod/serverless/modules/rp_scale.py

@@ -69,18 +69,21 @@ async def get_jobs(self, session):
                 break
 
             self.current_concurrency = self.concurrency_modifier(self.current_concurrency)
+            log.debug(f"Concurrency set to: {self.current_concurrency}")
 
-            tasks = [
-                asyncio.create_task(get_job(session, retry=False))
-                for _ in range(self.current_concurrency if job_list.get_job_list() else 1)
-            ]
+            log.debug(f"Jobs in progress: {job_list.get_job_count()}")
+            if job_list.get_job_count() < self.current_concurrency:
+                log.debug("Job list is less than concurrency, getting more jobs.")
 
-            for job_future in asyncio.as_completed(tasks):
-                job = await job_future
-                self.job_history.append(1 if job else 0)
-                if job:
-                    yield job
+                tasks = [
+                    asyncio.create_task(get_job(session, retry=False))
+                    for _ in range(self.current_concurrency if job_list.get_job_list() else 1)
+                ]
 
-            await asyncio.sleep(0)
+                for job_future in asyncio.as_completed(tasks):
+                    job = await job_future
+                    self.job_history.append(1 if job else 0)
+                    if job:
+                        yield job
 
-            log.debug(f"Concurrency set to: {self.current_concurrency}")
+            await asyncio.sleep(0)

runpod/serverless/modules/worker_state.py

@@ -83,3 +83,9 @@ def get_job_list(self):
         Returns the list of jobs as a string.
         '''
         return ','.join(str(job) for job in self.jobs) if self.jobs else None
+
+    def get_job_count(self):
+        '''
+        Returns the number of jobs.
+        '''
+        return len(self.jobs)

runpod/serverless/worker.py

@@ -42,7 +42,7 @@ async def _process_job(job, session, job_scaler, config):
     if rp_handler.is_generator(config["handler"]):
         is_stream = True
         generator_output = run_job_generator(config["handler"], job)
-        log.debug("Handler is a generator, streaming results.")
+        log.debug("Handler is a generator, streaming results.", job['id'])
 
         job_result = {'output': []}
         async for stream_output in generator_output:
@@ -66,13 +66,13 @@ async def _process_job(job, session, job_scaler, config):
     # If rp_debugger is set, debugger output will be returned.
     if config["rp_args"].get("rp_debugger", False) and isinstance(job_result, dict):
         job_result["output"]["rp_debugger"] = rp_debugger.get_debugger_output()
-        log.debug("rp_debugger | Flag set, returning debugger output.")
+        log.debug("rp_debugger | Flag set, returning debugger output.", job['id'])
 
         # Calculate ready delay for the debugger output.
         ready_delay = (config["reference_counter_start"] - REF_COUNT_ZERO) * 1000
         job_result["output"]["rp_debugger"]["ready_delay_ms"] = ready_delay
     else:
-        log.debug("rp_debugger | Flag not set, skipping debugger output.")
+        log.debug("rp_debugger | Flag not set, skipping debugger output.", job['id'])
         rp_debugger.clear_debugger_output()
 
     # Send the job result to SLS

tests/test_serverless/test_modules/test_job.py

@@ -296,7 +296,7 @@ async def test_run_job_generator_success(self):
         assert result == [{"output": "partial_output_1"}, {"output": "partial_output_2"}]
         assert mock_log.error.call_count == 0
         assert mock_log.info.call_count == 1
-        mock_log.info.assert_called_with('Finished', '123')
+        mock_log.info.assert_called_with('Finished running generator.', '123')
 
     async def test_run_job_generator_success_async(self):
         '''
@@ -311,7 +311,7 @@ async def test_run_job_generator_success_async(self):
         assert result == [{"output": "partial_output_1"}, {"output": "partial_output_2"}]
         assert mock_log.error.call_count == 0
         assert mock_log.info.call_count == 1
-        mock_log.info.assert_called_with('Finished', '123')
+        mock_log.info.assert_called_with('Finished running generator.', '123')
 
     async def test_run_job_generator_exception(self):
         '''
@@ -327,4 +327,4 @@ async def test_run_job_generator_exception(self):
         assert "error" in result[0]
         assert mock_log.error.call_count == 1
         assert mock_log.info.call_count == 1
-        mock_log.info.assert_called_with('Finished', '123')
+        mock_log.info.assert_called_with('Finished running generator.', '123')