parse cgroups

Signed-off-by: Vladimir Mandic <[email protected]>

launch.py

@@ -151,13 +151,9 @@ def run_extension_installer(ext_dir): # compatbility function
 
 
 def get_memory_stats():
-    import psutil
-    def gb(val: float):
-        return round(val / 1024 / 1024 / 1024, 2)
-    process = psutil.Process(os.getpid())
-    res = process.memory_info()
-    ram_total = 100 * res.rss / process.memory_percent()
-    return f'{gb(res.rss)}/{gb(ram_total)}'
+    from modules.memstats import ram_stats
+    res = ram_stats()
+    return f'{res["used"]}/{res["total"]}'
 
 
 def start_server(immediate=True, server=None):

modules/memstats.py

@@ -7,19 +7,42 @@
 
 fail_once = False
 mem = {}
+docker_limit = None
+runpod_limit = None
 
 
 def gb(val: float):
     return round(val / 1024 / 1024 / 1024, 2)
 
 
+def get_docker_limit():
+    global docker_limit # pylint: disable=global-statement
+    if docker_limit is not None:
+        return docker_limit
+    try:
+        with open('/sys/fs/cgroup/memory/memory.limit_in_bytes', 'r', encoding='utf8') as f:
+            docker_limit = float(f.read())
+    except Exception:
+        docker_limit = sys.float_info.max
+    return docker_limit
+
+
+def get_runpod_limit():
+    global runpod_limit # pylint: disable=global-statement
+    if runpod_limit is not None:
+        return runpod_limit
+    runpod_limit = float(os.environ.get('RUNPOD_MEM_GB', sys.float_info.max))
+    return runpod_limit
+
+
 def memory_stats():
     global fail_once # pylint: disable=global-statement
     mem.clear()
     try:
         process = psutil.Process(os.getpid())
         res = process.memory_info()
         ram_total = 100 * res.rss / process.memory_percent()
+        ram_total = min(ram_total, get_docker_limit(), get_runpod_limit())
         ram = { 'used': gb(res.rss), 'total': gb(ram_total) }
         mem.update({ 'ram': ram })
     except Exception as e:
@@ -54,6 +77,7 @@ def ram_stats():
         process = psutil.Process(os.getpid())
         res = process.memory_info()
         ram_total = 100 * res.rss / process.memory_percent()
+        ram_total = min(ram_total, docker_limit(), runpod_limit())
         ram = { 'used': gb(res.rss), 'total': gb(ram_total) }
         return ram
     except Exception:

modules/shared.py

@@ -7,9 +7,7 @@
 import threading
 import contextlib
 from types import SimpleNamespace
-from urllib.parse import urlparse
 from enum import Enum
-import psutil
 import requests
 import gradio as gr
 import fasteners
@@ -205,7 +203,7 @@ def default(obj):
 devices.backend = devices.get_backend(cmd_opts)
 devices.device = devices.get_optimal_device()
 mem_stat = memory_stats()
-cpu_memory = round(psutil.virtual_memory().total / 1024 / 1024 / 1024, 2)
+cpu_memory = mem_stat['ram']['total'] if "ram" in mem_stat else 0
 gpu_memory = mem_stat['gpu']['total'] if "gpu" in mem_stat else 0
 native = backend == Backend.DIFFUSERS
 if not files_cache.do_cache_folders:
@@ -307,6 +305,7 @@ def list_checkpoint_titles():
 
 
 def is_url(string):
+    from urllib.parse import urlparse
     parsed_url = urlparse(string)
     return all([parsed_url.scheme, parsed_url.netloc])