Remove try-except to ensure trace_link stops on exception

mlcommons · Jul 10, 2024 · 0102973 · 0102973
1 parent 7f8b892
commit 0102973
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/src/trace_link/trace_linker.py b/src/trace_link/trace_linker.py
@@ -422,11 +422,8 @@ def enforce_inter_thread_order(
 
             for future in as_completed(futures):
                 tid = futures[future]
-                try:
-                    future.result()
-                    logging.debug(f"Thread {tid} dependencies processed.")
-                except Exception as e:
-                    logging.error(f"Error processing thread {tid}: {e}")
+                future.result()
+                logging.debug(f"Thread {tid} dependencies processed.")
 
         return kineto_tid_cpu_ops_map
 

diff --git a/tests/trace_link/test_trace_linker.py b/tests/trace_link/test_trace_linker.py
@@ -425,6 +425,16 @@ def test_process_thread_inter_thread_order(mock_find_last, trace_linker):
     mock_find_last.assert_called()
 
 
+@patch("concurrent.futures.Future.result")
+@patch("chakra.src.trace_link.trace_linker.TraceLinker.process_thread_inter_thread_order")
+def test_enforce_inter_thread_order_exception(mock_process_thread, mock_future_result, trace_linker):
+    mock_future_result.side_effect = Exception("Test Exception")
+    kineto_tid_cpu_ops_map = {1: [MagicMock(spec=KinetoOperator)]}
+
+    with pytest.raises(Exception, match="Test Exception"):
+        trace_linker.enforce_inter_thread_order(kineto_tid_cpu_ops_map)
+
+
 def test_extract_pytorch_ops(trace_linker):
     mock_root_node = MagicMock(spec=PyTorchOperator)
     mock_child_node = MagicMock(spec=PyTorchOperator)