[Misc] Move device options to a single place (#8322)

2024-09-11 23:25:58 +03:00
parent 73202dbe77
commit 8baa454937
3 changed files with 23 additions and 20 deletions
--- a/vllm/engine/arg_utils.py
+++ b/vllm/engine/arg_utils.py
@@ -26,6 +26,16 @@ logger = init_logger(__name__)

 ALLOWED_DETAILED_TRACE_MODULES = ["model", "worker", "all"]

+DEVICE_OPTIONS = [
+    "auto",
+    "cuda",
+    "neuron",
+    "cpu",
+    "openvino",
+    "tpu",
+    "xpu",
+]
+

 def nullable_str(val: str):
    if not val or val == "None":
@@ -553,10 +563,7 @@ class EngineArgs:
        parser.add_argument("--device",
                            type=str,
                            default=EngineArgs.device,
-                            choices=[
-                                "auto", "cuda", "neuron", "cpu", "openvino",
-                                "tpu", "xpu"
-                            ],
+                            choices=DEVICE_OPTIONS,
                            help='Device type for vLLM execution.')
        parser.add_argument('--num-scheduler-steps',
                            type=int,