Fix command wrappers exiting after successful runs

setup_env.py

@@ -104,7 +104,7 @@ def run_command(command, shell=False, log_step=None):
             subprocess.run(command, shell=shell, check=True)
         except subprocess.CalledProcessError as e:
             logging.error(f"Error occurred while running command: {e}")
-        sys.exit(1)
+            sys.exit(1)
 
 def prepare_model():
     _, arch = system_info()

tests/test_command_wrappers.py

@@ -0,0 +1,49 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import setup_env
+from utils import e2e_benchmark
+
+
+def test_setup_env_run_command_does_not_exit_on_success_without_log_step():
+    with patch("setup_env.subprocess.run") as mock_run, patch("setup_env.sys.exit") as mock_exit:
+        setup_env.run_command(["echo", "ok"])
+
+    mock_run.assert_called_once_with(["echo", "ok"], shell=False, check=True)
+    mock_exit.assert_not_called()
+
+
+def test_setup_env_run_command_exits_on_failure_without_log_step():
+    error = setup_env.subprocess.CalledProcessError(1, ["echo", "fail"])
+    with patch("setup_env.subprocess.run", side_effect=error), patch(
+        "setup_env.sys.exit"
+    ) as mock_exit:
+        setup_env.run_command(["echo", "fail"])
+
+    mock_exit.assert_called_once_with(1)
+
+
+def test_e2e_benchmark_run_command_does_not_exit_on_success_without_log_step():
+    with patch("utils.e2e_benchmark.subprocess.run") as mock_run, patch(
+        "utils.e2e_benchmark.sys.exit"
+    ) as mock_exit:
+        e2e_benchmark.run_command(["echo", "ok"])
+
+    mock_run.assert_called_once_with(["echo", "ok"], shell=False, check=True)
+    mock_exit.assert_not_called()
+
+
+def test_e2e_benchmark_run_benchmark_uses_run_command_without_exiting_early():
+    e2e_benchmark.args = SimpleNamespace(
+        model="model.gguf",
+        n_token=32,
+        threads=4,
+        n_prompt=64,
+    )
+
+    with patch("utils.e2e_benchmark.platform.system", return_value="Linux"), patch(
+        "utils.e2e_benchmark.os.path.exists", return_value=True
+    ), patch("utils.e2e_benchmark.run_command") as mock_run:
+        e2e_benchmark.run_benchmark()
+
+    mock_run.assert_called_once()

utils/e2e_benchmark.py

@@ -20,7 +20,7 @@ def run_command(command, shell=False, log_step=None):
             subprocess.run(command, shell=shell, check=True)
         except subprocess.CalledProcessError as e:
             logging.error(f"Error occurred while running command: {e}")
-        sys.exit(1)
+            sys.exit(1)
 
 def run_benchmark():
     build_dir =  os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "build")
@@ -56,4 +56,4 @@ def parse_args():
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)
     args = parse_args()
-    run_benchmark()
+    run_benchmark()