diff --git a/tsl/profiler/lib/BUILD b/tsl/profiler/lib/BUILD index 9b6787baf..3ae0862d4 100644 --- a/tsl/profiler/lib/BUILD +++ b/tsl/profiler/lib/BUILD @@ -193,6 +193,7 @@ tsl_cc_test( "//tsl/profiler/protobuf:xplane_proto_cc", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/synchronization", "@com_google_absl//absl/time", "@com_google_googletest//:gtest_main", "@xla//xla/tsl/platform:test", diff --git a/tsl/profiler/lib/continuous_profiler_orchestrator.h b/tsl/profiler/lib/continuous_profiler_orchestrator.h index f00f6597a..9573273a1 100644 --- a/tsl/profiler/lib/continuous_profiler_orchestrator.h +++ b/tsl/profiler/lib/continuous_profiler_orchestrator.h @@ -101,6 +101,23 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { return status; } + std::vector SerializeChunks() { + std::vector chunks = PopBuffer(); + std::vector spaces; + spaces.reserve(chunks.size()); + tensorflow::profiler::XSpace space; + for (auto& chunk : chunks) { + space.Clear(); + absl::Status status = profiler_->Serialize(std::move(chunk), &space); + if (status.ok()) { + spaces.push_back(std::move(space)); + } else { + LOG(ERROR) << "Failed to serialize profiler chunk: " << status; + } + } + return spaces; + } + // Returns the current polling interval (primarily for testing). absl::Duration polling_interval() const { absl::MutexLock lock(mutex_); diff --git a/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc b/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc index 567ab3da7..c2aee2c06 100644 --- a/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc +++ b/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc @@ -18,10 +18,11 @@ limitations under the License. #include #include #include +#include #include "absl/status/status.h" #include "absl/status/statusor.h" -#include "absl/time/clock.h" +#include "absl/synchronization/notification.h" #include "absl/time/time.h" #include "xla/tsl/platform/test.h" #include "tsl/profiler/lib/profiler_interface.h" @@ -31,7 +32,6 @@ namespace tsl { namespace profiler { namespace { -using ::testing::Invoke; using ::testing::Pointee; using ::testing::Return; @@ -67,15 +67,19 @@ TEST(ContinuousProfilerOrchestratorTest, // Setup Consume mock to return sequential chunks with high sizes to shrink // the interval + absl::Notification consumed_enough; std::atomic consume_count(0); EXPECT_CALL(*mock, Consume()) - .WillRepeatedly(Invoke([&]() -> absl::StatusOr { + .WillRepeatedly([&]() -> absl::StatusOr { int count = ++consume_count; + if (count >= 4 && !consumed_enough.HasBeenNotified()) { + consumed_enough.Notify(); + } return ConsumeResult{ .data = std::any(count), .estimated_size_bytes = 1000 * 1024 * 1024 // 1000MB (>512MB) }; - })); + }); ContinuousProfilerOrchestrator orchestrator( std::move(mock_profiler)); @@ -83,13 +87,8 @@ TEST(ContinuousProfilerOrchestratorTest, // Start orchestrator (spawns background loop) ASSERT_OK(orchestrator.Start()); - // Wait until we have consumed at least 4 chunks. - // Due to high watermark scaling, interval shrinks from 1s -> 500ms -> 250ms - // -> 125ms -> 100ms. This will happen in less than 1.0 second of real-time - // sleep. - while (consume_count < 4) { - absl::SleepFor(absl::Milliseconds(50)); - } + // Wait until we have consumed at least 4 chunks using notification. + consumed_enough.WaitForNotification(); // Stop orchestrator (terminates background loop) ASSERT_OK(orchestrator.Stop()); @@ -122,13 +121,17 @@ TEST(ContinuousProfilerOrchestratorTest, DynamicIntervalLowWatermarkScaling) { EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); // Consume returns a very small chunk (1MB < 5MB low watermark) + absl::Notification first_consumed; EXPECT_CALL(*mock, Consume()) - .WillRepeatedly(Invoke([]() -> absl::StatusOr { + .WillRepeatedly([&]() -> absl::StatusOr { + if (!first_consumed.HasBeenNotified()) { + first_consumed.Notify(); + } return ConsumeResult{ .data = std::any(1), .estimated_size_bytes = 1 * 1024 * 1024 // 1MB (<5MB) }; - })); + }); ContinuousProfilerOrchestrator orchestrator( std::move(mock_profiler)); @@ -137,9 +140,8 @@ TEST(ContinuousProfilerOrchestratorTest, DynamicIntervalLowWatermarkScaling) { // Start ASSERT_OK(orchestrator.Start()); - // Wait a very short time for the first immediate consume to run and adjust - // the interval - absl::SleepFor(absl::Milliseconds(50)); + // Wait for the first immediate consume to run and adjust the interval + first_consumed.WaitForNotification(); // Stop immediately ASSERT_OK(orchestrator.Stop()); @@ -148,6 +150,42 @@ TEST(ContinuousProfilerOrchestratorTest, DynamicIntervalLowWatermarkScaling) { EXPECT_EQ(orchestrator.polling_interval(), absl::Seconds(2)); } +TEST(ContinuousProfilerOrchestratorTest, SerializeChunks) { + 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 consumed_enough; + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + int count = ++consume_count; + if (count >= 2 && !consumed_enough.HasBeenNotified()) { + consumed_enough.Notify(); + } + return ConsumeResult{.data = std::any(count), + .estimated_size_bytes = 10 * 1024 * 1024}; + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler)); + + ASSERT_OK(orchestrator.Start()); + consumed_enough.WaitForNotification(); + ASSERT_OK(orchestrator.Stop()); + + EXPECT_CALL(*mock, SerializeMock(::testing::_, ::testing::_)) + .WillRepeatedly(Return(absl::InternalError("serialization error"))); + EXPECT_CALL(*mock, SerializeMock(Pointee(AnyEqInt(1)), ::testing::_)) + .WillOnce(Return(absl::OkStatus())); + + std::vector spaces = + orchestrator.SerializeChunks(); + EXPECT_EQ(spaces.size(), 1); +} + } // namespace } // namespace profiler } // namespace tsl diff --git a/tsl/profiler/lib/profiler_controller.cc b/tsl/profiler/lib/profiler_controller.cc index d9c58717c..58ee8d738 100644 --- a/tsl/profiler/lib/profiler_controller.cc +++ b/tsl/profiler/lib/profiler_controller.cc @@ -14,9 +14,11 @@ limitations under the License. ==============================================================================*/ #include "tsl/profiler/lib/profiler_controller.h" +#include #include #include +#include "absl/status/status.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/logging.h" #include "tsl/profiler/lib/profiler_interface.h" @@ -85,5 +87,20 @@ absl::Status ProfilerController::CollectData( return status; } +absl::StatusOr ProfilerController::Consume() { + if (state_ != ProfilerState::kStart) { + return absl::AbortedError("Consume called in the wrong order."); + } + if (!status_.ok()) { + return absl::AbortedError("Previous call returned an error."); + } + return profiler_->Consume(); +} + +absl::Status ProfilerController::Serialize( + std::any data, tensorflow::profiler::XSpace* space) { + return profiler_->Serialize(std::move(data), space); +} + } // namespace profiler } // namespace tsl diff --git a/tsl/profiler/lib/profiler_controller.h b/tsl/profiler/lib/profiler_controller.h index cc0334e9d..28dbd3e9f 100644 --- a/tsl/profiler/lib/profiler_controller.h +++ b/tsl/profiler/lib/profiler_controller.h @@ -15,6 +15,7 @@ limitations under the License. #ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_CONTROLLER_H_ #define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_CONTROLLER_H_ +#include #include #include "absl/status/status.h" @@ -45,6 +46,11 @@ class ProfilerController : public ProfilerInterface { absl::Status CollectData(tensorflow::profiler::XSpace* space) override; + absl::StatusOr Consume() override; + + absl::Status Serialize(std::any data, + tensorflow::profiler::XSpace* space) override; + private: enum class ProfilerState { kInit = 0, diff --git a/tsl/profiler/lib/profiler_session.cc b/tsl/profiler/lib/profiler_session.cc index 2f5352f47..7eeb6853d 100644 --- a/tsl/profiler/lib/profiler_session.cc +++ b/tsl/profiler/lib/profiler_session.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include "absl/memory/memory.h" #include "absl/status/status.h" @@ -77,6 +78,53 @@ absl::Status ProfilerSession::Status() { return status_; } +bool ProfilerSession::IsContinuousProfilingEnabled() const { +#if defined(IS_MOBILE_PLATFORM) + return false; +#else + const auto& advanced_config = options_.advanced_configuration(); + auto it = advanced_config.find("enable_continuous_profiling"); + return (it != advanced_config.end()) && it->second.bool_value(); +#endif +} + +absl::Status ProfilerSession::Stop() { +#if !defined(IS_MOBILE_PLATFORM) + absl::MutexLock l(mutex_); + if (profilers_ != nullptr) { + auto status = profilers_->Stop(); + stop_time_ns_ = profiler::GetCurrentTimeNanos(); + return status; + } +#endif + return absl::OkStatus(); +} + +std::vector ProfilerSession::SerializeChunks() { + std::vector spaces; +#if !defined(IS_MOBILE_PLATFORM) + absl::MutexLock l(mutex_); + if (profilers_ != nullptr) { + auto* orchestrator = dynamic_cast< + profiler::ContinuousProfilerOrchestrator*>( + profilers_.get()); + if (orchestrator != nullptr) { + spaces = orchestrator->SerializeChunks(); + } + for (auto& space : spaces) { + profiler::SetXSpacePidIfNotSet(space, + tsl::Env::Default()->GetProcessId()); + profiler::PostProcessSingleHostXSpace(&space, start_time_ns_, + stop_time_ns_); + SetProfileOptionsIntoSpace(options_, &space); + } + profilers_.reset(); + } + profiler_lock_.ReleaseIfActive(); +#endif + return spaces; +} + #if !defined(IS_MOBILE_PLATFORM) absl::Status ProfilerSession::CollectDataInternal(XSpace* space) { absl::MutexLock l(mutex_); diff --git a/tsl/profiler/lib/profiler_session.h b/tsl/profiler/lib/profiler_session.h index 6644017f5..40a39161c 100644 --- a/tsl/profiler/lib/profiler_session.h +++ b/tsl/profiler/lib/profiler_session.h @@ -15,15 +15,12 @@ limitations under the License. #ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_SESSION_H_ #define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_SESSION_H_ -#include +#include #include #include #include "absl/status/status.h" #include "absl/synchronization/mutex.h" -#include "xla/tsl/platform/status.h" -#include "xla/tsl/platform/types.h" -#include "tsl/platform/platform.h" #include "tsl/platform/thread_annotations.h" #include "tsl/profiler/protobuf/profiler_options.pb.h" #include "tsl/profiler/protobuf/xplane.pb.h" @@ -67,6 +64,13 @@ class ProfilerSession { absl::Status CollectData(tensorflow::profiler::XSpace* space) TF_LOCKS_EXCLUDED(mutex_); + absl::Status Stop() TF_LOCKS_EXCLUDED(mutex_); + + std::vector SerializeChunks() + TF_LOCKS_EXCLUDED(mutex_); + + bool IsContinuousProfilingEnabled() const; + private: // Constructs an instance of the class and starts profiling explicit ProfilerSession(const tensorflow::ProfileOptions& options);