From 7f2e2bb276791fbdbb70593f41a0ea86e44501d1 Mon Sep 17 00:00:00 2001
From: Henry Tan <henrytan@google.com>
Date: Thu, 3 Sep 2020 18:32:24 -0700
Subject: [PATCH] Fix server address assignment when constructing TpuPodState.

PiperOrigin-RevId: 330032670
Change-Id: I686cb5d2ba665480acfc16f06ac54228e5b1177b
---
 tensorflow/core/tpu/kernels/tpu_pod_state.cc | 37 +++++++++++---------
 tensorflow/core/tpu/kernels/tpu_pod_state.h  |  2 ++
 2 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/tensorflow/core/tpu/kernels/tpu_pod_state.cc b/tensorflow/core/tpu/kernels/tpu_pod_state.cc
index e7f13a657ed..7b02998b343 100644
--- a/tensorflow/core/tpu/kernels/tpu_pod_state.cc
+++ b/tensorflow/core/tpu/kernels/tpu_pod_state.cc
@@ -28,23 +28,6 @@ namespace tensorflow {
 const char kTpuPodStateResourceName[] = "tpu_pod_state";
 
 namespace {
-Status GetServerAddressAndPort(std::string* server_address, int* serving_port) {
-  TF_Status* status = TF_NewStatus();
-  char* server_address_output = nullptr;
-  auto cleanup = xla::MakeCleanup([&status, &server_address_output]() {
-    TF_DeleteStatus(status);
-    tpu::ConfigApiFn()->TpuConfigurationApi_FreeCharArrayFn(
-        server_address_output);
-  });
-  size_t server_address_output_size;
-  *serving_port = -1;
-  tpu::ConfigApiFn()->TpuConfigurationApi_GetServerAddressAndPortFn(
-      &server_address_output_size, &server_address_output, serving_port,
-      status);
-  CHECK_NE(*serving_port, -1);
-  TF_RETURN_IF_ERROR(StatusFromTF_Status(status));
-  return Status::OK();
-}
 
 // Attempt to delete resource_name from resource_manager's default_container.
 // Returns OK if the deletion succeeded, or if the resource was not found. Else
@@ -86,6 +69,26 @@ ConstructCacheService(ResourceMgr* rmgr, int serving_port,
 }
 }  // namespace
 
+Status GetServerAddressAndPort(std::string* server_address, int* serving_port) {
+  TF_Status* status = TF_NewStatus();
+  char* server_address_output = nullptr;
+  auto cleanup = xla::MakeCleanup([&status, &server_address_output]() {
+    TF_DeleteStatus(status);
+    tpu::ConfigApiFn()->TpuConfigurationApi_FreeCharArrayFn(
+        server_address_output);
+  });
+  size_t server_address_output_size;
+  *serving_port = -1;
+  tpu::ConfigApiFn()->TpuConfigurationApi_GetServerAddressAndPortFn(
+      &server_address_output_size, &server_address_output, serving_port,
+      status);
+  TF_RETURN_IF_ERROR(StatusFromTF_Status(status));
+  *server_address =
+      std::string(server_address_output, server_address_output_size);
+  CHECK_NE(*serving_port, -1);
+  return Status::OK();
+}
+
 TpuPodState::TpuPodState(
     int service_port, std::unique_ptr<TpuCompilationCacheService> cache_service)
     : cache_service_(std::move(cache_service)), service_port_(service_port) {}
diff --git a/tensorflow/core/tpu/kernels/tpu_pod_state.h b/tensorflow/core/tpu/kernels/tpu_pod_state.h
index 07ad3bee553..9515a8ee8f5 100644
--- a/tensorflow/core/tpu/kernels/tpu_pod_state.h
+++ b/tensorflow/core/tpu/kernels/tpu_pod_state.h
@@ -56,6 +56,8 @@ Status ConstructTpuPodState(
     ResourceMgr* rmgr, const std::vector<int32_t>& num_devices_per_host,
     tpu::TpuCompilationCacheInterface* compilation_cache,
     std::string* host_config_proto);
+
+Status GetServerAddressAndPort(std::string* server_address, int* serving_port);
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_CORE_TPU_KERNELS_TPU_POD_STATE_H_