fix: limit number of concurrent tasks (#168)

Co-authored-by: Avram Tudor <[email protected]>

run.sh

@@ -3,7 +3,7 @@
 if nvcc --version
 then
   export CUDA_VISIBLE_DEVICES=0
-  export LLAMA_N_CTX=96000
+  export LLAMA_N_CTX=90000
 else
   export LLAMA_N_CTX=8192
 fi

skynet/env.py

@@ -105,6 +105,7 @@ def tobool(val: str | None):
 
 # jobs
 job_timeout = int(os.environ.get('JOB_TIMEOUT', 60 * 5))  # 5 minutes default
+max_concurrency = int(os.environ.get('MAX_CONCURRENCY', 30))
 
 # summaries
 summary_minimum_payload_length = int(os.environ.get('SUMMARY_MINIMUM_PAYLOAD_LENGTH', 100))

skynet/modules/ttt/openai_api/app.py

@@ -38,7 +38,7 @@ def initialize():
             '--model',
             llama_path,
             '--gpu_memory_utilization',
-            str(0.98),
+            str(0.95),
             '--max-model-len',
             str(llama_n_ctx),
             '--port',

skynet/modules/ttt/summaries/jobs.py

@@ -3,7 +3,7 @@
 import time
 import uuid
 
-from skynet.env import enable_batching, job_timeout, modules, redis_exp_seconds
+from skynet.env import enable_batching, job_timeout, max_concurrency, modules, redis_exp_seconds
 from skynet.logs import get_logger
 from skynet.modules.monitoring import (
     OPENAI_API_RESTART_COUNTER,
@@ -23,15 +23,15 @@
 
 log = get_logger(__name__)
 
-TIME_BETWEEN_JOBS_CHECK = 1
+TIME_BETWEEN_JOBS_CHECK = 0.3
 TIME_BETWEEN_JOBS_CHECK_ON_ERROR = 10
 
 PENDING_JOBS_KEY = "jobs:pending"
 RUNNING_JOBS_KEY = "jobs:running"
 ERROR_JOBS_KEY = "jobs:error"
 
 background_task = None
-current_tasks = set()
+current_tasks = set[asyncio.Task]()
 
 
 def restart():
@@ -48,11 +48,10 @@ def can_run_next_job() -> bool:
     if 'summaries:executor' not in modules:
         return False
 
-    # TODO: add limit even when batching is enabled.
     if enable_batching:
-        return True
+        return len(current_tasks) < max_concurrency
 
-    return all(t.done() for t in current_tasks)
+    return len(current_tasks) == 0
 
 
 async def update_summary_queue_metric() -> None: