1 /* Copyright 2018 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 // 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
33 namespace tensorflow {
34
35 namespace {
36 using strings::StrCat;
37
38 struct NameCounts {
39 mutex counts_mutex;
40 std::unordered_map<string, int> counts;
41 };
42
MakeUniqueFilename(string name,const string & suffix=".pbtxt")43 string 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
69 struct GraphDumperConfig {
70 mutex mu;
71
72 // The dumper and suffix configured.
73 struct Config {
IsSettensorflow::__anon7b34d4dc0111::GraphDumperConfig::Config74 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.
IsSettensorflow::__anon7b34d4dc0111::GraphDumperConfig83 bool IsSet() TF_LOCKS_EXCLUDED(mu) {
84 mutex_lock lock(mu);
85 return config.IsSet();
86 }
87 };
88
GetGraphDumperConfig()89 GraphDumperConfig& GetGraphDumperConfig() {
90 static GraphDumperConfig config;
91 return config;
92 }
93
94 // WritableFile that simply prints to stderr.
95 class StderrWritableFile : public WritableFile {
96 public:
StderrWritableFile()97 StderrWritableFile() {}
98
Append(StringPiece data)99 Status Append(StringPiece data) override {
100 fprintf(stderr, "%.*s", static_cast<int>(data.size()), data.data());
101 return Status::OK();
102 }
103
Close()104 Status Close() override { return Status::OK(); }
105
Flush()106 Status Flush() override {
107 fflush(stderr);
108 return Status::OK();
109 }
110
Name(StringPiece * result) const111 Status Name(StringPiece* result) const override {
112 *result = "stderr";
113 return Status::OK();
114 }
115
Sync()116 Status Sync() override { return Status::OK(); }
117
Tell(int64 * position)118 Status Tell(int64* position) override {
119 return errors::Unimplemented("Stream not seekable");
120 }
121 };
122
CreateWritableFile(Env * env,const string & dirname,const string & name,const string & suffix,string * filepath,std::unique_ptr<WritableFile> * file)123 Status 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 Status::OK();
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
WriteTextProtoToUniqueFile(const tensorflow::protobuf::Message & proto,WritableFile * file)162 Status 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 return file->Close();
170 }
171
WriteTextProtoToUniqueFile(const tensorflow::protobuf::MessageLite & proto,WritableFile * file)172 Status WriteTextProtoToUniqueFile(
173 const tensorflow::protobuf::MessageLite& proto, WritableFile* file) {
174 string s;
175 if (!SerializeToStringDeterministic(proto, &s)) {
176 return errors::Internal("Failed to serialize proto to string.");
177 }
178 TF_RETURN_IF_ERROR(file->Append(s));
179 return file->Close();
180 }
181
182 } // anonymous namespace
183
SetGraphDumper(std::function<Status (const Graph & graph,const FunctionLibraryDefinition * flib_def,WritableFile *)> dumper,string suffix)184 void SetGraphDumper(
185 std::function<Status(const Graph& graph,
186 const FunctionLibraryDefinition* flib_def,
187 WritableFile*)>
188 dumper,
189 string suffix) {
190 GraphDumperConfig& dumper_config = GetGraphDumperConfig();
191 mutex_lock lock(dumper_config.mu);
192 dumper_config.config.dumper = dumper;
193 dumper_config.config.suffix = suffix;
194 }
195
DumpGraphDefToFile(const string & name,GraphDef const & graph_def,const string & dirname)196 string DumpGraphDefToFile(const string& name, GraphDef const& graph_def,
197 const string& dirname) {
198 string filepath;
199 std::unique_ptr<WritableFile> file;
200 Status status = CreateWritableFile(Env::Default(), dirname, name, ".pbtxt",
201 &filepath, &file);
202 if (!status.ok()) {
203 return StrCat("(failed to create writable file: ", status.ToString(), ")");
204 }
205
206 status = WriteTextProtoToUniqueFile(graph_def, file.get());
207 if (!status.ok()) {
208 return StrCat("(failed to dump Graph to '", filepath,
209 "': ", status.ToString(), ")");
210 }
211 LOG(INFO) << "Dumped Graph to " << filepath;
212 return filepath;
213 }
214
DumpCostGraphDefToFile(const string & name,CostGraphDef const & graph_def,const string & dirname)215 string DumpCostGraphDefToFile(const string& name, CostGraphDef const& graph_def,
216 const string& dirname) {
217 string filepath;
218 std::unique_ptr<WritableFile> file;
219 Status status = CreateWritableFile(Env::Default(), dirname, name, ".pbtxt",
220 &filepath, &file);
221 if (!status.ok()) {
222 return StrCat("(failed to create writable file: ", status.ToString(), ")");
223 }
224
225 status = WriteTextProtoToUniqueFile(graph_def, file.get());
226 if (!status.ok()) {
227 return StrCat("(failed to dump Graph to '", filepath,
228 "': ", status.ToString(), ")");
229 }
230 LOG(INFO) << "Dumped Graph to " << filepath;
231 return filepath;
232 }
233
DumpGraphToFile(const string & name,Graph const & graph,const FunctionLibraryDefinition * flib_def,const string & dirname)234 string DumpGraphToFile(const string& name, Graph const& graph,
235 const FunctionLibraryDefinition* flib_def,
236 const string& dirname) {
237 auto& dumper_config = GetGraphDumperConfig();
238 if (dumper_config.IsSet()) {
239 GraphDumperConfig::Config config;
240 {
241 mutex_lock lock(dumper_config.mu);
242 config = dumper_config.config;
243 }
244 if (config.IsSet()) {
245 string filepath;
246 std::unique_ptr<WritableFile> file;
247 Status status = CreateWritableFile(Env::Default(), dirname, name,
248 config.suffix, &filepath, &file);
249 if (!status.ok()) {
250 return StrCat("(failed to create writable file: ", status.ToString(),
251 ")");
252 }
253 status = config.dumper(graph, flib_def, file.get());
254 if (!status.ok()) {
255 return StrCat("(failed to dump Graph to '", filepath,
256 "': ", status.ToString(), ")");
257 }
258 LOG(INFO) << "Dumped Graph to " << filepath;
259 return filepath;
260 }
261 }
262
263 GraphDef graph_def;
264 graph.ToGraphDef(&graph_def);
265 if (flib_def) {
266 *graph_def.mutable_library() = flib_def->ToProto();
267 }
268 return DumpGraphDefToFile(name, graph_def, dirname);
269 }
270
DumpFunctionDefToFile(const string & name,FunctionDef const & fdef,const string & dirname)271 string DumpFunctionDefToFile(const string& name, FunctionDef const& fdef,
272 const string& dirname) {
273 string filepath;
274 std::unique_ptr<WritableFile> file;
275 Status status = CreateWritableFile(Env::Default(), dirname, name, ".pbtxt",
276 &filepath, &file);
277 if (!status.ok()) {
278 return StrCat("(failed to create writable file: ", status.ToString(), ")");
279 }
280
281 status = WriteTextProtoToUniqueFile(fdef, file.get());
282 if (!status.ok()) {
283 return StrCat("(failed to dump FunctionDef to '", filepath,
284 "': ", status.ToString(), ")");
285 }
286 LOG(INFO) << "Dumped FunctionDef to " << filepath;
287 return filepath;
288 }
289
290 } // namespace tensorflow
291