measuring_memory_usage_in_pytest_xdist.py

import os

import psutil
import pytest


LIMIT = 5
SHARED_MEMORY_USAGE_INFO = "memory_usage"


def pytest_configure(config):
    """
    Defined appropriate plugins selection in pytest_configure hook

    Parameters
    ----------
    config : _pytest.config.Config
    """
    plugin = MemoryUsage(config)
    config.pluginmanager.register(plugin)


def is_master(config):
    """True if the code running the given pytest.config object is running in a xdist master
    node or not running xdist at all.
    """
    return not hasattr(config, 'workerinput')


def get_usage_memory():
    """
    Measures memory usage per Python process

    Returns
    -------
    memory_usage : float
    """
    process = psutil.Process(os.getpid())
    memory_use = process.memory_info()
    return memory_use.rss / 1024  # to KB


class MemoryUsage:
    def __init__(self, config):
        """
        Defined appropriate plugins selection in pytest_configure hook

        Parameters
        ----------
        config : _pytest.config.Config
        """
        self.config = config
        self.is_master = is_master(config)
        self.stats = {}

    def add(self, name):
        self.stats[name] = self.stats.get(name) or {}
        return self.stats[name]

    def pytest_runtest_setup(self, item):
        """Record maxrss for pre-setup."""
        self.add(item.nodeid)["setup"] = get_usage_memory()

    @pytest.hookimpl(hookwrapper=True)
    def pytest_runtest_call(self, item):
        """
        Track test memory

        Parameters
        ----------
        item : _pytest.main.Item
        """
        start = get_usage_memory()
        yield
        end = get_usage_memory()
        self.add(item.nodeid)["diff"] = end - start
        self.add(item.nodeid)["end"] = end
        self.add(item.nodeid)["start"] = start

    def pytest_terminal_summary(self, terminalreporter):
        tr = terminalreporter
        if self.stats:
            tr._tw.sep("=", "TOP {} tests which took most RAM".format(LIMIT), yellow=True)
            stats = sorted(self.stats.items(), key=lambda x: x[-1]["diff"], reverse=True)
            for test_name, info in stats[:LIMIT]:
                line = "setup({}KB) usage ({}KB) - {}".format(info["setup"], info["diff"], test_name)
                tr._tw.line(line)

    def pytest_testnodedown(self, node, error):
        """
        Get statistic about memory usage for test cases from xdist nodes
        and merge to master stats
        """
        node_stats = node.workeroutput[SHARED_MEMORY_USAGE_INFO]
        self.stats.update(node_stats)

    @pytest.hookimpl(hookwrapper=True, trylast=True)
    def pytest_sessionfinish(self, session, exitstatus):
        """
        Dump memory usage statistics to `workeroutput`
        Executed once per node if with xdist and will gen from mater node

        Parameters
        ----------
        session : _pytest.Session
        exitstatus : int
        """
        yield
        if not self.is_master:
            self.config.workeroutput[SHARED_MEMORY_USAGE_INFO] = self.stats

Hi @DKorytkin, this plugin seems very interesting but I am not sure how to enable it and use it. Could you please add some instructions?

Also some nitpicking, Kb should be note kB if you mean kilobyte. Lowercase b stands for bits (8x difference) which can be confusing.

Furthermore, some pedantic nitpick: 1 kB == 1000 B instead 1024 B. The convention for 1024 Byte is now 1 kibibyte noted KiB. https://en.wikipedia.org/wiki/Kilobyte

I assume the answer to my question is to add this code to a conftest.py file in the source folder.

@ogrisel, yes, you can copy/paste this code to your project. It should work

sorry for the delay

it was used in this article
https://korytkin.medium.com/how-to-get-data-from-pytest-xdist-nodes-2fbf2f0fe957