Increase timeout of collectives to handle the case where:
Each replica is busy compiling their clusters, and once one of the replica finishes compilation, it launches and waits at the collective barrier.
PiperOrigin-RevId: 398029227
Change-Id: Ia9b966044742dc3c76fdde962d07fd803d596f13
diff --git a/tensorflow/compiler/tf2xla/xla_helpers.cc b/tensorflow/compiler/tf2xla/xla_helpers.cc
index 235a6d5..adff95d 100644
--- a/tensorflow/compiler/tf2xla/xla_helpers.cc
+++ b/tensorflow/compiler/tf2xla/xla_helpers.cc
@@ -229,7 +229,8 @@
xla::ExecutableRunOptions& run_options,
xla::DeviceAssignment& device_assignment,
xla::gpu::GpuExecutableRunOptions& gpu_options) {
- static const int kTimeoutSeconds = 30;
+ // TODO(nnigania): workaround for b/199436990
+ static const int kTimeoutSeconds = 300;
if (!collective_reduce_info) {
// An empty device assignment is sufficient for the case where no
// collectives are present.
