add function ensureCollectTraceDone to wait and cleanup collectTraceT…

…hread
pytorch · Aug 6, 2024 · 1732581 · 1732581
1 parent 0f27758
commit 1732581
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/libkineto/src/CuptiActivityProfiler.cpp b/libkineto/src/CuptiActivityProfiler.cpp
@@ -1096,6 +1096,13 @@ void CuptiActivityProfiler::collectTrace(bool collection_done,
     UST_LOGGER_MARK_COMPLETED(kCollectionStage);
 }
 
+void CuptiActivityProfiler::ensureCollectTraceDone() {
+  if (collectTraceThread && collectTraceThread->joinable()) {
+    std::lock_guard<std::mutex> guard(mutex_);
+    collectTraceThread->join();
+    collectTraceThread.reset(nullptr);
+  }
+}
 void CuptiActivityProfiler::startTraceInternal(
     const time_point<system_clock>& now) {
   captureWindowStartTime_ = libkineto::timeSinceEpoch(now);
@@ -1251,11 +1258,7 @@ const time_point<system_clock> CuptiActivityProfiler::performRunLoopStep(
       }
 
       // 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);
-      }
+      ensureCollectTraceDone();
 
       // FIXME: Probably want to allow interruption here
       // for quickly handling trace request via synchronous API

diff --git a/libkineto/src/CuptiActivityProfiler.h b/libkineto/src/CuptiActivityProfiler.h
@@ -167,6 +167,8 @@ class CuptiActivityProfiler {
   // Collect CPU and GPU traces
   void collectTrace(bool collectionDone,  const std::chrono::time_point<std::chrono::system_clock>& now );
 
+  // Ensure collectTrace is done
+  void ensureCollectTraceDone();
   // Process CPU and GPU traces
   void processTrace(ActivityLogger& logger) {
     std::lock_guard<std::mutex> guard(mutex_);