server.py 28.1 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554
#!/usr/bin/env python3
"""H3 Console:给 gateway.py 加一层 Web UI + 上传 + 任务持久化。

不重做调度——gateway.py 已经做了选后端/分区路由/限流,这里只反代它的
/v1/videos* 接口,加三样它没有的东西:文件上传(返回可喂给 api.py 的 URL)、
SQLite 任务影子记录(gateway.py 的 JOBS 在后端进程内存里,进程重启就丢)、
静态前端。

弹性 partition 调度器(kill+restart 某个 GPU 后端切 fl2va/ref2va)只在
H3_ENABLE_SUPERVISOR=1 时启用——这一层直接操作服务器上的真实进程,本地
Mac 开发环境没有意义,见 DESIGN.md 「四点五」。
"""
import os, re, sqlite3, time, uuid, json, asyncio, subprocess
from pathlib import Path
from typing import Optional

import httpx
from fastapi import APIRouter, FastAPI, HTTPException, UploadFile, File, Request
from fastapi.responses import StreamingResponse, FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image

BASE = Path(__file__).resolve().parent
ASSETS = BASE / "assets"
ASSETS.mkdir(exist_ok=True)
DB_PATH = BASE / "console.sqlite3"

GATEWAY_URL = os.environ.get("H3_GATEWAY_URL", "http://192.168.110.67:8000")
CONSOLE_PUBLIC_BASE = os.environ.get("H3_CONSOLE_PUBLIC_BASE")  # 供后端 worker 回读上传图片用;见 README
MAX_UPLOAD_BYTES = 20 * 1024 * 1024  # 20MB,够单张参考图/首尾帧,防止占满磁盘
ENABLE_SUPERVISOR = os.environ.get("H3_ENABLE_SUPERVISOR") == "1"
# 部署在共享域名的子路径下时设这个(比如挂在 qwen.wudaima.com/console 而不是独立域名),
# 见 README「子路径部署」。留空就是原来的行为,挂在根路径,本地/局域网直连不受影响。
BASE_PATH = os.environ.get("H3_CONSOLE_BASE_PATH", "").rstrip("/")

# GPU/端口映射,对应服务器上 start-all.sh 写死的分配。仅 supervisor 用得到。
BACKEND_TOPOLOGY = [
    {"gpu": 1, "port": 8010},
    {"gpu": 2, "port": 8011},
    {"gpu": 0, "port": 8012},
    {"gpu": 3, "port": 8013},
]
H3_ROOT = Path("/data/models/minimax-h3")  # 仅服务器上存在,本地开发不会用到

app = FastAPI(title="H3 Console")
client = httpx.AsyncClient(timeout=30.0)
# 所有 /v1/* 接口挂在这个 router 上而不是直接挂在 app 上,是为了能整体套一层 BASE_PATH
# 前缀(app.include_router(router, prefix=BASE_PATH),见文件末尾)——BASE_PATH 留空时
# 前缀是空字符串,行为和直接挂在 app 上完全一样。
router = APIRouter()


def db():
    con = sqlite3.connect(DB_PATH)
    con.row_factory = sqlite3.Row
    return con


def init_db():
    with db() as con:
        con.execute("""
            CREATE TABLE IF NOT EXISTS jobs (
                id TEXT PRIMARY KEY,
                gateway_id TEXT,
                partition TEXT NOT NULL,
                backend TEXT,
                status TEXT NOT NULL,
                req_json TEXT NOT NULL,
                created_at REAL NOT NULL,
                updated_at REAL NOT NULL,
                last_known_json TEXT
            )
        """)
        # 老库升级:CREATE TABLE IF NOT EXISTS 不会给已存在的表加列,只能单独 ALTER。
        # 第二次启动起这行必然报"列已存在",吞掉就行,不影响其它列。
        try:
            con.execute("ALTER TABLE jobs ADD COLUMN deleted_at REAL")
        except sqlite3.OperationalError:
            pass


init_db()


# ---------------- 资产上传 ----------------
# api.py 的 GenReq 只认 URL 字段(first_frame_url / reference_image_urls),内部用
# urllib.urlopen 拉取。所以上传接口的活是「收文件 -> 落盘 -> 吐一个后端能拉到的 URL」,
# 不改 api.py 一行代码。
@router.post("/v1/assets")
async def upload_asset(file: UploadFile = File(...)):
    ext = Path(file.filename or "").suffix.lower()
    if ext not in (".jpg", ".jpeg", ".png", ".webp"):
        raise HTTPException(400, f"不支持的文件类型 {ext},只收图片(jpg/png/webp)")
    body = await file.read()
    if len(body) > MAX_UPLOAD_BYTES:
        raise HTTPException(413, f"文件超过 {MAX_UPLOAD_BYTES // 1024 // 1024}MB 上限")
    aid = uuid.uuid4().hex[:16]
    path = ASSETS / f"{aid}{ext}"
    path.write_bytes(body)
    try:
        Image.open(path).verify()  # 校验确实是能解码的图片,不是随便什么二进制
    except Exception:
        path.unlink(missing_ok=True)
        raise HTTPException(400, "文件不是有效图片")
    return {"id": aid, "url": f"{_public_base()}/assets/{aid}{ext}"}


def _public_base() -> str:
    if CONSOLE_PUBLIC_BASE:
        return CONSOLE_PUBLIC_BASE.rstrip("/")
    # 本地开发时没配这个环境变量会直接报错提示——见 README「已知限制」:
    # 上传图片的 URL 必须是 api.py 所在服务器能反向拉取到的地址,本地 Mac 默认猜不出这个地址。
    raise HTTPException(
        500, "未设置 H3_CONSOLE_PUBLIC_BASE:上传接口需要知道自己对外可访问的地址,"
             "这个地址会被服务器上的 api.py worker 反向拉取,见 README。")


app.mount(f"{BASE_PATH}/assets", StaticFiles(directory=str(ASSETS)), name="assets")


# ---------------- 生成任务:console 自己是排队的主人,gateway.py 只是"当前有空的执行者" ----------------
# 关键点:gateway.py 对没有就绪后端的 partition 直接 503(`没有就绪的 fl2va 实例`),不排队。
# 如果 console 只是把这个 503 原样透传给用户,"有几个任务排队就用几张卡跑 fl2va"就无从谈起——
# 请求在 503 那一步就死了,根本没机会被弹性调度器看到、更别说触发 flip。
# 所以 console 自己维护身份(console 生成的 id)和排队状态:提交时先落 sqlite 记一条 status=queued,
# 尝试提交给 gateway;提交不上(503/429)就留在本地队列里,交给下面的 reconciler_loop 每隔几秒重试,
# 弹性调度器看到的 pending 计数正是这些「还没提交上去」的本地排队任务。
def _partition_of(req: dict) -> str:
    # 和 gateway.py 里 `want = "fl2va" if (...) else "ref2va"` 保持完全一致的判断,
    # 不要各写一套、两边判断分岔的话 console 统计出来的 pending 计数就和 gateway 实际路由对不上。
    return "fl2va" if (req.get("first_frame_url") or req.get("last_frame_url")) else "ref2va"


def _row_view(row) -> dict:
    """把一行 sqlite 记录还原成对外的任务视图,用 console 自己的 id,不暴露 gateway 内部 id。"""
    if row["last_known_json"]:
        view = json.loads(row["last_known_json"])
    else:
        view = {}
    view["id"] = row["id"]
    view["status"] = row["status"]
    view["partition"] = row["partition"]
    view["backend"] = row["backend"]
    view.setdefault("req", json.loads(row["req_json"]))
    view["created_at"] = row["created_at"]
    view["deleted_at"] = row["deleted_at"]
    # gateway.py/api.py 原样带过来的 `poll`/`download` 字段用的是 gateway 内部 id(不是 console
    # 自己的 id),照抄会指向一个 console 自己的 sqlite 里根本查不到的路径,直接 404——这两个
    # 字段本来就该跟 `id` 一样,统一改成用 console 的 id 重新拼。
    if "poll" in view:
        view["poll"] = f"{BASE_PATH}/v1/videos/{row['id']}"
    if "download" in view:
        view["download"] = f"{BASE_PATH}/v1/videos/{row['id']}/file"
    return view


# create_video 首次提交、get_video 轮询时顺手提交、reconciler_loop 每 5 秒扫描,三条路都会在
# gateway_id 还是 NULL 的窗口里调 _try_submit(cid)。三个协程共享同一个事件循环,`await
# client.post(...)` 让出控制权的那段时间里,另外两条路完全可能也读到「还没提交过」,
# 于是同一个 cid 被 POST 给 gateway 两次——两张卡各烧一遍 GPU,console 只记得住最后一次
# UPDATE 写回的 gateway_id,另一条产物是真实存在但没人知道、也没人能停的孤儿任务。
# 2026-08-20 实测踩到过:一个 15s/832×480 的探测任务真的在两张卡上各 OOM 了一次。
# 修法:进程内占坑(不是数据库锁——server.py 单进程单事件循环,内存里一个 set 就够,
# 两行判断中间没有 await,天然原子,不会被其它协程插进来)。
_SUBMITTING: set[str] = set()


async def _try_submit(cid: str):
    """把一条本地排队的任务尝试提交给 gateway。给 create_video 首次调用和 reconciler_loop 复用。"""
    if cid in _SUBMITTING:
        return  # 已经有一次提交在飞行中,等它写完 gateway_id 再说,不要并发再发一次
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (cid,)).fetchone()
    if row is None or row["gateway_id"] is not None:
        return  # 任务不存在,或已经提交过
    _SUBMITTING.add(cid)
    try:
        req = json.loads(row["req_json"])
        try:
            resp = await client.post(f"{GATEWAY_URL}/v1/videos", json=req)
        except httpx.HTTPError:
            return  # gateway 暂时连不上,下一轮再试
        now = time.time()
        if resp.status_code < 400:
            job = resp.json()
            gateway_id = job["id"]
            job["id"] = cid  # 对外只暴露 console 自己的 id
            with db() as con:
                con.execute(
                    "UPDATE jobs SET gateway_id=?, backend=?, status=?, updated_at=?, last_known_json=? WHERE id=?",
                    (gateway_id, job.get("backend"), job.get("status", "queued"),
                     now, json.dumps(job, ensure_ascii=False), cid),
                )
        elif resp.status_code in (503, 429):
            return  # 没有就绪后端 / 队列满,留在本地队列,等 flip 或别的任务跑完再试
        else:
            # 真正的客户端错误(比如 422 参数越界、400 turbo+steps>8),重试也不会成功,直接判失败
            err_view = {"id": cid, "status": "failed", "error": resp.text,
                        "req": req, "partition": row["partition"], "created_at": row["created_at"]}
            with db() as con:
                con.execute("UPDATE jobs SET status='failed', updated_at=?, last_known_json=? WHERE id=?",
                            (now, json.dumps(err_view, ensure_ascii=False), cid))
    finally:
        _SUBMITTING.discard(cid)


async def _refresh_from_gateway(cid: str, gateway_id: str):
    """已经提交给 gateway 的任务,去问一下最新状态,刷新本地记录。"""
    try:
        resp = await client.get(f"{GATEWAY_URL}/v1/videos/{gateway_id}")
    except httpx.HTTPError:
        return
    now = time.time()
    if resp.status_code == 404:
        # gateway 完全不认识这个 gateway_id——通常是所在后端进程被重启(比如手动 flip 分区
        # 或崩溃),内存态清空,这条任务已经不会再有任何进展。不能像别的非 200 一样悄悄跳过,
        # 那样 status 会永远卡在 queued/running,界面看着像在跑,实际已经死了、永远不会再更新。
        with db() as con:
            row = con.execute("SELECT last_known_json FROM jobs WHERE id=?", (cid,)).fetchone()
        if row is None:
            return
        view = json.loads(row["last_known_json"]) if row["last_known_json"] else {"id": cid}
        if view.get("status") in ("done", "failed"):
            return  # 已经是终态,不要覆盖
        view["status"] = "failed"
        view["error"] = "任务在执行途中丢失:所在后端进程被重启,内存态被清空,不会再有进展"
        with db() as con:
            con.execute(
                "UPDATE jobs SET status='failed', updated_at=?, last_known_json=? WHERE id=?",
                (now, json.dumps(view, ensure_ascii=False), cid),
            )
        return
    if resp.status_code != 200:
        return  # 其它非 200(后端一时不可达等)多半是暂时的,保留原状态,下一轮再试
    job = resp.json()
    job["id"] = cid
    with db() as con:
        con.execute(
            "UPDATE jobs SET status=?, backend=?, updated_at=?, last_known_json=? WHERE id=?",
            (job.get("status"), job.get("backend"), now, json.dumps(job, ensure_ascii=False), cid),
        )


@router.post("/v1/videos", status_code=202)
async def create_video(req: dict):
    partition = _partition_of(req)
    cid = uuid.uuid4().hex[:12]
    now = time.time()
    synth = {"id": cid, "status": "queued", "req": req, "partition": partition, "created_at": now}
    with db() as con:
        con.execute(
            "INSERT INTO jobs (id, gateway_id, partition, backend, status, req_json, created_at, updated_at, last_known_json) "
            "VALUES (?,NULL,?,NULL,'queued',?,?,?,?)",
            (cid, partition, json.dumps(req, ensure_ascii=False), now, now, json.dumps(synth, ensure_ascii=False)),
        )
    await _try_submit(cid)  # 尽力立即提交一次;提交不上就留给 reconciler_loop
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (cid,)).fetchone()
    return _row_view(row)


@router.get("/v1/videos/{jid}")
async def get_video(jid: str):
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if row["gateway_id"]:
        await _refresh_from_gateway(jid, row["gateway_id"])
    else:
        await _try_submit(jid)  # 还在本地排队时,顺手趁轮询这一下也试着提交
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    return _row_view(row)


@router.get("/v1/videos/{jid}/file")
async def download_video(jid: str, request: Request):
    with db() as con:
        row = con.execute(
            "SELECT gateway_id, backend, status, last_known_json FROM jobs WHERE id=?", (jid,)
        ).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if not row["gateway_id"] or not row["backend"]:
        raise HTTPException(409, "任务还在本地排队,尚未分配到可用后端")

    # 不经过 gateway.py 转发:它自己那层 `StreamingResponse(urlopen(...))` 不认 Range 请求,
    # 探测 Range 一律回 200 而不是 206,还不带 Content-Length/Accept-Ranges——`<video>`
    # 标签播放前先探一次 Range 拿时长/可否拖动,探不到就一直卡在 readyState=0 播不出来
    # (下载能成功是因为普通 GET 不需要这些,只有网页里的播放器会被这个坑绊住)。
    # api.py 自己那层是 Starlette `FileResponse`,原生支持 Range;console 既然已经从
    # gateway 的响应里存了 `backend`(真实拥有这条任务的那个 GPU 后端),直接绕过 gateway
    # 这层转发、把 Range 头透传给 api.py,把它的状态码/头原样带回去,问题在这一层就解决,
    # 不用碰 gateway.py/api.py 一行代码。
    # 下载/播放到本地的文件名——不带这个的话浏览器只能从 URL 最后一段猜,我们的 URL 都是
    # 以 "/file" 结尾(没有扩展名),猜出来清一色叫 "file.mp4",多个任务的视频存到同一个
    # 下载目录会互相覆盖。用 console 自己的任务 id 拼一个每条任务都不同的文件名。
    download_filename = f"h3-{jid}.mp4"
    content_disposition = f'attachment; filename="{download_filename}"'

    upstream_url = f"{row['backend']}/v1/videos/{row['gateway_id']}/file"
    headers = {}
    if "range" in request.headers:
        headers["range"] = request.headers["range"]
    try:
        req = client.build_request("GET", upstream_url, headers=headers)
        upstream = await client.send(req, stream=True)
    except httpx.HTTPError:
        upstream = None

    if upstream is None or upstream.status_code >= 400:
        if upstream is not None:
            await upstream.aclose()
        # 后端进程重启会把 api.py 内存里的 JOBS 字典清空(这是它本来的设计,任务状态
        # 不落盘),导致它忘了这条任务——但视频文件本身在磁盘上没被动过。console 和
        # api.py 跑在同一台机器上,任务完成时已经把 `file` 路径存进了自己的 sqlite,
        # 后端忘了就直接读文件兜底,不用等它想起来(它也不会想起来,重启就是清空)。
        cached = json.loads(row["last_known_json"]) if row["last_known_json"] else {}
        file_path = cached.get("file")
        if file_path and os.path.exists(file_path):
            return FileResponse(file_path, media_type="video/mp4",
                                 headers={"Content-Disposition": content_disposition})
        raise HTTPException(404, "后端进程重启后忘了这条任务的记录,磁盘上也没找到对应文件")

    async def stream():
        async for chunk in upstream.aiter_bytes():
            yield chunk
        await upstream.aclose()

    passthrough = {k: v for k, v in upstream.headers.items()
                   if k.lower() in ("content-length", "content-range", "accept-ranges", "content-type")}
    passthrough["content-disposition"] = content_disposition
    return StreamingResponse(stream(), status_code=upstream.status_code, headers=passthrough)


@router.get("/v1/videos")
async def list_videos(limit: int = 50, trash: bool = False):
    # console 自己的 sqlite 是任务列表的唯一真相来源(本地排队中的任务,gateway 根本不知道它存在,
    # 没法从 gateway 反查出来)。实时进度由后台 reconciler_loop 持续刷新进 last_known_json,
    # 这里直接读,不在请求路径里再打一遍 gateway,避免列表接口被拖慢。
    # 软删除:deleted_at 非空的任务默认不出现在列表里,?trash=true 反过来只看这些——
    # 数据还在同一张表里,删除/恢复就是 UPDATE 一个时间戳,不用挪表、不用真删文件。
    where = "deleted_at IS NOT NULL" if trash else "deleted_at IS NULL"
    order = "deleted_at DESC" if trash else "created_at DESC"
    with db() as con:
        rows = con.execute(f"SELECT * FROM jobs WHERE {where} ORDER BY {order} LIMIT ?", (limit,)).fetchall()
    return {"jobs": [_row_view(r) for r in rows]}


@router.delete("/v1/videos/{jid}")
async def delete_video(jid: str):
    """软删除:只挪出主列表,不碰 gateway/api.py 上的任务本身(可能还在跑),不删磁盘文件。"""
    with db() as con:
        cur = con.execute("UPDATE jobs SET deleted_at=? WHERE id=? AND deleted_at IS NULL", (time.time(), jid))
    if cur.rowcount == 0:
        raise HTTPException(404, "无此任务,或已经在回收站里")
    return {"id": jid, "deleted": True}


@router.post("/v1/videos/{jid}/restore")
async def restore_video(jid: str):
    with db() as con:
        cur = con.execute("UPDATE jobs SET deleted_at=NULL WHERE id=? AND deleted_at IS NOT NULL", (jid,))
    if cur.rowcount == 0:
        raise HTTPException(404, "无此任务,或不在回收站里")
    return {"id": jid, "deleted": False}


@router.get("/v1/backends")
async def backends():
    resp = await client.get(f"{GATEWAY_URL}/health")
    resp.raise_for_status()
    return resp.json()


@router.get("/v1/models")
def models():
    # 静态信息,对应服务器上 api.py 里的 PARTITIONS 常量。只读展示,不做运行时切换。
    return {
        "partitions": {
            "fl2va": {"dit": "minimax-h3-fl2va-nf4.safetensors",
                      "lora": "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors",
                      "desc": "首末帧条件生成"},
            "ref2va": {"dit": "minimax-h3-ref2va-nf4.safetensors",
                       "lora": "minimax_h3_ref2v_turbo_4step_v0.1_comfyui_bf16.safetensors",
                       "desc": "文生视频 / 多参考图生成"},
        },
        "presets": [
            {"label": "标准档 832×480×124f(5.17s)", "width": 832, "height": 480, "frames": 124},
            {"label": "长时长档 448×256×362f(15s)", "width": 448, "height": 256, "frames": 362},
            {"label": "高分辨率档 1344×768×64f(2.7s)", "width": 1344, "height": 768, "frames": 64},
        ],
    }


# ---------------- 弹性 partition 调度器(仅服务器部署时启用) ----------------
def _pending_counts() -> dict:
    with db() as con:
        rows = con.execute(
            "SELECT partition, COUNT(*) c FROM jobs WHERE status IN ('queued','running') GROUP BY partition"
        ).fetchall()
    return {r["partition"]: r["c"] for r in rows}


async def _flip_backend(gpu: int, port: int, new_partition: str):
    # 和今天手动测试时用的手法一致:按端口查 PID、杀、用现成的 api-run.sh 重启到新分区。
    # 不新写进程管理代码,复用已经在跑的脚本。
    # 注意:ss 的输出里 pid 不是独立的空格分隔 token,是嵌在 `users:(("python",pid=NNNN,fd=14))`
    # 里面的——按空格 split 后拿 `tok.startswith("pid=")` 永远不会命中,得用正则从整段里抠出来
    # (手动测试时用的是 `grep -o "pid=[0-9]*"`,这里必须做等价的事,不能简化成 split)。
    pid_out = subprocess.run(["ss", "-lptn", f"sport = :{port}"], capture_output=True, text=True).stdout
    for pid in re.findall(r"pid=(\d+)", pid_out):
        subprocess.run(["kill", pid])
    subprocess.Popen(
        ["setsid", "nohup", "./api-run.sh", str(gpu), str(port), new_partition],
        cwd=str(H3_ROOT), stdout=open(H3_ROOT / "logs" / f"api-{port}.log", "a"),
        stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, start_new_session=True,
    )


async def supervisor_loop():
    idle_since: dict[int, float] = {}
    # 一个 flip 从 kill 到新进程报 ready 要 13-20 秒(api.py 冷启动实测),远长于 10 秒的 tick 间隔。
    # 如果不记「已经在切换中」,接下来两三轮 tick 会看到 fl2va_ready 还是 0(新进程没起来),
    # 误判成需求仍未满足,再多切几张卡——今天调试时真实撞到过:一个 pending=1 的请求把三张卡都切了
    # 过去(而且当时 kill 那步还有个 bug 没真的杀掉旧进程,新进程绑端口失败悄悄退出,见下)。
    flipping: dict[int, tuple[str, float]] = {}  # port -> (目标 partition, 发起时间)
    while True:
        await asyncio.sleep(10)
        try:
            health = (await client.get(f"{GATEWAY_URL}/health")).json()["backends"]
            pending = _pending_counts()
            now = time.time()

            for port in list(flipping):
                target, started = flipping[port]
                b = next((x for x in health if x["backend"].endswith(f":{port}")), None)
                if b and b.get("partition") == target and b.get("status") == "ok":
                    flipping.pop(port)  # 切换完成,新分区已经在 health 里报 ready 了
                elif now - started > 60:
                    flipping.pop(port)  # 60 秒还没起来,当作这次尝试失败,让下一轮按实际情况重新决定

            fl2va_ready = [b for b in health if b.get("partition") == "fl2va" and b.get("status") == "ok"]
            ref2va_ready = [b for b in health if b.get("partition") == "ref2va" and b.get("status") == "ok"]
            fl2va_inflight = sum(1 for t, _ in flipping.values() if t == "fl2va")
            ref2va_inflight = sum(1 for t, _ in flipping.values() if t == "ref2va")

            def idle_backends(group):
                out = []
                for b in group:
                    port = int(b["backend"].rsplit(":", 1)[1])
                    if port in flipping:
                        continue  # 已经在切换中,不能再选它
                    if b.get("queued", 0) == 0 and b.get("running", 0) == 0:
                        idle_since.setdefault(port, now)
                        if now - idle_since[port] >= 30:
                            out.append(port)
                    else:
                        idle_since.pop(port, None)
                return out

            fl2va_pending = pending.get("fl2va", 0)
            ref2va_pending = pending.get("ref2va", 0)
            print(f"[supervisor] tick: pending={pending} "
                  f"fl2va={len(fl2va_ready)}(+{fl2va_inflight} 切换中) "
                  f"ref2va={len(ref2va_ready)}(+{ref2va_inflight} 切换中)", flush=True)

            if fl2va_pending > len(fl2va_ready) + fl2va_inflight:
                candidates = idle_backends(ref2va_ready)
                if candidates:
                    port = candidates[0]
                    gpu = next(t["gpu"] for t in BACKEND_TOPOLOGY if t["port"] == port)
                    print(f"[supervisor] flip port={port} gpu={gpu} ref2va -> fl2va", flush=True)
                    flipping[port] = ("fl2va", now)
                    idle_since.pop(port, None)
                    await _flip_backend(gpu, port, "fl2va")
            elif ref2va_pending > len(ref2va_ready) + ref2va_inflight:
                candidates = idle_backends(fl2va_ready)
                if candidates:
                    port = candidates[0]
                    gpu = next(t["gpu"] for t in BACKEND_TOPOLOGY if t["port"] == port)
                    print(f"[supervisor] flip port={port} gpu={gpu} fl2va -> ref2va", flush=True)
                    flipping[port] = ("ref2va", now)
                    idle_since.pop(port, None)
                    await _flip_backend(gpu, port, "ref2va")
        except Exception as e:
            print(f"[supervisor] tick failed: {e}", flush=True)


# ---------------- 常驻对账循环(本地/服务器都跑,跟 GPU 无关) ----------------
# 两件事:本地排队中的任务(gateway_id 还是空)隔几秒重试提交一次;已提交的任务刷新最新状态。
# 这是任务列表能实时更新、以及「没后端时先排队而不是直接报错」这两件事真正生效的地方。
async def reconciler_loop():
    while True:
        await asyncio.sleep(5)
        try:
            with db() as con:
                pending = con.execute(
                    "SELECT id FROM jobs WHERE gateway_id IS NULL AND status='queued'"
                ).fetchall()
                in_flight = con.execute(
                    "SELECT id, gateway_id FROM jobs WHERE gateway_id IS NOT NULL AND status IN ('queued','running')"
                ).fetchall()
            for r in pending:
                await _try_submit(r["id"])
            for r in in_flight:
                await _refresh_from_gateway(r["id"], r["gateway_id"])
        except Exception as e:
            print(f"[reconciler] tick failed: {e}", flush=True)


@app.on_event("startup")
async def _startup():
    asyncio.create_task(reconciler_loop())
    if ENABLE_SUPERVISOR:
        asyncio.create_task(supervisor_loop())
        print("[console] elastic partition supervisor enabled", flush=True)
    else:
        print("[console] supervisor disabled (H3_ENABLE_SUPERVISOR!=1) — fl2va/ref2va allocation is whatever the backends currently report", flush=True)


# 任务详情/新建各自有一个真实的 URL,方便复制粘贴分享、刷新不丢状态——都是同一个单页应用
# (static/index.html 自己那份 JS 读 location.pathname 决定显示哪个视图),这三条路由只是让
# 浏览器地址栏里出现符合预期的路径,不是三个不同的页面文件。
#
# 页面不再直接 FileResponse 原文件,而是读进来做一次字符串替换:把 __H3_BASE_PATH__ 占位符
# 换成 BASE_PATH。部署在共享域名子路径下时(比如 qwen.wudaima.com/console),前端 JS 里所有
# fetch/history.pushState 用的绝对路径都得带上这个前缀,否则请求会打到域名根路径而不是我们
# 这个子路径——BASE_PATH 留空时替换成空字符串,跟原来的行为完全一样。
def render_index() -> HTMLResponse:
    html = (BASE / "static" / "index.html").read_text(encoding="utf-8")
    return HTMLResponse(html.replace("__H3_BASE_PATH__", BASE_PATH))


@router.get("/")
async def index_page():
    return render_index()


@router.get("/task.html/{tid}")
async def task_page(tid: str):
    return render_index()


@router.get("/newtask.html")
async def newtask_page():
    return render_index()


# BASE_PATH 留空时这个前缀是空字符串,router 上的路径原样挂到 app 根上,行为跟以前
# 直接用 @app.xxx 一模一样;非空时(比如 "/console")整体套一层前缀,一次性生效,
# 不用每个路由自己拼。
app.include_router(router, prefix=BASE_PATH)