mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
2ed47fecc5
multiprocessing.Queue relies on, among other things, background threads to send messages between processes. This works in the happy path but can cause issues if a process is exiting by bypassing atexit handlers or crashing because the writer to the Queue can terminate while the reader is blocked reading the queue. The reader sees the queue as non-empty yet even with a timeout will actually block forever. An example of a Queue deadlock is here: https://gist.github.com/chipturner/342f72341f087737befe9df84d0e41ce Since the error reporting case here is a simple one-shot message from the dying child to the parent, we can just use a file-based rendezvous. This eliminates the deadlock when a large traceback is still being flushed to the network when a child exits. Pull Request resolved: https://github.com/pytorch/pytorch/pull/114688 Approved by: https://github.com/suo, https://github.com/yifuwang
282 lines
10 KiB
Python
282 lines
10 KiB
Python
import logging
|
|
import multiprocessing
|
|
import multiprocessing.connection
|
|
import os
|
|
import pickle
|
|
import signal
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import warnings
|
|
from typing import Optional
|
|
|
|
from . import _prctl_pr_set_pdeathsig # type: ignore[attr-defined]
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class ProcessException(Exception):
|
|
__slots__ = ["error_index", "error_pid"]
|
|
|
|
def __init__(self, msg: str, error_index: int, pid: int):
|
|
super().__init__(msg)
|
|
self.msg = msg
|
|
self.error_index = error_index
|
|
self.pid = pid
|
|
|
|
def __reduce__(self):
|
|
return type(self), (self.msg, self.error_index, self.pid)
|
|
|
|
|
|
class ProcessRaisedException(ProcessException):
|
|
"""Exception raised when a process failed due to an exception raised by the code."""
|
|
|
|
def __init__(
|
|
self,
|
|
msg: str,
|
|
error_index: int,
|
|
error_pid: int,
|
|
):
|
|
super().__init__(msg, error_index, error_pid)
|
|
|
|
|
|
class ProcessExitedException(ProcessException):
|
|
"""Exception raised when a process failed due to signal or exited with a specific code."""
|
|
|
|
__slots__ = ["exit_code"]
|
|
|
|
def __init__(
|
|
self,
|
|
msg: str,
|
|
error_index: int,
|
|
error_pid: int,
|
|
exit_code: int,
|
|
signal_name: Optional[str] = None,
|
|
):
|
|
super().__init__(msg, error_index, error_pid)
|
|
self.exit_code = exit_code
|
|
self.signal_name = signal_name
|
|
|
|
def __reduce__(self):
|
|
return (
|
|
type(self),
|
|
(self.msg, self.error_index, self.pid, self.exit_code, self.signal_name),
|
|
)
|
|
|
|
|
|
def _wrap(fn, i, args, error_file):
|
|
# prctl(2) is a Linux specific system call.
|
|
# On other systems the following function call has no effect.
|
|
# This is set to ensure that non-daemonic child processes can
|
|
# terminate if their parent terminates before they do.
|
|
_prctl_pr_set_pdeathsig(signal.SIGINT)
|
|
|
|
try:
|
|
fn(i, *args)
|
|
except KeyboardInterrupt:
|
|
pass # SIGINT; Killed by parent, do nothing
|
|
except Exception:
|
|
# Propagate exception to parent process, keeping original traceback
|
|
import traceback
|
|
|
|
with open(error_file, "wb") as fh:
|
|
pickle.dump(traceback.format_exc(), fh)
|
|
sys.exit(1)
|
|
|
|
|
|
class ProcessContext:
|
|
def __init__(self, processes, error_files):
|
|
self.error_files = error_files
|
|
self.processes = processes
|
|
self.sentinels = {
|
|
process.sentinel: index for index, process in enumerate(processes)
|
|
}
|
|
|
|
def pids(self):
|
|
return [int(process.pid) for process in self.processes]
|
|
|
|
def join(self, timeout=None):
|
|
r"""Join one or more processes within spawn context.
|
|
|
|
Attempt to join one or more processes in this spawn context.
|
|
If one of them exited with a non-zero exit status, this function
|
|
kills the remaining processes and raises an exception with the cause
|
|
of the first process exiting.
|
|
|
|
Returns ``True`` if all processes have been joined successfully,
|
|
``False`` if there are more processes that need to be joined.
|
|
|
|
Args:
|
|
timeout (float): Wait this long before giving up on waiting.
|
|
"""
|
|
# Ensure this function can be called even when we're done.
|
|
if len(self.sentinels) == 0:
|
|
return True
|
|
|
|
# Wait for any process to fail or all of them to succeed.
|
|
ready = multiprocessing.connection.wait(
|
|
self.sentinels.keys(),
|
|
timeout=timeout,
|
|
)
|
|
|
|
error_index = None
|
|
for sentinel in ready:
|
|
index = self.sentinels.pop(sentinel)
|
|
process = self.processes[index]
|
|
process.join()
|
|
if process.exitcode != 0:
|
|
error_index = index
|
|
break
|
|
|
|
# Return if there was no error.
|
|
if error_index is None:
|
|
# Return whether or not all processes have been joined.
|
|
return len(self.sentinels) == 0
|
|
|
|
# Assume failure. Terminate processes that are still alive.
|
|
# Try SIGTERM then SIGKILL if the process isn't going down.
|
|
# The reason is related to python signal handling is limited
|
|
# to main thread and if that is in c/c++ land and stuck it won't
|
|
# to handle it. We have seen processes getting stuck not handling
|
|
# SIGTERM for the above reason.
|
|
timeout: int = 30
|
|
for process in self.processes:
|
|
if process.is_alive():
|
|
log.warning("Terminating process %s via signal SIGTERM", process.pid)
|
|
process.terminate()
|
|
end = time.monotonic() + timeout
|
|
for process in self.processes:
|
|
time_to_wait = max(0, end - time.monotonic())
|
|
process.join(time_to_wait)
|
|
for process in self.processes:
|
|
if process.is_alive():
|
|
log.warning(
|
|
"Unable to shutdown process %s via SIGTERM , forcefully exiting via SIGKILL",
|
|
process.pid,
|
|
)
|
|
process.kill()
|
|
process.join()
|
|
|
|
# The file will only be created if the process crashed.
|
|
failed_process = self.processes[error_index]
|
|
if not os.access(self.error_files[error_index], os.R_OK):
|
|
exitcode = self.processes[error_index].exitcode
|
|
if exitcode < 0:
|
|
try:
|
|
name = signal.Signals(-exitcode).name
|
|
except ValueError:
|
|
name = f"<Unknown signal {-exitcode}>"
|
|
raise ProcessExitedException(
|
|
"process %d terminated with signal %s" % (error_index, name),
|
|
error_index=error_index,
|
|
error_pid=failed_process.pid,
|
|
exit_code=exitcode,
|
|
signal_name=name,
|
|
)
|
|
else:
|
|
raise ProcessExitedException(
|
|
"process %d terminated with exit code %d" % (error_index, exitcode),
|
|
error_index=error_index,
|
|
error_pid=failed_process.pid,
|
|
exit_code=exitcode,
|
|
)
|
|
|
|
with open(self.error_files[error_index], "rb") as fh:
|
|
original_trace = pickle.load(fh)
|
|
msg = "\n\n-- Process %d terminated with the following error:\n" % error_index
|
|
msg += original_trace
|
|
raise ProcessRaisedException(msg, error_index, failed_process.pid)
|
|
|
|
|
|
class SpawnContext(ProcessContext):
|
|
def __init__(self, processes, error_files):
|
|
warnings.warn("SpawnContext is renamed to ProcessContext since 1.4 release.")
|
|
super().__init__(processes, error_files)
|
|
|
|
|
|
# Note: [start_processes]
|
|
# mp.start_processes handles both start_method='spawn' and 'fork'. It's supposed to be a
|
|
# more generalized API than mp.spawn. Currently we only document mp.spawn as it's the
|
|
# CUDA compatible start_method. However, in environments like Ipython notebooks, 'fork'
|
|
# works better than 'spawn'. Every helper function we created for mp.spawn is indeed
|
|
# general enough, and backends like XLA can reuse them in Colab notebooks as well.
|
|
# Currently we only add this API first, we can consider adding it to documentation as
|
|
# needed in the future.
|
|
def start_processes(
|
|
fn, args=(), nprocs=1, join=True, daemon=False, start_method="spawn"
|
|
):
|
|
mp = multiprocessing.get_context(start_method)
|
|
error_files = []
|
|
processes = []
|
|
for i in range(nprocs):
|
|
# Each process is assigned a file to write tracebacks to. We
|
|
# use the file being non-empty to indicate an exception
|
|
# occurred (vs an expected shutdown). Note: this previously
|
|
# used a multiprocessing.Queue but that can be prone to
|
|
# deadlocks, so we went with a simpler solution for a one-shot
|
|
# message between processes.
|
|
tf = tempfile.NamedTemporaryFile(
|
|
prefix="pytorch-errorfile-", suffix=".pickle", delete=False
|
|
)
|
|
tf.close()
|
|
os.unlink(tf.name)
|
|
process = mp.Process(
|
|
target=_wrap,
|
|
args=(fn, i, args, tf.name),
|
|
daemon=daemon,
|
|
)
|
|
process.start()
|
|
error_files.append(tf.name)
|
|
processes.append(process)
|
|
|
|
context = ProcessContext(processes, error_files)
|
|
if not join:
|
|
return context
|
|
|
|
# Loop on join until it returns True or raises an exception.
|
|
while not context.join():
|
|
pass
|
|
|
|
|
|
def spawn(fn, args=(), nprocs=1, join=True, daemon=False, start_method="spawn"):
|
|
r"""Spawns ``nprocs`` processes that run ``fn`` with ``args``.
|
|
|
|
If one of the processes exits with a non-zero exit status, the
|
|
remaining processes are killed and an exception is raised with the
|
|
cause of termination. In the case an exception was caught in the
|
|
child process, it is forwarded and its traceback is included in
|
|
the exception raised in the parent process.
|
|
|
|
Args:
|
|
fn (function): Function is called as the entrypoint of the
|
|
spawned process. This function must be defined at the top
|
|
level of a module so it can be pickled and spawned. This
|
|
is a requirement imposed by multiprocessing.
|
|
|
|
The function is called as ``fn(i, *args)``, where ``i`` is
|
|
the process index and ``args`` is the passed through tuple
|
|
of arguments.
|
|
|
|
args (tuple): Arguments passed to ``fn``.
|
|
nprocs (int): Number of processes to spawn.
|
|
join (bool): Perform a blocking join on all processes.
|
|
daemon (bool): The spawned processes' daemon flag. If set to True,
|
|
daemonic processes will be created.
|
|
start_method (str): (deprecated) this method will always use ``spawn``
|
|
as the start method. To use a different start method
|
|
use ``start_processes()``.
|
|
|
|
Returns:
|
|
None if ``join`` is ``True``,
|
|
:class:`~ProcessContext` if ``join`` is ``False``
|
|
|
|
"""
|
|
if start_method != "spawn":
|
|
msg = (
|
|
"This method only supports start_method=spawn (got: %s).\n"
|
|
"To use a different start_method use:\n\t\t"
|
|
" torch.multiprocessing.start_processes(...)" % start_method
|
|
)
|
|
warnings.warn(msg)
|
|
return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
|