Revert "Revert "[PROTON] Fix hook manager when a session is deactivated multiple times (#7743)"" (#4958)

This reverts commit
76f767f.

Author: whitneywhtsang

third_party/proton/proton/hooks/hook.py

@@ -77,8 +77,9 @@ def deactivate(session: Optional[int] = None) -> None:
         deactivated_hooks = set()
         for session in sessions:
             for hook in HookManager.session_hooks[session]:
+                if hook in HookManager.active_hooks:
+                    deactivated_hooks.add(hook)
                 HookManager.session_hooks[session][hook] = False
-                deactivated_hooks.add(hook)
 
         # Check if any other sessions rely on this hook
         for hook in deactivated_hooks:

third_party/proton/test/test_api.py

@@ -1,6 +1,7 @@
 import json
 import triton.profiler as proton
 import pathlib
+from triton.profiler.hooks.hook import HookManager
 
 
 def test_profile_single_session(tmp_path: pathlib.Path):
@@ -103,6 +104,14 @@ def test_hook(tmp_path: pathlib.Path):
     temp_file = tmp_path / "test_hook.hatchet"
     session_id0 = proton.start(str(temp_file.with_suffix("")), hook="triton")
     proton.activate(session_id0)
+    proton.activate(session_id0)
+    assert len(
+        HookManager.active_hooks) == 1, ("Activate a session multiple times should maintain a single instance of hook")
+    assert list(HookManager.session_hooks[session_id0].values())[0] is True
+    proton.deactivate(session_id0)
+    assert list(HookManager.session_hooks[session_id0].values())[0] is False
+    assert len(HookManager.active_hooks) == 0
+    # Deactivate a session multiple times should not raise an error
     proton.deactivate(session_id0)
     proton.finalize(None)
     assert temp_file.exists()