autograder.util.invoke

 1import multiprocessing
 2import sys
 3import time
 4import traceback
 5import typing
 6
 7REAP_TIME_SEC: float = 5
 8
 9_multiprocessing_initialized: bool = False  # pylint: disable=invalid-name
10
11def _init_multiprocessing() -> None:
12    """
13    Initialize Python multiprocessing.
14    Note that this should only be called on Linux code paths.
15    """
16
17    global _multiprocessing_initialized  # pylint: disable=global-statement
18
19    if (_multiprocessing_initialized):
20        return
21
22    multiprocessing.set_start_method('fork')
23    _multiprocessing_initialized = True
24
25def with_timeout(timeout: typing.Union[float, None], function: typing.Callable) -> typing.Tuple[bool, typing.Any]:
26    """
27    Run the given function in a different process with the given timeout.
28    If the timeout is None, then no timeout will be checked (and the code will be run on the same process).
29
30    Return: (success, function return value)
31    On timeout, success will be false and the value will be None.
32    On error, success will be false and value will be the string stacktrace.
33    On successful completion, success will be true and value may be None (if nothing was returned).
34    """
35
36    if ((timeout is None) or (not sys.platform.startswith('linux'))):
37        # Mac and Windows have some pickling issues with multiprocessing.
38        # Just run them without a timeout.
39        # Any autograder will be run on a Linux machine and will be safe.
40        start_time = time.time()
41        value = function()
42        runtime = time.time() - start_time
43
44        if ((timeout is not None) and (runtime > timeout)):
45            return (False, None)
46
47        return (True, value)
48
49    _init_multiprocessing()
50
51    result: multiprocessing.Queue = multiprocessing.Queue(1)
52
53    # Note that we use processes instead of threads so they can be more completely killed.
54    process = multiprocessing.Process(target = _invoke_helper, args = (result, function))
55    process.start()
56
57    # Wait for at most the timeout.
58    process.join(timeout)
59
60    # Check to see if the process is still running.
61    if (process.is_alive()):
62        # Kill the long-running process.
63        process.terminate()
64
65        # Try to reap the process once before just giving up on it.
66        process.join(REAP_TIME_SEC)
67
68        return (False, None)
69
70    # Check to see if the process explicitly existed (like via sys.exit()).
71    if (result.empty()):
72        return (False, 'Code explicitly exited (like via sys.exit()).')
73
74    value, error = result.get()
75
76    if (error is not None):
77        _, stacktrace = error
78        return (False, stacktrace)
79
80    return (True, value)
81
82def _invoke_helper(result: multiprocessing.Queue, function: typing.Callable) -> None:
83    """ A helper function for running the given function. """
84
85    value = None
86    error = None
87
88    try:
89        value = function()
90    except Exception as ex:
91        error = (ex, traceback.format_exc())
92
93    sys.stdout.flush()
94
95    result.put((value, error))
96    result.close()