From 04f033b847f950220c26cde8edd9cb3bfa28ed9b Mon Sep 17 00:00:00 2001
From: Henry Tan <henrytan@google.com>
Date: Thu, 11 Jun 2020 12:19:29 -0700
Subject: [PATCH] TPU library internal change.
PiperOrigin-RevId: 315952033
Change-Id: Ie8700a558f6db09ffa47f3e0b1dabff21dbd2301
---
.../core/tpu/kernels/tpu_compile_op_common.cc | 14 +-------------
tensorflow/core/tpu/kernels/tpu_util.h | 4 ++++
2 files changed, 5 insertions(+), 13 deletions(-)
diff --git a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
index d7fad26f4fb..ae090913dc7 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
+++ b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
@@ -352,19 +352,7 @@ Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
return;
}
-#ifdef PLATFORM_GOOGLE
- // TODO(henrytan): copybara this out.
- LOG(ERROR) << "Aborting process due to cancelled TPUCompileOp. This "
- "termination is to ensure a consistent state. If your job "
- "does not restart, modify the retries allowed. See "
- "b/71383512.";
-#else
- LOG(ERROR) << "Restarting TPU host due to cancelled compilation.";
-#endif
-
- // Exit code 42 tells the Borglet to restart the task and not report it as a
- // failure: http://g3doc/borg/slave/g3doc/users/task_features
- std::quick_exit(42);
+ LogAndExit(42);
}
/* static */ Status TPUCompileOpKernelCommon::GetDynamicShapes(
diff --git a/tensorflow/core/tpu/kernels/tpu_util.h b/tensorflow/core/tpu/kernels/tpu_util.h
index 0ca94d0af59..ae0f08ba656 100644
--- a/tensorflow/core/tpu/kernels/tpu_util.h
+++ b/tensorflow/core/tpu/kernels/tpu_util.h
@@ -61,6 +61,10 @@ Status DynamicShapesToTensorShapes(const InputList& dynamic_shapes,
xla::StatusOr<xla::Shape> TpuShapeRepresentation(const TensorShape& shape,
DataType type,
bool use_fast_memory);
+
+// A callback called on exit.
+void LogAndExit(int code);
+
} // namespace tpu
} // namespace tensorflow