1/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16// Helper functions for dumping Graphs, GraphDefs, and FunctionDefs to files for
17// debugging.
18
19#include "tensorflow/core/util/dump_graph.h"
20
21#include <memory>
22#include <unordered_map>
23
24#include "absl/strings/match.h"
25#include "absl/strings/str_cat.h"
26#include "tensorflow/core/lib/strings/proto_serialization.h"
27#include "tensorflow/core/platform/env.h"
28#include "tensorflow/core/platform/file_system.h"
29#include "tensorflow/core/platform/mutex.h"
30#include "tensorflow/core/platform/path.h"
31#include "tensorflow/core/platform/strcat.h"
32
33namespace tensorflow {
34
35namespace {
36using strings::StrCat;
37
38struct NameCounts {
39 mutex counts_mutex;
40 std::unordered_map<string, int> counts;
41};
42
43string MakeUniqueFilename(string name, const string& suffix = ".pbtxt") {
44 static NameCounts& instance = *new NameCounts;
45
46 // Remove illegal characters from `name`.
47 for (int i = 0; i < name.size(); ++i) {
48 char ch = name[i];
49 if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?' ||
50 ch == '\\') {
51 name[i] = '_';
52 }
53 }
54
55 int count;
56 {
57 mutex_lock lock(instance.counts_mutex);
58 count = instance.counts[name]++;
59 }
60
61 string filename = name;
62 if (count > 0) {
63 absl::StrAppend(&filename, "_", count);
64 }
65 absl::StrAppend(&filename, suffix);
66 return filename;
67}
68
69struct GraphDumperConfig {
70 mutex mu;
71
72 // The dumper and suffix configured.
73 struct Config {
74 bool IsSet() const { return dumper != nullptr; }
75 std::function<Status(const Graph& graph,
76 const FunctionLibraryDefinition* flib_def,
77 WritableFile*)>
78 dumper = nullptr;
79 string suffix = ".pbtxt";
80 } config TF_GUARDED_BY(mu);
81
82 // Returns whether a custom dumper is set.
83 bool IsSet() TF_LOCKS_EXCLUDED(mu) {
84 mutex_lock lock(mu);
85 return config.IsSet();
86 }
87};
88
89GraphDumperConfig& GetGraphDumperConfig() {
90 static GraphDumperConfig config;
91 return config;
92}
93
94// WritableFile that simply prints to stderr.
95class StderrWritableFile : public WritableFile {
96 public:
97 StderrWritableFile() {}
98
99 Status Append(StringPiece data) override {
100 fprintf(stderr, "%.*s", static_cast<int>(data.size()), data.data());
101 return OkStatus();
102 }
103
104 Status Close() override { return OkStatus(); }
105
106 Status Flush() override {
107 fflush(stderr);
108 return OkStatus();
109 }
110
111 Status Name(StringPiece* result) const override {
112 *result = "stderr";
113 return OkStatus();
114 }
115
116 Status Sync() override { return OkStatus(); }
117
118 Status Tell(int64_t* position) override {
119 return errors::Unimplemented("Stream not seekable");
120 }
121};
122
123Status CreateWritableFile(Env* env, const string& dirname, const string& name,
124 const string& suffix, string* filepath,
125 std::unique_ptr<WritableFile>* file) {
126 string dir;
127 if (!dirname.empty()) {
128 dir = dirname;
129 } else {
130 const char* prefix = getenv("TF_DUMP_GRAPH_PREFIX");
131 if (prefix != nullptr) dir = prefix;
132 }
133 if (dir.empty()) {
134 LOG(WARNING)
135 << "Failed to dump " << name << " because dump location is not "
136 << " specified through either TF_DUMP_GRAPH_PREFIX environment "
137 << "variable or function argument.";
138 return errors::InvalidArgument("TF_DUMP_GRAPH_PREFIX not specified");
139 }
140
141 if (absl::EqualsIgnoreCase(dir, "sponge") ||
142 absl::EqualsIgnoreCase(dir, "test_undeclared_outputs_dir")) {
143 if (!io::GetTestUndeclaredOutputsDir(&dir)) {
144 LOG(WARNING) << "TF_DUMP_GRAPH_PREFIX=sponge, but "
145 "TEST_UNDECLARED_OUTPUT_DIRS is not set, dumping to log";
146 dir = "-";
147 }
148 }
149
150 *filepath = "NULL";
151 if (dir == "-") {
152 *file = std::make_unique<StderrWritableFile>();
153 *filepath = "(stderr)";
154 return OkStatus();
155 }
156
157 TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(dir));
158 *filepath = io::JoinPath(dir, MakeUniqueFilename(name, suffix));
159 return env->NewWritableFile(*filepath, file);
160}
161
162Status WriteTextProtoToUniqueFile(const tensorflow::protobuf::Message& proto,
163 WritableFile* file) {
164 string s;
165 if (!::tensorflow::protobuf::TextFormat::PrintToString(proto, &s)) {
166 return errors::FailedPrecondition("Unable to convert proto to text.");
167 }
168 TF_RETURN_IF_ERROR(file->Append(s));
169 StringPiece name;
170 TF_RETURN_IF_ERROR(file->Name(&name));
171 VLOG(5) << name;
172 VLOG(5) << s;
173 return file->Close();
174}
175
176Status WriteTextProtoToUniqueFile(
177 const tensorflow::protobuf::MessageLite& proto, WritableFile* file) {
178 string s;
179 if (!SerializeToStringDeterministic(proto, &s)) {
180 return errors::Internal("Failed to serialize proto to string.");
181 }
182 StringPiece name;
183 TF_RETURN_IF_ERROR(file->Name(&name));
184 VLOG(5) << name;
185 VLOG(5) << s;
186 TF_RETURN_IF_ERROR(file->Append(s));
187 return file->Close();
188}
189
190string DumpToFile(const string& name, const string& dirname,
191 const string& suffix, const string& type_name,
192 std::function<Status(WritableFile*)> dumper) {
193 string filepath;
194 std::unique_ptr<WritableFile> file;
195 Status status = CreateWritableFile(Env::Default(), dirname, name, suffix,
196 &filepath, &file);
197 if (!status.ok()) {
198 return StrCat("(failed to create writable file: ", status.ToString(), ")");
199 }
200
201 status = dumper(file.get());
202 if (!status.ok()) {
203 return StrCat("(failed to dump ", type_name, " to '", filepath,
204 "': ", status.ToString(), ")");
205 }
206 LOG(INFO) << "Dumped " << type_name << " to " << filepath;
207 return filepath;
208}
209
210} // anonymous namespace
211
212void SetGraphDumper(
213 std::function<Status(const Graph& graph,
214 const FunctionLibraryDefinition* flib_def,
215 WritableFile*)>
216 dumper,
217 string suffix) {
218 GraphDumperConfig& dumper_config = GetGraphDumperConfig();
219 mutex_lock lock(dumper_config.mu);
220 dumper_config.config.dumper = dumper;
221 dumper_config.config.suffix = suffix;
222}
223
224string DumpGraphDefToFile(const string& name, GraphDef const& graph_def,
225 const string& dirname) {
226 return DumpToFile(name, dirname, ".pbtxt", "Graph", [&](WritableFile* file) {
227 return WriteTextProtoToUniqueFile(graph_def, file);
228 });
229}
230
231string DumpCostGraphDefToFile(const string& name, CostGraphDef const& graph_def,
232 const string& dirname) {
233 return DumpToFile(name, dirname, ".pbtxt", "Graph", [&](WritableFile* file) {
234 return WriteTextProtoToUniqueFile(graph_def, file);
235 });
236}
237
238string DumpGraphToFile(const string& name, Graph const& graph,
239 const FunctionLibraryDefinition* flib_def,
240 const string& dirname) {
241 auto& dumper_config = GetGraphDumperConfig();
242 if (dumper_config.IsSet()) {
243 GraphDumperConfig::Config config;
244 {
245 mutex_lock lock(dumper_config.mu);
246 config = dumper_config.config;
247 }
248 if (config.IsSet()) {
249 return DumpToFile(name, dirname, config.suffix, "Graph",
250 [&](WritableFile* file) {
251 return config.dumper(graph, flib_def, file);
252 });
253 }
254 }
255
256 GraphDef graph_def;
257 graph.ToGraphDef(&graph_def);
258 if (flib_def) {
259 *graph_def.mutable_library() = flib_def->ToProto();
260 }
261 return DumpGraphDefToFile(name, graph_def, dirname);
262}
263
264string DumpFunctionDefToFile(const string& name, FunctionDef const& fdef,
265 const string& dirname) {
266 return DumpToFile(name, dirname, ".pbtxt", "FunctionDef",
267 [&](WritableFile* file) {
268 return WriteTextProtoToUniqueFile(fdef, file);
269 });
270}
271
272} // namespace tensorflow
273