add function ensureCollectTraceDone to wait and cleanup collectTraceThread

Author: staugust

libkineto/src/CuptiActivityProfiler.cpp

@@ -1113,6 +1113,22 @@ void CuptiActivityProfiler::toggleCollectionDynamic(const bool enable){
 #endif
 }
 
+void CuptiActivityProfiler::ensureCollectTraceDone() {
+  if (collectTraceThread && collectTraceThread->joinable()) {
+    std::lock_guard<std::mutex> guard(mutex_);
+    collectTraceThread->join();
+    collectTraceThread.reset(nullptr);
+  }
+#endif
+#ifdef HAS_ROCTRACER
+  if (enable) {
+    cupti_.enableActivities(derivedConfig_->profileActivityTypes());
+  } else {
+    cupti_.disableActivities(derivedConfig_->profileActivityTypes());
+  }
+#endif
+}
+
 void CuptiActivityProfiler::ensureCollectTraceDone() {
   if (collectTraceThread && collectTraceThread->joinable()) {
     std::lock_guard<std::mutex> guard(mutex_);
@@ -1204,6 +1220,8 @@ const time_point<system_clock> CuptiActivityProfiler::performRunLoopStep(
         std::lock_guard<std::mutex> guard(mutex_);
         stopTraceInternal(now);
         resetInternal();
+        LOG(ERROR) << "State: Warmup stopped by CUPTI. (Buffer size configured is " << config_->activitiesMaxGpuBufferSize() / 1024 / 1024 << "MB)";
+        UST_LOGGER_MARK_COMPLETED(kWarmUpStage);
         VLOG(0) << "Warmup -> WaitForRequest";
         break;
       }
@@ -1298,6 +1316,12 @@ const time_point<system_clock> CuptiActivityProfiler::performRunLoopStep(
         return new_wakeup_time;
       }
 
+      // Before processing, we should wait for collectTrace thread to be done.
+      if (collectTraceThread && collectTraceThread->joinable()) {
+        std::lock_guard<std::mutex> guard(mutex_);
+        collectTraceThread->join();
+        collectTraceThread.reset(nullptr);
+      }
       // Before processing, we should wait for collectTrace thread to be done.
       if (collectTraceThread && collectTraceThread->joinable()) {
         std::lock_guard<std::mutex> guard(mutex_);