139 lines
4.6 KiB
Python
Executable File
139 lines
4.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""本地 GPU 推理统一入口 v3 — LLM (llama-server Vulkan 4B) + 向量嵌入
|
||
用法:
|
||
gpu_infer.py llm "问题" [max_tokens] # LLM 对话（4B 常驻服务 llama-server-4b）
|
||
gpu_infer.py embed "文本" # 向量嵌入（bge-embed 服务）
|
||
gpu_infer.py health # 健康检查
|
||
gpu_infer.py bench # 性能基准
|
||
|
||
2026-09-07 牧尘铁律: 3B/7B 禁止使用。本机只部署 llama-server-4b (Qwen3.5-4B Q4_K_M,
|
||
GPU -ngl 99: VRAM ~3.1GB + host ~480MB)。3B 曾 CPU 推理 RSS 7.6GB 内存黑洞。
|
||
"""
|
||
import os
|
||
import sys
|
||
import time
|
||
import json
|
||
import subprocess
|
||
import urllib.request
|
||
|
||
LLAMA_DIR = "/tmp/llama-vulkan/llama-b10679"
|
||
EMBED_URL = "http://localhost:8000/v1/embeddings"
|
||
LLM_URL = "http://127.0.0.1:8080/v1/chat/completions"
|
||
LLAMA_SVC = "llama-server-4b"
|
||
|
||
|
||
def llm_chat(prompt, max_tokens=100, model="4b"):
|
||
"""LLM 对话（走 llama-server OpenAI 兼容 API，4B 常驻服务）"""
|
||
data = json.dumps({
|
||
"model": "qwen" + model,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"max_tokens": max_tokens,
|
||
"temperature": 0.7,
|
||
}).encode()
|
||
req = urllib.request.Request(LLM_URL, data=data, headers={"Content-Type": "application/json"})
|
||
with urllib.request.urlopen(req, timeout=300) as resp:
|
||
result = json.loads(resp.read())
|
||
return result["choices"][0]["message"]["content"] if result.get("choices") else None
|
||
|
||
|
||
def embed(text):
|
||
"""向量嵌入（bge-embed 服务）"""
|
||
if isinstance(text, str):
|
||
text = [text]
|
||
data = json.dumps({"input": text}).encode()
|
||
req = urllib.request.Request(EMBED_URL, data=data, headers={"Content-Type": "application/json"})
|
||
with urllib.request.urlopen(req, timeout=60) as resp:
|
||
result = json.loads(resp.read())
|
||
return [d["embedding"] for d in result["data"]]
|
||
|
||
|
||
def health():
|
||
"""健康检查"""
|
||
print("=== GPU 健康检查 ===")
|
||
r = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.used,memory.total,utilization.gpu", "--format=csv,noheader"], capture_output=True, text=True)
|
||
print(f"GPU: {r.stdout.strip()}")
|
||
|
||
# bge-embed
|
||
try:
|
||
req = urllib.request.Request("http://localhost:8000/health")
|
||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||
d = json.loads(resp.read())
|
||
print(f"bge-embed: ✅ {d['providers']}")
|
||
except Exception as e:
|
||
print(f"bge-embed: ❌ {e}")
|
||
|
||
# llama-server-4b
|
||
try:
|
||
req = urllib.request.Request("http://127.0.0.1:8080/health")
|
||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||
d = json.loads(resp.read())
|
||
cur = "4B"
|
||
print(f"llama-server-4b {cur}: ✅ {d.get('status', 'ok')}")
|
||
except Exception as e:
|
||
print(f"llama-server-4b: ❌ {e}")
|
||
|
||
# 4B 服务状态
|
||
r = subprocess.run(["systemctl", "--user", "is-active", LLAMA_SVC], capture_output=True, text=True)
|
||
print(f"llama-server-4b service: {r.stdout.strip()}")
|
||
|
||
|
||
def bench():
|
||
"""性能基准"""
|
||
print("=== GPU 推理性能基准 ===\n")
|
||
print("--- LLM 4B (llama-server-4b) ---")
|
||
prompts = ["你好", "用一句话解释什么是过拟合"]
|
||
for p in prompts:
|
||
llm_chat(p, max_tokens=20)
|
||
start = time.time()
|
||
result = llm_chat(p, max_tokens=50)
|
||
elapsed = time.time() - start
|
||
print(f" {p[:15]}: {elapsed:.2f}s -> {result[:40]}...")
|
||
|
||
print("\n--- 向量嵌入 (bge-m3 INT8) ---")
|
||
start = time.time()
|
||
emb = embed("测试向量嵌入性能")
|
||
elapsed = time.time() - start
|
||
print(f" 维度 {len(emb[0])}: {elapsed:.3f}s")
|
||
|
||
|
||
def main():
|
||
if len(sys.argv) < 2:
|
||
print(__doc__)
|
||
return
|
||
cmd = sys.argv[1]
|
||
|
||
if cmd == "llm":
|
||
if len(sys.argv) < 3:
|
||
print("用法: gpu_infer.py llm \"问题\" [max_tokens]")
|
||
return
|
||
prompt = sys.argv[2]
|
||
max_tokens = int(sys.argv[3]) if len(sys.argv) > 3 else 100
|
||
start = time.time()
|
||
result = llm_chat(prompt, max_tokens)
|
||
elapsed = time.time() - start
|
||
if result:
|
||
print(f"输出: {result}")
|
||
print(f"耗时: {elapsed:.2f}s")
|
||
|
||
elif cmd == "embed":
|
||
if len(sys.argv) < 3:
|
||
print("用法: gpu_infer.py embed \"文本\"")
|
||
return
|
||
start = time.time()
|
||
result = embed(sys.argv[2])
|
||
elapsed = time.time() - start
|
||
print(f"维度: {len(result[0])}, 耗时: {elapsed:.3f}s")
|
||
|
||
elif cmd == "health":
|
||
health()
|
||
|
||
elif cmd == "bench":
|
||
bench()
|
||
|
||
else:
|
||
print(__doc__)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|