• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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