CLI param for benchmark timeout

SpareCores · Jan 27, 2025 · 50f973e · 50f973e
1 parent 665526d
commit 50f973e
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/images/benchmark-llm/benchmark.py b/images/benchmark-llm/benchmark.py
@@ -42,13 +42,16 @@
     default="/models",
     help="Directory to cache/store downloaded models.",
 )
+cli_parser.add_argument(
+    "--benchmark-timeout",
+    type=int,
+    default=60,
+    help="Timeout in seconds for a single benchmark.",
+)
 cli_args = cli_parser.parse_args()
 
 # #############################################################################
 
-# max number of seconds to wait for a benchmark to finish
-TIMEOUT = 60
-
 # default command for llama-bench
 COMMAND = [
     "./llama-bench",
@@ -192,7 +195,7 @@ def max_ngl(model: str):
             result = run(
                 COMMAND + ["-m", model, "-ngl", str(ngl), "-r", "1", "-t", "1"],
                 capture_output=True,
-                timeout=TIMEOUT,
+                timeout=cli_args.benchmark_timeout,
             )
             if result.returncode == 0:
                 return ngl
@@ -232,7 +235,7 @@ def max_ngl(model: str):
                     cmd
                     + [benchmark["iteration_param"], str(iteration)]
                     + benchmark["extra_params"],
-                    timeout=TIMEOUT,
+                    timeout=cli_args.benchmark_timeout,
                 )
             except Exception as e:
                 logger.error(f"Error: {e}")