Просмотр исходного кода

chore: move `device` keys to a constant (#1020)

AlpinDale 2 месяцев назад
Родитель
Сommit
fe01e2ded8

+ 11 - 3
aphrodite/engine/args_tools.py

@@ -26,6 +26,16 @@ if TYPE_CHECKING:
 
 APHRODITE_USE_RAY_SPMD_WORKER = envs.APHRODITE_USE_RAY_SPMD_WORKER
 
+DEVICE_OPTIONS = [
+    "auto",
+    "cuda",
+    "neuron",
+    "cpu",
+    "openvino",
+    "tpu",
+    "xpu",
+]
+
 def nullable_kvs(val: str) -> Optional[Mapping[str, int]]:
     if len(val) == 0:
         return None
@@ -357,9 +367,7 @@ class EngineArgs:
             "--device",
             type=str,
             default=EngineArgs.device,
-            choices=[
-                "auto", "cuda", "neuron", "cpu", "openvino", "tpu", "xpu"
-            ],
+            choices=DEVICE_OPTIONS,
             help=("Category: Model Options\n"
                   "Device to use for model execution."),
         )

+ 2 - 2
tests/benchmarks/engine/latency.py

@@ -11,7 +11,7 @@ from tqdm import tqdm
 
 from aphrodite import LLM, SamplingParams
 from aphrodite.common.utils import FlexibleArgumentParser
-from aphrodite.engine.args_tools import EngineArgs
+from aphrodite.engine.args_tools import DEVICE_OPTIONS, EngineArgs
 from aphrodite.inputs import PromptInputs
 from aphrodite.quantization import QUANTIZATION_METHODS
 
@@ -212,7 +212,7 @@ if __name__ == '__main__':
         "--device",
         type=str,
         default="auto",
-        choices=["auto", "cuda", "cpu", "openvino", "tpu", "xpu"],
+        choices=DEVICE_OPTIONS,
         help='device type for Aphrodite execution, supporting CUDA, OpenVINO '
         'and CPU.')
     parser.add_argument('--block-size',

+ 3 - 2
tests/benchmarks/engine/throughput.py

@@ -15,7 +15,8 @@ from aphrodite.common.utils import (FlexibleArgumentParser,
                                     merge_async_iterators)
 from aphrodite.endpoints.openai.api_server import (
     build_async_engine_client_from_engine_args)
-from aphrodite.engine.args_tools import AsyncEngineArgs, EngineArgs
+from aphrodite.engine.args_tools import (DEVICE_OPTIONS, AsyncEngineArgs,
+                                         EngineArgs)
 from aphrodite.quantization import QUANTIZATION_METHODS
 
 
@@ -456,7 +457,7 @@ if __name__ == "__main__":
         "--device",
         type=str,
         default="auto",
-        choices=["auto", "cuda", "cpu", "openvino", "tpu", "xpu"],
+        choices=DEVICE_OPTIONS,
         help='device type for Aphrodite execution, supporting CUDA, OpenVINO '
         'and CPU.')
     parser.add_argument(