From c5aac1cc5962c1f20c076b7c3ec0773cbc48efde Mon Sep 17 00:00:00 2001 From: Yi Situ Date: Fri, 28 Aug 2020 18:52:38 -0700 Subject: [PATCH] Move profiler services behind profiler namespace. PiperOrigin-RevId: 329046847 Change-Id: Ibc26db95ec314af01c6cf0a9efd4c9a796af66ca --- tensorflow/compiler/xla/python/xla.cc | 8 ++++---- tensorflow/core/data/service/server_lib.cc | 2 +- .../core/distributed_runtime/rpc/grpc_server_lib.cc | 2 +- tensorflow/core/profiler/rpc/profiler_server.cc | 2 ++ tensorflow/core/profiler/rpc/profiler_server.h | 2 ++ tensorflow/core/profiler/rpc/profiler_service_impl.cc | 7 +++++++ tensorflow/core/profiler/rpc/profiler_service_impl.h | 5 +++++ tensorflow/python/profiler/internal/profiler_wrapper.cc | 6 ++++-- tensorflow/tools/def_file_filter/symbols_pybind.txt | 4 ++-- 9 files changed, 28 insertions(+), 10 deletions(-) diff --git a/tensorflow/compiler/xla/python/xla.cc b/tensorflow/compiler/xla/python/xla.cc index d5977f4f0cf..06605660b63 100644 --- a/tensorflow/compiler/xla/python/xla.cc +++ b/tensorflow/compiler/xla/python/xla.cc @@ -171,13 +171,13 @@ class TraceMeWrapper : public tensorflow::profiler::TraceMeWrapper { void BuildProfilerSubmodule(py::module* m) { py::module profiler = m->def_submodule("profiler", "TensorFlow profiler integration"); - py::class_> + py::class_> profiler_server_class(profiler, "ProfilerServer"); profiler.def( "start_server", - [](int port) -> std::unique_ptr { - auto server = absl::make_unique(); + [](int port) -> std::unique_ptr { + auto server = absl::make_unique(); server->StartProfilerServer(port); return server; }, diff --git a/tensorflow/core/data/service/server_lib.cc b/tensorflow/core/data/service/server_lib.cc index 2b6110fe48a..83a6e67c584 100644 --- a/tensorflow/core/data/service/server_lib.cc +++ b/tensorflow/core/data/service/server_lib.cc @@ -82,7 +82,7 @@ int GrpcDataServerBase::BoundPort() { return bound_port(); } void GrpcDataServerBase::AddProfilerServiceToBuilder( ::grpc::ServerBuilder& builder) { - profiler_service_ = CreateProfilerService(); + profiler_service_ = profiler::CreateProfilerService(); builder.RegisterService(profiler_service_.get()); } diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc index 83e072559e9..71be10f69e5 100644 --- a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc +++ b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc @@ -249,7 +249,7 @@ Status GrpcServer::Init(const GrpcServerOptions& opts) { .release(); eager_service_ = new eager::GrpcEagerServiceImpl(&worker_env_, &builder); - profiler_service_ = CreateProfilerService(); + profiler_service_ = profiler::CreateProfilerService(); builder.RegisterService(profiler_service_.get()); // extra service: diff --git a/tensorflow/core/profiler/rpc/profiler_server.cc b/tensorflow/core/profiler/rpc/profiler_server.cc index 966a94a1116..cfff3fc05de 100644 --- a/tensorflow/core/profiler/rpc/profiler_server.cc +++ b/tensorflow/core/profiler/rpc/profiler_server.cc @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/core/profiler/rpc/profiler_service_impl.h" namespace tensorflow { +namespace profiler { void ProfilerServer::StartProfilerServer(int32 port) { std::string server_address = absl::StrCat("[::]:", port); @@ -54,4 +55,5 @@ ProfilerServer::~ProfilerServer() { } } +} // namespace profiler } // namespace tensorflow diff --git a/tensorflow/core/profiler/rpc/profiler_server.h b/tensorflow/core/profiler/rpc/profiler_server.h index b7148e7e686..45680e83b6c 100644 --- a/tensorflow/core/profiler/rpc/profiler_server.h +++ b/tensorflow/core/profiler/rpc/profiler_server.h @@ -22,6 +22,7 @@ limitations under the License. #include "tensorflow/core/profiler/profiler_service.grpc.pb.h" namespace tensorflow { +namespace profiler { class ProfilerServer { public: @@ -34,6 +35,7 @@ class ProfilerServer { std::unique_ptr<::grpc::Server> server_; }; +} // namespace profiler } // namespace tensorflow #endif // TENSORFLOW_CORE_PROFILER_RPC_PROFILER_SERVER_H_ diff --git a/tensorflow/core/profiler/rpc/profiler_service_impl.cc b/tensorflow/core/profiler/rpc/profiler_service_impl.cc index ba463813fc0..8eadd87bf77 100644 --- a/tensorflow/core/profiler/rpc/profiler_service_impl.cc +++ b/tensorflow/core/profiler/rpc/profiler_service_impl.cc @@ -34,6 +34,7 @@ limitations under the License. #include "tensorflow/core/profiler/protobuf/xplane.pb.h" namespace tensorflow { +namespace profiler { namespace { const absl::string_view kXPlanePb = "xplane.pb"; @@ -115,4 +116,10 @@ std::unique_ptr CreateProfilerService() { return absl::make_unique(); } +} // namespace profiler + +std::unique_ptr CreateProfilerService() { + return absl::make_unique(); +} + } // namespace tensorflow diff --git a/tensorflow/core/profiler/rpc/profiler_service_impl.h b/tensorflow/core/profiler/rpc/profiler_service_impl.h index 00a850acbf2..3960b33f5be 100644 --- a/tensorflow/core/profiler/rpc/profiler_service_impl.h +++ b/tensorflow/core/profiler/rpc/profiler_service_impl.h @@ -23,6 +23,11 @@ namespace tensorflow { std::unique_ptr CreateProfilerService(); +namespace profiler { + +std::unique_ptr CreateProfilerService(); + +} // namespace profiler } // namespace tensorflow #endif // TENSORFLOW_CORE_PROFILER_RPC_PROFILER_SERVICE_IMPL_H_ diff --git a/tensorflow/python/profiler/internal/profiler_wrapper.cc b/tensorflow/python/profiler/internal/profiler_wrapper.cc index 401201018d9..7ddce914010 100644 --- a/tensorflow/python/profiler/internal/profiler_wrapper.cc +++ b/tensorflow/python/profiler/internal/profiler_wrapper.cc @@ -16,6 +16,7 @@ limitations under the License. #include #include "absl/memory/memory.h" +#include "absl/strings/match.h" #include "absl/strings/numbers.h" #include "pybind11/pybind11.h" #include "pybind11/pytypes.h" @@ -50,7 +51,7 @@ tensorflow::Status ValidateHostPortPair(const std::string& host_port) { // Must be host:port, port must be a number, host must not contain a '/', // host also must not be empty. if (parts.size() != 2 || !absl::SimpleAtoi(parts[1], &port) || - parts[0].find("/") != std::string::npos || parts[0].empty()) { + absl::StrContains(parts[0], "/") || parts[0].empty()) { return tensorflow::errors::InvalidArgument( "Could not interpret \"", host_port, "\" as a host-port pair."); } @@ -123,7 +124,8 @@ PYBIND11_MODULE(_pywrap_profiler, m) { .def("export_to_tb", &ProfilerSessionWrapper::ExportToTensorBoard); m.def("start_server", [](int port) { - auto profiler_server = absl::make_unique(); + auto profiler_server = + absl::make_unique(); profiler_server->StartProfilerServer(port); // Intentionally release profiler server. Should transfer ownership to // caller instead. diff --git a/tensorflow/tools/def_file_filter/symbols_pybind.txt b/tensorflow/tools/def_file_filter/symbols_pybind.txt index 93305a0707c..ed4dadad5e1 100644 --- a/tensorflow/tools/def_file_filter/symbols_pybind.txt +++ b/tensorflow/tools/def_file_filter/symbols_pybind.txt @@ -337,8 +337,8 @@ tensorflow::ProfilerSession::Status tensorflow::ProfilerSession::~ProfilerSession [profiler_server_impl] # profiler -tensorflow::ProfilerServer::StartProfilerServer -tensorflow::ProfilerServer::~ProfilerServer +tensorflow::profiler::ProfilerServer::StartProfilerServer +tensorflow::profiler::ProfilerServer::~ProfilerServer [profiler_client_impl] # profiler tensorflow::profiler::ProfileGrpc