xref: /aosp_15_r20/external/tensorflow/tensorflow/core/profiler/rpc/client/save_profile.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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