Updates benchmarking CLI

Author: ppak10

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "additive-manufacturing"
-version = "0.0.30"
+version = "0.0.31"
 authors = [
   { name = "Peter Pak", email = "ppak10@gmail.com" },
 ]

src/am/benchmark/cli.py

@@ -32,6 +32,11 @@ def benchmark(
         num_runs: int = typer.Option(
             1, help="Number of benchmark runs (for error bar statistics)."
         ),
+        vllm_args: list[str] = typer.Option(
+            [],
+            "--vllm-arg",
+            help="Extra argument forwarded to 'vllm serve' (repeatable, e.g. --vllm-arg='--max-model-len=8192').",
+        ),
         workspace_option: WorkspaceOption = None,
         num_proc: NumProc = 1,
     ) -> None:
@@ -63,6 +68,7 @@ def benchmark(
                 out_path=workspace_folder_path,
                 proctor_url=proctor_url,
                 proctor_model=proctor_model,
+                vllm_extra_args=vllm_args or None,
             )
 
         except Exception as e:

src/am/benchmark/llm/__init__.py

@@ -85,6 +85,7 @@ def benchmark_llm(
     out_path: Path | None = None,
     proctor_url: str | None = None,
     proctor_model: str = PROCTOR_MODEL,
+    vllm_extra_args: list[str] | None = None,
 ) -> list[dict]:
 
     missing = [pkg for pkg in ("datasets",) if importlib.util.find_spec(pkg) is None]
@@ -147,7 +148,10 @@ def _with_proctor(runner, vision_runner, serve_model):
     else:
         from .server import managed_vllm_server
 
-        with managed_vllm_server(model, port=port) as (server_url, serve_model):
+        with managed_vllm_server(model, port=port, extra_args=vllm_extra_args) as (
+            server_url,
+            serve_model,
+        ):
             runner = lambda questions, max_tokens: _run_openai_compatible(
                 server_url, serve_model, questions, max_tokens
             )