Concurrency modifier

src/handler.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-
 from typing import Generator
 import runpod
 from utils import validate_and_convert_sampling_params, initialize_llm_engine, JobManager, ServerlessConfig
@@ -9,8 +8,8 @@ serverless_config = ServerlessConfig()
 job_manager = JobManager()
 llm = initialize_llm_engine()
 
-def concurrency_modifier() -> int:
-    return max(0, serverless_config.max_concurrency - job_manager.total_running_jobs)
+def concurrency_modifier(current_concurrency) -> int:
+    return max(0, serverless_config.max_concurrency - current_concurrency)
 
 async def handler(job: dict) -> Generator[dict, None, None]:
     job_input = job["input"]
@@ -22,7 +21,6 @@ async def handler(job: dict) -> Generator[dict, None, None]:
     validated_params = validate_and_convert_sampling_params(sampling_params)
     request_id = random_uuid()
     results_generator = llm.generate(prompt, validated_params, request_id)
-    job_manager.increment_job_count()
 
     batch, last_output_text = [], ""
     async for request_output in results_generator:
@@ -41,9 +39,7 @@ async def handler(job: dict) -> Generator[dict, None, None]:
 
     if batch:
         yield batch
-
-    job_manager.decrement_job_count()
-
+        
 runpod.serverless.start({
     "handler": handler,
     "concurrency_modifier": concurrency_modifier,

src/utils.py

@@ -51,15 +51,6 @@ def initialize_llm_engine() -> AsyncLLMEngine:
         logging.error(f"Error initializing vLLM engine: {e}")
         raise
 
-class JobManager:
-    def __init__(self):
-        self.total_running_jobs = 0
-
-    def increment_job_count(self):
-        self.total_running_jobs += 1
-
-    def decrement_job_count(self):
-        self.total_running_jobs -= 1
 
 def validate_and_convert_sampling_params(params: Dict[str, Any]) -> SamplingParams:
     validated_params = {}