TPU library internal change.

PiperOrigin-RevId: 315952033
Change-Id: Ie8700a558f6db09ffa47f3e0b1dabff21dbd2301

tensorflow/core/tpu/kernels/tpu_compile_op_common.cc

@@ -352,19 +352,7 @@ Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
     return;
   }
 
-#ifdef PLATFORM_GOOGLE
+  LogAndExit(42);
-  // 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);
 }
 
 /* static */ Status TPUCompileOpKernelCommon::GetDynamicShapes(

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