diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index 5a8856aba..bbfbd2ce8 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -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 = []