utils.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181
  1. import os
  2. import sys
  3. import time
  4. import errno
  5. import signal
  6. import warnings
  7. import subprocess
  8. import traceback
  9. try:
  10. import psutil
  11. except ImportError:
  12. psutil = None
  13. def kill_process_tree(process, use_psutil=True):
  14. """Terminate process and its descendants with SIGKILL"""
  15. if use_psutil and psutil is not None:
  16. _kill_process_tree_with_psutil(process)
  17. else:
  18. _kill_process_tree_without_psutil(process)
  19. def recursive_terminate(process, use_psutil=True):
  20. warnings.warn(
  21. "recursive_terminate is deprecated in loky 3.2, use kill_process_tree"
  22. "instead",
  23. DeprecationWarning,
  24. )
  25. kill_process_tree(process, use_psutil=use_psutil)
  26. def _kill_process_tree_with_psutil(process):
  27. try:
  28. descendants = psutil.Process(process.pid).children(recursive=True)
  29. except psutil.NoSuchProcess:
  30. return
  31. # Kill the descendants in reverse order to avoid killing the parents before
  32. # the descendant in cases where there are more processes nested.
  33. for descendant in descendants[::-1]:
  34. try:
  35. descendant.kill()
  36. except psutil.NoSuchProcess:
  37. pass
  38. try:
  39. psutil.Process(process.pid).kill()
  40. except psutil.NoSuchProcess:
  41. pass
  42. process.join()
  43. def _kill_process_tree_without_psutil(process):
  44. """Terminate a process and its descendants."""
  45. try:
  46. if sys.platform == "win32":
  47. _windows_taskkill_process_tree(process.pid)
  48. else:
  49. _posix_recursive_kill(process.pid)
  50. except Exception: # pragma: no cover
  51. details = traceback.format_exc()
  52. warnings.warn(
  53. "Failed to kill subprocesses on this platform. Please install"
  54. "psutil: https://github.com/giampaolo/psutil\n"
  55. f"Details:\n{details}"
  56. )
  57. # In case we cannot introspect or kill the descendants, we fall back to
  58. # only killing the main process.
  59. #
  60. # Note: on Windows, process.kill() is an alias for process.terminate()
  61. # which in turns calls the Win32 API function TerminateProcess().
  62. process.kill()
  63. process.join()
  64. def _windows_taskkill_process_tree(pid):
  65. # On windows, the taskkill function with option `/T` terminate a given
  66. # process pid and its children.
  67. try:
  68. subprocess.check_output(
  69. ["taskkill", "/F", "/T", "/PID", str(pid)], stderr=None
  70. )
  71. except subprocess.CalledProcessError as e:
  72. # In Windows, taskkill returns 128, 255 for no process found.
  73. if e.returncode not in [128, 255]:
  74. # Let's raise to let the caller log the error details in a
  75. # warning and only kill the root process.
  76. raise # pragma: no cover
  77. def _kill(pid):
  78. # Not all systems (e.g. Windows) have a SIGKILL, but the C specification
  79. # mandates a SIGTERM signal. While Windows is handled specifically above,
  80. # let's try to be safe for other hypothetic platforms that only have
  81. # SIGTERM without SIGKILL.
  82. kill_signal = getattr(signal, "SIGKILL", signal.SIGTERM)
  83. try:
  84. os.kill(pid, kill_signal)
  85. except OSError as e:
  86. # if OSError is raised with [Errno 3] no such process, the process
  87. # is already terminated, else, raise the error and let the top
  88. # level function raise a warning and retry to kill the process.
  89. if e.errno != errno.ESRCH:
  90. raise # pragma: no cover
  91. def _posix_recursive_kill(pid):
  92. """Recursively kill the descendants of a process before killing it."""
  93. try:
  94. children_pids = subprocess.check_output(
  95. ["pgrep", "-P", str(pid)], stderr=None, text=True
  96. )
  97. except subprocess.CalledProcessError as e:
  98. # `ps` returns 1 when no child process has been found
  99. if e.returncode == 1:
  100. children_pids = ""
  101. else:
  102. raise # pragma: no cover
  103. # Decode the result, split the cpid and remove the trailing line
  104. for cpid in children_pids.splitlines():
  105. cpid = int(cpid)
  106. _posix_recursive_kill(cpid)
  107. _kill(pid)
  108. def get_exitcodes_terminated_worker(processes):
  109. """Return a formatted string with the exitcodes of terminated workers.
  110. If necessary, wait (up to .25s) for the system to correctly set the
  111. exitcode of one terminated worker.
  112. """
  113. patience = 5
  114. # Catch the exitcode of the terminated workers. There should at least be
  115. # one. If not, wait a bit for the system to correctly set the exitcode of
  116. # the terminated worker.
  117. exitcodes = [
  118. p.exitcode for p in list(processes.values()) if p.exitcode is not None
  119. ]
  120. while not exitcodes and patience > 0:
  121. patience -= 1
  122. exitcodes = [
  123. p.exitcode
  124. for p in list(processes.values())
  125. if p.exitcode is not None
  126. ]
  127. time.sleep(0.05)
  128. return _format_exitcodes(exitcodes)
  129. def _format_exitcodes(exitcodes):
  130. """Format a list of exit code with names of the signals if possible"""
  131. str_exitcodes = [
  132. f"{_get_exitcode_name(e)}({e})" for e in exitcodes if e is not None
  133. ]
  134. return "{" + ", ".join(str_exitcodes) + "}"
  135. def _get_exitcode_name(exitcode):
  136. if sys.platform == "win32":
  137. # The exitcode are unreliable on windows (see bpo-31863).
  138. # For this case, return UNKNOWN
  139. return "UNKNOWN"
  140. if exitcode < 0:
  141. try:
  142. import signal
  143. return signal.Signals(-exitcode).name
  144. except ValueError:
  145. return "UNKNOWN"
  146. elif exitcode != 255:
  147. # The exitcode are unreliable on forkserver were 255 is always returned
  148. # (see bpo-30589). For this case, return UNKNOWN
  149. return "EXIT"
  150. return "UNKNOWN"