wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-03-22 17:43:33 +08:00
Code Issues Actions 10 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

81
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" if not os.path.exists(cgroup_mem_limit_path):
cgroup_mem_limit_path = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
# 兼容 cgroup v1 cgroup_mem_usage_path = "/sys/fs/cgroup/memory/memory.usage_in_bytes"
if not os.path.exists(cgroup_mem_limit_path):
cgroup_mem_limit_path = "/sys/fs/cgroup/memory/memory.limit_in_bytes" try:
cgroup_mem_usage_path = "/sys/fs/cgroup/memory/memory.usage_in_bytes" with open(cgroup_mem_limit_path, "r") as f:
raw_limit = f.read().strip()
# 读取内存限制（字节） if raw_limit == "max":
try: fallback_to_host_available = True
with open(cgroup_mem_limit_path, "r") as f: mem_limit = psutil.virtual_memory().total
mem_limit = int(f.read().strip()) else:
# cgroup v1 中 limit_in_bytes 为 9223372036854771712 表示无限制， fallback 到系统内存 mem_limit = int(raw_limit)
if mem_limit == 9223372036854771712: if mem_limit == 9223372036854771712:
mem_limit = psutil.virtual_memory().total fallback_to_host_available = True
except (FileNotFoundError, ValueError): mem_limit = psutil.virtual_memory().total
mem_limit = psutil.virtual_memory().total except (FileNotFoundError, ValueError):
fallback_to_host_available = True
# 读取已使用内存（字节） mem_limit = psutil.virtual_memory().total
try:
with open(cgroup_mem_usage_path, "r") as f: if fallback_to_host_available:
mem_used = int(f.read().strip()) return psutil.virtual_memory().available / (1024**3)
except (FileNotFoundError, ValueError):
mem_used = psutil.virtual_memory().total - psutil.virtual_memory().available try:
with open(cgroup_mem_usage_path, "r") as f:
# 计算可用内存（GB） mem_used = int(f.read().strip())
mem_available = mem_limit - mem_used except (FileNotFoundError, ValueError):
return mem_available / (1024**3) mem_used = psutil.virtual_memory().total - psutil.virtual_memory().available
# === 替换结束 ===
mem_available = max(0, mem_limit - mem_used)
if _ram_gb() > ram_headroom: return mem_available / (1024**3)
return
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 = []