wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-03-18 15:44:55 +08:00
Code Issues Actions 18 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 os
import gc
import itertools
import psutil
@ -382,45 +383,47 @@ class RAMPressureCache(LRUCache):
return super().get(node_id)
def poll(self, ram_headroom):
def _ram_gb():
"""读取容器可用内存（GB），优先从 cgroup v2/v1 读取"""
# === 容器内存检测核心逻辑 ===
# 1. 尝试读取 cgroup v2 内存限制（Docker/K8s 主流）
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):
cgroup_mem_limit_path = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
cgroup_mem_usage_path = "/sys/fs/cgroup/memory/memory.usage_in_bytes"
# 读取内存限制（字节）
try:
with open(cgroup_mem_limit_path, "r") as f:
mem_limit = int(f.read().strip())
# cgroup v1 中 limit_in_bytes 为 9223372036854771712 表示无限制， fallback 到系统内存
if mem_limit == 9223372036854771712:
mem_limit = psutil.virtual_memory().total
except (FileNotFoundError, ValueError):
mem_limit = psutil.virtual_memory().total
# 读取已使用内存（字节）
try:
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）
mem_available = mem_limit - mem_used
return mem_available / (1024**3)
# === 替换结束 ===
if _ram_gb() > ram_headroom:
return
gc.collect()
if _ram_gb() > ram_headroom:
return
def _ram_gb():
fallback_to_host_available = False
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_mem_usage_path = "/sys/fs/cgroup/memory/memory.usage_in_bytes"
try:
with open(cgroup_mem_limit_path, "r") as f:
raw_limit = f.read().strip()
if raw_limit == "max":
fallback_to_host_available = True
mem_limit = psutil.virtual_memory().total
else:
mem_limit = int(raw_limit)
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:
return psutil.virtual_memory().available / (1024**3)
try:
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
mem_available = max(0, mem_limit - mem_used)
return mem_available / (1024**3)
if _ram_gb() > ram_headroom:
return
gc.collect()
if _ram_gb() > ram_headroom:
return
clean_list = []