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