Files
vllm_agent_strategy/scripts/smoke_check.py
2026-07-21 16:04:18 +08:00

43 lines
1.4 KiB
Python

from __future__ import annotations
import argparse
import json
from app.clients.modelhub import ModelHubClient
from app.domain.gpu import expand_gpu_aliases
from app.domain.vllm_configs import gen_vllm_config
from app.settings import load_settings
def redact(payload: dict) -> dict:
clone = dict(payload)
if "contestApiToken" in clone:
clone["contestApiToken"] = "***"
return clone
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--no-submit", action="store_true", help="Only print a redacted payload")
parser.add_argument("--model-id", default="Qwen/Qwen2.5-0.5B-Instruct")
parser.add_argument("--gpu-alias", default="k100")
args = parser.parse_args()
settings = load_settings(require_secrets=not args.no_submit)
client = ModelHubClient(settings)
pairs = expand_gpu_aliases([args.gpu_alias])
if not pairs:
raise SystemExit(f"unsupported gpu alias for vLLM: {args.gpu_alias}")
_, gpu_type = pairs[0]
config = gen_vllm_config(gpu_type, args.model_id, settings.default_max_model_len)
payload = client.build_contest_payload(args.model_id, gpu_type, config)
print(json.dumps(redact(payload), ensure_ascii=False, indent=2))
if args.no_submit:
return
result = client.create_contest_task(args.model_id, gpu_type, config)
print(result)
if __name__ == "__main__":
main()