diff --git a/main.py b/main.py index 75566a9..ca4db28 100644 --- a/main.py +++ b/main.py @@ -243,16 +243,15 @@ def check_queue_available(gpu: str, token: str) -> int: def build_config_params(gpu: str) -> str: - """构建 YAML 配置""" + """构建 YAML 配置 - 完全匹配平台自动生成的格式""" config = GPU_CONFIGS.get(gpu, {}) framework = config.get('framework', 'vllm') - docker_image = config.get('docker_image', '') if framework == 'llama.cpp': + # hygon_k100-ai 使用 llama.cpp params = { 'framework': 'llama.cpp', - 'docker_image': docker_image, - 'nv_docker_image': docker_image, + 'nv_framework': 'llama.cpp', 'api': 'completion', 'max_tokens': 1024, 'temperature': 0.7, @@ -260,22 +259,21 @@ def build_config_params(gpu: str) -> str: 'top_p': 0.9, 'lang': 'zh', 'max_model_len': 4096, - 'modelFormat': 'GGUF', 'sut_config': { 'gpu_num': 1, 'values': { 'command': [ - '/bin/bash', '-ic', - 'llama-server --model /model --alias llm --threads 20 ' - '--n-gpu-layers 999 --prio 3 --min_p 0.01 ' - '--ctx-size 4096 --host 0.0.0.0 --port 8000 --jinja --flash-attn off' + 'llama-server', '--model', '/model', '--alias', 'llm', + '--threads', '20', '--n-gpu-layers', '999', '--prio', '3', + '--min_p', '0.01', '--ctx-size', '4096', + '--host', '0.0.0.0', '--port', '8000', + '--jinja', '--flash-attn', 'off', ] } }, 'ref_config': { 'gpu_num': 1, 'values': { - 'cpu_num': 2, 'command': [ 'llama-server', '--model', '/model', '--alias', 'llm', '--threads', '20', '--n-gpu-layers', '999', @@ -283,13 +281,12 @@ def build_config_params(gpu: str) -> str: ] } }, - 'model': 'llm', } else: + # vllm (P800, Ascend, MetaX) params = { 'framework': 'vllm', - 'docker_image': docker_image, - 'nv_docker_image': docker_image, + 'nv_framework': 'vllm', 'api': 'completion', 'max_tokens': 1024, 'temperature': 0.7, @@ -297,20 +294,9 @@ def build_config_params(gpu: str) -> str: 'top_p': 0.9, 'lang': 'zh', 'max_model_len': 2048, - 'modelFormat': 'HuggingFace', 'sut_config': { 'gpu_num': 1, 'values': { - 'command': [ - '/bin/bash', '-ic', - 'vllm serve /model --port 8000 --served-model-name llm --max-model-len 2048 --dtype auto --gpu-memory-utilization 0.95 -tp 1 --enforce-eager --trust-remote-code' - ] - } - }, - 'ref_config': { - 'gpu_num': 1, - 'values': { - 'cpu_num': 2, 'command': [ 'vllm', 'serve', '/model', '--port', '8000', '--served-model-name', 'llm', '--max-model-len', '2048', @@ -319,7 +305,16 @@ def build_config_params(gpu: str) -> str: ] } }, - 'model': 'llm', + 'ref_config': { + 'gpu_num': 1, + 'values': { + 'command': [ + 'vllm', 'serve', '/model', '--port', '80', + '--served-model-name', 'llm', '--max-model-len', '4096', + '--enforce-eager', '--trust-remote-code', '-tp', '1', + ] + } + }, } return yaml.dump(params, default_flow_style=False, allow_unicode=True, width=1000)