[PROTON] Fix proton's support for multiple profiling sessions (#5140)

Also change the default behavior or `activate` and `deactivate` to apply
for all sessions but not only session 0

Author: Jokeren

third_party/proton/csrc/Proton.cpp

@@ -26,10 +26,16 @@ void initProton(pybind11::module &&m) {
     SessionManager::instance().activateSession(sessionId);
   });
 
+  m.def("activate_all",
+        []() { SessionManager::instance().activateAllSessions(); });
+
   m.def("deactivate", [](size_t sessionId) {
     SessionManager::instance().deactivateSession(sessionId);
   });
 
+  m.def("deactivate_all",
+        []() { SessionManager::instance().deactivateAllSessions(); });
+
   m.def("finalize", [](size_t sessionId, const std::string &outputFormat) {
     auto outputFormatEnum = parseOutputFormat(outputFormat);
     SessionManager::instance().finalizeSession(sessionId, outputFormatEnum);

third_party/proton/csrc/include/Profiler/Profiler.h

@@ -27,10 +27,9 @@ class Profiler {
   /// If the profiler is already started, this function does nothing.
   Profiler *start() {
     std::unique_lock<std::shared_mutex> lock(mutex);
-    if (this->isInitialized)
-      return this;
-    this->doStart();
-    this->isInitialized = true;
+    if (this->initializedCount == 0)
+      this->doStart();
+    this->initializedCount++;
     return this;
   }
 
@@ -45,10 +44,11 @@ class Profiler {
   /// Stop the profiler.
   Profiler *stop() {
     std::unique_lock<std::shared_mutex> lock(mutex);
-    if (!this->isInitialized)
+    if (this->initializedCount == 0)
       return this;
-    this->doStop();
-    this->isInitialized = false;
+    this->initializedCount--;
+    if (this->initializedCount == 0)
+      this->doStop();
     return this;
   }
 
@@ -80,7 +80,9 @@ class Profiler {
 
   mutable std::shared_mutex mutex;
   std::set<Data *> dataSet;
-  bool isInitialized{false};
+
+private:
+  int initializedCount{};
 };
 
 } // namespace proton

third_party/proton/csrc/include/Session/Session.h

@@ -77,8 +77,12 @@ class SessionManager : public Singleton<SessionManager> {
 
   void activateSession(size_t sessionId);
 
+  void activateAllSessions();
+
   void deactivateSession(size_t sessionId);
 
+  void deactivateAllSessions();
+
   void enterScope(const Scope &scope);
 
   void exitScope(const Scope &scope);

third_party/proton/csrc/lib/Session/Session.cpp

@@ -84,11 +84,25 @@ void SessionManager::activateSession(size_t sessionId) {
   activateSessionImpl(sessionId);
 }
 
+void SessionManager::activateAllSessions() {
+  std::unique_lock<std::shared_mutex> lock(mutex);
+  for (auto iter : sessionActive) {
+    activateSessionImpl(iter.first);
+  }
+}
+
 void SessionManager::deactivateSession(size_t sessionId) {
   std::unique_lock<std::shared_mutex> lock(mutex);
   deActivateSessionImpl(sessionId);
 }
 
+void SessionManager::deactivateAllSessions() {
+  std::unique_lock<std::shared_mutex> lock(mutex);
+  for (auto iter : sessionActive) {
+    deActivateSessionImpl(iter.first);
+  }
+}
+
 void SessionManager::activateSessionImpl(size_t sessionId) {
   throwIfSessionNotInitialized(sessions, sessionId);
   if (sessionActive[sessionId])
@@ -116,6 +130,7 @@ void SessionManager::removeSession(size_t sessionId) {
   }
   auto path = sessions[sessionId]->path;
   sessionPaths.erase(path);
+  sessionActive.erase(sessionId);
   sessions.erase(sessionId);
 }
 

third_party/proton/proton/profile.py

@@ -85,36 +85,42 @@ def start(
     return libproton.start(name, context, data, backend)
 
 
-def activate(session: Optional[int] = 0) -> None:
+def activate(session: Optional[int] = None) -> None:
     """
     Activate the specified session.
     The profiling session will be active and data will be recorded.
 
     Args:
-        session (int): The session ID of the profiling session. Defaults to 0 (the first session started.)
+        session (int): The session ID of the profiling session. Defaults to None (all sessions)
 
     Returns:
         None
     """
     if is_command_line() and session != 0:
         raise ValueError("Only one session can be activated when running from the command line.")
-    libproton.activate(session)
+    if session is None:
+        libproton.activate_all()
+    else:
+        libproton.activate(session)
 
 
-def deactivate(session: Optional[int] = 0) -> None:
+def deactivate(session: Optional[int] = None) -> None:
     """
     Stop the specified session.
     The profiling session's data will still be in the memory, but no more data will be recorded.
 
     Args:
-        session (int): The session ID of the profiling session. Defaults to 0 (the first session started.)
+        session (int): The session ID of the profiling session. Defaults to None (all sessions)
 
     Returns:
         None
     """
     if is_command_line() and session != 0:
         raise ValueError("Only one session can be deactivated when running from the command line.")
-    libproton.deactivate(session)
+    if session is None:
+        libproton.deactivate_all()
+    else:
+        libproton.deactivate(session)
 
 
 def finalize(session: Optional[int] = None, output_format: str = "hatchet") -> None:

third_party/proton/test/test_api.py

@@ -3,7 +3,7 @@
 import pathlib
 
 
-def test_profile(tmp_path: pathlib.Path):
+def test_profile_single_session(tmp_path: pathlib.Path):
     temp_file0 = tmp_path / "test_profile0.hatchet"
     session_id0 = proton.start(str(temp_file0.with_suffix("")))
     proton.activate()
@@ -29,6 +29,28 @@ def test_profile(tmp_path: pathlib.Path):
     pathlib.Path("test.hatchet").unlink()
 
 
+def test_profile_multiple_sessions(tmp_path: pathlib.Path):
+    temp_file0 = tmp_path / "test_profile0.hatchet"
+    proton.start(str(temp_file0.with_suffix("")))
+    temp_file1 = tmp_path / "test_profile1.hatchet"
+    proton.start(str(temp_file1.with_suffix("")))
+    proton.activate()
+    proton.deactivate()
+    proton.finalize()
+    assert temp_file0.exists()
+    assert temp_file1.exists()
+
+    temp_file2 = tmp_path / "test_profile2.hatchet"
+    session_id2 = proton.start(str(temp_file2.with_suffix("")))
+    temp_file3 = tmp_path / "test_profile3.hatchet"
+    session_id3 = proton.start(str(temp_file3.with_suffix("")))
+    proton.deactivate(session_id2)
+    proton.deactivate(session_id3)
+    proton.finalize()
+    assert temp_file2.exists()
+    assert temp_file3.exists()
+
+
 def test_profile_decorator(tmp_path: pathlib.Path):
     temp_file = tmp_path / "test_profile_decorator.hatchet"
 

third_party/proton/test/test_profile.py

@@ -257,3 +257,23 @@ def test_deactivate(tmp_path: pathlib.Path):
     assert "device_id" not in data[0]["metrics"]
     assert len(data[0]["children"]) == 1
     assert "device_id" in data[0]["children"][0]["metrics"]
+
+
+def test_multiple_sessions(tmp_path: pathlib.Path):
+    temp_file0 = tmp_path / "test_multiple_sessions0.hatchet"
+    temp_file1 = tmp_path / "test_multiple_sessions1.hatchet"
+    session_id0 = proton.start(str(temp_file0.with_suffix("")))
+    session_id1 = proton.start(str(temp_file1.with_suffix("")))
+    torch.randn((10, 10), device="cuda")
+    torch.randn((10, 10), device="cuda")
+    proton.deactivate(session_id0)
+    proton.finalize(session_id0)
+    torch.randn((10, 10), device="cuda")
+    proton.finalize(session_id1)
+    # kernel has been invokved twice in session 0 and three times in session 1
+    with temp_file0.open() as f:
+        data = json.load(f)
+    assert int(data[0]["children"][0]["metrics"]["count"]) == 2
+    with temp_file1.open() as f:
+        data = json.load(f)
+    assert int(data[0]["children"][0]["metrics"]["count"]) == 3