gateway.py
6.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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
"""四个 H3 实例前面的调度层。
api.py 是单 worker 串行，一个进程独占一张卡，所以并发的唯一办法是往不同端口投。
这一层把「投哪个端口」收掉：对外只有一个地址，接口和 api.py 逐字相同，
调用方原来怎么写现在还怎么写，只把端口换成 8000。
自己不存任何状态。任务 id 归哪个后端，问一遍就知道——四个后端一次扫描的代价可以忽略，
换来的是这个进程随时能重启，不丢任何东西（api.py 的 JOBS 本来也在内存里，
调度层再存一份 id→后端 的映射，就是第二个会丢的地方）。
"""
import os, json, threading, urllib.request, urllib.error
from concurrent.futures import ThreadPoolExecutor
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
BACKENDS = [f"http://127.0.0.1:{port.strip()}" for port in
os.environ.get("H3_BACKENDS", "8010,8011,8012,8013").split(",")]
POOL = ThreadPoolExecutor(max_workers=16)
# 最闲的那个后端已经压了这么多条时，再来的提交直接 429 顶回去。
# 没有这个上限的话队列是无界的：投 1000 条全部返回 202，最后一条两天后才开始跑，
# 调用方却收到「已接受」。宁可当场说满了。默认 4 条 ≈ 45 分钟积压。
MAX_QUEUE = int(os.environ.get("H3_MAX_QUEUE", "4"))
# 选后端是「读一遍各家队列深度，再往最闲的那个投」。这中间不能插进另一个提交，
# 否则同时来的请求会读到同一份快照、一起挑中同一个后端，全堆在一张卡上。
# 锁只圈住选 + 投这一段（约 30 毫秒），任务本身要跑十几分钟，序列化的代价可以忽略。
PICK = threading.Lock()
def fetch(url, data=None, timeout=30):
req = urllib.request.Request(url, data=data,
headers={"Content-Type": "application/json"} if data else {})
with urllib.request.urlopen(req, timeout=timeout) as resp:
return json.load(resp)
def health_all():
def one(backend):
try:
return dict(fetch(backend + "/health", timeout=5), backend=backend)
except Exception as e:
return {"backend": backend, "status": "down", "error": f"{type(e).__name__}: {e}"}
return list(POOL.map(one, BACKENDS))
def find_owner(jid):
"""任务归谁，挨个问。后端对陌生 id 返回 404，urlopen 抛异常，接住当没有。"""
def one(backend):
try:
return backend, fetch(f"{backend}/v1/videos/{jid}", timeout=5)
except Exception:
return backend, None
for backend, job in POOL.map(one, BACKENDS):
if job is not None:
return backend, job
raise HTTPException(404, "无此任务")
app = FastAPI(title="MiniMax-H3 Gateway", version="1.0")
@app.get("/health")
def health():
backends = health_all()
ready = [b for b in backends if b.get("status") == "ok"]
return {"status": "ok" if ready else "unavailable",
"ready": len(ready), "total": len(backends),
"queued": sum(b.get("queued", 0) for b in ready),
"running": sum(b.get("running", 0) for b in ready),
"backends": backends}
# 同步 def，让 FastAPI 把它丢进线程池。写成 async def 会让下面这些阻塞的 HTTP 调用
# 卡住事件循环——那样确实也能歪打正着地把提交排成队，但任何一个后端卡住，
# 5 秒超时就会连轮询一起堵死。要串行就用锁明说。
@app.post("/v1/videos", status_code=202)
def create(req: dict):
# 给了首末帧就只能走 fl2va；其余都走 ref2va——ref2va 不给参考图就退回纯文生视频，
# 所以「没有首末帧」这一类不需要再细分。
want = "fl2va" if (req.get("first_frame_url") or req.get("last_frame_url")) else "ref2va"
with PICK:
candidates = [b for b in health_all()
if b.get("status") == "ok" and b.get("partition") == want]
if not candidates:
raise HTTPException(503, f"没有就绪的 {want} 实例")
# 排队最短的那个。queued 是还在队列里的，running 是正在跑的那条，两个都要算。
pick = min(candidates, key=lambda b: b["queued"] + b["running"])
ahead = pick["queued"] + pick["running"]
if ahead >= MAX_QUEUE:
raise HTTPException(429, f"队列满：{len(candidates)} 个 {want} 实例每个都压着 "
f"{ahead} 条（上限 {MAX_QUEUE}）。等一会儿再投，"
f"或调大 H3_MAX_QUEUE")
try:
job = fetch(pick["backend"] + "/v1/videos", data=json.dumps(req).encode())
except urllib.error.HTTPError as e:
raise HTTPException(e.code, e.read().decode())
# eta_seconds 是这一条自己的执行时间，不含排队。前面还压着 ahead 条，
# 真正的等待还要加上它们——它们规格可能不同，这里不替调用方瞎乘。
return dict(job, backend=pick["backend"], queue_ahead=ahead, poll=f"/v1/videos/{job['id']}")
@app.get("/v1/videos/{jid}")
def get(jid: str):
backend, job = find_owner(jid)
return dict(job, backend=backend)
@app.get("/v1/videos/{jid}/file")
def download(jid: str):
backend, job = find_owner(jid)
if job["status"] != "done":
raise HTTPException(409, f"任务尚未完成（当前 {job['status']}）")
resp = urllib.request.urlopen(f"{backend}/v1/videos/{jid}/file", timeout=300)
return StreamingResponse(resp, media_type="video/mp4",
headers={"Content-Disposition": f'attachment; filename="{jid}.mp4"'})
@app.get("/v1/videos")
def listjobs(limit: int = 50):
def one(backend):
try:
return [dict(j, backend=backend)
for j in fetch(f"{backend}/v1/videos?limit={limit}", timeout=5)["jobs"]]
except Exception:
return []
jobs = [job for part in POOL.map(one, BACKENDS) for job in part]
jobs.sort(key=lambda j: j["created_at"], reverse=True)
return {"jobs": jobs[:limit]}