Source code for sfepy.base.mem_usage

"""
Memory usage functions.
"""
from __future__ import absolute_import
import sys
import collections

import numpy as nm
import scipy.sparse as sp

from sfepy.base.base import basestr, Struct, Output
import six

[docs] def get_mem_usage(obj, usage=None, name=None, traversal_order=None, level=0): """ Get lower bound of memory usage of an object. Takes into account strings, numpy arrays and scipy CSR sparse matrices, descends into sequences, mappings and objects. Parameters ---------- obj : any object The object to be measured. usage : dict The dict with memory usage records, serving also as a cache of already traversed objects. name : str The name to be given to the object in its record. traversal_order : list, internal The traversal order of the object. level : int, internal The recurrence level. Returns ------- usage : int The object's lower bound of memory usage. """ if usage is None: usage = {} if name is None: name = getattr(obj, 'name', '-') if traversal_order is None: traversal_order = [0] to = traversal_order key = id(obj) if key in usage: usage[key].nrefs += 1 return 0 else: record = usage.setdefault(key, Struct(name=name, kind=type(obj).__name__, usage=0, nrefs=1, traversal_order=to[0], level=level)) level += 1 if isinstance(obj, nm.ndarray): record.usage = obj.nbytes elif isinstance(obj, sp.csr_matrix): record.usage = (get_mem_usage(obj.data, usage, name='data', traversal_order=to, level=level) + get_mem_usage(obj.indices, usage, name='indices', traversal_order=to, level=level) + get_mem_usage(obj.indptr, usage, name='indptr', traversal_order=to, level=level)) elif isinstance(obj, basestr): record.usage = len(obj) elif isinstance(obj, Struct): for subname, sub in six.iteritems(obj.__dict__): to[0] += 1 record.usage += get_mem_usage(sub, usage, name='attribute %s of %s' % (subname, getattr(obj, 'name', record.kind)), traversal_order=to, level=level) elif isinstance(obj, collections.Mapping): try: for subname, sub in six.iteritems(obj): to[0] += 1 record.usage += get_mem_usage(sub, usage, name='item %s of %s' % (subname, record.kind), traversal_order=to, level=level) except: pass elif isinstance(obj, collections.Sequence): for ii, sub in enumerate(obj): to[0] += 1 record.usage += get_mem_usage(sub, usage, name='item %d of %s' % (ii, record.kind), traversal_order=to, level=level) else: record.usage = sys.getsizeof(obj) return record.usage
[docs] def raise_if_too_large(size, factor=1.0): """ Raise MemoryError if the total system memory is lower than `size` times safety `factor`. Use `factor=None` for skipping the memory check. """ if factor is None: return import psutil mem = psutil.virtual_memory() if factor * size > mem.total: mb = 1000**2 raise MemoryError('insufficent memory {} MB to allocate {} MB' ' with safety factor {}' .format(mem.total/mb, size/mb, factor))