diff --git a/tests/pytorch/distributed/conftest.py b/tests/pytorch/distributed/conftest.py index 799d6b4ca..92aafcfdb 100644 --- a/tests/pytorch/distributed/conftest.py +++ b/tests/pytorch/distributed/conftest.py @@ -31,26 +31,38 @@ _LAUNCHERS = ("torchrun", "mpirun") @cache -def _terminate_timeout_seconds(): +def _terminate_timeout_seconds(run_timeout=None): """Grace between SIGTERM and the SIGKILL backstop.""" try: - return str(max(1, int(os.environ.get("TE_DIST_LAUNCH_KILL_AFTER", "60")))) + explicit = int(os.environ.get("TE_DIST_LAUNCH_KILL_AFTER", "60")) + if run_timeout and run_timeout < 2*explicit: + return max(1, run_timeout // 2) + return max(1, explicit) except ValueError: - return "60" + return 60 -def _launch_timeout_seconds(): - """Inner per-launch bound, in seconds, as a string for the coreutil.""" - explicit = os.environ.get("TE_DIST_LAUNCH_TIMEOUT") - if explicit: - return explicit +def _launch_timeout_seconds(run_timeout=None, terminate_timeout=None): + """Inner per-launch bound, in seconds""" # Fire a bit before the outer pytest-timeout so the child is reaped cleanly # rather than orphaned when the watchdog calls os._exit(). try: - outer = int(os.environ.get("PYTEST_TIMEOUT", "1200")) + explicit = os.environ.get("TE_DIST_LAUNCH_TIMEOUT") + if explicit: + return int(explicit) + outer = run_timeout + if run_timeout: + # extra 10 seconds for timeout to kill the process before timeout itself is killed + outer = run_timeout - 10 + else: + outer = int(os.environ.get("PYTEST_TIMEOUT", "1200")) except ValueError: - return "1200" - return str(max(60, outer - 60 - int(_terminate_timeout_seconds()))) + return 1200 + if terminate_timeout is None: + terminate_timeout = _terminate_timeout_seconds(run_timeout) + # Give at least 1 minute for the child to run to avoid too fast process termination + # if pytest timeout or subprocess.run timeout are set to a small value. + return max(60, outer - terminate_timeout) def _is_launcher(cmd): @@ -62,8 +74,10 @@ def _is_launcher(cmd): return head in _LAUNCHERS -def _wrap(cmd): - return ["timeout", "-k", _terminate_timeout_seconds(), "-v", _launch_timeout_seconds(), *cmd] +def _wrap(cmd, timeout): + terminate_timeout = _terminate_timeout_seconds(timeout) + launch_timeout = _launch_timeout_seconds(timeout, terminate_timeout) + return ["timeout", "-k", str(terminate_timeout), "-v", str(launch_timeout), *cmd] # Patch at import (collection) time so it is active for every test in this dir. @@ -72,7 +86,7 @@ def _wrap(cmd): def _run_with_timeout(cmd, *args, **kwargs): if _is_launcher(cmd): - cmd = _wrap(cmd) + cmd = _wrap(cmd, kwargs.get("timeout", 0)) return _orig_run(cmd, *args, **kwargs) subprocess.run = _run_with_timeout