多处理功能上的超时装饰器

Question

这个装饰器直接取自我在网上找到的一个例子：

class TimedOutExc(Exception):
    pass


def timeout(timeout):
    def decorate(f):
        def handler(signum, frame):
            raise TimedOutExc()

        def new_f(*args, **kwargs):

            old = signal.signal(signal.SIGALRM, handler)
            signal.alarm(timeout)

            try:
                result = f(*args, **kwargs)
            except TimedOutExc:
                return None
            finally:
                signal.signal(signal.SIGALRM, old)
            signal.alarm(0)
            return result

        new_f.func_name = f.func_name
        return new_f

    return decorate

如果 f 函数超时，它会抛出异常。

好吧，它可以工作，但是当我在多处理函数上使用这个装饰器并由于超时而停止时，它不会终止计算中涉及的进程。我该怎么做？

我不想启动异常并停止程序。基本上我想要的是当 f 超时时，让它返回 None 然后终止所涉及的进程。

Answer 1

虽然我同意亚伦回答的要点，但我想详细说明一下。

multiprocessing 启动的进程必须在要修饰的函数中停止；我不认为这可以简单地从装饰器本身来完成（装饰函数是唯一知道它启动了什么计算的实体）。

除了让修饰函数捕获SIGALARM，您还可以捕获自定义TimedOutExc 异常——这可能更灵活。您的示例将变为：

import signal
import functools

class TimedOutExc(Exception):
    """
    Raised when a timeout happens
    """

def timeout(timeout):
    """
    Return a decorator that raises a TimedOutExc exception
    after timeout seconds, if the decorated function did not return.
    """

    def decorate(f):

        def handler(signum, frame):
            raise TimedOutExc()

        @functools.wraps(f)  # Preserves the documentation, name, etc.
        def new_f(*args, **kwargs):

            old_handler = signal.signal(signal.SIGALRM, handler)
            signal.alarm(timeout)

            result = f(*args, **kwargs)  # f() always returns, in this scheme

            signal.signal(signal.SIGALRM, old_handler)  # Old signal handler is restored
            signal.alarm(0)  # Alarm removed

            return result

        return new_f

    return decorate

@timeout(10)
def function_that_takes_a_long_time():
    try:
        # ... long, parallel calculation ...
    except TimedOutExc:
        # ... Code that shuts down the processes ...
        # ...
        return None  # Or exception raised, which means that the calculation is not complete

Answer 2

我怀疑装饰器可以做到这一点：装饰器是函数的包装器；该函数是一个黑匣子。装饰器和它包装的函数之间没有通信。

您需要做的是重写您的函数代码以使用SIGALRM 处理程序来终止它已启动的任何进程。