mirror of
https://github.com/vladmandic/automatic
synced 2026-08-25 22:20:46 +02:00
CalamitousFelicitousness
d38f2c97b2
Move the per-component placement walk from the api server into memstats next to ram and gpu stats, and add it to the --monitor tick so offload placement is visible in a log rather than only over the api. - guard the walk internally: the supervisor loop logs its monitor line unguarded, and the walk can race a reload or an offload rewrap - monitor reports gb and merges into a fresh dict, since memory_stats returns a module global that the per-generation log also prints - endpoint keeps reporting raw bytes
235 lines
7.6 KiB
Python
235 lines
7.6 KiB
Python
import re
|
|
import sys
|
|
import os
|
|
import types
|
|
from collections import deque
|
|
import psutil
|
|
import torch
|
|
from modules import shared, errors, devices
|
|
from modules.logger import log
|
|
|
|
|
|
fail_once = False
|
|
ram = {}
|
|
gpu = {}
|
|
mem = {}
|
|
process = None
|
|
docker_limit = None
|
|
runpod_limit = None
|
|
|
|
|
|
def gb(val: float):
|
|
return round(val / 1024 / 1024 / 1024, 2)
|
|
|
|
|
|
def get_docker_limit():
|
|
global docker_limit # pylint: disable=global-statement
|
|
if docker_limit is not None:
|
|
return docker_limit
|
|
try:
|
|
with open('/sys/fs/cgroup/memory/memory.limit_in_bytes', encoding='utf8') as f:
|
|
docker_limit = float(f.read())
|
|
except Exception:
|
|
docker_limit = sys.float_info.max
|
|
if docker_limit == 0:
|
|
docker_limit = sys.float_info.max
|
|
return docker_limit
|
|
|
|
|
|
def get_runpod_limit():
|
|
global runpod_limit # pylint: disable=global-statement
|
|
if runpod_limit is not None:
|
|
return runpod_limit
|
|
runpod_limit = float(os.environ.get('RUNPOD_MEM_GB', 0)) * 1024 * 1024 * 1024
|
|
if runpod_limit == 0:
|
|
runpod_limit = sys.float_info.max
|
|
return runpod_limit
|
|
|
|
|
|
def ram_stats():
|
|
global process, fail_once # pylint: disable=global-statement
|
|
try:
|
|
if process is None:
|
|
process = psutil.Process(os.getpid())
|
|
res = process.memory_info()
|
|
if 'total' not in ram:
|
|
process = psutil.Process(os.getpid())
|
|
mem_percent = process.memory_percent()
|
|
if mem_percent > 0:
|
|
ram_total = 100 * res.rss / mem_percent
|
|
else:
|
|
ram_total = res.rss
|
|
ram_total = min(ram_total, get_docker_limit(), get_runpod_limit())
|
|
ram['total'] = gb(ram_total)
|
|
ram['rss'] = gb(res.rss)
|
|
except Exception as e:
|
|
ram['total'] = 0
|
|
ram['rss'] = 0
|
|
ram['error'] = str(e)
|
|
if not fail_once:
|
|
log.error(f'RAM stats: {e}')
|
|
errors.display(e, 'RAM stats')
|
|
fail_once = True
|
|
try:
|
|
vmem = psutil.virtual_memory()
|
|
ram['used'] = gb(vmem.used) if hasattr(vmem, 'used') else 0
|
|
ram['free'] = gb(vmem.free) if hasattr(vmem, 'free') else 0
|
|
ram['avail'] = gb(vmem.available) if hasattr(vmem, 'available') else 0
|
|
ram['buffers'] = gb(vmem.buffers) if hasattr(vmem, 'buffers') else 0
|
|
ram['cached'] = gb(vmem.cached) if hasattr(vmem, 'cached') else 0
|
|
except Exception as e:
|
|
ram['used'] = 0
|
|
ram['free'] = 0
|
|
ram['avail'] = 0
|
|
ram['buffers'] = 0
|
|
ram['cached'] = 0
|
|
ram['error'] = str(e)
|
|
if not fail_once:
|
|
log.error(f'RAM stats: {e}')
|
|
errors.display(e, 'RAM stats')
|
|
fail_once = True
|
|
return ram
|
|
|
|
|
|
def gpu_stats():
|
|
global fail_once # pylint: disable=global-statement
|
|
try:
|
|
free, total = torch.cuda.mem_get_info()
|
|
gpu['used'] = gb(total - free)
|
|
gpu['total'] = gb(total)
|
|
stats = dict(torch.cuda.memory_stats())
|
|
if stats.get('num_ooms', 0) > 0:
|
|
shared.state.oom = True
|
|
gpu['active'] = gb(stats.get('active_bytes.all.current', 0))
|
|
gpu['peak'] = gb(stats.get('active_bytes.all.peak', 0))
|
|
gpu['retries'] = stats.get('num_alloc_retries', 0)
|
|
gpu['oom'] = stats.get('num_ooms', 0)
|
|
except Exception as e:
|
|
gpu['total'] = 0
|
|
gpu['used'] = 0
|
|
gpu['error'] = str(e)
|
|
if not fail_once:
|
|
# log.warning(f'GPU stats: {e}')
|
|
# errors.display(e, 'GPU stats')
|
|
fail_once = True
|
|
return gpu
|
|
|
|
|
|
def model_stats(as_gb: bool = False):
|
|
"""Loaded-model bytes per component and device, so resident weights can be told from offloaded ones."""
|
|
try:
|
|
from modules.modeldata import model_data
|
|
pipe = model_data.sd_model # raw slot: the shared.sd_model property can trigger a model load
|
|
if pipe is None:
|
|
return {}
|
|
components = getattr(pipe, 'components', None) or ({ 'model': pipe } if isinstance(pipe, torch.nn.Module) else {})
|
|
placement = {}
|
|
seen = set()
|
|
for name, component in components.items():
|
|
if not isinstance(component, torch.nn.Module):
|
|
continue
|
|
devmap = {}
|
|
for tensors in (component.parameters(), component.buffers()):
|
|
for t in tensors:
|
|
ptr = 0 if t.is_meta else t.untyped_storage().data_ptr()
|
|
if ptr:
|
|
if ptr in seen: # tied weights and offload rewraps share one storage across tensors
|
|
continue
|
|
seen.add(ptr)
|
|
devmap[t.device.type] = devmap.get(t.device.type, 0) + t.numel() * t.element_size()
|
|
if devmap:
|
|
placement[name] = { d: gb(v) for d, v in devmap.items() } if as_gb else devmap
|
|
return placement
|
|
except Exception as err: # walk can race a reload or an offload rewrap; every caller is a diagnostic that must not take its caller down
|
|
return { 'error': f'{err}' }
|
|
|
|
|
|
def memory_stats():
|
|
mem['ram'] = ram_stats()
|
|
mem['gpu'] = gpu_stats()
|
|
mem['job'] = shared.state.job
|
|
try:
|
|
mem['gpu']['swap'] = round(mem['gpu']['active'] - mem['gpu']['used']) if mem['gpu']['active'] > mem['gpu']['used'] else 0
|
|
except Exception:
|
|
mem['gpu']['swap'] = 0
|
|
return mem
|
|
|
|
|
|
def reset_stats():
|
|
try:
|
|
torch.cuda.reset_memory_stats()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class Object:
|
|
pattern = r"'(.*?)'"
|
|
|
|
def get_size(self, obj, seen=None):
|
|
size = sys.getsizeof(obj)
|
|
if seen is None:
|
|
seen = set()
|
|
obj_id = id(obj)
|
|
if obj_id in seen:
|
|
return 0 # Avoid double counting
|
|
seen.add(obj_id)
|
|
if isinstance(obj, dict):
|
|
size += sum(self.get_size(k, seen) + self.get_size(v, seen) for k, v in obj.items())
|
|
elif isinstance(obj, (list, tuple, set, frozenset, deque)):
|
|
size += sum(self.get_size(i, seen) for i in obj)
|
|
return size
|
|
|
|
def __init__(self, name, obj):
|
|
self.id = id(obj)
|
|
self.name = name
|
|
self.fn = sys._getframe(2).f_code.co_name
|
|
self.refcount = sys.getrefcount(obj)
|
|
if torch.is_tensor(obj):
|
|
self.type = obj.dtype
|
|
self.size = obj.element_size() * obj.nelement()
|
|
else:
|
|
self.type = re.findall(self.pattern, str(type(obj)))[0]
|
|
self.size = self.get_size(obj)
|
|
def __str__(self):
|
|
return f'{self.fn}.{self.name} type={self.type} size={self.size} ref={self.refcount}'
|
|
|
|
|
|
def get_objects(gcl=None, threshold:int=1024*1024):
|
|
devices.torch_gc(force=True)
|
|
if gcl is None:
|
|
# gcl = globals()
|
|
gcl = {}
|
|
log.trace(f'Memory: modules={len(sys.modules)}')
|
|
for _module_name, module in sys.modules.items():
|
|
try:
|
|
if not isinstance(module, types.ModuleType):
|
|
continue
|
|
namespace = vars(module)
|
|
gcl.update(namespace)
|
|
except Exception:
|
|
pass # Some modules may not allow introspection
|
|
objects = []
|
|
seen = []
|
|
|
|
log.trace(f'Memory: items={len(gcl)} threshold={threshold}')
|
|
for name, obj in gcl.items():
|
|
if id(obj) in seen:
|
|
continue
|
|
seen.append(id(obj))
|
|
if name == '__name__':
|
|
name = obj
|
|
elif name.startswith('__'):
|
|
continue
|
|
try:
|
|
o = Object(name, obj)
|
|
if o.size >= threshold:
|
|
objects.append(o)
|
|
except Exception:
|
|
pass
|
|
|
|
objects = sorted(objects, key=lambda x: x.size, reverse=True)
|
|
for obj in objects:
|
|
log.trace(f'Memory: {obj}')
|
|
|
|
return objects
|