experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add support for specifying a custom exit code in WorkerHeartbeatRequest.

Browse Source 
PiperOrigin-RevId: 264909269
This commit is contained in:
A. Unique TensorFlower 2019-08-22 13:41:06 -07:00 committed by TensorFlower Gardener
parent 6f737e0dd6
commit dbf5399ae9
2 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/core/util/event.proto
View File

@ -106,9 +106,14 @@ message WatchdogConfig {
int64 timeout_ms = 1;
}
message RequestedExitCode {
int32 exit_code = 1;
}
message WorkerHeartbeatRequest {
WorkerShutdownMode shutdown_mode = 1;
WatchdogConfig watchdog_config = 2;
RequestedExitCode exit_code = 3;
}
message WorkerHeartbeatResponse {

6
tensorflow/python/tpu/session_support.py
View File

@ -145,12 +145,14 @@ class WorkerHeartbeatManager(object):
# Default timeout is set to allow other shutdown triggered operations (log
# flushing etc) to finish before terminating the worker.
def shutdown(self, wait_time_in_ms=60000):
def shutdown(self, wait_time_in_ms=60000, exit_code=None):
"""Shutdown all workers after `shutdown_timeout_secs`."""
logging.info('Shutting down %s.', self)
req = event_pb2.WorkerHeartbeatRequest(
watchdog_config=event_pb2.WatchdogConfig(timeout_ms=wait_time_in_ms),
shutdown_mode=event_pb2.SHUTDOWN_AFTER_TIMEOUT)
shutdown_mode=event_pb2.SHUTDOWN_AFTER_TIMEOUT,
exit_code=event_pb2.RequestedExitCode(
exit_code=exit_code) if exit_code is not None else None)
self.configure(req)
# Wait for workers to shutdown.