From a291f6a904d7a554937413c812c8bcd60accb164 Mon Sep 17 00:00:00 2001 From: Subham Soni Date: Thu, 1 Jan 2026 23:57:19 -0800 Subject: [PATCH] Add named arguments to the pywrap_profiler_plugin functions. PiperOrigin-RevId: 851212029 --- tsl/profiler/lib/profiler_session.cc | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tsl/profiler/lib/profiler_session.cc b/tsl/profiler/lib/profiler_session.cc index ebf2c81f2..64beb4569 100644 --- a/tsl/profiler/lib/profiler_session.cc +++ b/tsl/profiler/lib/profiler_session.cc @@ -65,24 +65,36 @@ absl::Status ProfilerSession::Status() { absl::Status ProfilerSession::CollectDataInternal(XSpace* space) { absl::MutexLock l(mutex_); TF_RETURN_IF_ERROR(status_); - LOG(INFO) << "Profiler session collecting data."; + LOG(INFO) << "[SERVER] ProfilerSession::CollectDataInternal started."; if (profilers_ != nullptr) { + LOG(INFO) << "[SERVER] Stopping profiler collection."; profilers_->Stop().IgnoreError(); stop_time_ns_ = profiler::GetCurrentTimeNanos(); + LOG(INFO) << "[SERVER] Collecting data from profilers."; profilers_->CollectData(space).IgnoreError(); profilers_.reset(); // data has been collected. + LOG(INFO) << "[SERVER] Finished collecting data from profilers."; + } else { + LOG(INFO) << "[SERVER] No active profilers to collect data from."; } // Allow another session to start. profiler_lock_.ReleaseIfActive(); + LOG(INFO) << "[SERVER] ProfilerSession::CollectDataInternal finished."; return absl::OkStatus(); } #endif absl::Status ProfilerSession::CollectData(XSpace* space) { #if !defined(IS_MOBILE_PLATFORM) - space->add_hostnames(port::Hostname()); + std::string hostname = port::Hostname(); + LOG(INFO) << "[SERVER] ProfilerSession::CollectData started for host: " + << hostname; + space->add_hostnames(hostname); TF_RETURN_IF_ERROR(CollectDataInternal(space)); + LOG(INFO) << "[SERVER] Post-processing XSpace data for host: " << hostname; profiler::PostProcessSingleHostXSpace(space, start_time_ns_, stop_time_ns_); + LOG(INFO) << "[SERVER] ProfilerSession::CollectData finished for host: " + << hostname; #endif return absl::OkStatus(); }