Take a snapshot in a function.

take_inner_snapshot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.append(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tracemalloc_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def foo(x, y, z):\n",
    "    dataset1 = np.random.uniform(low=-1., high=1., size=100).astype(np.float64)\n",
    "    print('x', x)\n",
    "    dataset2 = np.random.uniform(low=-1., high=1., size=1000).astype(np.float64)\n",
    "\n",
    "    if x == 0:\n",
    "        return 0\n",
    "    elif x == 1:\n",
    "        return 1\n",
    "\n",
    "    dataset3 = np.random.uniform(low=-1., high=1., size=3000).astype(np.float64)\n",
    "    return 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x 20\n",
      "Top 10 lines\n",
      "#1: <tracemalloc_utils-src>:11: 23.5 KiB\n",
      "    dataset3 = np.random.uniform(low=-1., high=1., size=3000).astype(np.float64)\n",
      "#2: <tracemalloc_utils-src>:4: 7.9 KiB\n",
      "    dataset2 = np.random.uniform(low=-1., high=1., size=1000).astype(np.float64)\n",
      "#3: <tracemalloc_utils-src>:2: 1.6 KiB\n",
      "    dataset1 = np.random.uniform(low=-1., high=1., size=100).astype(np.float64)\n",
      "#4: <tracemalloc_utils-src>:3: 0.3 KiB\n",
      "    print('x', x)\n",
      "Total allocated size: 33.4 KiB\n"
     ]
    }
   ],
   "source": [
    "snapshot, source = tracemalloc_utils.take_inner_snapshot(\n",
    "    function=foo,\n",
    "    function_args=dict(x=20, y=11, z=12),\n",
    "    setup='import numpy as np',\n",
    ")\n",
    "tracemalloc_utils.display_top(snapshot, source=source)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以下は通常の呼び出し"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tracemalloc\n",
    "import linecache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_top(snapshot, key_type='lineno', limit=10):\n",
    "    snapshot = snapshot.filter_traces((\n",
    "        tracemalloc.Filter(False, \"<frozen importlib._bootstrap>\"),\n",
    "        tracemalloc.Filter(False, \"<unknown>\"),\n",
    "    ))\n",
    "    top_stats = snapshot.statistics(key_type)\n",
    "\n",
    "    print(\"Top %s lines\" % limit)\n",
    "    for index, stat in enumerate(top_stats[:limit], 1):\n",
    "        frame = stat.traceback[0]\n",
    "        # replace \"/path/to/module/file.py\" with \"module/file.py\"\n",
    "        filename = os.sep.join(frame.filename.split(os.sep)[-2:])\n",
    "        print(\"#%s: %s:%s: %.1f KiB\"\n",
    "              % (index, filename, frame.lineno, stat.size / 1024))\n",
    "        line = linecache.getline(frame.filename, frame.lineno).strip()\n",
    "        if line:\n",
    "            print('    %s' % line)\n",
    "\n",
    "    other = top_stats[limit:]\n",
    "    if other:\n",
    "        size = sum(stat.size for stat in other)\n",
    "        print(\"%s other: %.1f KiB\" % (len(other), size / 1024))\n",
    "    total = sum(stat.size for stat in top_stats)\n",
    "    print(\"Total allocated size: %.1f KiB\" % (total / 1024))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "tracemalloc.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x 20\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "foo(x=20, y=11, z=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top 10 lines\n",
      "#1: lib\\codeop.py:133: 6.2 KiB\n",
      "    codeob = compile(source, filename, symbol, self.flags, 1)\n",
      "#2: core\\compilerop.py:99: 5.0 KiB\n",
      "    return compile(source, filename, symbol, self.flags | PyCF_ONLY_AST, 1)\n",
      "#3: traitlets\\traitlets.py:600: 4.8 KiB\n",
      "    elif hasattr(obj, '_%s_validate' % self.name):\n",
      "#4: core\\inputsplitter.py:636: 4.7 KiB\n",
      "    out = _flush(t, out)\n",
      "#5: config\\configurable.py:121: 4.4 KiB\n",
      "    cfgs.append(self.parent._find_my_config(cfg))\n",
      "#6: traitlets\\traitlets.py:459: 4.1 KiB\n",
      "    self.metadata = self.metadata.copy()\n",
      "#7: config\\configurable.py:122: 3.6 KiB\n",
      "    my_config = Config()\n",
      "#8: traitlets\\traitlets.py:1102: 3.5 KiB\n",
      "    self.notify_change = hold\n",
      "#9: ipykernel\\jsonutil.py:191: 3.4 KiB\n",
      "    out[unicode_type(k)] = json_clean(v)\n",
      "#10: traitlets\\traitlets.py:234: 3.3 KiB\n",
      "    value = getattr(object, key)\n",
      "206 other: 95.4 KiB\n",
      "Total allocated size: 138.3 KiB\n"
     ]
    }
   ],
   "source": [
    "snapshot = tracemalloc.take_snapshot()\n",
    "display_top(snapshot)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}

tracemalloc_utils.py

import os
import inspect
import ast
import types
import linecache
from tracemalloc import start, take_snapshot, Filter


DUMMY_SRC_NAME = '<tracemalloc_utils-src>'


class Transformer(ast.NodeTransformer):
    '''Add tracemalloc functions.'''
    def __init__(self, result_id):
        self._result_id = result_id

    def visit_FunctionDef(self, node):
    
        pre_hook_expr = ast.Expr(
            value=ast.Call(
                func=ast.Name(id='start', ctx=ast.Load()),
                args=[],
                keywords=[]
            )
        )

        finalbody = [
            ast.Global(names=[self._result_id]),
            ast.Assign(
                targets=[ast.Name(id=self._result_id, ctx=ast.Store())],
                value=ast.Call(
                    func=ast.Name(id='take_snapshot', ctx=ast.Load()),
                    args=[],
                    keywords=[]
                )
            )
        ]
        
        body_elems = [pre_hook_expr]
        body_elems.extend([elem for elem in node.body])
        node.body.clear()
        node.body.append(
            ast.Try(
                body=body_elems,
                handlers=[],
                orelse=[],
                finalbody=finalbody
            )
        )
        return ast.fix_missing_locations(node)    


def take_inner_snapshot(
    function,
    function_args,
    setup='pass'
):
    '''Take a inner snapshot.'''
    if not isinstance(function, types.FunctionType):
        raise TypeError

    if not isinstance(function_args, dict):
        raise TypeError
    
    if not isinstance(setup, str):
        raise TypeError

    # Add modules temporarily.
    temp = {'SNAPSHOT': None}
    code = compile(setup, DUMMY_SRC_NAME, 'exec')
    exec(code, globals(), temp)

    globals().update(temp)
    global SNAPSHOT

    # Modify the function.
    source = inspect.getsource(function)
    node = ast.parse(source)
    node = Transformer(result_id='SNAPSHOT').visit(node)

    # Take the snapshot.
    _locals = {}
    code = compile(node, DUMMY_SRC_NAME, 'exec')
    exec(code, globals(), _locals)
    dst_function = _locals[function.__name__]
    dst_function(**function_args)
    snapshot = SNAPSHOT
    
    # Restore.
    for key in temp.keys():
        globals().pop(key, None)

    return snapshot, source


def display_top(
    snapshot,
    source=None,
    key_type='lineno',
    limit=10
):
    if source is not None:
        source_lines = source.split(sep='\n')

    snapshot = snapshot.filter_traces([Filter(True, DUMMY_SRC_NAME),])
    
    top_stats = snapshot.statistics(key_type)

    print('Top {} lines'.format(limit))
    for index, stat in enumerate(top_stats[:limit], 1):
        frame = stat.traceback[0]
        
        # replace "/path/to/module/file.py" with "module/file.py"
        filename = os.sep.join(frame.filename.split(os.sep)[-2:])

        fmt = '#{index}: {filename}:{lineno}: {size:.1f} KiB'.format(
            index=index,
            filename=filename,
            lineno=frame.lineno,
            size=stat.size / 1024
        )
        print(fmt)

        line = ''
        if frame.filename == DUMMY_SRC_NAME:
            if source is not None:
                line = source_lines[frame.lineno-1].strip()
        else:
            line = linecache.getline(frame.filename, frame.lineno).strip()
            
        if line:
            print('    %s' % line)

    other = top_stats[limit:]
    if other:
        size = sum(stat.size for stat in other)
        fmt = '{length} other: {size:.1f} KiB'.format(
            length=len(other),
            size=size / 1024
        )
        print(fmt)

    total = sum(stat.size for stat in top_stats)
    fmt = 'Total allocated size: {size:.1f} KiB'.format(size=total / 1024)
    print(fmt)