Merge branch 'main' into zejun/kineto_for_pti_0.11

Author: zejun-chen

libkineto/CMakeLists.txt

@@ -172,7 +172,7 @@ set_target_properties(kineto_base kineto_api PROPERTIES
 set(KINETO_COMPILE_OPTIONS "-DKINETO_NAMESPACE=libkineto")
 list(APPEND KINETO_COMPILE_OPTIONS "-DFMT_HEADER_ONLY")
 list(APPEND KINETO_COMPILE_OPTIONS "-DENABLE_IPC_FABRIC")
-list(APPEND KINETO_COMPILE_OPTIONS "-DTMP_IMPL_MEMORY_PROFILING_ON_DEMAND")
+list(APPEND KINETO_COMPILE_OPTIONS)
 if(NOT MSVC)
   list(APPEND KINETO_COMPILE_OPTIONS "-std=c++17")
 else()

libkineto/include/ClientInterface.h

@@ -18,11 +18,9 @@ class ClientInterface {
   virtual void prepare(bool, bool, bool, bool, bool) = 0;
   virtual void start() = 0;
   virtual void stop() = 0;
-#ifdef TMP_IMPL_MEMORY_PROFILING_ON_DEMAND
   virtual void start_memory_profile() = 0;
   virtual void stop_memory_profile() = 0;
   virtual void export_memory_profile(const std::string&) = 0;
-#endif
 };
 
 } // namespace libkineto

libkineto/include/output_base.h

@@ -63,6 +63,8 @@ class ActivityLogger {
     handleTraceStart(std::unordered_map<std::string, std::string>(), "");
   }
 
+  virtual void finalizeMemoryTrace(const std::string&, const Config&) = 0;
+
   virtual void finalizeTrace(
       const Config& config,
       std::unique_ptr<ActivityBuffers> buffers,

libkineto/src/ActivityProfilerController.cpp

@@ -245,13 +245,33 @@ void ActivityProfilerController::profilerLoop() {
 }
 
 void ActivityProfilerController::memoryProfilerLoop() {
-  const std::string& path = asyncRequestConfig_->activitiesLogFile();
+  std::string path = asyncRequestConfig_->activitiesLogFile();
   auto profile_time = asyncRequestConfig_->profileMemoryDuration();
-  profiler_->performMemoryLoop(path, profile_time);
+  std::unique_ptr<Config> config = asyncRequestConfig_->clone();
+  while (!stopRunloop_) {
+    // Perform Double-checked locking to reduce overhead of taking lock.
+    if (asyncRequestConfig_ && !profiler_->isActive()) {
+      std::lock_guard<std::mutex> lock(asyncConfigLock_);
+      if (asyncRequestConfig_ && !profiler_->isActive()) {
+        logger_ = makeLogger(*asyncRequestConfig_);
+        path = asyncRequestConfig_->activitiesLogFile();
+        profile_time = asyncRequestConfig_->profileMemoryDuration();
+        config = asyncRequestConfig_->clone();
+        asyncRequestConfig_ = nullptr;
+      } else {
+        continue;
+      }
+    } else {
+      continue;
+    }
+
+    profiler_->performMemoryLoop(path, profile_time, logger_.get(), *config);
+  }
 }
 
 void ActivityProfilerController::step() {
-  // Do not remove this copy to currentIter. Otherwise count is not guaranteed.
+  // Do not remove this copy to currentIter. Otherwise count is not
+  // guaranteed.
   int64_t currentIter = ++iterationCount_;
   VLOG(0) << "Step called , iteration  = " << currentIter;
 
@@ -264,7 +284,6 @@ void ActivityProfilerController::step() {
       activateConfig(now);
     }
   }
-
   if (profiler_->isActive()) {
     auto now = system_clock::now();
     auto next_wakeup_time = now + Config::kControllerIntervalMsecs;

libkineto/src/Config.cpp

@@ -399,8 +399,9 @@ bool Config::handleOption(const std::string& name, std::string& val) {
     perThreadBufferEnabled_ = toBool(val);
   } else if (!name.compare(kProfileMemory)) {
     memoryProfilerEnabled_ = toBool(val);
-    activitiesLogFile_ = defaultMemoryTraceFileName();
-    activitiesLogUrl_ = fmt::format("file://{}", activitiesLogFile_);
+    if (memoryProfilerEnabled_) {
+      activitiesLogFile_ = defaultMemoryTraceFileName();
+    }
   } else if (!name.compare(kProfileMemoryDuration)) {
     profileMemoryDuration_ = toInt32(val);
   } else if (!name.compare(kActivitiesLogFileKey)) {

libkineto/src/CuptiActivityProfiler.cpp

@@ -1421,16 +1421,21 @@ const time_point<system_clock> CuptiActivityProfiler::performRunLoopStep(
 
 const void CuptiActivityProfiler::performMemoryLoop(
     const string& path,
-    uint32_t profile_time) {
-#ifdef TMP_IMPL_MEMORY_PROFILING_ON_DEMAND
+    uint32_t profile_time,
+    ActivityLogger* logger,
+    Config& config) {
+  currentRunloopState_ = RunloopState::CollectTrace;
   if (libkineto::api().client()) {
     libkineto::api().client()->start_memory_profile();
     LOG(INFO) << "Running memory profiling for " << profile_time << " ms";
     std::this_thread::sleep_for(std::chrono::milliseconds(profile_time));
+    LOG(INFO) << "Exporting memory profiling results to " << path;
     libkineto::api().client()->export_memory_profile(path);
     libkineto::api().client()->stop_memory_profile();
+    LOG(INFO) << "Finalizing trace";
+    logger->finalizeMemoryTrace(path, config);
   }
-#endif
+  currentRunloopState_ = RunloopState::WaitForRequest;
 }
 
 void CuptiActivityProfiler::finalizeTrace(

libkineto/src/CuptiActivityProfiler.h

@@ -140,7 +140,11 @@ class CuptiActivityProfiler {
       const std::chrono::time_point<std::chrono::system_clock>& nextWakeupTime,
       int64_t currentIter = -1);
 
-  const void performMemoryLoop(const std::string& path, uint32_t profile_time);
+  const void performMemoryLoop(
+      const std::string& path,
+      uint32_t profile_time,
+      ActivityLogger* logger,
+      Config& config);
 
   // Used for async requests
   void setLogger(ActivityLogger* logger) {

libkineto/src/DeviceProperties.cpp

@@ -70,7 +70,7 @@ static const std::vector<gpuDeviceProp>& deviceProps() {
 static const std::string createDevicePropertiesJson(
     size_t id,
     const gpuDeviceProp& props) {
-  std::string gpuSpecific = "";
+  std::string gpuSpecific;
 #if defined(HAS_CUPTI)
   gpuSpecific = fmt::format(
       R"JSON(

libkineto/src/output_json.cpp

@@ -160,6 +160,10 @@ void ChromeTraceLogger::openTraceFile() {
   }
 }
 
+void ChromeTraceLogger::finalizeMemoryTrace(const std::string&, const Config&) {
+  LOG(INFO) << "finalizeMemoryTrace not implemented for ChromeTraceLogger";
+}
+
 ChromeTraceLogger::ChromeTraceLogger(const std::string& traceFileName) {
   fileName_ = traceFileName.empty() ? defaultFileName() : traceFileName;
   traceOf_.clear(std::ios_base::badbit);
@@ -399,7 +403,7 @@ void ChromeTraceLogger::handleActivity(const libkineto::ITraceActivity& op) {
       external_id = op.correlationId();
     }
   }
-  std::string arg_values = "";
+  std::string arg_values;
   if (external_id != 0) {
     arg_values.append(fmt::format("\"External id\": {}", external_id));
   }
@@ -526,7 +530,7 @@ void ChromeTraceLogger::handleActivity(const libkineto::ITraceActivity& op) {
     pgMap.insert({processGroupName, pg_config});
   }
 
-  std::string args = "";
+  std::string args;
   if (!arg_values.empty()) {
     args = fmt::format(
         R"JSON(,

libkineto/src/output_json.h

@@ -35,21 +35,21 @@ class Config;
 
 struct pgConfig {
   pgConfig() = default;
-  std::string pg_name{""};
-  std::string pg_desc{""};
-  std::string backend_config{""};
-  std::string pg_size{""};
-  std::string ranks{""};
+  std::string pg_name;
+  std::string pg_desc;
+  std::string backend_config;
+  std::string pg_size;
+  std::string ranks;
 };
 
 struct DistributedInfo {
   DistributedInfo() = default;
 
-  std::string backend{""};
-  std::string rank{""};
-  std::string world_size{""};
-  std::string pg_count{""};
-  std::string nccl_version{""};
+  std::string backend;
+  std::string rank;
+  std::string world_size;
+  std::string pg_count;
+  std::string nccl_version;
   bool distInfo_present_{false};
 };
 
@@ -81,6 +81,8 @@ class ChromeTraceLogger : public libkineto::ActivityLogger {
       std::unordered_map<std::string, std::vector<std::string>>& metadata)
       override;
 
+  void finalizeMemoryTrace(const std::string&, const Config&) override;
+
   std::string traceFileName() const {
     return fileName_;
   }