WebHost: Set Generator memory limit to 4GiB (#4319)

* WebHost: Set Generator memory limit to 4GiB

* WebHost: make generator memory limit configurable, better naming

* Update WebHostLib/__init__.py

Co-authored-by: Fabian Dill <Berserker66@users.noreply.github.com>

* Update docs/webhost configuration sample.yaml

---------

Co-authored-by: Fabian Dill <Berserker66@users.noreply.github.com>

WebHostLib/__init__.py

@@ -39,6 +39,8 @@ app.config["SECRET_KEY"] = bytes(socket.gethostname(), encoding="utf-8")
 app.config["JOB_THRESHOLD"] = 1
 # after what time in seconds should generation be aborted, freeing the queue slot. Can be set to None to disable.
 app.config["JOB_TIME"] = 600
+# memory limit for generator processes in bytes
+app.config["GENERATOR_MEMORY_LIMIT"] = 4294967296
 app.config['SESSION_PERMANENT'] = True
 
 # waitress uses one thread for I/O, these are for processing of views that then get sent

WebHostLib/autolauncher.py

@@ -6,6 +6,7 @@ import multiprocessing
 import typing
 from datetime import timedelta, datetime
 from threading import Event, Thread
+from typing import Any
 from uuid import UUID
 
 from pony.orm import db_session, select, commit
@@ -53,7 +54,21 @@ def launch_generator(pool: multiprocessing.pool.Pool, generation: Generation):
         generation.state = STATE_STARTED
 
 
-def init_db(pony_config: dict):
+def init_generator(config: dict[str, Any]) -> None:
+    try:
+        import resource
+    except ModuleNotFoundError:
+        pass  # unix only module
+    else:
+        # set soft limit for memory to from config (default 4GiB)
+        soft_limit = config["GENERATOR_MEMORY_LIMIT"]
+        old_limit, hard_limit = resource.getrlimit(resource.RLIMIT_AS)
+        if soft_limit != old_limit:
+            resource.setrlimit(resource.RLIMIT_AS, (soft_limit, hard_limit))
+            logging.debug(f"Changed AS mem limit {old_limit} -> {soft_limit}")
+        del resource, soft_limit, hard_limit
+
+    pony_config = config["PONY"]
     db.bind(**pony_config)
     db.generate_mapping()
 
@@ -105,8 +120,8 @@ def autogen(config: dict):
         try:
             with Locker("autogen"):
 
-                with multiprocessing.Pool(config["GENERATORS"], initializer=init_db,
+                with multiprocessing.Pool(config["GENERATORS"], initializer=init_generator,
-                                          initargs=(config["PONY"],), maxtasksperchild=10) as generator_pool:
+                                          initargs=(config,), maxtasksperchild=10) as generator_pool:
                     with db_session:
                         to_start = select(generation for generation in Generation if generation.state == STATE_STARTED)
 

docs/webhost configuration sample.yaml

@@ -27,8 +27,14 @@
 # If you wish to deploy, uncomment the following line and set it to something not easily guessable.
 # SECRET_KEY: "Your secret key here"
 
-# TODO
+# Slot limit to post a generation to Generator process pool instead of rolling directly in WebHost process
-#JOB_THRESHOLD: 2
+#JOB_THRESHOLD: 1
+
+# After what time in seconds should generation be aborted, freeing the queue slot. Can be set to None to disable.
+#JOB_TIME: 600
+
+# Memory limit for Generator processes in bytes, -1 for unlimited. Currently only works on Linux.
+#GENERATOR_MEMORY_LIMIT: 4294967296
 
 # waitress uses one thread for I/O, these are for processing of view that get sent
 #WAITRESS_THREADS: 10