def maybe_cloud_tpu():
"""If FLAGS.cloud_tpu is set, setup Cloud instances."""
if not FLAGS.cloud_tpu:
yield
return
tf.logging.info("Running on Cloud TPU")
if (not FLAGS.data_dir.startswith("gs://") or
not FLAGS.output_dir.startswith("gs://")):
raise ValueError("To run on Cloud TPUs, data_dir and output_dir need to "
"be gs:// paths, i.e. on Google Cloud Storage.")
FLAGS.use_tpu = True
with cloud_tpu.cloud_tpu(
FLAGS.cloud_vm_name,
FLAGS.cloud_tpu_name,
delete_on_done=FLAGS.cloud_delete_on_done) as tpu_master:
FLAGS.master = tpu_master
yield
def maybe_cloud_tpu():
"""If FLAGS.cloud_tpu is set, setup Cloud instances."""
if not FLAGS.cloud_tpu:
yield
return
tf.logging.info("Running on Cloud TPU")
if (not FLAGS.data_dir.startswith("gs://") or
not FLAGS.output_dir.startswith("gs://")):
raise ValueError("To run on Cloud TPUs, data_dir and output_dir need to "
"be gs:// paths, i.e. on Google Cloud Storage.")
FLAGS.use_tpu = True
with cloud_tpu.cloud_tpu(
FLAGS.cloud_vm_name,
FLAGS.cloud_tpu_name,
delete_on_done=FLAGS.cloud_delete_on_done) as tpu_master:
FLAGS.master = tpu_master
yield
