improved test load

2025-10-09 19:37:39 -07:00
parent 74bd932327
commit e251afda2b
3 changed files with 78 additions and 17 deletions
@@ -292,12 +292,12 @@ def test_load_cmd(
    args = arg_parser.parse_args()
    if hasattr(args, "comfy_model"):
        os.environ["COMFY_MODEL"] = args.comfy_model
-    server_url = dict(
-        prod="https://run.vast.ai",
-        alpha="https://run-alpha.vast.ai",
-        candidate="https://run-candidate.vast.ai",
-        local="http://localhost:8080",
-    )[args.instance]
+    server_url = {
+        "prod": "https://run.vast.ai",
+        "alpha": "https://run-alpha.vast.ai",
+        "candidate": "https://run-candidate.vast.ai",
+        "local": "http://localhost:8080",
+    }.get(args.instance, "http://localhost:8080")
    run_test(
        num_requests=args.num_requests,
        requests_per_second=args.requests_per_second,