Add build target to initialize a TPU model server.

PiperOrigin-RevId: 357796214
Change-Id: I5ee082fe9be39ee46953b6a1628170a8805aa8c6

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,
+)

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

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_