From eef699b9f8c158d51a00f127b1a2854d29c6751b Mon Sep 17 00:00:00 2001
From: Russell Power <power@google.com>
Date: Wed, 2 Jan 2019 14:22:57 -0800
Subject: [PATCH] Rename shutdown mode field to better fit behavior.

PiperOrigin-RevId: 227581084
---
 .../contrib/tpu/python/tpu/session_support.py    | 16 ++++++++++------
 tensorflow/core/util/event.proto                 |  2 +-
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/tensorflow/contrib/tpu/python/tpu/session_support.py b/tensorflow/contrib/tpu/python/tpu/session_support.py
index 3e463823c82..f5735cecc38 100644
--- a/tensorflow/contrib/tpu/python/tpu/session_support.py
+++ b/tensorflow/contrib/tpu/python/tpu/session_support.py
@@ -185,7 +185,8 @@ def all_worker_devices(session):
   """Return a list of devices for each worker in the system."""
   devices = session.list_devices()
   return [
-      device.name for device in devices
+      device.name
+      for device in devices
       if ':CPU:' in device.name and 'coordinator' not in device.name
   ]
 
@@ -255,12 +256,14 @@ class WatchdogManager(threading.Thread):
     self._worker_manager.configure(
         event_pb2.WorkerHeartbeatRequest(
             watchdog_config=event_pb2.WatchdogConfig(
-                timeout_ms=self.shutdown_timeout * 1000,)))
+                timeout_ms=self.shutdown_timeout * 1000,),
+            shutdown_mode=event_pb2.WAIT_FOR_COORDINATOR))
 
   def configure_and_run(self):
-    logging.info('Enabling watchdog timer with %d second timeout '
-                 'and %d second ping interval.',
-                 self.shutdown_timeout, self.ping_interval)
+    logging.info(
+        'Enabling watchdog timer with %d second timeout '
+        'and %d second ping interval.', self.shutdown_timeout,
+        self.ping_interval)
     self._reset_manager()
     self._running = True
     self.start()
@@ -269,7 +272,8 @@ class WatchdogManager(threading.Thread):
     logging.info('Stopping worker watchdog.')
     self._worker_manager.configure(
         event_pb2.WorkerHeartbeatRequest(
-            watchdog_config=event_pb2.WatchdogConfig(timeout_ms=-1,)))
+            watchdog_config=event_pb2.WatchdogConfig(timeout_ms=-1,),
+            shutdown_mode=event_pb2.NOT_CONFIGURED))
     self._running = False
     self.join()
 
diff --git a/tensorflow/core/util/event.proto b/tensorflow/core/util/event.proto
index 9ce85be5511..2d3ae627773 100644
--- a/tensorflow/core/util/event.proto
+++ b/tensorflow/core/util/event.proto
@@ -95,7 +95,7 @@ enum WorkerHealth {
 // signal is received.
 enum WorkerShutdownMode {
   DEFAULT = 0;
-  SHUTDOWN_IMMEDIATELY = 1;
+  NOT_CONFIGURED = 1;
   WAIT_FOR_COORDINATOR = 2;
 }