diff --git a/libkineto/src/AsyncActivityProfilerHandler.cpp b/libkineto/src/AsyncActivityProfilerHandler.cpp index f5af85154..01b614442 100644 --- a/libkineto/src/AsyncActivityProfilerHandler.cpp +++ b/libkineto/src/AsyncActivityProfilerHandler.cpp @@ -266,6 +266,7 @@ void AsyncActivityProfilerHandler::activateConfig( std::chrono::time_point now) { logger_ = ActivityProfilerController::makeLogger(*asyncRequestConfig_); profiler_.setLogger(logger_.get()); + Logger::resetLoggerObservers(); LOGGER_OBSERVER_SET_TRIGGER_ON_DEMAND(); profiler_.configure(*asyncRequestConfig_, now); asyncRequestConfig_ = nullptr; diff --git a/libkineto/src/GenericActivityProfiler.cpp b/libkineto/src/GenericActivityProfiler.cpp index 896d96440..063dcbfe0 100644 --- a/libkineto/src/GenericActivityProfiler.cpp +++ b/libkineto/src/GenericActivityProfiler.cpp @@ -143,6 +143,7 @@ const std::unordered_set& getLoggerMedataAllowList() { } // namespace void GenericActivityProfiler::processTraceInternal(ActivityLogger& logger) { + UST_LOGGER_STAGE_SCOPE(kPostProcessingStage); LOG(INFO) << "Processing " << traceBuffers_->cpu.size() << " CPU buffers"; VLOG(0) << "Profile time range: " << captureWindowStartTime_ << " - " << captureWindowEndTime_; @@ -974,7 +975,6 @@ time_point GenericActivityProfiler::performRunLoopStep( // for quickly handling trace request via synchronous API std::lock_guard guard(mutex_); processTraceInternal(*logger_); - UST_LOGGER_MARK_COMPLETED(kPostProcessingStage); resetInternal(); VLOG(0) << "ProcessTrace -> WaitForRequest"; break; diff --git a/libkineto/src/Logger.cpp b/libkineto/src/Logger.cpp index a81914920..7e2438cdb 100644 --- a/libkineto/src/Logger.cpp +++ b/libkineto/src/Logger.cpp @@ -157,6 +157,13 @@ void Logger::setLoggerObserverOnDemand() { } } +void Logger::resetLoggerObservers() { + std::lock_guard guard(loggerObserversMutex()); + for (auto observer : loggerObservers()) { + observer->reset(); + } +} + void Logger::addLoggerObserverAddMetadata( const std::string& key, const std::string& value) { @@ -166,6 +173,13 @@ void Logger::addLoggerObserverAddMetadata( } } +USTLoggerStageGuard::~USTLoggerStageGuard() { + try { + UST_LOGGER_MARK_COMPLETED(stage_); + } catch (...) { + } +} + } // namespace KINETO_NAMESPACE #endif // USE_GOOGLE_LOG diff --git a/libkineto/src/Logger.h b/libkineto/src/Logger.h index 1f7b88d80..b4b9be4af 100644 --- a/libkineto/src/Logger.h +++ b/libkineto/src/Logger.h @@ -27,6 +27,7 @@ #define LOGGER_OBSERVER_SET_TRIGGER_ON_DEMAND() #define LOGGER_OBSERVER_ADD_METADATA(key, value) #define UST_LOGGER_MARK_COMPLETED(stage) +#define UST_LOGGER_STAGE_SCOPE(stage) #define USDT_LOGGER_EMIT_MESSAGE(usdt_type) #define USDT_EMIT_START_TRACE() #define USDT_EMIT_STOP_TRACE() @@ -129,6 +130,8 @@ class Logger { static void setLoggerObserverOnDemand(); + static void resetLoggerObservers(); + static void addLoggerObserverAddMetadata(const std::string& key, const std::string& value); private: @@ -155,6 +158,23 @@ class VoidLogger { void operator&(std::ostream&) {} }; +// RAII helper used to ensure a UST stage row is emitted on every exit from a scope. +// Bucketed LOG(ERROR) / LOG(WARNING) calls within the scope are picked up by the +// emitted row as usual. +class USTLoggerStageGuard { + public: + explicit USTLoggerStageGuard(const std::string& stage) : stage_(stage) {} + ~USTLoggerStageGuard(); + + USTLoggerStageGuard(const USTLoggerStageGuard&) = delete; + USTLoggerStageGuard& operator=(const USTLoggerStageGuard&) = delete; + USTLoggerStageGuard(USTLoggerStageGuard&&) = delete; + USTLoggerStageGuard& operator=(USTLoggerStageGuard&&) = delete; + + private: + std::string stage_; +}; + } // namespace KINETO_NAMESPACE #ifdef LOG // Undefine in case these are already defined (quite likely) @@ -254,6 +274,9 @@ struct __to_constant__ { // UST Logger Semantics to describe when a stage is complete. #define UST_LOGGER_MARK_COMPLETED(stage) LOG(libkineto::LoggerOutputType::STAGE) << "Completed Stage: " << stage +// RAII helper that fires UST_LOGGER_MARK_COMPLETED(stage) on scope exit. +#define UST_LOGGER_STAGE_SCOPE(stage) libkineto::USTLoggerStageGuard LOCAL_VARNAME(ust_stage_guard)(stage) + #define USDT_LOGGER_EMIT_MESSAGE(usdt_type) LOG(libkineto::LoggerOutputType::USDT) << usdt_type #define USDT_EMIT_START_TRACE() USDT_LOGGER_EMIT_MESSAGE("profiler_start") #define USDT_EMIT_STOP_TRACE() USDT_LOGGER_EMIT_MESSAGE("profiler_stop") diff --git a/libkineto/src/LoggerCollector.h b/libkineto/src/LoggerCollector.h index e50bf64b3..102647a0f 100644 --- a/libkineto/src/LoggerCollector.h +++ b/libkineto/src/LoggerCollector.h @@ -41,6 +41,8 @@ class LoggerCollector : public ILoggerObserver { } void reset() override { + buckets_.clear(); + devices.clear(); trace_duration_ms = 0; event_count = 0; destinations.clear(); diff --git a/libkineto/src/SyncActivityProfilerHandler.cpp b/libkineto/src/SyncActivityProfilerHandler.cpp index ca27dc9aa..e21d40a3b 100644 --- a/libkineto/src/SyncActivityProfilerHandler.cpp +++ b/libkineto/src/SyncActivityProfilerHandler.cpp @@ -44,6 +44,7 @@ void SyncActivityProfilerHandler::prepareTrace(const Config& config) { profiler_.reset(); } + Logger::resetLoggerObservers(); profiler_.configure(config, now); } @@ -62,7 +63,6 @@ std::unique_ptr SyncActivityProfilerHandler:: profiler_.processTrace(*logger); // Will follow up with another patch for logging URLs when ActivityTrace // is moved. - UST_LOGGER_MARK_COMPLETED(kPostProcessingStage); // Logger Metadata contains a map of LOGs collected in Kineto // logger_level -> List of log lines