Merge pull request #273 from runpod/refresh-worker-datarace

Refresh worker datarace

Author: justinmerrell

runpod/serverless/modules/rp_scale.py

@@ -64,15 +64,12 @@ async def get_jobs(self, session):
         Returns:
             List[Any]: A list of job data retrieved from the server.
         """
-        while True:
-            if not self.is_alive():
-                break
-
+        while self.is_alive():
             self.current_concurrency = self.concurrency_modifier(self.current_concurrency)
             log.debug(f"Concurrency set to: {self.current_concurrency}")
 
             log.debug(f"Jobs in progress: {job_list.get_job_count()}")
-            if job_list.get_job_count() < self.current_concurrency:
+            if job_list.get_job_count() < self.current_concurrency and self.is_alive():
                 log.debug("Job list is less than concurrency, getting more jobs.")
 
                 tasks = [

runpod/serverless/worker.py

@@ -54,7 +54,7 @@ async def _process_job(job, session, job_scaler, config):
 
             await stream_result(session, stream_output, job)
     else:
-        is_stream = 0
+        is_stream = False
         job_result = await run_job(config["handler"], job)
 
     # If refresh_worker is set, pod will be reset after job is complete.

tests/test_serverless/test_worker.py

@@ -459,59 +459,6 @@ def mock_is_alive():
         with patch("runpod.serverless.modules.rp_scale.JobScaler.is_alive", wraps=mock_is_alive):
             runpod.serverless.start(config)
 
-    @patch("runpod.serverless.modules.rp_scale.get_job")
-    @patch("runpod.serverless.worker.run_job")
-    @patch("runpod.serverless.worker.send_result")
-    async def test_run_worker_multi_processing_availability_ratio(
-            self, mock_send_result, mock_run_job, mock_get_job):
-        '''
-        Test run_worker with multi processing enabled, the scale-up and
-        scale-down behavior with availability ratio.
-
-        Args:
-            mock_send_result (_type_): _description_
-            mock_stream_result (_type_): _description_
-            mock_run_job (_type_): _description_
-            mock_get_job (_type_): _description_
-            mock_session (_type_): _description_
-        '''
-
-        # Let the test be a long running one so we can capture the scale-up and scale-down.
-        config = {
-            "handler": MagicMock(),
-            "refresh_worker": False,
-            "rp_args": {
-                "rp_debugger": True,
-                "rp_log_level": "DEBUG"
-            }
-        }
-
-        # Let's stop after the 20th call.
-        scale_behavior = {
-            'counter': 0
-        }
-
-        def mock_is_alive():
-            res = scale_behavior['counter'] < 10
-            scale_behavior['counter'] += 1
-
-            # Let's oscillate between upscaling, downscaling, upscaling, downscaling, ...
-            if scale_behavior['counter'] % 2 == 0:
-                mock_get_job.return_value = {
-                    "id": "123", "input": {"number": 1}}
-            else:
-                mock_get_job.return_value = None
-            return res
-
-        # Define the mock behaviors
-        mock_run_job.return_value = {"output": {"result": "odd"}}
-        with patch("runpod.serverless.modules.rp_scale.JobScaler.is_alive", wraps=mock_is_alive):
-            runpod.serverless.start(config)
-
-        # 5 calls with actual jobs
-        assert mock_run_job.call_count == 5
-        assert mock_send_result.call_count == 5
-
     # Test with sls-core
     async def test_run_worker_with_sls_core(self):
         '''