oneclick_gsqrco_serve: one-command GSQ-RCO serving (tunnel-aware)

oneclick_gsqrco_serve.py

@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""oneclick_gsqrco_serve.py — SAMAI-9B 最优量化模型一键 serving（GSQ-RCO-lite, R14）
+
+用法（Kaggle/Colab T4 容器内）:
+    HF_TOKEN=hf_xxx python3 oneclick_gsqrco_serve.py
+    python3 oneclick_gsqrco_serve.py --token hf_xxx
+
+自动完成:
+  1. HF 拉权重: m13_r10_gsqrco_v1.gguf(6.48G, ppl 7.99 vs 旧 mixbit 8.11)
+              + mmproj_m11.gguf(921M) + llama.cpp CUDA tar(sm_75)
+  2. llama-server: 端口 8101, key=1234, ctx 梯子 131072(双卡)/65536(单卡)→32768,
+                   -fa on -ctk q8_0 -ctv q8_0 -t 2 -ngl 99
+  3. 公网隧道（二选一，自动识别）:
+     a. 容器有 kaggle_server.py → 打 /v1 代理补丁后重启它（隧道原生出 OpenAI API,
+        单隧道双功能; 坑#67 同机单身份/#68 监督环——勿抢端口, 补丁本体）
+     b. 无 → pip 装 aitun 并起免 token 隧道 -p 8101, 解析打印 aitun.cc URL
+  4. 健康检查 + 打印调用示例
+
+日志: /tmp/k8b/oneclick.log   旗标: ONECLICK_DONE / ONECLICK_FAIL
+"""
+import json, os, re, subprocess, sys, time, traceback
+
+K8B, LCF = "/tmp/k8b", "/tmp/lcfin"
+REPO = "tchbcb/samai-9b"
+GGUF = K8B + "/m13_r10_gsqrco_v1.gguf"
+MMP = K8B + "/mmproj_m11.gguf"
+LOG = K8B + "/oneclick.log"
+os.makedirs(K8B, exist_ok=True)
+os.makedirs(LCF, exist_ok=True)
+log = open(LOG, "a", buffering=1)
+
+
+def P(m):
+    log.write("[%s] %s\n" % (time.strftime("%m-%d %H:%M:%S"), m))
+
+
+def sh(c, t=3600):
+    p = subprocess.run(c, shell=True, capture_output=True, text=True,
+                       timeout=t, errors="replace")
+    return (p.stdout or "") + (p.stderr or "")
+
+
+def get_token(argv):
+    if "--token" in argv:
+        return argv[argv.index("--token") + 1]
+    if os.environ.get("HF_TOKEN"):
+        return os.environ["HF_TOKEN"]
+    p = os.path.expanduser("~/.cache/huggingface/token")
+    if os.path.exists(p):
+        return open(p).read().strip()
+    raise SystemExit("缺少 HF token: python3 oneclick_gsqrco_serve.py --token hf_xxx")
+
+
+def hfget(tok, fn, dst, minb):
+    if os.path.exists(dst) and os.path.getsize(dst) >= minb:
+        P("have %s (%.1fG)" % (fn, os.path.getsize(dst) / 1e9))
+        return
+    from huggingface_hub import hf_hub_download
+    P("pull %s ..." % fn)
+    t0 = time.time()
+    p = hf_hub_download(REPO, fn, token=tok)
+    sh("cp -f '%s' '%s'" % (p, dst), 1800)
+    assert os.path.getsize(dst) >= minb, "pull short: " + fn
+    P("pulled %s %.1fG %.0fs" % (fn, os.path.getsize(dst) / 1e9, time.time() - t0))
+
+
+def lc_setup(tok):
+    b = LCF + "/llama-server"
+    for tar_name in ("llama_cpp_cuda_colab_t4.tar.gz", "llama_cpp_cuda_t4_v2.tar.gz"):
+        if os.path.exists(b) and os.path.getsize(b) > 5e6:
+            v = sh("LD_LIBRARY_PATH=%s %s --version 2>&1 | head -1" % (LCF, b), 60)
+            if "version" in v.lower() or "build" in v.lower():
+                P("lcfin ok: %s" % v[-90:])
+                return True
+        TAR = K8B + "/lc_tar.tar.gz"
+        try:
+            hfget(tok, tar_name, TAR, 3e7)
+        except Exception as e:
+            P("tar %s fail: %s" % (tar_name, repr(e)[:100]))
+            continue
+        sh("rm -rf %s/* && tar xzf %s -C %s" % (LCF, TAR, LCF), 900)
+        sh('LS=$(find %s -name llama-server -type f | head -1); [ -n "$LS" ] && '
+           'cp -f "$LS" %s/; find %s -name "*.so*" -type f -exec cp -f {} %s/ \\; '
+           '2>/dev/null; rm -f %s/libcuda.so* %s/libnvidia-*'
+           % (LCF, LCF, LCF, LCF, LCF, LCF), 300)
+        for cand in ("/usr/lib64-nvidia/libcuda.so.1",
+                     "/usr/lib/x86_64-linux-gnu/libcuda.so.1",
+                     "/usr/local/nvidia/lib64/libcuda.so.1"):
+            if os.path.exists(cand):
+                sh("ln -sf %s %s/libcuda.so.1" % (cand, LCF), 30)
+                break
+    return os.path.exists(b) and os.path.getsize(b) > 5e6
+
+
+def n_gpu():
+    try:
+        return len([l for l in sh("nvidia-smi -L").strip().splitlines() if "GPU" in l])
+    except Exception:
+        return 1
+
+
+def serve():
+    sh("pkill -f llama-server; sleep 2", 60)
+    ng = n_gpu()
+    env = ("LD_LIBRARY_PATH=%s:/usr/local/nvidia/lib64:/usr/lib64-nvidia:"
+           "/usr/local/cuda/lib64" % LCF)
+    import urllib.request
+    for c in (131072 if ng >= 2 else 65536, 32768):
+        cmd = ("%s %s/llama-server -m %s --mmproj %s --host 0.0.0.0 --port 8101 "
+               "--api-key 1234 -c %d -fa on -ctk q8_0 -ctv q8_0 -t 2 -ngl 99"
+               % (env, LCF, GGUF, MMP, c))
+        P("fire ctx=%d (ngpu=%d)" % (c, ng))
+        subprocess.Popen(["bash", "-c", cmd + " > /tmp/k8b/serve_oneclick.log 2>&1"],
+                         stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT,
+                         stdin=subprocess.DEVNULL, start_new_session=True)
+        for i in range(66):
+            time.sleep(5)
+            try:
+                h = urllib.request.urlopen("http://127.0.0.1:8101/health",
+                                           timeout=4).read().decode()
+                if '"ok"' in h:
+                    P("SERVE_OK ctx=%d" % c)
+                    return True
+            except Exception:
+                pass
+            if i in (12, 30):
+                t = sh("tail -3 /tmp/k8b/serve_oneclick.log")[-300:].lower()
+                if ("out of memory" in t) or ("cuda error" in t):
+                    break
+        sh("pkill -f llama-server; sleep 2", 60)
+    return False
+
+
+# ---- /v1 代理补丁文本（kaggle_server 模式; 坑#67/#68）----
+V1_PATCH = '''
+# ==== R14 BEGIN: /v1 proxy -> llama-server 8101 (oneclick) ====
+import requests as _r14rq
+_R14_LLM = "http://127.0.0.1:8101"
+
+@app.route('/v1', methods=['GET', 'POST', 'OPTIONS'])
+@app.route('/v1/<path:_r14p>', methods=['GET', 'POST', 'OPTIONS', 'PUT', 'PATCH', 'DELETE'])
+def _r14_v1(_r14p=''):
+    _up = _R14_LLM + '/v1' + ('' if _r14p == '' else '/' + _r14p)
+    _h = {k: v for k, v in request.headers.items()
+          if k.lower() not in ('host', 'connection', 'accept-encoding', 'content-length')}
+    try:
+        _r = _r14rq.request(request.method, _up, params=request.args,
+                            data=request.get_data(cache=False), headers=_h,
+                            stream=True, timeout=(10, 600))
+    except Exception as _e:
+        return jsonify(error='llm upstream unreachable: %s' % repr(_e)[:120]), 502
+    _hd = {k: v for k, v in _r.headers.items()
+           if k.lower() not in ('content-length', 'transfer-encoding', 'connection',
+                                'content-encoding', 'keep-alive')}
+    if request.method == 'OPTIONS':
+        _r.close()
+        return ('', _r.status_code, _hd)
+    _ct = _r.headers.get('Content-Type', '')
+    if 'text/event-stream' in _ct or _r.headers.get('Transfer-Encoding', '').lower() == 'chunked':
+        def _r14gen():
+            try:
+                for _c in _r.iter_content(chunk_size=4096):
+                    if _c:
+                        yield _c
+            finally:
+                _r.close()
+        return Response(_r14gen(), status=_r.status_code, headers=_hd)
+    _body = _r.content
+    _r.close()
+    return Response(_body, status=_r.status_code, headers=_hd)
+# ==== R14 END: /v1 proxy ====
+'''
+
+
+def find_ks():
+    for c in ("/kaggle/working/kaggle_server.py",
+              os.path.join(os.getcwd(), "kaggle_server.py")):
+        if os.path.exists(c):
+            return c
+    return None
+
+
+def patch_kaggle_server(ks):
+    src = open(ks).read()
+    if "R14 BEGIN" in src:
+        P("kaggle_server already patched")
+    else:
+        m = re.search(r"^(\w+)\s*=\s*Flask\(", src, re.M)
+        assert m, "Flask app object not found"
+        anchor = "if __name__ == '__main__':"
+        assert anchor in src, "main anchor not found"
+        open(ks + ".bak_r14", "w").write(src)
+        src = src.replace(anchor, V1_PATCH + "\n\n" + anchor, 1)
+        compile(src, ks, "exec")
+        open(ks, "w").write(src)
+        P("kaggle_server patched (backup .bak_r14)")
+    return True
+
+
+def restart_ksguarded(ks):
+    """杀 5000 持有者; 监督环优先重生, 8s 无重生则兜底自启（坑#68）"""
+    d = os.path.dirname(ks)
+    sh("KPID=$(ss -tlnp | grep ':5000 ' | grep -oP 'pid=\\K[0-9]+' | head -1); "
+       "[ -n \"$KPID\" ] && kill $KPID; sleep 6", 40)
+    for _ in range(4):
+        if ":5000 " in sh("ss -tlnp", 15):
+            break
+        time.sleep(3)
+        if not (":5000 " in sh("ss -tlnp", 15)):
+            subprocess.Popen(["bash", "-c",
+                              "cd %s && nohup python3 kaggle_server.py "
+                              ">> /tmp/k8b/oneclick_ks.log 2>&1 &" % d],
+                             stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT,
+                             stdin=subprocess.DEVNULL, start_new_session=True)
+            time.sleep(5)
+    P("kaggle_server restarted: %s" % (":5000 " in sh("ss -tlnp", 15)))
+
+
+def tunnel_url_from(path):
+    try:
+        m = re.search(r"Proxy URL: (https?://\S+?)/?\s", open(path).read() + "\n")
+        return m.group(1) if m else None
+    except Exception:
+        return None
+
+
+def main():
+    t0 = time.time()
+    P("==== oneclick start ====")
+    try:
+        tok = get_token(sys.argv)
+        hfhome = os.path.expanduser("~/.cache/huggingface")
+        os.makedirs(hfhome, exist_ok=True)
+        open(hfhome + "/token", "w").write(tok)
+        os.chmod(hfhome + "/token", 0o600)
+        hfget(tok, "m13_r10_gsqrco_v1.gguf", GGUF, 5e9)
+        hfget(tok, "mmproj_m11.gguf", MMP, 5e8)
+        if not lc_setup(tok):
+            raise RuntimeError("llama-server unavailable")
+        if not serve():
+            raise RuntimeError("serve failed at all ctx")
+        ks = find_ks()
+        pub = None
+        if ks:
+            patch_kaggle_server(ks)
+            restart_ksguarded(ks)
+            for lg in ("/tmp/k8b/aitun_main.log", "/tmp/k8b/aitun_8101.log"):
+                pub = tunnel_url_from(lg)
+                if pub:
+                    break
+            print("\n[MODE] kaggle_server /v1 ���理（单隧道双功能）")
+            print("[API] <你的隧道URL>/v1  （沿用现有 aitun 隧道码）")
+        else:
+            sh("which aitun || pip install -q aitun || pip install -q "
+               "git+https://github.com/ctz168/aitun.git", 600)
+            sh("pkill -f 'aitun-clien[t].* -p 8101'", 20)
+            subprocess.Popen(["bash", "-c",
+                              "nohup aitun -p 8101 > /tmp/k8b/aitun_api.log 2>&1 &"],
+                             stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT,
+                             stdin=subprocess.DEVNULL, start_new_session=True)
+            for _ in range(10):
+                time.sleep(3)
+                pub = tunnel_url_from("/tmp/k8b/aitun_api.log")
+                if pub:
+                    break
+            print("\n[MODE] aitun 直连隧道 (-p 8101)")
+            print("[API] %s/v1" % (pub or "https://aitun.cc/<隧道码>"))
+        open(K8B + "/ONECLICK_DONE", "w").write("%.0fs" % (time.time() - t0))
+        P("==== ONECLICK_DONE %.0fs ====" % (time.time() - t0))
+        print("[OK] serving 上线 127.0.0.1:8101  key=1234  ctx 见 serve_oneclick.log")
+        print("[调用示例]")
+        print("  curl '%s/v1/chat/completions' \\" % (pub or "https://aitun.cc/<隧道码>"))
+        print("    -H 'Authorization: Bearer 1234' -H 'Content-Type: application/json' \\")
+        print('    -d \'{"model":"m13_r10_gsqrco_v1","messages":'
+              '[{"role":"user","content":"你好"}],"max_tokens":400}\'')
+        print("[openai python]")
+        print("  from openai import OpenAI")
+        print("  client = OpenAI(base_url='%s/v1', api_key='1234')"
+              % (pub or "https://aitun.cc/<隧道码>"))
+        print("  r = client.chat.completions.create(model='m13_r10_gsqrco_v1', "
+              "messages=[{'role':'user','content':'你好'}])")
+        print("[本地] http://127.0.0.1:8101/v1")
+    except Exception as e:
+        traceback.print_exc(file=log)
+        open(K8B + "/ONECLICK_FAIL", "w").write(repr(e)[:300])
+        P("==== ONECLICK_FAIL: %s ====" % repr(e)[:200])
+        print("FAIL:", repr(e)[:200], "— 日志:", LOG)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()