Merge pull request #129417 from vishwakt:fix-debug-events-writer-reinit

PiperOrigin-RevId: 1005360689

tensorflow/core/util/debug_events_writer.cc

@@ -39,9 +39,10 @@ void MaybeSetDebugEventTimestamp(DebugEvent* debug_event, Env* env) {
 }  // namespace
 
 SingleDebugEventFileWriter::SingleDebugEventFileWriter(
-    const std::string& file_path)
+    const std::string& file_path, bool append)
     : env_(Env::Default()),
       file_path_(file_path),
+      append_(append),
       num_outstanding_events_(0),
       writer_mu_() {}
 
@@ -56,9 +57,15 @@ absl::Status SingleDebugEventFileWriter::InitLocked() {
   // Flush() and Close() call have access to writable_file_.
   record_writer_.reset();
 
-  TF_RETURN_WITH_CONTEXT_IF_ERROR(
-      env_->NewWritableFile(file_path_, &writable_file_),
-      "Creating writable file ", file_path_);
+  if (append_) {
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(
+        env_->NewAppendableFile(file_path_, &writable_file_),
+        "Opening appendable file ", file_path_);
+  } else {
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(
+        env_->NewWritableFile(file_path_, &writable_file_),
+        "Creating writable file ", file_path_);
+  }
   record_writer_ = std::make_unique<io::RecordWriter>(writable_file_.get());
   if (record_writer_ == nullptr) {
     return absl::UnknownError(
@@ -165,7 +172,18 @@ absl::Status DebugEventsWriter::LookUpDebugEventsWriter(
 
 absl::Status DebugEventsWriter::Init() {
   mutex_lock l(initialization_mu_);
+  return InitLocked();
+}
 
+absl::Status DebugEventsWriter::InitUnlessClosed() {
+  mutex_lock l(initialization_mu_);
+  if (is_closed_) {
+    return absl::OkStatus();
+  }
+  return InitLocked();
+}
+
+absl::Status DebugEventsWriter::InitLocked() {
   // TODO(cais): We currently don't check for file deletion. When the need
   // arises, check and fix file deletion.
   if (is_initialized_) {
@@ -177,23 +195,30 @@ absl::Status DebugEventsWriter::Init() {
                                     "Failed to create directory ", dump_root_);
   }
 
-  int64_t time_in_seconds = env_->NowMicros() / 1e6;
-  file_prefix_ = io::JoinPath(
-      dump_root_, absl::StrFormat("%s.%010lld.%s", kFileNamePrefix,
-                                  static_cast<long long>(time_in_seconds),
-                                  port::Hostname().c_str()));
+  // After Close(), the files are reopened for appending, so that everything
+  // written under the dump root stays in a single file set.
+  if (!is_closed_) {
+    int64_t time_in_seconds = env_->NowMicros() / 1e6;
+    file_prefix_ = io::JoinPath(
+        dump_root_, absl::StrFormat("%s.%010lld.%s", kFileNamePrefix,
+                                    static_cast<long long>(time_in_seconds),
+                                    port::Hostname().c_str()));
+  }
   TF_RETURN_IF_ERROR(InitNonMetadataFile(SOURCE_FILES));
   TF_RETURN_IF_ERROR(InitNonMetadataFile(STACK_FRAMES));
   TF_RETURN_IF_ERROR(InitNonMetadataFile(GRAPHS));
 
   // The metadata file should be created.
   std::string metadata_filename = GetFileNameInternal(METADATA);
-  auto metadata_writer =
-      std::make_unique<SingleDebugEventFileWriter>(metadata_filename);
+  auto metadata_writer = std::make_unique<SingleDebugEventFileWriter>(
+      metadata_filename, is_closed_);
   if (metadata_writer == nullptr) {
     return absl::UnknownError(
         "Could not create debug event metadata file writer");
   }
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(metadata_writer->Init(),
+                                  "Initializing debug event writer at path ",
+                                  metadata_filename);
 
   {
     mutex_lock l(writers_mu_);
@@ -218,54 +243,64 @@ absl::Status DebugEventsWriter::Init() {
   TF_RETURN_IF_ERROR(InitNonMetadataFile(EXECUTION));
   TF_RETURN_IF_ERROR(InitNonMetadataFile(GRAPH_EXECUTION_TRACES));
   is_initialized_ = true;
+  is_closed_ = false;
   return absl::OkStatus();
 }
 
 absl::Status DebugEventsWriter::WriteSourceFile(SourceFile* source_file) {
+  std::unique_ptr<SourceFile> file(source_file);
+  TF_RETURN_IF_ERROR(InitUnlessClosed());
   DebugEvent debug_event;
-  debug_event.set_allocated_source_file(source_file);
+  debug_event.set_allocated_source_file(file.release());
   return SerializeAndWriteDebugEvent(&debug_event, SOURCE_FILES);
 }
 
 absl::Status DebugEventsWriter::WriteStackFrameWithId(
     StackFrameWithId* stack_frame_with_id) {
+  std::unique_ptr<StackFrameWithId> frame(stack_frame_with_id);
+  TF_RETURN_IF_ERROR(InitUnlessClosed());
   DebugEvent debug_event;
-  debug_event.set_allocated_stack_frame_with_id(stack_frame_with_id);
+  debug_event.set_allocated_stack_frame_with_id(frame.release());
   return SerializeAndWriteDebugEvent(&debug_event, STACK_FRAMES);
 }
 
 absl::Status DebugEventsWriter::WriteGraphOpCreation(
     GraphOpCreation* graph_op_creation) {
+  std::unique_ptr<GraphOpCreation> op_creation(graph_op_creation);
+  TF_RETURN_IF_ERROR(InitUnlessClosed());
   DebugEvent debug_event;
-  debug_event.set_allocated_graph_op_creation(graph_op_creation);
+  debug_event.set_allocated_graph_op_creation(op_creation.release());
   return SerializeAndWriteDebugEvent(&debug_event, GRAPHS);
 }
 
 absl::Status DebugEventsWriter::WriteDebuggedGraph(
     DebuggedGraph* debugged_graph) {
+  std::unique_ptr<DebuggedGraph> graph(debugged_graph);
+  TF_RETURN_IF_ERROR(InitUnlessClosed());
   DebugEvent debug_event;
-  debug_event.set_allocated_debugged_graph(debugged_graph);
+  debug_event.set_allocated_debugged_graph(graph.release());
   return SerializeAndWriteDebugEvent(&debug_event, GRAPHS);
 }
 
 absl::Status DebugEventsWriter::WriteExecution(Execution* execution) {
+  std::unique_ptr<Execution> exec(execution);
+  TF_RETURN_IF_ERROR(InitUnlessClosed());
   if (circular_buffer_size_ <= 0) {
     // No cyclic-buffer behavior.
     DebugEvent debug_event;
-    debug_event.set_allocated_execution(execution);
+    debug_event.set_allocated_execution(exec.release());
     return SerializeAndWriteDebugEvent(&debug_event, EXECUTION);
   } else {
     // Circular buffer behavior.
     DebugEvent debug_event;
     MaybeSetDebugEventTimestamp(&debug_event, env_);
-    debug_event.set_allocated_execution(execution);
+    debug_event.set_allocated_execution(exec.release());
     std::string serialized;
     debug_event.SerializeToString(&serialized);
 
     tf_shared_lock writers_lock(writers_mu_);
     if (execution_writer_ == nullptr) {
-      return absl::FailedPreconditionError(
-          "Debug events execution writer is not initialized");
+      return absl::OkStatus();
     }
     mutex_lock l(execution_buffer_mu_);
     execution_buffer_.emplace_back(std::move(serialized));
@@ -278,24 +313,24 @@ absl::Status DebugEventsWriter::WriteExecution(Execution* execution) {
 
 absl::Status DebugEventsWriter::WriteGraphExecutionTrace(
     GraphExecutionTrace* graph_execution_trace) {
-  TF_RETURN_IF_ERROR(Init());
+  std::unique_ptr<GraphExecutionTrace> trace(graph_execution_trace);
+  TF_RETURN_IF_ERROR(InitUnlessClosed());
   if (circular_buffer_size_ <= 0) {
     // No cyclic-buffer behavior.
     DebugEvent debug_event;
-    debug_event.set_allocated_graph_execution_trace(graph_execution_trace);
+    debug_event.set_allocated_graph_execution_trace(trace.release());
     return SerializeAndWriteDebugEvent(&debug_event, GRAPH_EXECUTION_TRACES);
   } else {
     // Circular buffer behavior.
     DebugEvent debug_event;
     MaybeSetDebugEventTimestamp(&debug_event, env_);
-    debug_event.set_allocated_graph_execution_trace(graph_execution_trace);
+    debug_event.set_allocated_graph_execution_trace(trace.release());
     std::string serialized;
     debug_event.SerializeToString(&serialized);
 
     tf_shared_lock writers_lock(writers_mu_);
     if (graph_execution_traces_writer_ == nullptr) {
-      return absl::FailedPreconditionError(
-          "Debug events graph execution traces writer is not initialized");
+      return absl::OkStatus();
     }
     mutex_lock l(graph_execution_trace_buffer_mu_);
     graph_execution_trace_buffer_.emplace_back(std::move(serialized));
@@ -328,6 +363,7 @@ absl::Status DebugEventsWriter::WriteGraphExecutionTrace(
 
 void DebugEventsWriter::WriteSerializedNonExecutionDebugEvent(
     const std::string& debug_event_str, DebugEventFileType type) {
+  InitUnlessClosed().IgnoreError();
   tf_shared_lock l(writers_mu_);
   std::unique_ptr<SingleDebugEventFileWriter>* writer = nullptr;
   SelectWriter(type, &writer);
@@ -338,6 +374,7 @@ void DebugEventsWriter::WriteSerializedNonExecutionDebugEvent(
 
 void DebugEventsWriter::WriteSerializedExecutionDebugEvent(
     const std::string& debug_event_str, DebugEventFileType type) {
+  InitUnlessClosed().IgnoreError();
   tf_shared_lock writers_lock(writers_mu_);
   const std::unique_ptr<SingleDebugEventFileWriter>* writer = nullptr;
   std::deque<std::string>* buffer = nullptr;
@@ -408,7 +445,7 @@ absl::Status DebugEventsWriter::FlushNonExecutionFilesLocked() {
 }
 
 absl::Status DebugEventsWriter::FlushNonExecutionFiles() {
-  TF_RETURN_IF_ERROR(Init());
+  TF_RETURN_IF_ERROR(InitUnlessClosed());
   tf_shared_lock l(writers_mu_);
   return FlushNonExecutionFilesLocked();
 }
@@ -443,13 +480,13 @@ absl::Status DebugEventsWriter::FlushExecutionFilesLocked() {
 }
 
 absl::Status DebugEventsWriter::FlushExecutionFiles() {
-  TF_RETURN_IF_ERROR(Init());
+  TF_RETURN_IF_ERROR(InitUnlessClosed());
   tf_shared_lock l(writers_mu_);
   return FlushExecutionFilesLocked();
 }
 
 std::string DebugEventsWriter::FileName(DebugEventFileType type) {
-  Init().IgnoreError();
+  InitUnlessClosed().IgnoreError();
   mutex_lock l(initialization_mu_);
   return GetFileNameInternal(type);
 }
@@ -503,11 +540,13 @@ absl::Status DebugEventsWriter::Close() {
     }
     graph_execution_traces_writer_.reset();
   }
+  is_initialized_ = false;
+  is_closed_ = true;
 
   if (failed_to_close_files.empty()) {
     return absl::OkStatus();
   }
-  return absl::FailedPreconditionError(absl::StrCat(
+  return absl::FailedPreconditionError(absl::StrFormat(
       "Failed to close %d debug-events files associated with tfdbg",
       failed_to_close_files.size()));
 }
@@ -528,6 +567,7 @@ DebugEventsWriter::DebugEventsWriter(const std::string& dump_root,
       dump_root_(dump_root),
       tfdbg_run_id_(tfdbg_run_id),
       is_initialized_(false),
+      is_closed_(false),
       initialization_mu_(),
       circular_buffer_size_(circular_buffer_size),
       execution_buffer_(),
@@ -539,7 +579,8 @@ DebugEventsWriter::DebugEventsWriter(const std::string& dump_root,
 
 absl::Status DebugEventsWriter::InitNonMetadataFile(DebugEventFileType type) {
   const std::string filename = GetFileNameInternal(type);
-  auto new_writer = std::make_unique<SingleDebugEventFileWriter>(filename);
+  auto new_writer =
+      std::make_unique<SingleDebugEventFileWriter>(filename, is_closed_);
   if (new_writer == nullptr) {
     return absl::UnknownError(absl::StrCat(
         "Could not create debug event file writer for ", filename));
@@ -571,12 +612,8 @@ absl::Status DebugEventsWriter::SerializeAndWriteDebugEvent(
   SelectWriter(type, &writer);
   if (writer != nullptr && *writer != nullptr) {
     (*writer)->WriteSerializedDebugEvent(str);
-    return absl::OkStatus();
   }
-
-  return absl::InternalError(absl::StrCat(
-      "Unable to find debug events file writer for DebugEventsFileType ",
-      type));
+  return absl::OkStatus();
 }
 
 void DebugEventsWriter::SelectWriter(

tensorflow/core/util/debug_events_writer.h

@@ -49,7 +49,10 @@ enum DebugEventFileType {
 // TFRecord files, and hence utilizes multiple objects of this helper class.
 class SingleDebugEventFileWriter {
  public:
-  explicit SingleDebugEventFileWriter(const std::string& file_path);
+  // If `append` is true, an existing file is appended to instead of being
+  // overwritten.
+  explicit SingleDebugEventFileWriter(const std::string& file_path,
+                                      bool append = false);
 
   absl::Status Init();
 
@@ -66,6 +69,7 @@ class SingleDebugEventFileWriter {
 
   Env* env_;
   const std::string file_path_;
+  const bool append_;
   std::atomic_int_fast32_t num_outstanding_events_;
 
   std::unique_ptr<WritableFile> writable_file_ TF_GUARDED_BY(writer_mu_);
@@ -129,6 +133,8 @@ class DebugEventsWriter {
   // Idempotent: if the metadata file exists and is open, this is a no-op.
   // If on the other hand the file was opened, but has since disappeared (e.g.
   // deleted by another process), this will open a new file.
+  // After Close(), only a direct call to Init() reopens the files, which are
+  // then appended to.
   absl::Status Init();
 
   // The four DebugEvent fields below are written _without_ the circular
@@ -212,7 +218,8 @@ class DebugEventsWriter {
   absl::Status FlushExecutionFiles();
 
   // Close() calls FlushNonExecutionFiles() and FlushExecutionFiles()
-  // and then closes the current debug events files.
+  // and then closes the current debug events files. Events written after
+  // Close() are dropped until Init() is called again.
   absl::Status Close();
 
  private:
@@ -233,6 +240,12 @@ class DebugEventsWriter {
   // suffix.
   std::string FileName(DebugEventFileType type);
 
+  // Same as Init(), but a no-op if Close() has been called since the files
+  // were last opened.
+  absl::Status InitUnlessClosed();
+
+  absl::Status InitLocked() TF_EXCLUSIVE_LOCKS_REQUIRED(initialization_mu_);
+
   // Initialize the TFRecord writer for non-metadata file type.
   absl::Status InitNonMetadataFile(DebugEventFileType type)
       TF_EXCLUSIVE_LOCKS_REQUIRED(initialization_mu_);
@@ -258,6 +271,7 @@ class DebugEventsWriter {
 
   std::string file_prefix_ TF_GUARDED_BY(initialization_mu_);
   bool is_initialized_ TF_GUARDED_BY(initialization_mu_);
+  bool is_closed_ TF_GUARDED_BY(initialization_mu_);
   mutex initialization_mu_;
 
   const int64_t circular_buffer_size_;

tensorflow/core/util/debug_events_writer_test.cc

@@ -263,6 +263,77 @@ TEST_F(DebugEventsWriterTest, InitTwiceDoesNotCreateNewMetadataFile) {
   EXPECT_GE(actuals[0].debug_metadata().file_version().size(), 0);
 }
 
+TEST_F(DebugEventsWriterTest, InitAfterCloseAppendsToTheSameFiles) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+  GraphExecutionTrace* trace = new GraphExecutionTrace();
+  trace->set_tfdbg_context_id("graph_0");
+  TF_ASSERT_OK(writer->WriteGraphExecutionTrace(trace));
+  TF_ASSERT_OK(writer->Close());
+  const std::string traces_path =
+      GetDebugEventFileName(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES);
+
+  TF_ASSERT_OK(writer->Init());
+  EXPECT_EQ(
+      GetDebugEventFileName(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES),
+      traces_path);
+  trace = new GraphExecutionTrace();
+  trace->set_tfdbg_context_id("graph_1");
+  TF_ASSERT_OK(writer->WriteGraphExecutionTrace(trace));
+  TF_ASSERT_OK(writer->Close());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES,
+                       &actuals);
+  ASSERT_EQ(actuals.size(), 2);
+  EXPECT_EQ(actuals[0].graph_execution_trace().tfdbg_context_id(), "graph_0");
+  EXPECT_EQ(actuals[1].graph_execution_trace().tfdbg_context_id(), "graph_1");
+
+  ReadDebugEventProtos(writer, DebugEventFileType::METADATA, &actuals);
+  EXPECT_EQ(actuals.size(), 2);
+}
+
+TEST_F(DebugEventsWriterTest, FlushAfterCloseDoesNotReopenFiles) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  TF_ASSERT_OK(writer->Init());
+  TF_ASSERT_OK(writer->Close());
+  TF_ASSERT_OK(writer->FlushNonExecutionFiles());
+  TF_ASSERT_OK(writer->FlushExecutionFiles());
+
+  SourceFile* source_file = new SourceFile();
+  source_file->set_file_path("/home/tf_programs/main.py");
+  TF_ASSERT_OK(writer->WriteSourceFile(source_file));
+  GraphExecutionTrace* trace = new GraphExecutionTrace();
+  trace->set_tfdbg_context_id("graph_0");
+  TF_ASSERT_OK(writer->WriteGraphExecutionTrace(trace));
+  TF_ASSERT_OK(writer->FlushNonExecutionFiles());
+  TF_ASSERT_OK(writer->FlushExecutionFiles());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+  ReadDebugEventProtos(writer, DebugEventFileType::GRAPH_EXECUTION_TRACES,
+                       &actuals);
+  EXPECT_EQ(actuals.size(), 0);
+}
+
+TEST_F(DebugEventsWriterTest, WriteSourceFileWithoutPreviousInitCall) {
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+
+  SourceFile* source_file = new SourceFile();
+  source_file->set_file_path("/home/tf_programs/main.py");
+  TF_ASSERT_OK(writer->WriteSourceFile(source_file));
+  TF_ASSERT_OK(writer->Close());
+
+  std::vector<DebugEvent> actuals;
+  ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
+  ASSERT_EQ(actuals.size(), 1);
+  EXPECT_EQ(actuals[0].source_file().file_path(), "/home/tf_programs/main.py");
+}
+
 TEST_F(DebugEventsWriterTest, WriteSourceFile) {
   DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
       dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);

tensorflow/python/debug/lib/dumping_callback_test.py

@@ -1150,6 +1150,49 @@ class DumpingCallbackTest(
       self.assertEqual(reader.num_graph_execution_traces(), 0)
       self.assertFalse(reader.outermost_graphs())
 
+  def testDisableTracingKeepsTracedFunctions(self):
+
+    @def_function.function
+    def times_two(x):
+      return x * 2.0
+
+    x = constant_op.constant([1.0, 2.0])
+    dumping_callback.enable_dump_debug_info(
+        self.dump_root, tensor_debug_mode="FULL_HEALTH"
+    )
+    self.evaluate(times_two(x))
+    dumping_callback.disable_dump_debug_info()
+    self.assertAllClose(self.evaluate(times_two(x)), [2.0, 4.0])
+
+  def testReenablingWithSameDumpRootKeepsDataFromBothRuns(self):
+
+    @def_function.function
+    def times_two(x):
+      return x * 2.0
+
+    @def_function.function
+    def plus_one(x):
+      return x + 1.0
+
+    x = constant_op.constant([1.0, 2.0])
+    for func in (times_two, plus_one):
+      writer = dumping_callback.enable_dump_debug_info(
+          self.dump_root, tensor_debug_mode="FULL_HEALTH"
+      )
+      self.evaluate(func(x))
+      writer.FlushNonExecutionFiles()
+      writer.FlushExecutionFiles()
+      dumping_callback.disable_dump_debug_info()
+
+    with debug_events_reader.DebugDataReader(self.dump_root) as reader:
+      reader.update()
+      executed_op_types = [
+          trace.op_type
+          for trace in reader.graph_execution_traces()
+          if trace.op_type not in ["Const", "Placeholder"]
+      ]
+      self.assertCountEqual(executed_op_types, ["Mul", "AddV2"])
+
   @parameterized.named_parameters(
       ("NoTensor", "NO_TENSOR"),
       ("CurtHealth", "CURT_HEALTH"),