api_client.py

"""七个空间共用的可恢复任务客户端；仅 URL 和请求 JSON 随空间变化。"""

import argparse
import json
import os
import time
import uuid
from pathlib import Path

import httpx


def main():
    """提交或恢复任务，轮询终态后通过 result_url 保存结果。

    Returns:
        下载成功后结束；状态文件保留幂等键和任务 ID，便于断线后继续。
    """
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", default="https://idlecloudx-moss-image-base.hf.space")
    parser.add_argument("--request", type=Path, required=True)
    parser.add_argument("--state", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--timeout", type=int, default=900)
    args = parser.parse_args()
    api_key = os.environ.get("JOB_API_KEY")
    if not api_key:
        raise SystemExit("请设置 JOB_API_KEY；HF_TOKEN 可选，游客无需设置。")
    payload = json.loads(args.request.read_text(encoding="utf-8-sig"))
    base_url = args.url.rstrip("/")
    if args.state.exists():
        state = json.loads(args.state.read_text(encoding="utf-8"))
        if state["payload"] != payload or state["base_url"] != base_url:
            raise SystemExit("状态文件属于另一请求，请为新任务使用新的 --state 路径。")
    else:
        state = {"base_url": base_url, "payload": payload, "idempotency_key": uuid.uuid4().hex}
        # 先保存幂等键，提交响应丢失时仍能安全恢复同一个任务。
        args.state.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8")
    headers = {"X-API-Key": api_key, "Idempotency-Key": state["idempotency_key"]}
    hf_token = os.environ.get("HF_TOKEN", "").strip()
    if hf_token:
        headers["Authorization"] = f"Bearer {hf_token}"
    deadline = time.monotonic() + args.timeout
    # 下载只使用 URL 自带的令牌，不把任务密钥或 HF_TOKEN 发送到结果地址。
    with httpx.Client(headers=headers, timeout=30) as client, httpx.Client(timeout=120) as downloader:
        while time.monotonic() < deadline:
            downloading = "result_url" in state
            try:
                if downloading:
                    response = downloader.get(state["result_url"])
                elif "job_id" not in state:
                    response = client.post(f"{base_url}/api/jobs", json=payload)
                else:
                    response = client.get(f"{base_url}/api/jobs/{state['job_id']}")
                is_json = "application/json" in response.headers.get("content-type", "")
                # 业务失败是终态；只对限流、重启或传输故障重试。
                if not downloading and is_json and response.status_code == 500:
                    failure = response.json()
                    if failure.get("status") == "failed":
                        raise SystemExit(json.dumps(failure["error"], ensure_ascii=False))
                if response.status_code in {408, 429, 500, 502, 503, 504}:
                    time.sleep(min(30, max(2, int(response.headers.get("Retry-After", "5")))))
                    continue
                response.raise_for_status()
                # HF 启动等待页可能以 HTTP 200 返回 HTML。
                if "text/html" in response.headers.get("content-type", ""):
                    time.sleep(5)
                    continue
            except httpx.TransportError:
                time.sleep(3)
                continue
            if downloading:
                args.output.write_bytes(response.content)
                print(f"已保存：{args.output.resolve()}")
                return
            body = response.json()
            if "job_id" not in state:
                state["job_id"] = body["job_id"]
                print(f"任务已提交：{state['job_id']}", flush=True)
            elif body["status"] in {"queued", "running"}:
                time.sleep(int(response.headers.get("Retry-After", "2")))
                continue
            elif body["status"] == "failed":
                raise SystemExit(json.dumps(body["error"], ensure_ascii=False))
            else:
                state["result_url"] = body["result_url"]
            args.state.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8")
    raise SystemExit("等待超时；用相同命令和状态文件继续，不会重复生成。")


if __name__ == "__main__":
    main()