1 /* Copyright 2017 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/client/save_profile.h"
17
18 #include <memory>
19 #include <sstream>
20 #include <string>
21 #include <vector>
22
23 #include "absl/strings/match.h"
24 #include "absl/strings/str_cat.h"
25 #include "absl/strings/str_replace.h"
26 #include "absl/strings/string_view.h"
27 #include "absl/strings/strip.h"
28 #include "absl/time/clock.h"
29 #include "absl/time/time.h"
30 #include "tensorflow/core/lib/io/zlib_compression_options.h"
31 #include "tensorflow/core/lib/io/zlib_outputbuffer.h"
32 #include "tensorflow/core/platform/env.h"
33 #include "tensorflow/core/platform/errors.h"
34 #include "tensorflow/core/platform/file_system.h"
35 #include "tensorflow/core/platform/logging.h"
36 #include "tensorflow/core/platform/status.h"
37 #include "tensorflow/core/profiler/profiler_service.pb.h"
38 #include "tensorflow/core/profiler/utils/file_system_utils.h"
39
40 // Windows.h #defines ERROR, but it is also used in
41 // tensorflow/core/util/event.proto
42 #undef ERROR
43 #include "tensorflow/core/util/events_writer.h"
44
45 namespace tensorflow {
46 namespace profiler {
47 namespace {
48
49
50 constexpr char kProtoTraceFileName[] = "trace";
51 constexpr char kTfStatsHelperSuffix[] = "tf_stats_helper_result";
52
DumpToolData(absl::string_view run_dir,absl::string_view host,const ProfileToolData & tool,std::ostream * os)53 Status DumpToolData(absl::string_view run_dir, absl::string_view host,
54 const ProfileToolData& tool, std::ostream* os) {
55 // Don't save the intermediate results for combining the per host tool data.
56 if (absl::EndsWith(tool.name(), kTfStatsHelperSuffix)) return OkStatus();
57 std::string host_prefix = host.empty() ? "" : absl::StrCat(host, ".");
58 std::string path =
59 ProfilerJoinPath(run_dir, absl::StrCat(host_prefix, tool.name()));
60 TF_RETURN_IF_ERROR(WriteStringToFile(Env::Default(), path, tool.data()));
61 if (os) {
62 *os << "Dumped tool data for " << tool.name() << " to " << path << '\n';
63 }
64 return OkStatus();
65 }
66
WriteGzippedDataToFile(const std::string & filepath,const std::string & data)67 Status WriteGzippedDataToFile(const std::string& filepath,
68 const std::string& data) {
69 std::unique_ptr<WritableFile> file;
70 TF_RETURN_IF_ERROR(Env::Default()->NewWritableFile(filepath, &file));
71 io::ZlibCompressionOptions options = io::ZlibCompressionOptions::GZIP();
72 io::ZlibOutputBuffer buffer(file.get(), options.input_buffer_size,
73 options.output_buffer_size, options);
74 TF_RETURN_IF_ERROR(buffer.Init());
75 TF_RETURN_IF_ERROR(buffer.Append(data));
76 TF_RETURN_IF_ERROR(buffer.Close());
77 TF_RETURN_IF_ERROR(file->Close());
78 return OkStatus();
79 }
80
GetOrCreateRunDir(const std::string & repository_root,const std::string & run,std::string * run_dir,std::ostream * os)81 Status GetOrCreateRunDir(const std::string& repository_root,
82 const std::string& run, std::string* run_dir,
83 std::ostream* os) {
84 // Creates a directory to <repository_root>/<run>/.
85 *run_dir = ProfilerJoinPath(repository_root, run);
86 *os << "Creating directory: " << *run_dir << '\n';
87 TF_RETURN_IF_ERROR(Env::Default()->RecursivelyCreateDir(*run_dir));
88 return OkStatus();
89 }
90 } // namespace
91
GetTensorBoardProfilePluginDir(const std::string & logdir)92 std::string GetTensorBoardProfilePluginDir(const std::string& logdir) {
93 constexpr char kPluginName[] = "plugins";
94 constexpr char kProfileName[] = "profile";
95 return ProfilerJoinPath(logdir, kPluginName, kProfileName);
96 }
97
MaybeCreateEmptyEventFile(const std::string & logdir)98 Status MaybeCreateEmptyEventFile(const std::string& logdir) {
99 // Suffix for an empty event file. it should be kept in sync with
100 // _EVENT_FILE_SUFFIX in tensorflow/python/eager/profiler.py.
101 constexpr char kProfileEmptySuffix[] = ".profile-empty";
102 TF_RETURN_IF_ERROR(Env::Default()->RecursivelyCreateDir(logdir));
103
104 std::vector<std::string> children;
105 TF_RETURN_IF_ERROR(Env::Default()->GetChildren(logdir, &children));
106 for (const std::string& child : children) {
107 if (absl::EndsWith(child, kProfileEmptySuffix)) {
108 return OkStatus();
109 }
110 }
111 EventsWriter event_writer(ProfilerJoinPath(logdir, "events"));
112 return event_writer.InitWithSuffix(kProfileEmptySuffix);
113 }
114
SaveProfile(const std::string & repository_root,const std::string & run,const std::string & host,const ProfileResponse & response,std::ostream * os)115 Status SaveProfile(const std::string& repository_root, const std::string& run,
116 const std::string& host, const ProfileResponse& response,
117 std::ostream* os) {
118 if (response.tool_data().empty()) return OkStatus();
119 std::string run_dir;
120 TF_RETURN_IF_ERROR(GetOrCreateRunDir(repository_root, run, &run_dir, os));
121 // Windows file names do not support colons.
122 std::string hostname = absl::StrReplaceAll(host, {{":", "_"}});
123 for (const auto& tool_data : response.tool_data()) {
124 TF_RETURN_IF_ERROR(DumpToolData(run_dir, hostname, tool_data, os));
125 }
126 return OkStatus();
127 }
128
SaveGzippedToolData(const std::string & repository_root,const std::string & run,const std::string & host,const std::string & tool_name,const std::string & data)129 Status SaveGzippedToolData(const std::string& repository_root,
130 const std::string& run, const std::string& host,
131 const std::string& tool_name,
132 const std::string& data) {
133 std::string run_dir;
134 std::stringstream ss;
135 Status status = GetOrCreateRunDir(repository_root, run, &run_dir, &ss);
136 LOG(INFO) << ss.str();
137 TF_RETURN_IF_ERROR(status);
138 std::string host_prefix = host.empty() ? "" : absl::StrCat(host, ".");
139 std::string path =
140 ProfilerJoinPath(run_dir, absl::StrCat(host_prefix, tool_name));
141 TF_RETURN_IF_ERROR(WriteGzippedDataToFile(path, data));
142 LOG(INFO) << "Dumped gzipped tool data for " << tool_name << " to " << path;
143 return OkStatus();
144 }
145
GetCurrentTimeStampAsString()146 std::string GetCurrentTimeStampAsString() {
147 return absl::FormatTime("%E4Y_%m_%d_%H_%M_%S", absl::Now(),
148 absl::LocalTimeZone());
149 }
150
151 } // namespace profiler
152 } // namespace tensorflow
153