gateway-test.py
2.02 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
"""验调度层：往 :8000 投 5 条（比后端多一条），看是不是每卡一条、多出来那条排队。
用短任务（22 帧）只是为了快，路由逻辑和长任务完全一样。
"""
import json, time, urllib.request
GW = "http://127.0.0.1:8000"
SPEC = {"prompt": "一只橘猫趴在洒满阳光的窗台上，慢慢眨眼，尾巴轻轻摆动",
"width": 832, "height": 480, "frames": 22, "steps": 4, "turbo": True, "seed": 7}
def call(path, data=None, timeout=60):
req = urllib.request.Request(GW + path, data=data,
headers={"Content-Type": "application/json"} if data else {})
with urllib.request.urlopen(req, timeout=timeout) as r:
return json.load(r)
ids = []
for i in range(5):
job = call("/v1/videos", json.dumps(dict(SPEC, seed=7 + i)).encode())
ids.append(job["id"])
print(f"第{i+1}条 -> {job['backend']} id={job['id']} 队列位置={job['queue_position']}", flush=True)
print("\n--- 派发结果")
used = {}
for jid in ids:
b = call(f"/v1/videos/{jid}")["backend"]
used.setdefault(b, []).append(jid)
for b, js in sorted(used.items()):
print(f"{b}: {len(js)} 条 {js}")
print(f"用上了 {len(used)}/4 个后端")
print("\n--- 轮询到全部完成")
while True:
time.sleep(30)
states = [call(f"/v1/videos/{jid}") for jid in ids]
print(" " + " ".join(f"{s['id'][:6]}:{s['status']}" for s in states), flush=True)
if all(s["status"] in ("done", "failed") for s in states):
break
print("\n--- 结果")
for s in states:
took = (s["finished_at"] - s["started_at"]) if s.get("started_at") else 0
print(f"{s['id']} {s['status']} {took:.0f}s {s.get('size_bytes')}字节 {s['backend']}")
ok = [s for s in states if s["status"] == "done"]
if ok:
with urllib.request.urlopen(f"{GW}/v1/videos/{ok[0]['id']}/file", timeout=300) as r:
data = r.read()
print(f"\n经调度层下载 {ok[0]['id']}: {len(data)} 字节 "
f"(后端记录 {ok[0]['size_bytes']}) 一致={len(data) == ok[0]['size_bytes']}")