Handle psutil virtual_memory failures

comfy/model_management.py

@@ -17,7 +17,6 @@
 """
 from __future__ import annotations
 
-import psutil
 import logging
 from enum import Enum
 from comfy.cli_args import args, PerformanceFeature
@@ -30,6 +29,7 @@ import gc
 import os
 from contextlib import contextmanager, nullcontext
 import comfy.memory_management
+import comfy.psutil_utils
 import comfy.utils
 import comfy.quant_ops
 import comfy_aimdo.host_buffer
@@ -317,7 +317,7 @@ def get_total_memory(dev=None, torch_total_too=False):
         dev = get_torch_device()
 
     if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
-        mem_total = psutil.virtual_memory().total
+        mem_total = comfy.psutil_utils.virtual_memory_total()
         mem_total_torch = mem_total
     else:
         if directml_enabled:
@@ -360,7 +360,7 @@ def mac_version():
         return None
 
 total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
-total_ram = psutil.virtual_memory().total / (1024 * 1024)
+total_ram = comfy.psutil_utils.virtual_memory_total() / (1024 * 1024)
 logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
 
 try:
@@ -648,7 +648,7 @@ def ensure_pin_budget(size, evict_active=False):
     if args.fast_disk:
         shortfall = TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY
     else:
-        shortfall = size + max(comfy.memory_management.RAM_CACHE_HEADROOM / 2, 2048 * 1024 ** 2) - psutil.virtual_memory().available
+        shortfall = size + max(comfy.memory_management.RAM_CACHE_HEADROOM / 2, 2048 * 1024 ** 2) - comfy.psutil_utils.virtual_memory_available()
     if shortfall <= 0:
         return True
 
@@ -1656,7 +1656,7 @@ def get_free_memory(dev=None, torch_free_too=False):
         dev = get_torch_device()
 
     if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
-        mem_free_total = psutil.virtual_memory().available
+        mem_free_total = comfy.psutil_utils.virtual_memory_available()
         mem_free_torch = mem_free_total
     else:
         if directml_enabled:

comfy/psutil_utils.py

@@ -0,0 +1,37 @@
+import logging
+import os
+from collections import namedtuple
+
+import psutil
+
+_FallbackVMem = namedtuple("svmem", ["total", "available"])
+_virtual_memory_warned = False
+
+
+def _fallback_total_memory():
+    try:
+        return os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
+    except (AttributeError, OSError, ValueError):
+        return 0
+
+
+def virtual_memory(default_total=None, default_available=None):
+    global _virtual_memory_warned
+    try:
+        return psutil.virtual_memory()
+    except RuntimeError as e:
+        if not _virtual_memory_warned:
+            logging.warning("psutil.virtual_memory() failed; using fallback memory values: %s", e)
+            _virtual_memory_warned = True
+
+        total = default_total if default_total is not None else _fallback_total_memory()
+        available = default_available if default_available is not None else total
+        return _FallbackVMem(total=total, available=available)
+
+
+def virtual_memory_available(default=None):
+    return virtual_memory(default_available=default).available
+
+
+def virtual_memory_total(default=None):
+    return virtual_memory(default_total=default).total

comfy_execution/caching.py

@@ -1,11 +1,11 @@
 import asyncio
 import bisect
 import itertools
-import psutil
 import time
 import torch
 from typing import Sequence, Mapping, Dict
 from comfy.model_patcher import ModelPatcher
+import comfy.psutil_utils
 from comfy_execution.graph import DynamicPrompt
 from abc import ABC, abstractmethod
 
@@ -525,7 +525,7 @@ class RAMPressureCache(LRUCache):
         super().set_local(node_id, value)
 
     def ram_release(self, target, free_active=False):
-        if psutil.virtual_memory().available >= target:
+        if comfy.psutil_utils.virtual_memory_available(default=target) >= target:
             return
 
         clean_list = []
@@ -555,7 +555,7 @@ class RAMPressureCache(LRUCache):
             #break OOM score ties on the last touch timestamp (pure LRU)
             bisect.insort(clean_list, (oom_score, self.timestamps[key], key))
 
-        while psutil.virtual_memory().available < target and clean_list:
+        while comfy.psutil_utils.virtual_memory_available(default=target) < target and clean_list:
             _, _, key = clean_list.pop()
             del self.cache[key]
             self.used_generation.pop(key, None)

execution.py

@@ -2,7 +2,6 @@ import copy
 import heapq
 import inspect
 import logging
-import psutil
 import sys
 import threading
 import time
@@ -15,6 +14,7 @@ import torch
 
 from comfy.cli_args import args
 import comfy.memory_management
+import comfy.psutil_utils
 import comfy.model_management
 import comfy.model_prefetch
 import comfy_aimdo.model_vbar
@@ -793,7 +793,7 @@ class PromptExecutor:
 
                     if self.cache_type == CacheType.RAM_PRESSURE:
                         ram_release_callback(ram_inactive_headroom)
-                        ram_shortfall = ram_headroom - psutil.virtual_memory().available
+                        ram_shortfall = ram_headroom - comfy.psutil_utils.virtual_memory_available(default=ram_headroom)
                         freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
                         if freed < ram_shortfall:
                             if freed > 64 * (1024 ** 2):