From 632b2acccf1430db91c4255dae8fa9ac1a80513c Mon Sep 17 00:00:00 2001
From: Allen Wang <allencwang@google.com>
Date: Tue, 16 Feb 2021 13:45:37 -0800
Subject: [PATCH] Add build target to initialize a TPU model server.
PiperOrigin-RevId: 357796214
Change-Id: I5ee082fe9be39ee46953b6a1628170a8805aa8c6
---
tensorflow/core/tpu/BUILD | 20 +++++
.../core/tpu/tpu_model_server_initializer.cc | 78 +++++++++++++++++++
.../core/tpu/tpu_model_server_initializer.h | 33 ++++++++
3 files changed, 131 insertions(+)
create mode 100644 tensorflow/core/tpu/tpu_model_server_initializer.cc
create mode 100644 tensorflow/core/tpu/tpu_model_server_initializer.h
diff --git a/tensorflow/core/tpu/BUILD b/tensorflow/core/tpu/BUILD
index 5d15e278762..a9faa1871e9 100644
--- a/tensorflow/core/tpu/BUILD
+++ b/tensorflow/core/tpu/BUILD
@@ -319,3 +319,23 @@ cc_library(
"//tensorflow/core/lib/core:status",
],
)
+
+cc_library(
+ name = "tpu_model_server_initializer",
+ srcs = ["tpu_model_server_initializer.cc"],
+ hdrs = ["tpu_model_server_initializer.h"],
+ visibility = ["//visibility:public"],
+ deps = [
+ ":libtftpu_header",
+ ":tpu_api",
+ ":tpu_api_dlsym_set_fn",
+ ":tpu_executor_init_fns",
+ ":tpu_initializer_helper",
+ ":tpu_library_init_fns",
+ ":tpu_ops_c_api_hdrs",
+ "//tensorflow/core:lib",
+ "//tensorflow/stream_executor/tpu:tpu_executor",
+ "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
+ ],
+ alwayslink = True,
+)
diff --git a/tensorflow/core/tpu/tpu_model_server_initializer.cc b/tensorflow/core/tpu/tpu_model_server_initializer.cc
new file mode 100644
index 00000000000..554d11586d2
--- /dev/null
+++ b/tensorflow/core/tpu/tpu_model_server_initializer.cc
@@ -0,0 +1,78 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tpu/tpu_model_server_initializer.h"
+
+#include <dlfcn.h>
+
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tpu/tpu_api_dlsym_set_fn.h"
+
+#if !defined(PLATFORM_GOOGLE)
+#include "tensorflow/core/tpu/tpu_api.h"
+#include "tensorflow/core/tpu/tpu_initializer_helper.h"
+#include "tensorflow/stream_executor/tpu/tpu_platform.h"
+#endif
+
+namespace tensorflow {
+namespace tpu {
+
+
+#if defined(PLATFORM_GOOGLE)
+Status InitializeTpuModelServer(void* library_handle) {
+ return errors::Unimplemented("You must statically link in a TPU library.");
+}
+#else // PLATFORM_GOOGLE
+#include "tensorflow/core/tpu/tpu_library_init_fns.inc"
+
+Status InitializeTpuModelServer(void* library_handle) {
+ Status s = InitializeTpuStructFns(library_handle);
+
+ // Retrieve arguments from environment if applicable
+ std::pair<std::vector<std::string>, std::vector<const char*> > args =
+ GetLibTpuInitArguments();
+
+ // TPU platform registration must only be performed after the library is
+ // loaded. We do not want to register a TPU platform in XLA without the
+ // supporting library providing the necessary APIs.
+ if (s.ok()) {
+ void (*initialize_fn)(bool init_library, int num_args, const char** args);
+ initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
+ dlsym(library_handle, "TfTpu_Initialize"));
+ (*initialize_fn)(/*init_library=*/true, args.second.size(),
+ args.second.data());
+
+ RegisterTpuPlatform();
+ }
+
+ OpsApiFn()->TfTpu_InitializeTpuModelServerFn();
+ return s;
+}
+
+bool FindAndLoadTpuModelServer() {
+ if (!TryAcquireTpuLock()) return false;
+ void* library = dlopen("libtpu.so", RTLD_NOW);
+ if (library) {
+ InitializeTpuModelServer(library);
+ }
+ return true;
+}
+
+static bool tpu_library_finder = FindAndLoadTpuModelServer();
+#endif // PLATFORM_GOOGLE
+
+} // namespace tpu
+} // namespace tensorflow
diff --git a/tensorflow/core/tpu/tpu_model_server_initializer.h b/tensorflow/core/tpu/tpu_model_server_initializer.h
new file mode 100644
index 00000000000..d85786cf1c4
--- /dev/null
+++ b/tensorflow/core/tpu/tpu_model_server_initializer.h
@@ -0,0 +1,33 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TPU_TPU_MODEL_SERVER_INITIALIZER_H_
+#define TENSORFLOW_CORE_TPU_TPU_MODEL_SERVER_INITIALIZER_H_
+
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tpu/libtftpu.h"
+#include "tensorflow/core/tpu/tpu_ops_c_api.h"
+#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
+
+
+namespace tensorflow {
+namespace tpu {
+
+Status InitializeTpuModelServer(void* library_handle);
+
+} // namespace tpu
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CORE_TPU_TPU_MODEL_SERVER_INITIALIZER_H_