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 <[email protected]> * Update docs/webhost configuration sample.yaml --------- Co-authored-by: Fabian Dill <[email protected]> (cherry picked from commit 4a5ba75)
ArchipelagoMW · Dec 10, 2024 · b934061 · b934061
1 parent 4938b5f
commit b934061
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 5 deletions.
diff --git a/WebHostLib/__init__.py b/WebHostLib/__init__.py
@@ -39,6 +39,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

diff --git a/WebHostLib/autolauncher.py b/WebHostLib/autolauncher.py
@@ -6,6 +6,7 @@
 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 keep_running():
         try:
             with Locker("autogen"):
 
-                with multiprocessing.Pool(config["GENERATORS"], initializer=init_db,
-                                          initargs=(config["PONY"],), maxtasksperchild=10) as generator_pool:
+                with multiprocessing.Pool(config["GENERATORS"], initializer=init_generator,
+                                          initargs=(config,), maxtasksperchild=10) as generator_pool:
                     with db_session:
                         to_start = select(generation for generation in Generation if generation.state == STATE_STARTED)
 

diff --git a/docs/webhost configuration sample.yaml b/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
-#JOB_THRESHOLD: 2
+# Slot limit to post a generation to Generator process pool instead of rolling directly in WebHost process
+#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