fix(ml): better model unloading (#3340)

* restart process on inactivity

* formatting

* always update `last_called`

* load models sequentially

* renamed variable, updated docs

* formatting

* made poll env name consistent with model ttl env

machine-learning/app/main.py

@@ -1,5 +1,9 @@
 import asyncio
+import gc
+import os
+import sys
 import threading
+import time
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 from zipfile import BadZipFile
@@ -34,7 +38,10 @@ def init_state() -> None:
     )
     # asyncio is a huge bottleneck for performance, so we use a thread pool to run blocking code
     app.state.thread_pool = ThreadPoolExecutor(settings.request_threads) if settings.request_threads > 0 else None
-    app.state.locks = {model_type: threading.Lock() for model_type in ModelType}
+    app.state.lock = threading.Lock()
+    app.state.last_called = None
+    if settings.model_ttl > 0 and settings.model_ttl_poll_s > 0:
+        asyncio.ensure_future(idle_shutdown_task())
     log.info(f"Initialized request thread pool with {settings.request_threads} threads.")
 
 
@@ -79,9 +86,9 @@ async def predict(
 
 
 async def run(model: InferenceModel, inputs: Any) -> Any:
+    app.state.last_called = time.time()
     if app.state.thread_pool is None:
         return model.predict(inputs)
-
     return await asyncio.get_running_loop().run_in_executor(app.state.thread_pool, model.predict, inputs)
 
 
@@ -90,7 +97,7 @@ async def load(model: InferenceModel) -> InferenceModel:
         return model
 
     def _load() -> None:
-        with app.state.locks[model.model_type]:
+        with app.state.lock:
             model.load()
 
     loop = asyncio.get_running_loop()
@@ -113,3 +120,27 @@ async def load(model: InferenceModel) -> InferenceModel:
         else:
             await loop.run_in_executor(app.state.thread_pool, _load)
         return model
+
+
+async def idle_shutdown_task() -> None:
+    while True:
+        log.debug("Checking for inactivity...")
+        if app.state.last_called is not None and time.time() - app.state.last_called > settings.model_ttl:
+            log.info("Shutting down due to inactivity.")
+            loop = asyncio.get_running_loop()
+            for task in asyncio.all_tasks(loop):
+                if task is not asyncio.current_task():
+                    try:
+                        task.cancel()
+                    except asyncio.CancelledError:
+                        pass
+            sys.stderr.close()
+            sys.stdout.close()
+            sys.stdout = sys.stderr = open(os.devnull, "w")
+            try:
+                await app.state.model_cache.cache.clear()
+                gc.collect()
+                loop.stop()
+            except asyncio.CancelledError:
+                pass
+        await asyncio.sleep(settings.model_ttl_poll_s)