wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-03-23 01:53:32 +08:00
Code Issues Actions 8 Packages Projects Releases Wiki Activity

Adapt to obtaining resources within Docker.

Browse Source
This commit is contained in:
root 2026-03-09 15:45:31 +08:00
parent 3fb4fe9a2b
commit 6fdde2271a
1 changed files with 42 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
comfy_execution/caching.py
View File

@ -1,4 +1,5 @@
import bisect import bisect
import os
import gc import gc
import itertools import itertools
import psutil import psutil
@ -382,45 +383,47 @@ class RAMPressureCache(LRUCache):
return super().get(node_id) return super().get(node_id)
def poll(self, ram_headroom): def poll(self, ram_headroom):
def _ram_gb(): def _ram_gb():
"""读取容器可用内存（GB），优先从 cgroup v2/v1 读取""" fallback_to_host_available = False
# === 容器内存检测核心逻辑 === cgroup_mem_limit_path = "/sys/fs/cgroup/memory.max"
# 1. 尝试读取 cgroup v2 内存限制（Docker/K8s 主流） cgroup_mem_usage_path = "/sys/fs/cgroup/memory.current"
cgroup_mem_limit_path = "/sys/fs/cgroup/memory.max"
cgroup_mem_usage_path = "/sys/fs/cgroup/memory.current"
# 兼容 cgroup v1 if not os.path.exists(cgroup_mem_limit_path):
if not os.path.exists(cgroup_mem_limit_path): cgroup_mem_limit_path = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
cgroup_mem_limit_path = "/sys/fs/cgroup/memory/memory.limit_in_bytes" cgroup_mem_usage_path = "/sys/fs/cgroup/memory/memory.usage_in_bytes"
cgroup_mem_usage_path = "/sys/fs/cgroup/memory/memory.usage_in_bytes"
# 读取内存限制（字节） try:
try: with open(cgroup_mem_limit_path, "r") as f:
with open(cgroup_mem_limit_path, "r") as f: raw_limit = f.read().strip()
mem_limit = int(f.read().strip()) if raw_limit == "max":
# cgroup v1 中 limit_in_bytes 为 9223372036854771712 表示无限制， fallback 到系统内存 fallback_to_host_available = True
if mem_limit == 9223372036854771712: mem_limit = psutil.virtual_memory().total
mem_limit = psutil.virtual_memory().total else:
except (FileNotFoundError, ValueError): mem_limit = int(raw_limit)
mem_limit = psutil.virtual_memory().total if mem_limit == 9223372036854771712:
fallback_to_host_available = True
mem_limit = psutil.virtual_memory().total
except (FileNotFoundError, ValueError):
fallback_to_host_available = True
mem_limit = psutil.virtual_memory().total
# 读取已使用内存（字节） if fallback_to_host_available:
try: return psutil.virtual_memory().available / (1024**3)
with open(cgroup_mem_usage_path, "r") as f:
mem_used = int(f.read().strip())
except (FileNotFoundError, ValueError):
mem_used = psutil.virtual_memory().total - psutil.virtual_memory().available
# 计算可用内存（GB） try:
mem_available = mem_limit - mem_used with open(cgroup_mem_usage_path, "r") as f:
return mem_available / (1024**3) mem_used = int(f.read().strip())
# === 替换结束 === except (FileNotFoundError, ValueError):
mem_used = psutil.virtual_memory().total - psutil.virtual_memory().available
if _ram_gb() > ram_headroom: mem_available = max(0, mem_limit - mem_used)
return return mem_available / (1024**3)
gc.collect()
if _ram_gb() > ram_headroom: if _ram_gb() > ram_headroom:
return return
gc.collect()
if _ram_gb() > ram_headroom:
return
clean_list = [] clean_list = []