Adds timer to avoid saturating CPU usage

juanroesel · juanroesel · commit 9c0820d279dc · 2024-06-28T16:54:00.000-07:00
diff --git a/llama_cpp/_utils.py b/llama_cpp/_utils.py
@@ -153,11 +153,10 @@ async def monitor_task_queue(status_dict: Dict[str, Union[int, float]]):
     a shared status dictionary with the number of tasks that have not
     started and the number of tasks that are currently running.
     It recursively calls itself to continuously monitor the task queue.
-    NOTE: There will always be 4 tasks running in the task queue:
+    NOTE: There will always be 3 tasks running in the task queue:
     - LifespanOn.main: Main application coroutine
     - Server.serve: Server coroutine
     - monitor_task_queue: Task queue monitoring coroutine
-    - RequestReponseCycle.run_asgi: ASGI single cycle coroutine
     Any upcoming requests will be added to the task queue in the form of
     another RequestReponseCycle.run_asgi coroutine.
     """
@@ -175,6 +174,8 @@ async def monitor_task_queue(status_dict: Dict[str, Union[int, float]]):
         for task in all_tasks
     }
 
+    await asyncio.sleep(5)  # adds a delay of 5 seconds to avoid overloading the CPU
+
     asyncio.create_task(
         monitor_task_queue(status_dict)
     )  # pass status_dict to the next task
diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py
@@ -246,22 +246,20 @@ async def authenticate(
     response_model=HealthMetrics,
     summary="Server's health",
 )
-async def check_health(
-    
-):
-    # 4 running tasks + new scheduled request
-    if 0 <= task_queue_status.get("running_tasks_count", 0) <= 5:
+async def check_health():
+    # 3 running tasks + new scheduled request
+    if 0 <= task_queue_status.get("running_tasks_count", 0) <= 4:
         return JSONResponse(
             content={"status": "OK", "task_queue_status": task_queue_status}
         )
-    # 1 - 6 scheduled requests
-    elif 5 < task_queue_status.get("running_tasks_count", 0) <= 10:
+    # 2 - 6 scheduled requests
+    elif 4 < task_queue_status.get("running_tasks_count", 0) < 10:
         return JSONResponse(
             content={"status": "Warning", "task_queue_status": task_queue_status}
         )
     # 7+ scheduled requests
     # TODO: Evaluate if in this case we should manually stop the execution of certain tasks to clear the queue
-    elif task_queue_status.get("running_tasks_count", 0) > 10:
+    elif task_queue_status.get("running_tasks_count", 0) >= 10:
         return JSONResponse(
             content={"status": "Critical", "task_queue_status": task_queue_status}
         )
