diff --git a/tsl/profiler/lib/continuous_profiler_orchestrator.h b/tsl/profiler/lib/continuous_profiler_orchestrator.h index 58954be37..a35e1be66 100644 --- a/tsl/profiler/lib/continuous_profiler_orchestrator.h +++ b/tsl/profiler/lib/continuous_profiler_orchestrator.h @@ -28,6 +28,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" +#include "absl/time/clock.h" #include "absl/time/time.h" #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/errors.h" // IWYU pragma: keep @@ -53,12 +54,19 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { static constexpr absl::Duration kDefaultPollingInterval = absl::Seconds(2); static constexpr absl::Duration kMinPollingInterval = absl::Seconds(2); static constexpr absl::Duration kMaxPollingInterval = absl::Seconds(5); + // Upper bound on how long Flush() will block waiting for the ingestion + // thread. Generous relative to kMaxPollingInterval (which bounds how long + // the ingestion thread sleeps) but small enough that a wedged Consume() + // cannot hold an RPC handler hostage. + static constexpr absl::Duration kFlushTimeout = absl::Seconds(30); explicit ContinuousProfilerOrchestrator( std::unique_ptr profiler, - size_t max_buffer_bytes = kDefaultMaxBufferBytes) + size_t max_buffer_bytes = kDefaultMaxBufferBytes, + absl::Duration flush_timeout = kFlushTimeout) : profiler_(std::move(profiler)), max_buffer_bytes_(max_buffer_bytes), + flush_timeout_(flush_timeout), is_running_(false), polling_interval_(kDefaultPollingInterval) {} @@ -132,6 +140,41 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { return spaces; } + // Asks the ingestion thread to run a Consume() cycle now and waits for it to + // finish, so that everything recorded before this call is in the buffer. + // Callers that need the most recent data should call Flush() before + // SerializeChunks(); otherwise SerializeChunks() only returns what the + // ingestion loop has already harvested, which can lag by up to one polling + // interval. + // + // No-op if the orchestrator is not running. The wait is bounded by the + // `flush_timeout` passed to the constructor: a Consume() implementation that + // wedges must not be able to block the caller (typically an RPC handler) + // indefinitely. On timeout this simply returns, and a following + // SerializeChunks() sees whatever was already buffered. + // + // Note this can cost up to two Consume() cycles in the worst case: a + // Consume() that was already in flight when the flush was requested cannot + // be credited against it, because it may have started before the data the + // caller cares about was recorded. + void Flush() { + absl::MutexLock lock(mutex_); + if (!is_running_) return; + uint64_t target_seq = ++flush_requested_seq_; + cv_.SignalAll(); + const absl::Time deadline = absl::Now() + flush_timeout_; + while (is_running_ && flush_completed_seq_ < target_seq) { + if (cv_.WaitWithDeadline(&mutex_, deadline)) { + if (flush_completed_seq_ >= target_seq) break; + LOG_EVERY_N_SEC(WARNING, 30) + << "ContinuousProfilerOrchestrator::Flush timed out after " + << flush_timeout_ + << " waiting for the ingestion thread; returning partial data."; + return; + } + } + } + // Returns the current polling interval (primarily for testing). absl::Duration polling_interval() const { absl::MutexLock lock(mutex_); @@ -225,9 +268,11 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { void IngestionLoop() { LOG(INFO) << "ContinuousProfilerOrchestrator::IngestionLoop started"; while (true) { + uint64_t seq_before_consume = 0; { absl::MutexLock lock(mutex_); if (!is_running_) break; + seq_before_consume = flush_requested_seq_; } absl::StatusOr result = profiler_->Consume(); @@ -238,9 +283,14 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { AdjustIntervalLocked(chunk_size); } + if (seq_before_consume > flush_completed_seq_) { + flush_completed_seq_ = seq_before_consume; + cv_.SignalAll(); + } if (!is_running_) break; - - // Wait using absl::CondVar on absl::Mutex + if (flush_requested_seq_ > flush_completed_seq_) { + continue; // Immediate consume requested while Consume() was in flight. + } cv_.WaitWithTimeout(&mutex_, polling_interval_); if (!is_running_) break; } @@ -272,6 +322,7 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { std::unique_ptr profiler_; const size_t max_buffer_bytes_; + const absl::Duration flush_timeout_; mutable absl::Mutex mutex_; absl::CondVar cv_; @@ -284,6 +335,8 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { size_t total_buffered_bytes_ ABSL_GUARDED_BY(mutex_) = 0; uint64_t dropped_chunks_count_ ABSL_GUARDED_BY(mutex_) = 0; uint64_t dropped_bytes_count_ ABSL_GUARDED_BY(mutex_) = 0; + uint64_t flush_requested_seq_ ABSL_GUARDED_BY(mutex_) = 0; + uint64_t flush_completed_seq_ ABSL_GUARDED_BY(mutex_) = 0; }; } // namespace profiler diff --git a/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc b/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc index 4341e5eeb..d17db04dd 100644 --- a/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc +++ b/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc @@ -17,12 +17,14 @@ limitations under the License. #include #include #include +#include // NOLINT(build/c++11) #include #include #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/synchronization/notification.h" +#include "absl/time/clock.h" #include "absl/time/time.h" #include "xla/tsl/platform/test.h" #include "tsl/profiler/lib/profiler_interface.h" @@ -494,6 +496,168 @@ TEST(ContinuousProfilerOrchestratorTest, StopDrainAccounting) { EXPECT_EQ(std::any_cast(chunks[1]), 2); } +TEST(ContinuousProfilerOrchestratorTest, FlushForcesAnImmediateConsume) { + // The polling interval is >= 2s. Flush() must not wait for it: it asks the + // ingestion thread for a Consume() now. + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + return ConsumeResult{.data = std::any(++consume_count), + .estimated_size_bytes = 10 * 1024 * 1024}; + }); + EXPECT_CALL(*mock, SerializeMock(::testing::_, ::testing::_)) + .WillRepeatedly(Return(absl::OkStatus())); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler)); + + ASSERT_OK(orchestrator.Start()); + + absl::Time start = absl::Now(); + orchestrator.Flush(); + absl::Duration elapsed = absl::Now() - start; + std::vector spaces = + orchestrator.SerializeChunks(); + + EXPECT_FALSE(spaces.empty()); + // Two Consume() cycles is the documented worst case (an in-flight Consume() + // cannot be credited against the flush), so allow for it but not for a full + // polling interval of idling on top. + EXPECT_LT(elapsed, absl::Seconds(5)) + << "Flush() should not have waited out the polling interval"; + EXPECT_GT(consume_count.load(), 0); + + ASSERT_OK(orchestrator.Stop()); +} + +TEST(ContinuousProfilerOrchestratorTest, SerializeChunksDoesNotFlush) { + // SerializeChunks() only drains what is already buffered. With a wedged + // Consume() and the default (30s) flush timeout, it must still return + // right away. + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + absl::Notification release_consume; + absl::Notification consume_entered; + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + if (!consume_entered.HasBeenNotified()) consume_entered.Notify(); + release_consume.WaitForNotification(); + return absl::OutOfRangeError("End of stream"); + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler)); + ASSERT_OK(orchestrator.Start()); + consume_entered.WaitForNotification(); + + absl::Time start = absl::Now(); + std::vector spaces = + orchestrator.SerializeChunks(); + absl::Duration elapsed = absl::Now() - start; + + EXPECT_TRUE(spaces.empty()); + EXPECT_LT(elapsed, absl::Seconds(10)); + + release_consume.Notify(); + ASSERT_OK(orchestrator.Stop()); +} + +TEST(ContinuousProfilerOrchestratorTest, FlushIsBoundedWhenConsumeWedges) { + // A Consume() that never returns must not hold Flush() (and hence the + // calling RPC handler) forever. + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + absl::Notification release_consume; + absl::Notification consume_entered; + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + if (!consume_entered.HasBeenNotified()) consume_entered.Notify(); + release_consume.WaitForNotification(); + return absl::OutOfRangeError("End of stream"); + }); + + constexpr absl::Duration kTestFlushTimeout = absl::Milliseconds(200); + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler), kDefaultMaxBufferBytes, kTestFlushTimeout); + ASSERT_OK(orchestrator.Start()); + consume_entered.WaitForNotification(); + + absl::Time start = absl::Now(); + orchestrator.Flush(); + absl::Duration elapsed = absl::Now() - start; + + EXPECT_GE(elapsed, kTestFlushTimeout - absl::Milliseconds(50)); + EXPECT_LT(elapsed, kTestFlushTimeout + absl::Seconds(5)); + + release_consume.Notify(); + ASSERT_OK(orchestrator.Stop()); +} + +TEST(ContinuousProfilerOrchestratorTest, FlushOnAStoppedOrchestratorReturns) { + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly(Return(absl::OutOfRangeError("End of stream"))); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler)); + + // Never started: Flush() must be a no-op rather than waiting for an ingestion + // thread that does not exist. + absl::Time start = absl::Now(); + orchestrator.Flush(); + EXPECT_LT(absl::Now() - start, absl::Seconds(5)); + EXPECT_TRUE(orchestrator.SerializeChunks().empty()); +} + +TEST(ContinuousProfilerOrchestratorTest, ConcurrentFlushesAllComplete) { + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + return ConsumeResult{.data = std::any(++consume_count), + .estimated_size_bytes = 1024}; + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler)); + ASSERT_OK(orchestrator.Start()); + + // The flush sequence counters are shared; several waiters must all be woken. + std::atomic completed(0); + std::vector threads; + threads.reserve(4); + for (int i = 0; i < 4; ++i) { + threads.emplace_back([&] { + orchestrator.Flush(); + completed.fetch_add(1); + }); + } + for (auto& t : threads) t.join(); + EXPECT_EQ(completed.load(), 4); + + ASSERT_OK(orchestrator.Stop()); +} + } // namespace } // namespace profiler } // namespace tsl