diff --git a/libkineto/include/ActivityProfilerInterface.h b/libkineto/include/ActivityProfilerInterface.h index 9098c09bf..2279722e6 100644 --- a/libkineto/include/ActivityProfilerInterface.h +++ b/libkineto/include/ActivityProfilerInterface.h @@ -81,6 +81,10 @@ class ActivityProfilerInterface { // Maps kernel thread id -> pthread id for CPU ops. // Client must record any new kernel thread where the activity has occured. virtual void recordThreadInfo(pid_t tid, pthread_t pthreadId) {} + + // Record trace metadata, currently supporting only string key and values, + // values with the same key are overwritten + virtual void addMetadata(const std::string& key, const std::string& value) = 0; }; } // namespace libkineto diff --git a/libkineto/src/ActivityProfiler.cpp b/libkineto/src/ActivityProfiler.cpp index 96e39a43d..7385180ae 100644 --- a/libkineto/src/ActivityProfiler.cpp +++ b/libkineto/src/ActivityProfiler.cpp @@ -127,6 +127,7 @@ void ActivityProfiler::processTraceInternal(ActivityLogger& logger) { << " CPU buffers"; VLOG(0) << "Profile time range: " << captureWindowStartTime_ << " - " << captureWindowEndTime_; + logger.handleTraceStart(metadata_); for (auto& cpu_trace : traceBuffers_->cpu) { string trace_name = cpu_trace->span.name; VLOG(0) << "Processing CPU buffer for " << trace_name << " (" @@ -696,6 +697,7 @@ void ActivityProfiler::finalizeTrace(const Config& config, ActivityLogger& logge } } } + // Thread info for (auto pair : threadInfo_) { const auto& thread_info = pair.second; @@ -728,6 +730,7 @@ void ActivityProfiler::resetTraceData() { clientActivityTraceMap_.clear(); disabledTraceSpans_.clear(); traceBuffers_ = nullptr; + metadata_.clear(); } diff --git a/libkineto/src/ActivityProfiler.h b/libkineto/src/ActivityProfiler.h index 094faef4c..1e495d4de 100644 --- a/libkineto/src/ActivityProfiler.h +++ b/libkineto/src/ActivityProfiler.h @@ -111,6 +111,11 @@ class ActivityProfiler { } } + void addMetadata(const std::string& key, const std::string& value) { + std::lock_guard guard(mutex_); + metadata_[key] = value; + } + private: class ExternalEventMap { public: @@ -366,6 +371,9 @@ class ActivityProfiler { // Buffers where trace data is stored std::unique_ptr traceBuffers_; + // Trace metadata + std::unordered_map metadata_; + }; } // namespace KINETO_NAMESPACE diff --git a/libkineto/src/ActivityProfilerController.cpp b/libkineto/src/ActivityProfilerController.cpp index e85421442..1e3d09a2a 100644 --- a/libkineto/src/ActivityProfilerController.cpp +++ b/libkineto/src/ActivityProfilerController.cpp @@ -147,4 +147,9 @@ std::unique_ptr ActivityProfilerController::stopTrace() return std::make_unique(std::move(logger), CuptiActivityInterface::singleton()); } +void ActivityProfilerController::addMetadata( + const std::string& key, const std::string& value) { + profiler_->addMetadata(key, value); +} + } // namespace KINETO_NAMESPACE diff --git a/libkineto/src/ActivityProfilerController.h b/libkineto/src/ActivityProfilerController.h index c14e61466..6e48449a3 100644 --- a/libkineto/src/ActivityProfilerController.h +++ b/libkineto/src/ActivityProfilerController.h @@ -58,9 +58,11 @@ class ActivityProfilerController { } void recordThreadInfo(pid_t tid, pthread_t pthreadId) { - profiler_->recordThreadInfo(tid, pthreadId); + profiler_->recordThreadInfo(tid, pthreadId); } + void addMetadata(const std::string& key, const std::string& value); + private: void profilerLoop(); diff --git a/libkineto/src/ActivityProfilerProxy.cpp b/libkineto/src/ActivityProfilerProxy.cpp index 304b225fd..f9f268651 100644 --- a/libkineto/src/ActivityProfilerProxy.cpp +++ b/libkineto/src/ActivityProfilerProxy.cpp @@ -85,7 +85,12 @@ bool ActivityProfilerProxy::enableForRegion(const std::string& match) { } void ActivityProfilerProxy::recordThreadInfo(pid_t tid, pthread_t pthreadId) { - controller_->recordThreadInfo(tid, pthreadId); + controller_->recordThreadInfo(tid, pthreadId); +} + +void ActivityProfilerProxy::addMetadata( + const std::string& key, const std::string& value) { + controller_->addMetadata(key, value); } } // namespace libkineto diff --git a/libkineto/src/ActivityProfilerProxy.h b/libkineto/src/ActivityProfilerProxy.h index c6a6d18f6..17cb698a9 100644 --- a/libkineto/src/ActivityProfilerProxy.h +++ b/libkineto/src/ActivityProfilerProxy.h @@ -60,6 +60,8 @@ class ActivityProfilerProxy : public ActivityProfilerInterface { bool enableForRegion(const std::string& match) override; + void addMetadata(const std::string& key, const std::string& value) override; + private: bool cpuOnly_{true}; ActivityProfilerController* controller_{nullptr}; diff --git a/libkineto/src/output_base.h b/libkineto/src/output_base.h index 14775b5b8..49a323e22 100644 --- a/libkineto/src/output_base.h +++ b/libkineto/src/output_base.h @@ -67,6 +67,13 @@ class ActivityLogger { const GpuActivity& activity) = 0; #endif // HAS_CUPTI + virtual void handleTraceStart( + const std::unordered_map& metadata) = 0; + + void handleTraceStart() { + handleTraceStart(std::unordered_map()); + } + virtual void finalizeTrace( const KINETO_NAMESPACE::Config& config, std::unique_ptr buffers, diff --git a/libkineto/src/output_json.cpp b/libkineto/src/output_json.cpp index 784dba644..46046283c 100644 --- a/libkineto/src/output_json.cpp +++ b/libkineto/src/output_json.cpp @@ -31,28 +31,48 @@ namespace KINETO_NAMESPACE { static constexpr int kSchemaVersion = 1; -static void writeHeader(std::ofstream& stream) { - stream << fmt::format(R"JSON( +void ChromeTraceLogger::handleTraceStart( + const std::unordered_map& metadata) { + traceOf_ << fmt::format(R"JSON( {{ "schemaVersion": {}, - "traceEvents": [ )JSON", kSchemaVersion); + + if (!metadata.empty()) { + traceOf_ << R"JSON( + "metadata": { + )JSON"; + bool first = true; + for (const auto& kv : metadata) { + if (!first) { + traceOf_ << ",\n"; + } + traceOf_ << fmt::format(R"( "{}": "{}")", kv.first, kv.second); + first = false; + } + traceOf_ << R"JSON( + }, + )JSON"; + } + + traceOf_ << R"JSON( + "traceEvents": [ + )JSON"; } -static void openTraceFile(std::string& name, std::ofstream& stream) { - stream.open(name, std::ofstream::out | std::ofstream::trunc); - if (!stream) { - PLOG(ERROR) << "Failed to open '" << name << "'"; +void ChromeTraceLogger::openTraceFile() { + traceOf_.open(fileName_, std::ofstream::out | std::ofstream::trunc); + if (!traceOf_) { + PLOG(ERROR) << "Failed to open '" << fileName_ << "'"; } else { - LOG(INFO) << "Tracing to " << name; - writeHeader(stream); + LOG(INFO) << "Tracing to " << fileName_; } } ChromeTraceLogger::ChromeTraceLogger(const std::string& traceFileName, int smCount) : fileName_(traceFileName), pid_(getpid()) { traceOf_.clear(std::ios_base::badbit); - openTraceFile(fileName_, traceOf_); + openTraceFile(); #ifdef HAS_CUPTI smCount_ = CuptiActivityInterface::singleton().smCount(); #endif diff --git a/libkineto/src/output_json.h b/libkineto/src/output_json.h index 84ac673c9..b7de0e463 100644 --- a/libkineto/src/output_json.h +++ b/libkineto/src/output_json.h @@ -60,6 +60,9 @@ class ChromeTraceLogger : public libkineto::ActivityLogger { void handleGpuActivity(const GpuActivity& activity) override; #endif // HAS_CUPTI + void handleTraceStart( + const std::unordered_map& metadata) override; + void finalizeTrace( const Config& config, std::unique_ptr buffers, @@ -73,6 +76,8 @@ class ChromeTraceLogger : public libkineto::ActivityLogger { void handleLinkEnd(const TraceActivity& e); #endif // HAS_CUPTI + void openTraceFile(); + std::string fileName_; std::ofstream traceOf_; diff --git a/libkineto/src/output_membuf.h b/libkineto/src/output_membuf.h index 3e3c76e64..be29bb9dc 100644 --- a/libkineto/src/output_membuf.h +++ b/libkineto/src/output_membuf.h @@ -87,6 +87,11 @@ class MemoryTraceLogger : public ActivityLogger { } #endif // HAS_CUPTI + void handleTraceStart( + const std::unordered_map& metadata) override { + metadata_ = metadata; + } + void finalizeTrace( const Config& config, std::unique_ptr buffers, @@ -100,6 +105,7 @@ class MemoryTraceLogger : public ActivityLogger { } void log(ActivityLogger& logger) { + logger.handleTraceStart(metadata_); for (auto& activity : activities_) { activity->log(logger); } @@ -151,6 +157,7 @@ class MemoryTraceLogger : public ActivityLogger { std::vector traceSpanList_; std::vector iterationList_; std::unique_ptr buffers_; + std::unordered_map metadata_; int64_t endTime_{0}; };