Measure function resource usage (cpu time, max memory consumption)

resource_usage.py

import platform
from contextlib import closing
from multiprocessing import Process, Pipe

import math
import resource
from functools import wraps
from multiprocessing.connection import Connection
from pydantic import BaseModel
from typing import Callable, Any


class ResultWithResources(BaseModel):
    real_result: Any
    cpu_time: int
    max_memory: int


class CommandResult(BaseModel):
    real_time: int
    cpu_time: int
    max_memory: int
    code: int


def _record_resource_usage(func: Callable, args: tuple, kwargs: dict, w: Connection) -> None:
    """Wrapper function, called in Process to record child resource usage"""
    rusage_before = resource.getrusage(resource.RUSAGE_CHILDREN)
    func_result = func(*args, **kwargs)
    rusage_after = resource.getrusage(resource.RUSAGE_CHILDREN)

    print(f'Before resource usage: {rusage_before}')
    print(f'After resource usage: {rusage_after}')

    cpu_time = rusage_after.ru_utime - rusage_before.ru_utime + rusage_after.ru_stime - rusage_before.ru_stime
    max_memory = rusage_after.ru_maxrss

    # mac os shows result in bytes, normal OSes -- in kilobytes
    if platform.platform() != 'darwin':
        max_memory *= 1024

    w.send(
        ResultWithResources(
            cpu_time=math.ceil(cpu_time * 1000),
            max_memory=math.ceil(max_memory),
            real_result=func_result,
        ),
    )


def measure_children_resource_usage(func: Callable) -> Callable[..., ResultWithResources]:
    @wraps(func)
    def wrapper(*args, **kwargs) -> ResultWithResources:
        r, w = Pipe(duplex=False)

        with closing(r), closing(w):
            p = Process(
                target=_record_resource_usage,
                kwargs={
                    'func': func,
                    'args': args,
                    'kwargs': kwargs,
                    'w': w,
                }
            )
            p.start()
            p.join()

            result: ResultWithResources = r.recv()

        return result

    return wrapper