From 498e5b4f6db80df13b54c44cbd657a2750067564 Mon Sep 17 00:00:00 2001
From: Rick Chao <rchao@google.com>
Date: Mon, 13 Apr 2020 18:02:56 -0700
Subject: [PATCH] Adjust the timeout secs in multi_process_runner.

PiperOrigin-RevId: 306348917
Change-Id: Ibe5253d0c8204bacf31a958b40b9d1d63ca05392
---
 tensorflow/python/distribute/multi_process_runner.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/distribute/multi_process_runner.py b/tensorflow/python/distribute/multi_process_runner.py
index e091f0da5e2..e258d98fb8c 100644
--- a/tensorflow/python/distribute/multi_process_runner.py
+++ b/tensorflow/python/distribute/multi_process_runner.py
@@ -89,6 +89,10 @@ BARRIER = 'barrier'
 
 _DEFAULT_MAX_SUBPROCESS_COUNT = 20
 
+# Default time out sec is selected so that it's handled before the default
+# "medium" timeout of the test runs.
+_DEFAULT_TIMEOUT_SEC = 200
+
 # Next pipe index to be global so that pipes are not reused across multiple
 # MultiProcessRunner usages.
 # TODO(rchao): Investigate possibility to remove this variable.
@@ -373,7 +377,7 @@ class MultiProcessRunner(object):
         break
     return list_to_return
 
-  def join(self, timeout=250):
+  def join(self, timeout=_DEFAULT_TIMEOUT_SEC):
     """Joins all the processes with timeout.
 
     Args:
@@ -621,7 +625,7 @@ def run(proc_func,
         grpc_fail_fast=None,
         stream_stdout=True,
         list_stdout=False,
-        timeout=250,
+        timeout=_DEFAULT_TIMEOUT_SEC,
         args=None,
         kwargs=None):  # pylint: disable=g-doc-args
   """Runs functions in local child processes.