From abffd2ddbba945adcd0200cb14f36b2d0045d076 Mon Sep 17 00:00:00 2001 From: Sannidhya Chauhan Date: Fri, 28 Aug 2026 04:34:17 -0700 Subject: [PATCH] Allow ProfilerController to consume data in the kCollectData state. PiperOrigin-RevId: 972507321 --- tsl/profiler/lib/profiler_collection.cc | 2 +- tsl/profiler/lib/profiler_controller.cc | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tsl/profiler/lib/profiler_collection.cc b/tsl/profiler/lib/profiler_collection.cc index 41f1cc1a2..490c96fff 100644 --- a/tsl/profiler/lib/profiler_collection.cc +++ b/tsl/profiler/lib/profiler_collection.cc @@ -70,7 +70,7 @@ absl::StatusOr ProfilerCollection::Consume() { data_vector.push_back(std::move(result->data)); total_estimated_size_bytes += result->estimated_size_bytes; } else if (absl::IsUnimplemented(result.status())) { - LOG(WARNING) << "Profiler consume not implemented: " << result.status(); + VLOG(1) << "Profiler consume not implemented: " << result.status(); data_vector.push_back(std::any()); } else { LOG(ERROR) << "Profiler consume failed: " << result.status(); diff --git a/tsl/profiler/lib/profiler_controller.cc b/tsl/profiler/lib/profiler_controller.cc index 2bf124318..fd311ccd7 100644 --- a/tsl/profiler/lib/profiler_controller.cc +++ b/tsl/profiler/lib/profiler_controller.cc @@ -88,7 +88,7 @@ absl::Status ProfilerController::CollectData( } absl::StatusOr ProfilerController::Consume() { - if (state_ != ProfilerState::kStart && state_ != ProfilerState::kStop) { + if (state_ == ProfilerState::kInit) { return absl::AbortedError("Consume called in the wrong order."); } if (!status_.ok()) {