Source code for schrodinger.test.performance.general

"""
General tools for performance tests.

"""

import functools
import subprocess
import time
import timeit

import psutil

#Default maximum time is 5 minutes.
MAX_TIME = 5 * 60

#Directory where input files are kept, added by test runner.
# INPUTFILE_DIR = None


def get_process_name(process):
    """Pretty-print a psutil.Process object."""
    cmd = subprocess.list2cmdline(process.cmdline)
    if not cmd:
        cmd = f'[{process.name}]'
    return f'PID{process.pid}: {cmd}'


def wait_for_successful_completion(process, timeout=MAX_TIME):
    """
    Wait until `process` is complete.

    If recursive is true, also wait for all child processes to complete. If the
    process is still running after `timeout`, kill it and all its children and
    raise an exception. If the exit code is not 0, raise an exception.

    :type process: psutil.Process or pid
    :param process: Process to watch, waiting for completion.
    :type timeout: int
    :param timeout: Duration (in seconds) to wait for the process.
    :type recursive: bool
    :param recursive: Should the child processes of `process` be monitored as
                      well?

    """
    if not isinstance(process, psutil.Process):
        process = psutil.Process(process)

    try:
        return_code = process.wait(timeout=timeout)
    except psutil._error.TimeoutExpired:
        # If there is an exception, kill all the children.
        for child in reversed(process.get_children(recursive=True)):
            print("killing", child.pid)
            child.terminate()
        process.terminate()
        raise
    if return_code:
        name = get_process_name(process)
        raise Exception('{} exited with return code {}'.format(
            name, return_code))


def call_in_loop(func, iterations, *args, **kwargs):
    """
    Calls function repeatedly to normalize test data.
    """
    testable = functools.partial(func, *args, **kwargs)
    duration = timeit.timeit(testable, number=iterations)
    return duration / iterations


def wait_and_monitor_usage(process, timeout=MAX_TIME, timeslice=0.1):
    '''
    Wait until `process` is complete. Query maxrss and cpu time
    used by the process and its children every so often (let the
    process to run for `timeslice` seconds between observations).

    If the process is still running after `timeout`, kill it and all
    its children and raise an exception. If the exit code is not 0, raise
    an exception.

    Returns maxrss (in megabytes) and cpu time (in seconds).

    :type process: psutil.Process or pid
    :param process: Process to watch, waiting for completion.

    :type timeout: float
    :param timeout: Duration (in seconds) to wait for the process.

    :type timeslice: float
    :param timeslice: Interval (in seconds) between observations.

    :rtype: (int, float)
    :return: Max RSS (in megabytes) and CPU time (user, in seconds).
    '''

    if not isinstance(process, psutil.Process):
        process = psutil.Process(process)

    maxrss = 0
    cputime = 0

    cmdline = process.cmdline()
    start_time = time.time()

    while time.time() - start_time < timeout:
        try:
            process.wait(timeslice)
            if process.returncode:
                raise subprocess.CalledProcessError(process.returncode, cmdline)
            return (maxrss / 1024.0 / 1024.0, cputime)
        except psutil.TimeoutExpired:
            rss, user = 0, 0
            for child in process.children(recursive=True):
                try:
                    rss += child.memory_info().rss
                    user += child.cpu_times().user
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass
            maxrss = max(maxrss, rss)
            cputime = max(cputime, user)

    for child in process.children(recursive=True):
        child.terminate()

    process.terminate()

    raise psutil.TimeoutExpired(timeout, process.pid)