From c83525a1887ac3d7c03d4d25351e421cd90069a4 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Thu, 26 Jul 2018 11:32:54 -0700
Subject: Ergonomic fixes to tensorflow::dump_graph::*

 - Teach the routines to create the dump directories if they don't exist.  I
   routinely forget to create the dump directory before launching the model.
 - Print out the file locations on LOG(INFO) and the error messages on
   LOG(WARNING).

PiperOrigin-RevId: 206192822
---
 tensorflow/compiler/tf2xla/dump_graph.cc | 53 ++++++++++++++++++--------------
 1 file changed, 30 insertions(+), 23 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/dump_graph.cc b/tensorflow/compiler/tf2xla/dump_graph.cc
index 03603ee9ba..24616c01c7 100644
--- a/tensorflow/compiler/tf2xla/dump_graph.cc
+++ b/tensorflow/compiler/tf2xla/dump_graph.cc
@@ -33,7 +33,7 @@ struct NameCounts {
   std::unordered_map<string, int> counts;
 };
 
-string MakeUniquePath(string name) {
+string MakeUniqueFilename(string name) {
   static NameCounts& instance = *new NameCounts;
 
   // Remove illegal characters from `name`.
@@ -50,26 +50,41 @@ string MakeUniquePath(string name) {
     count = instance.counts[name]++;
   }
 
-  legacy_flags::DumpGraphFlags* flags = legacy_flags::GetDumpGraphFlags();
-  string path = strings::StrCat(flags->tf_dump_graph_prefix, "/", name);
+  string filename = name;
   if (count > 0) {
-    strings::StrAppend(&path, "_", count);
+    strings::StrAppend(&filename, "_", count);
   }
-  strings::StrAppend(&path, ".pbtxt");
-  return path;
+  strings::StrAppend(&filename, ".pbtxt");
+  return filename;
+}
+
+string WriteTextProtoToUniqueFile(
+    Env* env, const string& name, const char* proto_type,
+    const ::tensorflow::protobuf::Message& proto) {
+  const string& dirname =
+      legacy_flags::GetDumpGraphFlags()->tf_dump_graph_prefix;
+  Status status = env->RecursivelyCreateDir(dirname);
+  if (!status.ok()) {
+    LOG(WARNING) << "Failed to create " << dirname << " for dumping "
+                 << proto_type << ": " << status;
+    return "(unavailable)";
+  }
+  string filepath = strings::StrCat(dirname, "/", MakeUniqueFilename(name));
+  status = WriteTextProto(Env::Default(), filepath, proto);
+  if (!status.ok()) {
+    LOG(WARNING) << "Failed to dump " << proto_type << " to file: " << filepath
+                 << " : " << status;
+    return "(unavailable)";
+  }
+  LOG(INFO) << "Dumped " << proto_type << " to " << filepath;
+  return filepath;
 }
 
 }  // anonymous namespace
 
 string DumpGraphDefToFile(const string& name, GraphDef const& graph_def) {
-  string path = MakeUniquePath(name);
-  Status status = WriteTextProto(Env::Default(), path, graph_def);
-  if (!status.ok()) {
-    VLOG(1) << "Failed to dump GraphDef to file: " << path << " : " << status;
-    path.clear();
-    path = "(unavailable)";
-  }
-  return path;
+  return WriteTextProtoToUniqueFile(Env::Default(), name, "GraphDef",
+                                    graph_def);
 }
 
 string DumpGraphToFile(const string& name, Graph const& graph,
@@ -83,15 +98,7 @@ string DumpGraphToFile(const string& name, Graph const& graph,
 }
 
 string DumpFunctionDefToFile(const string& name, FunctionDef const& fdef) {
-  string path = MakeUniquePath(name);
-  Status status = WriteTextProto(Env::Default(), path, fdef);
-  if (!status.ok()) {
-    VLOG(1) << "Failed to dump FunctionDef to file: " << path << " : "
-            << status;
-    path.clear();
-    path = "(unavailable)";
-  }
-  return path;
+  return WriteTextProtoToUniqueFile(Env::Default(), name, "FunctionDef", fdef);
 }
 
 }  // namespace dump_graph
-- 
cgit v1.2.3