Inflate RAM usage when max session is reached (#24)

Author: alejoe91

entrypoint.py

@@ -1,8 +1,10 @@
 import os
 import shutil
+import threading
 from argparse import ArgumentParser
 
 import requests
+import numpy as np
 import psutil
 import panel as pn
 from tornado.web import RequestHandler
@@ -12,6 +14,10 @@
 from aind_ephys_portal.panel.logging import list_gui_sessions, get_max_number_of_gui_sessions, LOG_DIR  # noqa: F401
 
 
+TARGET_MEMORY_TRIGGER_PERCENT = 75
+TARGET_CLEAR_TMP_ARR_SECONDS = 30
+
+
 if LOG_DIR.is_dir():
     print(f"Cleaning up old log files in {LOG_DIR}...")
     shutil.rmtree(LOG_DIR)
@@ -31,20 +37,57 @@ def get_ecs_task_id():
 MAX_GUI_SESSIONS_PER_TASK = get_max_number_of_gui_sessions()
 print(f"Max GUI sessions per task: {MAX_GUI_SESSIONS_PER_TASK}")
 
+_tmp_array_triggered = False
+_tmp_arr = None
+_tmp_arr_timer = None
+
+
+def _clear_tmp_arr():
+    global _tmp_arr, _tmp_arr_timer
+    print("Clearing temporary array to free up memory...")
+    _tmp_arr = None
+    _tmp_arr_timer = None
+    print(f"tmp_arr cleared after {TARGET_CLEAR_TMP_ARR_SECONDS}s")
+
 # 2. Health Check & Index Redirect
 class HealthHandler(RequestHandler):
     def get(self):
+        global _tmp_arr, _tmp_arr_timer, _tmp_array_triggered
         mem = psutil.virtual_memory()
         # count number of GUI app sessions
         gui_sessions = list_gui_sessions()
         task_id = get_ecs_task_id()
-        if mem.percent > 70 or len(gui_sessions) > MAX_GUI_SESSIONS_PER_TASK:
-            self.set_status(503)
+        if mem.percent > TARGET_MEMORY_TRIGGER_PERCENT:
+            self.set_status(200)
+            self.write(
+                f"Busy (RAM Usage):\nMemory at {mem.percent}% - Num sessions: {len(gui_sessions)} "
+                f"(max {MAX_GUI_SESSIONS_PER_TASK}) Task ID: {task_id}"
+            )
+        elif len(gui_sessions) > MAX_GUI_SESSIONS_PER_TASK:
+            self.set_status(200)
+            busy_msg = "(MAX SESSIONS EXCEEDED)"
+            if _tmp_arr is not None:
+                busy_msg += " (inflating memory)"
+            elif _tmp_array_triggered:
+                busy_msg += " (inflated memory released)"
             self.write(
-                f"Busy:\nMemory at {mem.percent}% - Num sessions: {len(gui_sessions)} "
+                f"Busy {busy_msg}:\nMemory at {mem.percent}% - Num sessions: {len(gui_sessions)} "
                 f"(max {MAX_GUI_SESSIONS_PER_TASK}) Task ID: {task_id}"
             )
+            # inflate RAM once per threshold-exceeded event so ECS spawns a new task
+            if _tmp_arr is None and not _tmp_array_triggered:
+                _tmp_array_triggered = True
+                total_memory = psutil.virtual_memory().total
+                used_memory = psutil.virtual_memory().used
+                target_memory = total_memory * TARGET_MEMORY_TRIGGER_PERCENT / 100
+                array_size = int((target_memory - used_memory) / 8)  # assuming float64 (8 bytes)
+                _tmp_arr = np.ones(array_size, dtype=np.float64)
+                print(f"Inflated memory with array of size {array_size} to trigger ECS scaling")
+                _tmp_arr_timer = threading.Timer(TARGET_CLEAR_TMP_ARR_SECONDS, _clear_tmp_arr)
+                _tmp_arr_timer.daemon = True
+                _tmp_arr_timer.start()
         else:
+            _tmp_array_triggered = False
             self.set_status(200)
             self.write(
                 f"Healthy:\nMemory at {mem.percent}% - Num sessions: {len(gui_sessions)} "

src/aind_ephys_portal/panel/logging.py

@@ -85,7 +85,7 @@ def list_gui_sessions():
 
 def get_max_number_of_gui_sessions():
     # Estimate number of sessions per worker for health check.
-    SESSION_AVG_RAM_USAGE_GB = 2.5
+    SESSION_AVG_RAM_USAGE_GB = 2
     TOTAL_RAM_GB = psutil.virtual_memory().total / (1024**3)
     MAX_SESSIONS_PER_WORKER = int(np.floor(TOTAL_RAM_GB / SESSION_AVG_RAM_USAGE_GB))