Deprecated JobScaler.concurrency_modifier (#331)

Author: deanq

examples/serverless/concurrent_handler.py

@@ -11,19 +11,6 @@ async def async_generator_handler(job):
     return job
 
 
-# --------------------------- Concurrency Modifier --------------------------- #
-def concurrency_modifier(current_concurrency=1):
-    '''
-    Concurrency modifier.
-    '''
-    desired_concurrency = current_concurrency
-
-    # Do some logic to determine the desired concurrency.
-
-    return desired_concurrency
-
-
 runpod.serverless.start({
     "handler": async_generator_handler,
-    "concurrency_modifier": concurrency_modifier
 })

runpod/serverless/__init__.py

@@ -111,7 +111,6 @@ def start(config: Dict[str, Any]):
     config (Dict[str, Any]): Configuration parameters for the worker.
 
     config["handler"] (Callable): The handler function to run.
-    config["concurrency_modifier"] (Callable): Concurrency modifier function to run.
 
     config["rp_args"] (Dict[str, Any]): Arguments for the worker, populated by runtime arguments.
     """

runpod/serverless/modules/rp_scale.py

@@ -4,8 +4,7 @@
 '''
 
 import asyncio
-import typing
-
+import warnings
 from runpod.serverless.modules.rp_logger import RunPodLogger
 from .rp_job import get_job
 from .worker_state import Jobs
@@ -14,35 +13,20 @@
 job_list = Jobs()
 
 
-def _default_concurrency_modifier(current_concurrency: int) -> int:
-    """
-    Default concurrency modifier.
-
-    This function returns the current concurrency without any modification.
-
-    Args:
-        current_concurrency (int): The current concurrency.
-
-    Returns:
-        int: The current concurrency.
-    """
-    return current_concurrency
-
-
 class JobScaler():
     """
     Job Scaler. This class is responsible for scaling the number of concurrent requests.
     """
 
-    def __init__(self, concurrency_modifier: typing.Any):
-        if concurrency_modifier is None:
-            self.concurrency_modifier = _default_concurrency_modifier
-        else:
-            self.concurrency_modifier = concurrency_modifier
-
-        self.background_get_job_tasks = set()
-        self.job_history = []
-        self.current_concurrency = 1
+    def __init__(self, concurrency_modifier = None):
+        if concurrency_modifier:
+            warnings.warn(
+                "JobScaler(concurrency_modifier) is deprecated ",
+                "and will be removed in a future version. "
+                "Please remove `concurrency_modifier` parameter.",
+                DeprecationWarning,
+                stacklevel=2
+            )
         self._is_alive = True
 
     def is_alive(self):
@@ -65,22 +49,12 @@ async def get_jobs(self, session):
             List[Any]: A list of job data retrieved from the server.
         """
         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 and self.is_alive():
-                log.debug("Job list is less than concurrency, getting more jobs.")
-
-                tasks = [
-                    asyncio.create_task(get_job(session, retry=False))
-                    for _ in range(self.current_concurrency if job_list.get_job_list() else 1)
-                ]
 
-                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))
+            ]
 
-            await asyncio.sleep(0)
+            for job_future in asyncio.as_completed(tasks):
+                if job := await job_future:
+                    yield job

runpod/serverless/worker.py

@@ -87,9 +87,7 @@ async def run_worker(config: Dict[str, Any]) -> None:
     client_session = AsyncClientSession()
 
     async with client_session as session:
-        job_scaler = rp_scale.JobScaler(
-            concurrency_modifier=config.get('concurrency_modifier', None)
-        )
+        job_scaler = rp_scale.JobScaler()
 
         while job_scaler.is_alive():
 

tests/test_serverless/test_worker.py

@@ -316,45 +316,6 @@ async def test_run_worker_generator_aggregate_handler(
         _, args, _ = mock_send_result.mock_calls[0]
         assert args[1] == {'output': ['test1', 'test2'], 'stopPod': True}
 
-    @patch("runpod.serverless.worker.AsyncClientSession")
-    @patch("runpod.serverless.modules.rp_scale.get_job")
-    @patch("runpod.serverless.worker.run_job")
-    @patch("runpod.serverless.worker.stream_result")
-    @patch("runpod.serverless.worker.send_result")
-    # pylint: disable=too-many-arguments
-    async def test_run_worker_concurrency(
-            self, mock_send_result, mock_stream_result, mock_run_job, mock_get_job, mock_session):
-        '''
-        Test run_worker with synchronous handler.
-
-        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_
-        '''
-        # Define the mock behaviors
-        mock_get_job.return_value = {"id": "123", "input": {"number": 1}}
-        mock_run_job.return_value = {"output": {"result": "odd"}}
-
-        def concurrency_modifier(current_concurrency):
-            return current_concurrency + 1
-
-        config_with_concurrency = self.config.copy()
-        config_with_concurrency['concurrency_modifier'] = concurrency_modifier
-
-        # Call the function
-        runpod.serverless.start(config_with_concurrency)
-
-        # Make assertions about the behaviors
-        mock_get_job.assert_called_once()
-        mock_run_job.assert_called_once()
-        mock_send_result.assert_called_once()
-
-        assert mock_stream_result.called is False
-        assert mock_session.called
-
     @patch("runpod.serverless.worker.AsyncClientSession")
     @patch("runpod.serverless.modules.rp_scale.get_job")
     @patch("runpod.serverless.worker.run_job")