Skip to content

Commit eda16d0

Browse files
committed
Add short/long name mapping of vllm engine args used in model config
1 parent aadb954 commit eda16d0

File tree

2 files changed

+15
-2
lines changed

2 files changed

+15
-2
lines changed

vec_inf/client/_client_vars.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,12 @@
6767
"error": "err_file",
6868
}
6969

70+
# vLLM engine args mapping between short and long names
71+
VLLM_SHORT_TO_LONG_MAP = {
72+
"-tp": "--tensor-parallel-size",
73+
"-pp": "--pipeline-parallel-size",
74+
"-O": "--compilation-config",
75+
}
7076

7177
# Slurm script templates
7278
class ShebangConfig(TypedDict):

vec_inf/client/_helper.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
KEY_METRICS,
2020
REQUIRED_FIELDS,
2121
SRC_DIR,
22+
VLLM_SHORT_TO_LONG_MAP,
2223
)
2324
from vec_inf.client._exceptions import (
2425
MissingRequiredFieldsError,
@@ -156,9 +157,15 @@ def _process_vllm_args(self, arg_string: str) -> dict[str, Any]:
156157
for arg in arg_string.split(","):
157158
if "=" in arg:
158159
key, value = arg.split("=")
159-
vllm_args[key] = value
160+
if key.strip() in VLLM_SHORT_TO_LONG_MAP:
161+
key = VLLM_SHORT_TO_LONG_MAP[key.strip()]
162+
vllm_args[key.strip()] = value.strip()
160163
else:
161-
vllm_args[arg] = True
164+
if "-O" in arg.strip():
165+
key = VLLM_SHORT_TO_LONG_MAP["-O"]
166+
vllm_args[key] = arg.strip()[2:].strip()
167+
else:
168+
vllm_args[arg.strip()] = True
162169
return vllm_args
163170

164171
def _get_launch_params(self) -> dict[str, Any]:

0 commit comments

Comments
 (0)