1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/profiler/rpc/profiler_service_impl.h"
17
18 #include "grpcpp/support/status.h"
19 #include "absl/container/flat_hash_set.h"
20 #include "absl/strings/str_cat.h"
21 #include "absl/strings/str_join.h"
22 #include "absl/strings/string_view.h"
23 #include "tensorflow/core/lib/core/errors.h"
24 #include "tensorflow/core/platform/env.h"
25 #include "tensorflow/core/platform/env_time.h"
26 #include "tensorflow/core/profiler/convert/op_stats_to_input_pipeline_analysis.h"
27 #include "tensorflow/core/profiler/convert/op_stats_to_overview_page.h"
28 #include "tensorflow/core/profiler/convert/op_stats_to_tf_stats.h"
29 #include "tensorflow/core/profiler/convert/xplane_to_op_stats.h"
30 #include "tensorflow/core/profiler/convert/xplane_to_trace_events.h"
31 #include "tensorflow/core/profiler/lib/profiler_session.h"
32 #include "tensorflow/core/profiler/protobuf/hardware_types.pb.h"
33 #include "tensorflow/core/profiler/protobuf/input_pipeline.pb.h"
34 #include "tensorflow/core/profiler/protobuf/op_stats.pb.h"
35 #include "tensorflow/core/profiler/protobuf/overview_page.pb.h"
36 #include "tensorflow/core/profiler/protobuf/tf_stats.pb.h"
37 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
38 #include "tensorflow/core/profiler/utils/group_events.h"
39 #include "tensorflow/core/protobuf/trace_events.pb.h"
40 #include "tensorflow/core/util/ptr_util.h"
41
42 namespace tensorflow {
43 namespace {
44
45 const absl::string_view kTensorflowStats = "tensorflow_stats";
46 const absl::string_view kInputPipeline = "input_pipeline";
47 const absl::string_view kOverviewPage = "overview_page";
48
HardwareTypeFromRunEnvironment(const profiler::RunEnvironment & run_env)49 profiler::HardwareType HardwareTypeFromRunEnvironment(
50 const profiler::RunEnvironment& run_env) {
51 if (run_env.device_type() == "GPU") return profiler::HardwareType::GPU;
52 if (run_env.device_type() == "CPU") return profiler::HardwareType::CPU_ONLY;
53 return profiler::HardwareType::UNKNOWN_HARDWARE;
54 }
55
56 template <typename Proto>
AddToolData(absl::string_view tool_name,const Proto & tool_output,ProfileResponse * response)57 void AddToolData(absl::string_view tool_name, const Proto& tool_output,
58 ProfileResponse* response) {
59 auto* tool_data = response->add_tool_data();
60 tool_data->set_name(string(tool_name));
61 tool_output.SerializeToString(tool_data->mutable_data());
62 }
63
64 // Returns the tool name with extension.
ToolName(absl::string_view tool)65 std::string ToolName(absl::string_view tool) {
66 return absl::StrCat(tool, ".pb");
67 }
68
CollectDataToResponse(const ProfileRequest & req,ProfilerSession * profiler,uint64 start_time_ns,ProfileResponse * response)69 Status CollectDataToResponse(const ProfileRequest& req,
70 ProfilerSession* profiler, uint64 start_time_ns,
71 ProfileResponse* response) {
72 profiler::XSpace xspace;
73 TF_RETURN_IF_ERROR(profiler->CollectData(&xspace));
74 GroupTfEvents(&xspace, /*event_group_name_map=*/nullptr);
75 {
76 uint64 end_time_ns = EnvTime::NowNanos();
77 profiler::Trace trace;
78 profiler::ConvertXSpaceToTraceEvents(start_time_ns, end_time_ns, xspace,
79 &trace);
80 trace.SerializeToString(response->mutable_encoded_trace());
81 }
82 absl::flat_hash_set<absl::string_view> tools(req.tools().begin(),
83 req.tools().end());
84 if (!tools.empty()) {
85 profiler::OpStats op_stats = profiler::ConvertXSpaceToOpStats(xspace);
86 profiler::HardwareType hw_type =
87 HardwareTypeFromRunEnvironment(op_stats.run_environment());
88 if (tools.contains(kOverviewPage)) {
89 profiler::OverviewPage overview_page_db =
90 profiler::ConvertOpStatsToOverviewPage(op_stats, hw_type);
91 AddToolData(ToolName(kOverviewPage), overview_page_db, response);
92 }
93 if (tools.contains(kInputPipeline)) {
94 profiler::InputPipelineAnalysisResult input_pipeline_analysis =
95 profiler::ConvertOpStatsToInputPipelineAnalysis(op_stats, hw_type);
96 AddToolData(ToolName(kInputPipeline), input_pipeline_analysis, response);
97 }
98 if (tools.contains(kTensorflowStats)) {
99 profiler::TfStatsDatabase tf_stats_db =
100 profiler::ConvertOpStatsToTfStats(op_stats);
101 AddToolData(ToolName(kTensorflowStats), tf_stats_db, response);
102 }
103 }
104 return Status::OK();
105 }
106
107 class ProfilerServiceImpl : public grpc::ProfilerService::Service {
108 public:
Monitor(::grpc::ServerContext * ctx,const MonitorRequest * req,MonitorResponse * response)109 ::grpc::Status Monitor(::grpc::ServerContext* ctx, const MonitorRequest* req,
110 MonitorResponse* response) override {
111 return ::grpc::Status(::grpc::StatusCode::UNIMPLEMENTED, "unimplemented.");
112 }
113
Profile(::grpc::ServerContext * ctx,const ProfileRequest * req,ProfileResponse * response)114 ::grpc::Status Profile(::grpc::ServerContext* ctx, const ProfileRequest* req,
115 ProfileResponse* response) override {
116 LOG(INFO) << "Received a profile request: " << req->DebugString();
117 uint64 start_time_ns = EnvTime::NowNanos();
118 std::unique_ptr<ProfilerSession> profiler = ProfilerSession::Create();
119 Status status = profiler->Status();
120 if (!status.ok()) {
121 return ::grpc::Status(::grpc::StatusCode::INTERNAL,
122 status.error_message());
123 }
124
125 Env* env = Env::Default();
126 for (size_t i = 0; i < req->duration_ms(); ++i) {
127 env->SleepForMicroseconds(EnvTime::kMillisToMicros);
128 if (ctx->IsCancelled()) {
129 return ::grpc::Status::CANCELLED;
130 }
131 }
132
133 status =
134 CollectDataToResponse(*req, profiler.get(), start_time_ns, response);
135 if (!status.ok()) {
136 return ::grpc::Status(::grpc::StatusCode::INTERNAL,
137 status.error_message());
138 }
139
140 return ::grpc::Status::OK;
141 }
142 };
143
144 } // namespace
145
CreateProfilerService()146 std::unique_ptr<grpc::ProfilerService::Service> CreateProfilerService() {
147 return MakeUnique<ProfilerServiceImpl>();
148 }
149
150 } // namespace tensorflow
151