reset-tpus.py

import tensorflow as tf
import sys
from tensorflow.contrib.tpu.python.tpu import session_support
 
def reset_tpu(name):
  print('Resetting: %s' % name) 
  resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu=name, job_name='worker')
  address = resolver.master()
  cluster_def = resolver.cluster_spec().as_cluster_def()
 
  print('Master %s.  Cluster %s' % (address, cluster_def))
  session = tf.Session(address, config=tf.ConfigProto(cluster_def=cluster_def))
  m = session_support.WorkerHeartbeatManager.from_devices(
      session, session_support.all_worker_devices(session)
  )
  print('Resetting workers: %s', m)
  m.shutdown(timeout_ms=5000)
 
if __name__ == '__main__':
  reset_tpu(sys.argv[1])