experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix server address assignment when constructing TpuPodState.

Browse Source 
PiperOrigin-RevId: 330032670
Change-Id: I686cb5d2ba665480acfc16f06ac54228e5b1177b
This commit is contained in:
Henry Tan 2020-09-03 18:32:24 -07:00 committed by TensorFlower Gardener
parent a982edc004
commit 7f2e2bb276
2 changed files with 22 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
tensorflow/core/tpu/kernels/tpu_pod_state.cc
View File

@ -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) {}

2
tensorflow/core/tpu/kernels/tpu_pod_state.h
View File

@ -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_