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

Add ProfilerService to TF workers to start by default.

Browse Source 
PiperOrigin-RevId: 324036243
Change-Id: I55f4c5b1ad66a70f2530ff74e4802927da9721d4
This commit is contained in:
A. Unique TensorFlower 2020-07-30 10:51:11 -07:00 committed by TensorFlower Gardener
parent 4138a1c45a
commit cb3a22b732
13 changed files with 82 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/compiler/xla/python/BUILD
View File

@ -400,7 +400,7 @@ pybind_extension(
"//tensorflow/core:lib_internal_impl", # buildcleaner: keep
"//tensorflow/core/profiler/lib:profiler_backends",
"//tensorflow/core/profiler/lib:profiler_session",
"//tensorflow/core/profiler/rpc:profiler_server",
"//tensorflow/core/profiler/rpc:profiler_server_impl",
"//tensorflow/python/profiler/internal:traceme_wrapper",
"//tensorflow/stream_executor:device_memory_allocator",
"//tensorflow/stream_executor:platform",

4
tensorflow/core/distributed_runtime/rpc/BUILD
View File

@ -17,6 +17,7 @@ load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency") # buildifier: disa
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_kernel_tests_linkstatic",
"tf_protos_profiler_service",
)
load(
"//tensorflow/core/platform:build_config_root.bzl",
@ -335,9 +336,10 @@ cc_library(
"//tensorflow/core/distributed_runtime:worker_cache_wrapper",
"//tensorflow/core/distributed_runtime:worker_env",
"//tensorflow/core/distributed_runtime/rpc/eager:grpc_eager_service_impl",
"//tensorflow/core/profiler/rpc:profiler_service_impl",
tf_grpc_dependency(),
tf_grpc_cc_dependency(),
],
] + tf_protos_profiler_service(),
alwayslink = 1,
)

4
tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc
View File

@ -51,6 +51,7 @@ limitations under the License.
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/mem.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/profiler/rpc/profiler_service_impl.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/core/util/env_var.h"
@ -248,6 +249,9 @@ Status GrpcServer::Init(const GrpcServerOptions& opts) {
.release();
eager_service_ = new eager::GrpcEagerServiceImpl(&worker_env_, &builder);
profiler_service_ = CreateProfilerService();
builder.RegisterService(profiler_service_.get());
// extra service:
if (opts.service_func != nullptr) {
opts.service_func(&worker_env_, &builder);

4
tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h
View File

@ -36,6 +36,7 @@ limitations under the License.
#include "tensorflow/core/framework/collective.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/profiler/profiler_service.grpc.pb.h"
namespace tensorflow {
@ -182,6 +183,9 @@ class GrpcServer : public ServerInterface {
std::unique_ptr<Thread> eager_thread_ TF_GUARDED_BY(mu_);
std::shared_ptr<WorkerSession> worker_session_;
// TensorFlow profiler service implementation.
std::unique_ptr<grpc::ProfilerService::Service> profiler_service_ = nullptr;
// The overall server configuration.
ServerDef server_def_ TF_GUARDED_BY(mu_);

4
tensorflow/core/platform/build_config.bzl
View File

@ -28,6 +28,7 @@ load(
_tf_platform_deps = "tf_platform_deps",
_tf_portable_deps_no_runtime = "tf_portable_deps_no_runtime",
_tf_portable_proto_lib = "tf_portable_proto_lib",
_tf_profiler_client_deps = "tf_profiler_client_deps",
_tf_proto_library = "tf_proto_library",
_tf_proto_library_cc = "tf_proto_library_cc",
_tf_proto_library_py = "tf_proto_library_py",
@ -38,6 +39,7 @@ load(
_tf_protos_grappler = "tf_protos_grappler",
_tf_protos_grappler_impl = "tf_protos_grappler_impl",
_tf_protos_profiler_impl = "tf_protos_profiler_impl",
_tf_protos_profiler_service = "tf_protos_profiler_service",
_tf_py_clif_cc = "tf_py_clif_cc",
_tf_pyclif_proto_library = "tf_pyclif_proto_library",
_tf_resource_deps = "tf_resource_deps",
@ -80,6 +82,8 @@ tf_protos_all_impl = _tf_protos_all_impl
tf_protos_grappler = _tf_protos_grappler
tf_protos_grappler_impl = _tf_protos_grappler_impl
tf_protos_profiler_impl = _tf_protos_profiler_impl
tf_protos_profiler_service = _tf_protos_profiler_service
tf_profiler_client_deps = _tf_profiler_client_deps
tf_py_clif_cc = _tf_py_clif_cc
tf_pyclif_proto_library = _tf_pyclif_proto_library
tf_resource_deps = _tf_resource_deps

10
tensorflow/core/platform/default/build_config.bzl
View File

@ -597,6 +597,16 @@ def tf_protos_profiler_impl():
clean_dep("//tensorflow/core/profiler:profiler_options_proto_cc_impl"),
]
def tf_protos_profiler_service():
return [
clean_dep("//tensorflow/core/profiler:profiler_analysis_proto_cc_impl"),
clean_dep("//tensorflow/core/profiler:profiler_service_proto_cc_impl"),
clean_dep("//tensorflow/core/profiler:profiler_service_monitor_result_proto_cc_impl"),
]
def tf_profiler_client_deps():
return [clean_dep("//tensorflow/core/profiler/rpc/client:profiler_client_headers")]
def tf_protos_grappler_impl():
return [clean_dep("//tensorflow/core/grappler/costs:op_performance_data_cc_impl")]

3
tensorflow/core/profiler/BUILD
View File

@ -25,6 +25,7 @@ tf_proto_library(
name = "profiler_service_monitor_result_proto",
srcs = ["profiler_service_monitor_result.proto"],
cc_api_version = 2,
make_default_target_header_only = True,
visibility = ["//visibility:public"],
)
@ -48,6 +49,7 @@ tf_proto_library(
has_services = 1,
cc_api_version = 2,
cc_grpc_version = 1,
make_default_target_header_only = True,
protodeps = [
":profiler_options_proto",
":profiler_service_monitor_result_proto",
@ -62,6 +64,7 @@ tf_proto_library(
has_services = 1,
cc_api_version = 2,
cc_grpc_version = 1,
make_default_target_header_only = True,
protodeps = [":profiler_service_proto"],
use_grpc_namespace = True,
visibility = ["//visibility:public"],

1
tensorflow/core/profiler/convert/BUILD
View File

@ -321,6 +321,7 @@ tf_cc_test(
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/distributed_runtime/rpc:grpc_server_lib",
"//tensorflow/core/profiler:profiler_service_proto_cc",
"//tensorflow/core/profiler/protobuf:input_pipeline_proto_cc",
"//tensorflow/core/profiler/protobuf:overview_page_proto_cc",

18
tensorflow/core/profiler/rpc/BUILD
View File

@ -1,5 +1,6 @@
load("//tensorflow:tensorflow.bzl", "tf_external_workspace_visible") # buildifier: disable=same-origin-load
load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency") # buildifier: disable=same-origin-load
load("//tensorflow:tensorflow.bzl", "tf_pybind_cc_library_wrapper") # buildifier: disable=same-origin-load
package(
licenses = ["notice"], # Apache 2.0
@ -11,7 +12,10 @@ cc_library(
hdrs = ["profiler_service_impl.h"],
features = ["-layering_check"],
visibility = tf_external_workspace_visible(
["//tensorflow_serving/model_servers:__pkg__"],
[
"//tensorflow/core/distributed_runtime/rpc:__pkg__",
"//tensorflow_serving/model_servers:__pkg__",
],
),
deps = [
"//tensorflow/core:lib",
@ -24,11 +28,12 @@ cc_library(
)
cc_library(
name = "profiler_server",
name = "profiler_server_impl",
srcs = ["profiler_server.cc"],
hdrs = ["profiler_server.h"],
visibility = [
"//tensorflow/compiler/xla/python:__pkg__",
"//tensorflow/python:__pkg__",
"//tensorflow/python/profiler/internal:__pkg__",
],
deps = [
@ -38,4 +43,13 @@ cc_library(
"@com_google_absl//absl/strings",
tf_grpc_cc_dependency(),
],
alwayslink = True,
)
tf_pybind_cc_library_wrapper(
name = "profiler_server_headers",
visibility = [
"//tensorflow/python/profiler/internal:__pkg__",
],
deps = [":profiler_server_impl"],
)

27
tensorflow/core/profiler/rpc/client/BUILD
View File

@ -1,4 +1,11 @@
load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
load("//tensorflow:tensorflow.bzl", "tf_pybind_cc_library_wrapper") # buildifier: disable=same-origin-load
# For platform specific build config
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_profiler_client_deps",
)
package(
licenses = ["notice"], # Apache 2.0
@ -8,17 +15,18 @@ cc_library(
name = "capture_profile",
srcs = ["capture_profile.cc"],
hdrs = ["capture_profile.h"],
visibility = ["//tensorflow/python/profiler/internal:__pkg__"],
visibility = [
"//tensorflow/python/profiler/internal:__pkg__",
],
deps = [
":profiler_client",
":save_profile",
"@com_google_absl//absl/strings",
"//tensorflow/core:lib",
"//tensorflow/core/profiler:profiler_analysis_proto_cc",
"//tensorflow/core/profiler:profiler_options_proto_cc",
"//tensorflow/core/profiler:profiler_service_proto_cc",
"//tensorflow/core/profiler/convert:xplane_to_profile_response",
"@com_google_absl//absl/strings",
],
] + tf_profiler_client_deps(),
)
cc_library(
@ -37,10 +45,10 @@ cc_library(
)
cc_library(
name = "profiler_client",
name = "profiler_client_impl",
srcs = ["profiler_client.cc"],
hdrs = ["profiler_client.h"],
visibility = ["//tensorflow/python/profiler/internal:__pkg__"],
visibility = ["//tensorflow/python:__pkg__"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
@ -48,4 +56,11 @@ cc_library(
"//tensorflow/core/profiler:profiler_service_proto_cc",
tf_grpc_cc_dependency(),
],
alwayslink = True,
)
tf_pybind_cc_library_wrapper(
name = "profiler_client_headers",
visibility = ["//tensorflow/python/profiler/internal:__pkg__"],
deps = [":profiler_client_impl"],
)

4
tensorflow/python/BUILD
View File

@ -6077,6 +6077,8 @@ pywrap_tensorflow_macro(
"//tensorflow/core/debug",
"//tensorflow/core/distributed_runtime:server_lib",
"//tensorflow/core/platform:stacktrace_handler",
"//tensorflow/core/profiler/rpc:profiler_server_impl",
"//tensorflow/core/profiler/rpc/client:profiler_client_impl",
"//tensorflow/core/profiler/internal:print_model_analysis",
"//tensorflow/lite/delegates/flex:delegate",
"//tensorflow/core/profiler/internal/cpu:python_tracer",
@ -6150,6 +6152,8 @@ filegroup(
"//tensorflow/core/profiler/internal:print_model_analysis", # tfprof
"//tensorflow/core/profiler/internal:traceme_recorder_impl", # profiler
"//tensorflow/core/profiler/lib:profiler_session_impl", # profiler
"//tensorflow/core/profiler/rpc:profiler_server_impl", # profiler
"//tensorflow/core/profiler/rpc/client:profiler_client_impl", # profiler
"//tensorflow/core/util:port", # util_port
"//tensorflow/core/util/tensor_bundle", # checkpoint_reader
"//tensorflow/lite/toco/python:toco_python_api", # toco

4
tensorflow/python/profiler/internal/BUILD
View File

@ -118,16 +118,16 @@ tf_python_pybind_extension(
],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core/profiler:profiler_service_proto_cc",
"//tensorflow/core/profiler/convert:op_stats_to_tf_stats",
"//tensorflow/core/profiler/convert:xplane_to_op_stats",
"//tensorflow/core/profiler/convert:xplane_to_profile_response",
"//tensorflow/core/profiler/convert:xplane_to_trace_events",
"//tensorflow/core/profiler/lib:profiler_session_headers",
"//tensorflow/core/profiler/rpc:profiler_server",
"//tensorflow/core/profiler/rpc:profiler_server_headers",
"//tensorflow/core/profiler/rpc/client:capture_profile",
"//tensorflow/core/profiler/rpc/client:save_profile",
"//tensorflow/python:pybind11_status",
"@com_github_grpc_grpc//:grpc++_public_hdrs",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@pybind11",

9
tensorflow/tools/def_file_filter/symbols_pybind.txt
View File

@ -335,6 +335,15 @@ tensorflow::ProfilerSession::SerializeToString
tensorflow::ProfilerSession::Status
tensorflow::ProfilerSession::~ProfilerSession
[profiler_server_impl] # profiler
tensorflow::ProfilerServer::StartProfilerServer
tensorflow::ProfilerServer::~ProfilerServer
[profiler_client_impl] # profiler
tensorflow::profiler::ProfileGrpc
tensorflow::profiler::NewSessionGrpc
tensorflow::profiler::MonitorGrpc
[status_macros] # tfcompile
xla::status_macros::MakeErrorStream::Impl::Impl
xla::status_macros::MakeErrorStream::Impl::~Impl