utils: Add profile.py to diagnose memory usage

Signed-off-by: Aaron LI <aly@aaronly.me>

1 files changed, 97 insertions, 0 deletions

diff --git a/fg21sim/utils/profile.py b/fg21sim/utils/profile.py
new file mode 100644
index 0000000..1ccbdb5
--- /dev/null
+++ b/fg21sim/utils/profile.py
@@ -0,0 +1,97 @@
+# Copyright (c) 2017 Weitian LI <weitian@aaronly.me>
+# MIT license
+
+"""
+Profiling utilities to help analyze the code.
+"""
+
+import sys
+import resource
+import gc
+from numbers import Number
+from collections import Set, Mapping, deque, defaultdict
+
+
+def getsize(obj):
+    """
+    Recursively iterate to sum size of object & members.
+
+    Returns
+    -------
+    size : int
+        The size of the object in units of "Bytes".
+
+    Credit
+    ------
+    * How do I determine the size of an object in Python?
+      https://stackoverflow.com/a/30316760/4856091
+    """
+    zero_depth_bases = (str, bytes, Number, range, bytearray)
+
+    def inner(obj, _seen_ids=set()):
+        obj_id = id(obj)
+        if obj_id in _seen_ids:
+            return 0
+
+        _seen_ids.add(obj_id)
+        size = sys.getsizeof(obj)
+        if isinstance(obj, zero_depth_bases):
+            # bypass remaining control flow and return
+            pass
+        elif isinstance(obj, (tuple, list, Set, deque)):
+            size += sum(inner(i) for i in obj)
+        elif isinstance(obj, Mapping) or hasattr(obj, "items"):
+            size += sum(inner(k) + inner(v) for k, v in obj.items())
+        # Check for custom object instances - may subclass above too
+        if hasattr(obj, "__dict__"):
+            size += inner(vars(obj))
+        if hasattr(obj, "__slots__"):
+            # can have ``__slots__`` with ``__dict__``
+            size += sum(inner(getattr(obj, s))
+                        for s in obj.__slots__ if hasattr(obj, s))
+        return size
+
+    return inner(obj)
+
+
+def get_objcounts():
+    """
+    Get the current objects and counts as an dictionary.
+
+    Credit
+    ------
+    * Working around memory leaks
+      https://stackoverflow.com/a/1641280/4856091
+    """
+    objcounts = defaultdict(int)
+    for obj in gc.get_objects():
+        objcounts[type(obj)] += 1
+    return objcounts
+
+
+def diff_objcounts(objc, objc_ref):
+    """
+    Compare the ``objc1`` to ``objc_ref`` and return the differences
+    in a list of ``(type, counts)``.
+
+    Credit
+    ------
+    * Working around memory leaks
+      https://stackoverflow.com/a/1641280/4856091
+    """
+    diff = [(k, objc[k]-objc_ref[k])
+            for k in objc
+            if (objc[k]-objc_ref[k])]
+    return diff
+
+
+def mem_usage(MiB=True):
+    """
+    Get the current memory usage.
+    """
+    usage = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss  # [KiB]
+    unit = "KiB"
+    if MiB:
+        usage /= 1024
+        unit = "MiB"
+    return (usage, unit)