TPUEstimator: Fix shutdown behavior after preemption.

PiperOrigin-RevId: 299005210
Change-Id: I925e5308c8470dce0e071361441266856278a391
This commit is contained in:
Russell Power 2020-03-04 20:09:00 -08:00 committed by TensorFlower Gardener
parent 30f796f327
commit ea95f4a740

View File

@ -418,7 +418,7 @@ class ResetComputation(object):
def __call__(self, run_context, all_workers, lame_workers):
del run_context, lame_workers
all_workers.shutdown()
all_workers.shutdown(exit_code=42)
logging.info('Resetting coordinator.')
raise CoordinatorResetError()
@ -435,7 +435,7 @@ class ShutdownLameWorkers(object):
pass
def __call__(self, run_context, all_workers, lame_workers):
lame_workers.shutdown()
lame_workers.shutdown(exit_code=42)
class ShutdownAllWorkers(object):
@ -449,4 +449,4 @@ class ShutdownAllWorkers(object):
pass
def __call__(self, run_context, all_workers, lame_workers):
all_workers.shutdown()
all_workers.shutdown(exit_code=42)