Added /terminate to free VRAM

Author: PrinsINT

base/Dockerfile

@@ -9,4 +9,4 @@ RUN pip install --no-deps --no-cache-dir uv && uv --no-cache pip install --syste
 
 EXPOSE 8080
 
-CMD ["./start.sh"]
+CMD ["python", "-u", "webservice.py"]

base/start.sh

@@ -31,17 +31,17 @@
 
 CALLBACK_SERVER: str = os.getenv("CALLBACK_SERVER") or ""
 NUM_WORKERS = int(os.getenv("NUM_WORKERS") or 1)
+# Needs to be defined as a reference object, e.g. dict.
+_global = {"pool": None}
 
 
 def run_pending_tasks() -> None:
     """
     Send a new task to the pool if there is no busy task.
     If there is no pool running, start a new pool.
     """
-    global pool
-
     # One task at a time.
-    tasks_in_queue = pool._taskqueue.qsize()
+    tasks_in_queue = _global["pool"]._taskqueue.qsize()
     if tasks_in_queue > 0:
         return
 
@@ -54,11 +54,11 @@ def run_pending_tasks() -> None:
             if sl.get_status()["busy"] is False:
                 sl.busy("Parsing file")  # Sets busy true
                 # Extra None check for typing
-                if (not is_pool_running(pool)) or pool is None:
+                if (not is_pool_running()) or _global["pool"] is None:
                     # Spawn pool if not running
-                    pool = mp.Pool(processes=NUM_WORKERS, initializer=process.init)
+                    _global["pool"] = mp.Pool(processes=NUM_WORKERS)
                 # Perform task at running pool
-                pool.apply_async(process_file, args=(sl.filename,))
+                _global["pool"].apply_async(process_file, args=(sl.filename,))
 
 
 def process_file(filename: str):
@@ -183,28 +183,31 @@ def send_error_to_callback_server(filename: str, out_path: str, message: str) ->
     keep_or_delete_file(r, out_path)
 
 
-def is_pool_running(pool: Optional[Pool]) -> bool:
+def is_pool_running() -> bool:
     """
     Check if the pool is running by trying to execute a dummy function.
     """
-    if pool is None:
+    if _global["pool"] is None:
         return False
     try:
-        pool.apply_async(lambda: None)
+        _global["pool"].apply_async(lambda: None)
     except ValueError as e:
         if str(e) == "Pool not running":
             return False
     return True
 
 
-# Pool needs to be defined after the functions it will execute.
-# https://stackoverflow.com/questions/41385708/multiprocessing-example-giving-attributeerror#comment101561695_42383397
-pool = None
+def terminate_pool() -> None:
+    """
+    Terminate the pool if it is running.
+    """
+    _global["pool"].terminate()
+
 
-if __name__ == "__main__":
+def run_background():
     # Can't use fork with the gpu.
     mp.set_start_method("spawn", force=True)
-    pool = mp.Pool(processes=NUM_WORKERS, initializer=process.init)
+    _global["pool"] = mp.Pool(processes=NUM_WORKERS)
 
     # It is ugly, but it is also used here:
     # https://pypi.org/project/schedule/

base/tagger_worker.py

@@ -13,6 +13,7 @@
 import os
 import subprocess
 import uuid
+import threading
 
 # Third-party
 import bottle
@@ -23,6 +24,7 @@
 from shared import OUTPUT_FOLDER, UPLOAD_FOLDER, ERROR_FOLDER
 from statuslogger import StatusLogger
 from process import OUTPUT_EXTENSION, PROCESSING_SPEED
+from tagger_worker import run_background, terminate_pool
 
 app = application = bottle.default_app()
 
@@ -163,4 +165,15 @@ def delete_file(id: str):
         return HTTPResponse("File " + id + " deleted", 200)
 
 
-app.run(host="0.0.0.0", port=8080)
+@post("/terminate")
+def terminate():
+    """
+    Terminate the worker pool immediately, freeing RAM (& VRAM).
+    """
+    terminate_pool()
+    return HTTPResponse("Terminated the worker pool", 200)
+
+
+if __name__ == "__main__":
+    threading.Thread(target=run_background, daemon=True).start()
+    app.run(host="0.0.0.0", port=8080)

base/webservice.py

@@ -34,6 +34,8 @@ def process(in_file: str, out_file: str) -> None:
     """
     Process the file at path "in_file" and write the result to path "out_file".
     """
+    if tagger is None or splitter is None:
+        init()
     with open(out_file, "w+", encoding="utf-8") as f_out:
         with open(in_file, "r", encoding="utf-8") as f_in:
             doc = f_in.read()

flair/process.py

@@ -50,6 +50,8 @@ def process(in_file: str, out_file: str) -> None:
     Pie outputs to the same directory as the input file.
     We process the file in a temporary directory so we don't polute /input.
     """
+    if tagger is None:
+        init()
     with tempfile.TemporaryDirectory() as temp_dir:
         # input
         temp_in_file = os.path.join(temp_dir, "file.txt")

pie/base/process.py

@@ -55,6 +55,8 @@ def process(in_file: str, out_file: str) -> None:
     """
     Process the file at path "in_file" and write the result to path "out_file".
     """
+    if nlp is None:
+        init()
     with open(out_file, "w+", encoding="utf-8") as f_out:
         with open(in_file, "r", encoding="utf-8") as f_in:
             is_xml = is_file_xml(in_file)

spacy/process.py

@@ -61,6 +61,8 @@ def process(in_file: str, out_file: str) -> None:
     """
     Process the file at path "in_file" and write the result to path "out_file".
     """
+    if xml_nlp is None or txt_nlp is None:
+        init()
     with open(out_file, "w+", encoding="utf-8") as f_out:
         with open(in_file, "r", encoding="utf-8") as f_in:
             is_xml = is_file_xml(in_file)