utils.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172
  1. import os
  2. import sys
  3. import time
  4. import errno
  5. import signal
  6. import warnings
  7. import threading
  8. import subprocess
  9. try:
  10. import psutil
  11. except ImportError:
  12. psutil = None
  13. WIN32 = sys.platform == "win32"
  14. def _flag_current_thread_clean_exit():
  15. """Put a ``_clean_exit`` flag on the current thread"""
  16. thread = threading.current_thread()
  17. thread._clean_exit = True
  18. def recursive_terminate(process, use_psutil=True):
  19. if use_psutil and psutil is not None:
  20. _recursive_terminate_with_psutil(process)
  21. else:
  22. _recursive_terminate_without_psutil(process)
  23. def _recursive_terminate_with_psutil(process, retries=5):
  24. try:
  25. children = psutil.Process(process.pid).children(recursive=True)
  26. except psutil.NoSuchProcess:
  27. return
  28. # Kill the children in reverse order to avoid killing the parents before
  29. # the children in cases where there are more processes nested.
  30. for child in children[::-1]:
  31. try:
  32. child.kill()
  33. except psutil.NoSuchProcess:
  34. pass
  35. process.terminate()
  36. process.join()
  37. def _recursive_terminate_without_psutil(process):
  38. """Terminate a process and its descendants.
  39. """
  40. try:
  41. _recursive_terminate(process.pid)
  42. except OSError as e:
  43. warnings.warn("Failed to kill subprocesses on this platform. Please"
  44. "install psutil: https://github.com/giampaolo/psutil")
  45. # In case we cannot introspect the children, we fall back to the
  46. # classic Process.terminate.
  47. process.terminate()
  48. process.join()
  49. def _recursive_terminate(pid):
  50. """Recursively kill the descendants of a process before killing it.
  51. """
  52. if sys.platform == "win32":
  53. # On windows, the taskkill function with option `/T` terminate a given
  54. # process pid and its children.
  55. try:
  56. subprocess.check_output(
  57. ["taskkill", "/F", "/T", "/PID", str(pid)],
  58. stderr=None)
  59. except subprocess.CalledProcessError as e:
  60. # In windows, taskkill return 1 for permission denied and 128, 255
  61. # for no process found.
  62. if e.returncode not in [1, 128, 255]:
  63. raise
  64. elif e.returncode == 1:
  65. # Try to kill the process without its descendants if taskkill
  66. # was denied permission. If this fails too, with an error
  67. # different from process not found, let the top level function
  68. # raise a warning and retry to kill the process.
  69. try:
  70. os.kill(pid, signal.SIGTERM)
  71. except OSError as e:
  72. if e.errno != errno.ESRCH:
  73. raise
  74. else:
  75. try:
  76. children_pids = subprocess.check_output(
  77. ["pgrep", "-P", str(pid)],
  78. stderr=None
  79. )
  80. except subprocess.CalledProcessError as e:
  81. # `ps` returns 1 when no child process has been found
  82. if e.returncode == 1:
  83. children_pids = b''
  84. else:
  85. raise
  86. # Decode the result, split the cpid and remove the trailing line
  87. children_pids = children_pids.decode().split('\n')[:-1]
  88. for cpid in children_pids:
  89. cpid = int(cpid)
  90. _recursive_terminate(cpid)
  91. try:
  92. os.kill(pid, signal.SIGTERM)
  93. except OSError as e:
  94. # if OSError is raised with [Errno 3] no such process, the process
  95. # is already terminated, else, raise the error and let the top
  96. # level function raise a warning and retry to kill the process.
  97. if e.errno != errno.ESRCH:
  98. raise
  99. def get_exitcodes_terminated_worker(processes):
  100. """Return a formated string with the exitcodes of terminated workers.
  101. If necessary, wait (up to .25s) for the system to correctly set the
  102. exitcode of one terminated worker.
  103. """
  104. patience = 5
  105. # Catch the exitcode of the terminated workers. There should at least be
  106. # one. If not, wait a bit for the system to correctly set the exitcode of
  107. # the terminated worker.
  108. exitcodes = [p.exitcode for p in list(processes.values())
  109. if p.exitcode is not None]
  110. while len(exitcodes) == 0 and patience > 0:
  111. patience -= 1
  112. exitcodes = [p.exitcode for p in list(processes.values())
  113. if p.exitcode is not None]
  114. time.sleep(.05)
  115. return _format_exitcodes(exitcodes)
  116. def _format_exitcodes(exitcodes):
  117. """Format a list of exit code with names of the signals if possible"""
  118. str_exitcodes = ["{}({})".format(_get_exitcode_name(e), e)
  119. for e in exitcodes if e is not None]
  120. return "{" + ", ".join(str_exitcodes) + "}"
  121. def _get_exitcode_name(exitcode):
  122. if sys.platform == "win32":
  123. # The exitcode are unreliable on windows (see bpo-31863).
  124. # For this case, return UNKNOWN
  125. return "UNKNOWN"
  126. if exitcode < 0:
  127. try:
  128. import signal
  129. if sys.version_info > (3, 5):
  130. return signal.Signals(-exitcode).name
  131. # construct an inverse lookup table
  132. for v, k in signal.__dict__.items():
  133. if (v.startswith('SIG') and not v.startswith('SIG_') and
  134. k == -exitcode):
  135. return v
  136. except ValueError:
  137. return "UNKNOWN"
  138. elif exitcode != 255:
  139. # The exitcode are unreliable on forkserver were 255 is always returned
  140. # (see bpo-30589). For this case, return UNKNOWN
  141. return "EXIT"
  142. return "UNKNOWN"