1 /* Copyright 2020 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 #include "tensorflow/core/profiler/convert/xplane_to_profile_response.h"
16
17 #include <string>
18
19 #include "absl/container/flat_hash_set.h"
20 #include "absl/strings/str_cat.h"
21 #include "absl/strings/string_view.h"
22 #include "tensorflow/core/platform/errors.h"
23 #include "tensorflow/core/platform/status.h"
24 #include "tensorflow/core/platform/types.h"
25 #include "tensorflow/core/profiler/convert/op_stats_to_input_pipeline_analysis.h"
26 #include "tensorflow/core/profiler/convert/op_stats_to_overview_page.h"
27 #include "tensorflow/core/profiler/convert/op_stats_to_tf_stats.h"
28 #include "tensorflow/core/profiler/convert/trace_events_to_json.h"
29 #include "tensorflow/core/profiler/convert/xplane_to_memory_profile.h"
30 #include "tensorflow/core/profiler/convert/xplane_to_op_stats.h"
31 #include "tensorflow/core/profiler/convert/xplane_to_trace_events.h"
32 #include "tensorflow/core/profiler/profiler_service.pb.h"
33 #include "tensorflow/core/profiler/protobuf/input_pipeline.pb.h"
34 #include "tensorflow/core/profiler/protobuf/kernel_stats.pb.h"
35 #include "tensorflow/core/profiler/protobuf/memory_profile.pb.h"
36 #include "tensorflow/core/profiler/protobuf/overview_page.pb.h"
37 #include "tensorflow/core/profiler/protobuf/tf_stats.pb.h"
38 #include "tensorflow/core/profiler/protobuf/trace_events.pb.h"
39 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
40 #include "tensorflow/core/profiler/rpc/client/save_profile.h"
41 #include "tensorflow/core/profiler/utils/xplane_schema.h"
42
43 namespace tensorflow {
44 namespace profiler {
45 namespace {
46
47 const absl::string_view kTraceViewer = "trace_viewer";
48 const absl::string_view kTensorflowStats = "tensorflow_stats";
49 const absl::string_view kInputPipeline = "input_pipeline";
50 const absl::string_view kOverviewPage = "overview_page";
51 const absl::string_view kKernelStats = "kernel_stats";
52 const absl::string_view kMemoryProfile = "memory_profile";
53 const absl::string_view kXPlanePb = "xplane.pb";
54
55 template <typename Proto>
AddToolData(absl::string_view tool_name,const Proto & tool_output,ProfileResponse * response)56 void AddToolData(absl::string_view tool_name, const Proto& tool_output,
57 ProfileResponse* response) {
58 auto* tool_data = response->add_tool_data();
59 tool_data->set_name(string(tool_name));
60 tool_output.SerializeToString(tool_data->mutable_data());
61 }
62
63 // Returns the tool name with extension.
ToolName(absl::string_view tool)64 std::string ToolName(absl::string_view tool) {
65 if (tool == kTraceViewer) return "trace.json.gz";
66 if (tool == kMemoryProfile) return "memory_profile.json.gz";
67 return absl::StrCat(tool, ".pb");
68 }
69
70 } // namespace
71
ConvertXSpaceToProfileResponse(const XSpace & xspace,const ProfileRequest & req,ProfileResponse * response)72 Status ConvertXSpaceToProfileResponse(const XSpace& xspace,
73 const ProfileRequest& req,
74 ProfileResponse* response) {
75 absl::flat_hash_set<absl::string_view> tools(req.tools().begin(),
76 req.tools().end());
77 if (tools.empty()) return Status::OK();
78 if (tools.contains(kXPlanePb)) {
79 AddToolData(kXPlanePb, xspace, response);
80 }
81 if (tools.contains(kTraceViewer)) {
82 Trace trace;
83 ConvertXSpaceToTraceEvents(xspace, &trace);
84 if (trace.trace_events().empty()) {
85 response->set_empty_trace(true);
86 return Status::OK();
87 }
88 TF_RETURN_IF_ERROR(SaveGzippedToolData(
89 req.repository_root(), req.session_id(), req.host_name(),
90 ToolName(kTraceViewer), TraceEventsToJson(trace)));
91 // Trace viewer is the only tool, skip OpStats conversion.
92 if (tools.size() == 1) return Status::OK();
93 }
94
95 OpStatsOptions options;
96 options.generate_kernel_stats_db = true;
97 options.generate_op_metrics_db = true;
98 options.generate_step_db = true;
99 options.maybe_drop_incomplete_steps = true;
100 OpStats op_stats = ConvertXSpaceToOpStats(xspace, options);
101 if (tools.contains(kOverviewPage)) {
102 OverviewPage overview_page_db = ConvertOpStatsToOverviewPage(op_stats);
103 AddToolData(ToolName(kOverviewPage), overview_page_db, response);
104 if (tools.contains(kInputPipeline)) {
105 AddToolData(ToolName(kInputPipeline), overview_page_db.input_analysis(),
106 response);
107 }
108 } else if (tools.contains(kInputPipeline)) {
109 AddToolData(ToolName(kInputPipeline),
110 ConvertOpStatsToInputPipelineAnalysis(op_stats), response);
111 }
112 if (tools.contains(kTensorflowStats)) {
113 TfStatsDatabase tf_stats_db = ConvertOpStatsToTfStats(op_stats);
114 AddToolData(ToolName(kTensorflowStats), tf_stats_db, response);
115 }
116 if (tools.contains(kKernelStats)) {
117 AddToolData(ToolName(kKernelStats), op_stats.kernel_stats_db(), response);
118 }
119 if (tools.contains(kMemoryProfile)) {
120 std::string json_output;
121 TF_RETURN_IF_ERROR(ConvertXSpaceToMemoryProfileJson(xspace, &json_output));
122 TF_RETURN_IF_ERROR(SaveGzippedToolData(
123 req.repository_root(), req.session_id(), req.host_name(),
124 ToolName(kMemoryProfile), json_output));
125 }
126 return Status::OK();
127 }
128
129 } // namespace profiler
130 } // namespace tensorflow
131