OLD | NEW |
---|---|
(Empty) | |
1 # Copyright 2016 The Chromium Authors. All rights reserved. | |
2 # Use of this source code is governed by a BSD-style license that can be | |
3 # found in the LICENSE file. | |
4 | |
5 import multiprocessing | |
6 import os | |
7 import Queue | |
8 import resource | |
9 import signal | |
10 | |
11 import psutil | |
12 | |
13 | |
14 def _LimitMemory(memory_share): | |
15 """Limits the memory available to this process, to avoid OOM issues. | |
16 | |
17 Args: | |
18 memory_share: (float) Share coefficient of the total physical memory that | |
19 the process can use. | |
20 """ | |
21 total_memory = psutil.virtual_memory().total | |
22 memory_limit = memory_share * total_memory | |
23 resource.setrlimit(resource.RLIMIT_AS, (memory_limit, -1L)) | |
24 | |
25 | |
26 def _MultiprocessingWrapper(queue, memory_share, function, args): | |
27 """Helper function that sets a memory limit on the current process, then | |
28 calls |function| on |args| and write the results to |queue|. | |
29 | |
30 Args: | |
31 queue: (multiprocessing.Queue) Queue where the results of the wrapped | |
32 function are written. | |
33 memory_share: (float) Share coefficient of the total physical memory that | |
34 the process can use. | |
35 function: The wrapped function. | |
36 args: (list) Arguments for the wrapped function. | |
37 """ | |
38 try: | |
39 if memory_share: | |
40 _LimitMemory(memory_share) | |
41 | |
42 queue.put(function(*args)) | |
43 except Exception: | |
44 queue.put(None) | |
45 | |
46 | |
47 def RunOutOfProcess(function, args, logger, timeout_seconds, memory_share=None): | |
blundell
2016/06/06 14:12:00
nit: maybe RunInSeparateProcess()?
| |
48 """Runs a function in a separate process, and kills it after the timeout is | |
49 reached. | |
50 | |
51 Args: | |
52 function: The function to run. | |
53 args: (list) Arguments for the wrapped function. | |
54 timeout_seconds: (float) Timeout in seconds after which the subprocess is | |
55 terminated. | |
56 memory_share: (float) Set this parameter to limit the memory available to | |
57 the spawned subprocess. This is a ratio of the total system | |
58 memory (between 0 and 1). | |
59 Returns: | |
60 The result of the wrapped function, or None if the call failed. | |
61 """ | |
62 queue = multiprocessing.Queue() | |
63 process = multiprocessing.Process(target=_MultiprocessingWrapper, | |
64 args=(queue, memory_share, function, args)) | |
65 process.daemon = True | |
66 process.start() | |
67 | |
68 result = None | |
69 | |
70 try: | |
71 logger.info('Wait for result.') | |
72 # Note: If the subprocess somehow crashes (e.g. Python crashing), this | |
73 # process will wait the full timeout. Could be avoided but probably not | |
74 # worth the extra complexity. | |
75 result = queue.get(block=True, timeout=timeout_seconds) | |
76 except Queue.Empty: | |
77 logger.warning('Subprocess timeout.') | |
78 process.terminate() | |
79 | |
80 logger.info('Wait for process to terminate.') | |
81 process.join(timeout=5) | |
82 | |
83 if process.is_alive(): | |
84 logger.warning('Process still alive, hard killing now.') | |
85 os.kill(process.pid, signal.SIGKILL) | |
86 | |
87 return result | |
OLD | NEW |