timeout_pool.py

# -*- coding: utf-8 -*-

import os
import sys
from multiprocessing import Process
import subprocess
import time
import psutil
import tqdm
from utils import get_logger_for_file
import logging

_logger = get_logger_for_file(logging.ERROR, 'timeout_pool')


class TimeoutPool:

    def __init__(self, process, timeout, memory_limit):
        self.ps = []
        self.process = process
        self.timeout = timeout
        self.memory_limit = memory_limit

    @staticmethod
    def meta_print(pid, meta):
        _logger.debug("run %d : %s" % (pid, meta['args']), file=sys.stdout)

    @staticmethod
    def meta_timeout_print(pid, meta):
        _logger.warning("timeout %d : %s" % (pid, meta['args']), file=sys.stderr)

    @staticmethod
    def meta_too_much_mem_print(pid, meta):
        rss = TimeoutPool.get_memory_rss(pid)
        _logger.warning("too much mem %d (mem=%d) : %s" % (pid, rss, meta['args']), file=sys.stderr)

    @staticmethod
    def meta_builder(pid, args):
        return {'pid': pid, 'args': args}

    @staticmethod
    def kill_children(pid):
        try:
            child_process_list = psutil.Process(pid).children(recursive=True)
        except psutil.NoSuchProcess:
            child_process_list = []
        for c in child_process_list:
            try:
                c.terminate()
            except psutil.NoSuchProcess:
                pass

    @staticmethod
    def kill_process(p, wait_for_end=1.0, with_children=True):
        if with_children:
            TimeoutPool.kill_children(p.pid)
        p.terminate()
        # p.kill()
        p.join(wait_for_end)
        if p.exitcode is None:
            # force to kill this process
            os.system('kill %d' % p.pid)

    @staticmethod
    def get_memory_rss(pid, with_children=True):
        total_mem = 0
        try:
            total_mem = psutil.Process(pid).memory_full_info().rss
            for c in psutil.Process(pid).children(recursive=True):
                try:
                    total_mem += c.memory_full_info().rss
                except psutil.NoSuchProcess:
                    pass
        except psutil.NoSuchProcess:
            pass
        return total_mem

    def check_memory_usage(self, meta_too_much_mem_print, wait_for_end):
        new_ps = []
        for p_info in self.ps:
            tmp_p, tmp_start, tmp_meta = p_info
            if tmp_p.is_alive():
                try:
                    tmp_rss = TimeoutPool.get_memory_rss(tmp_p.pid)
                    if self.memory_limit > 0 and tmp_rss > self.memory_limit:
                        # kill process with pid and update self.ps
                        if meta_too_much_mem_print:
                            meta_too_much_mem_print(tmp_p.pid, tmp_meta)
                        TimeoutPool.kill_process(p_info[0], wait_for_end)
                    else:
                        new_ps.append(p_info)
                except psutil.AccessDenied:
                    pass
            else:
                # this process has end
                pass
        self.ps = new_ps

    def check_timeout_process(self, meta_timeout_print, wait_for_end):
        now = time.time()
        new_ps = []
        for p_info in self.ps:
            tmp_p, tmp_start, tmp_meta = p_info
            if tmp_p.is_alive():
                if self.timeout > 0 and now - tmp_start > self.timeout:
                    if meta_timeout_print:
                        meta_timeout_print(tmp_p.pid, tmp_meta)
                    TimeoutPool.kill_process(tmp_p, wait_for_end)
                else:
                    new_ps.append(p_info)
            else:
                # this process has end
                pass
        self.ps = new_ps

    @staticmethod
    def _run_cmd_wrapper(cmd, out_file, err_file):
        """
        close the output stream after finishing this command
        """
        if out_file is not None:
            out_stream = open(out_file, 'w')
        else:
            out_stream = subprocess.DEVNULL
        if err_file is not None:
            err_stream = open(err_file, 'w')
        else:
            err_stream = subprocess.DEVNULL
        # cmd_args = cmd.split()
        # subprocess.run(cmd_args, stdout=out_stream, stderr=err_stream)
        # when the command line has space and we cannot find the executable and args
        # use the following interface
        # then we cannot seperate stdin and stdout, everything will be pushed into stdout
        ret, all_output = subprocess.getstatusoutput(cmd)
        if out_file is not None:
            out_stream.write(all_output)

        if out_file is not None:
            out_stream.close()
        if err_file is not None:
            err_stream.close()
        if ret != 0:
            _logger.error(f'Fail run :{cmd}')

    def map(self,
            target,
            args_list,
            meta_builder=None,
            meta_print=None,
            meta_timeout_print=None,
            meta_too_much_mem_print=None,
            wait_for_end=1.0,
            check_freq=0.1,
            enable_pbar=True):
        if meta_builder is None:
            meta_builder = TimeoutPool.meta_builder
            meta_print = TimeoutPool.meta_print
            meta_timeout_print = TimeoutPool.meta_timeout_print
            meta_toomuchmem_print = TimeoutPool.meta_too_much_mem_print
        else:
            meta_builder = meta_builder
            meta_print = meta_print
            meta_timeout_print = meta_timeout_print
            meta_toomuchmem_print = meta_too_much_mem_print

        pbar = None
        step = max(len(args_list) // 10000, 1)
        if enable_pbar:
            pbar = tqdm.tqdm(total=len(args_list), desc='Progress: ', leave=False)
        count = 0
        for args in args_list:
            while len(self.ps) >= self.process:
                self.check_memory_usage(meta_toomuchmem_print, wait_for_end)
                if len(self.ps) < self.process:
                    break

                self.check_timeout_process(meta_timeout_print, wait_for_end)
                if len(self.ps) < self.process and check_freq > 0:
                    break
                else:
                    time.sleep(check_freq)

            # new process
            p = Process(target=target, args=args)
            p_start = time.time()
            p.start()
            p_meta = meta_builder(p.pid, args)
            self.ps.append((p, p_start, p_meta))
            meta_print(p.pid, p_meta)
            # update
            if pbar:
                count += 1
                if count == step:
                    pbar.update(step)
                    count = 0

        # finish all processes
        while len(self.ps) > 0:
            self.check_memory_usage(meta_toomuchmem_print, wait_for_end)
            if len(self.ps) == 0:
                break

            self.check_timeout_process(meta_timeout_print, wait_for_end)
            if len(self.ps) == 0:
                break
            time.sleep(check_freq)
        # close pbar if it's enabled
        if pbar:
            pbar.close()