feat: 基于 LLaMA-Factory 实现模型推理引擎
利用容器内已有的 LLaMA-Factory ChatModel (huggingface 后端) 实现真实 模型推理,无需额外安装 vLLM。 Compute 端新增: - compute/engines/llama_factory/inference.py InferenceSession: 模型加载/卸载/对话/流式输出 支持 base model 和 LoRA adapter,线程安全 - compute/api/main.py 新增 5 个推理端点: POST /inference/load - 加载模型 POST /inference/unload - 卸载释放 GPU 显存 GET /inference/status - 查询会话状态 POST /inference/chat - 非流式对话 POST /inference/chat/stream - SSE 流式对话 后端新增: - platform.py 推理代理端点(local/chat, local/chat/stream, local/preload, local/unload, local/status, trained/preload) - ComputeNodeClient._request 通用请求方法 - _select_first_online_node 自动选择在线节点 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -6,7 +6,9 @@ from pathlib import Path
|
|||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from fastapi import APIRouter, BackgroundTasks, Body, File, HTTPException, Query, UploadFile
|
from fastapi import APIRouter, BackgroundTasks, Body, File, HTTPException, Query, UploadFile
|
||||||
from fastapi.responses import PlainTextResponse
|
from fastapi.responses import PlainTextResponse, StreamingResponse
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
from app.core.config import get_settings
|
from app.core.config import get_settings
|
||||||
from app.db.platform_store import get_platform_store
|
from app.db.platform_store import get_platform_store
|
||||||
@@ -20,6 +22,15 @@ def ok(data: Any = None, message: str = "ok") -> dict[str, Any]:
|
|||||||
return {"code": 0, "message": message, "data": data}
|
return {"code": 0, "message": message, "data": data}
|
||||||
|
|
||||||
|
|
||||||
|
def _select_first_online_node(store: Any) -> dict[str, Any] | None:
|
||||||
|
"""Select the first online compute node for inference."""
|
||||||
|
nodes = store.compute_nodes()
|
||||||
|
for node in nodes:
|
||||||
|
if node.get("enabled") and node.get("scheduler_status") == "online":
|
||||||
|
return node
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def fail(status_code: int, message: str) -> HTTPException:
|
def fail(status_code: int, message: str) -> HTTPException:
|
||||||
return HTTPException(status_code=status_code, detail={"code": status_code, "message": message, "data": None})
|
return HTTPException(status_code=status_code, detail={"code": status_code, "message": message, "data": None})
|
||||||
|
|
||||||
@@ -1102,17 +1113,99 @@ async def model_chat_batch(payload: dict[str, Any] = Body(...)) -> dict[str, Any
|
|||||||
|
|
||||||
@router.post("/model-chat/local/chat")
|
@router.post("/model-chat/local/chat")
|
||||||
async def model_chat_local(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
|
async def model_chat_local(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
|
||||||
return ok({"response": "local chat adapter is not connected yet", "request": payload})
|
"""Proxy chat to the compute node running the inference model."""
|
||||||
|
store = get_platform_store()
|
||||||
|
node = _select_first_online_node(store)
|
||||||
|
if not node:
|
||||||
|
return ok({"response": "no online compute node available for inference", "request": payload})
|
||||||
|
try:
|
||||||
|
client = ComputeNodeClient(node["api_base_url"])
|
||||||
|
result = await client._request("POST", "/inference/chat", json_data=payload)
|
||||||
|
return ok(result)
|
||||||
|
except Exception as exc:
|
||||||
|
return ok({"response": f"inference failed: {exc}", "request": payload})
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/model-chat/local/chat/stream")
|
||||||
|
async def model_chat_local_stream(payload: dict[str, Any] = Body(...)) -> StreamingResponse:
|
||||||
|
"""Stream chat from the compute node."""
|
||||||
|
store = get_platform_store()
|
||||||
|
node = _select_first_online_node(store)
|
||||||
|
if not node:
|
||||||
|
return StreamingResponse(
|
||||||
|
iter(['data: {"error": "no online compute node"}\n\n']),
|
||||||
|
media_type="text/event-stream",
|
||||||
|
)
|
||||||
|
client = ComputeNodeClient(node["api_base_url"])
|
||||||
|
|
||||||
|
async def stream_proxy():
|
||||||
|
async with httpx.AsyncClient(timeout=300) as http:
|
||||||
|
url = f"{node['api_base_url'].rstrip('/')}/modelTF/inference/chat/stream"
|
||||||
|
async with http.stream("POST", url, json=payload, headers=client.headers()) as resp:
|
||||||
|
async for chunk in resp.aiter_bytes():
|
||||||
|
yield chunk
|
||||||
|
|
||||||
|
return StreamingResponse(stream_proxy(), media_type="text/event-stream")
|
||||||
|
|
||||||
|
|
||||||
@router.post("/model-chat/local/preload")
|
@router.post("/model-chat/local/preload")
|
||||||
async def model_chat_local_preload(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
|
async def model_chat_local_preload(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
|
||||||
return ok({"loaded": True, "request": payload})
|
"""Load a model on the compute node for inference."""
|
||||||
|
store = get_platform_store()
|
||||||
|
node = _select_first_online_node(store)
|
||||||
|
if not node:
|
||||||
|
return ok({"loaded": False, "error": "no online compute node"})
|
||||||
|
try:
|
||||||
|
client = ComputeNodeClient(node["api_base_url"])
|
||||||
|
result = await client._request("POST", "/inference/load", json_data=payload)
|
||||||
|
return ok(result)
|
||||||
|
except Exception as exc:
|
||||||
|
return ok({"loaded": False, "error": str(exc)})
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/model-chat/local/unload")
|
||||||
|
async def model_chat_local_unload() -> dict[str, Any]:
|
||||||
|
"""Unload the inference model from the compute node."""
|
||||||
|
store = get_platform_store()
|
||||||
|
node = _select_first_online_node(store)
|
||||||
|
if not node:
|
||||||
|
return ok({"unloaded": False, "error": "no online compute node"})
|
||||||
|
try:
|
||||||
|
client = ComputeNodeClient(node["api_base_url"])
|
||||||
|
result = await client._request("POST", "/inference/unload", json_data={})
|
||||||
|
return ok(result)
|
||||||
|
except Exception as exc:
|
||||||
|
return ok({"unloaded": False, "error": str(exc)})
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/model-chat/local/status")
|
||||||
|
async def model_chat_local_status() -> dict[str, Any]:
|
||||||
|
"""Get inference session status from compute node."""
|
||||||
|
store = get_platform_store()
|
||||||
|
node = _select_first_online_node(store)
|
||||||
|
if not node:
|
||||||
|
return ok({"loaded": False, "error": "no online compute node"})
|
||||||
|
try:
|
||||||
|
client = ComputeNodeClient(node["api_base_url"])
|
||||||
|
result = await client._request("GET", "/inference/status")
|
||||||
|
return ok(result)
|
||||||
|
except Exception as exc:
|
||||||
|
return ok({"loaded": False, "error": str(exc)})
|
||||||
|
|
||||||
|
|
||||||
@router.post("/model-chat/trained/preload")
|
@router.post("/model-chat/trained/preload")
|
||||||
async def model_chat_trained_preload(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
|
async def model_chat_trained_preload(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
|
||||||
return ok({"loaded": True, "request": payload})
|
"""Load a trained model (base + adapter) on the compute node for inference."""
|
||||||
|
store = get_platform_store()
|
||||||
|
node = _select_first_online_node(store)
|
||||||
|
if not node:
|
||||||
|
return ok({"loaded": False, "error": "no online compute node"})
|
||||||
|
try:
|
||||||
|
client = ComputeNodeClient(node["api_base_url"])
|
||||||
|
result = await client._request("POST", "/inference/load", json_data=payload)
|
||||||
|
return ok(result)
|
||||||
|
except Exception as exc:
|
||||||
|
return ok({"loaded": False, "error": str(exc)})
|
||||||
|
|
||||||
|
|
||||||
@router.get("/compute/nodes")
|
@router.get("/compute/nodes")
|
||||||
|
|||||||
@@ -182,6 +182,17 @@ class ComputeNodeClient:
|
|||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
return _unwrap_dict(response.json())
|
return _unwrap_dict(response.json())
|
||||||
|
|
||||||
|
async def _request(self, method: str, path: str, json_data: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||||
|
"""Generic request method for compute API endpoints."""
|
||||||
|
url = _join_url(self.api_base_url, f"{self.route_prefix}{path}")
|
||||||
|
async with httpx.AsyncClient(timeout=300, headers=self.headers()) as client:
|
||||||
|
if method.upper() == "GET":
|
||||||
|
response = await client.get(url)
|
||||||
|
else:
|
||||||
|
response = await client.post(url, json=json_data)
|
||||||
|
response.raise_for_status()
|
||||||
|
return _unwrap_dict(response.json())
|
||||||
|
|
||||||
async def upload_file(
|
async def upload_file(
|
||||||
self,
|
self,
|
||||||
filename: str,
|
filename: str,
|
||||||
|
|||||||
@@ -10,10 +10,11 @@ from pathlib import Path
|
|||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
|
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
|
||||||
from fastapi.responses import FileResponse, JSONResponse
|
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
|
||||||
|
|
||||||
from compute.agent.process_manager import ProcessManager
|
from compute.agent.process_manager import ProcessManager
|
||||||
from compute.engines.llama_factory.adapter import build_command, parse_log_line, prepare_runtime_files
|
from compute.engines.llama_factory.adapter import build_command, parse_log_line, prepare_runtime_files
|
||||||
|
from compute.engines.llama_factory.inference import get_inference_session
|
||||||
|
|
||||||
|
|
||||||
def create_app() -> FastAPI:
|
def create_app() -> FastAPI:
|
||||||
@@ -666,6 +667,85 @@ def create_app() -> FastAPI:
|
|||||||
metrics = [parse_log_line(line) for line in window["content"].splitlines()]
|
metrics = [parse_log_line(line) for line in window["content"].splitlines()]
|
||||||
return {"job_id": job_id, **window, "metrics": [m for m in metrics if m]}
|
return {"job_id": job_id, **window, "metrics": [m for m in metrics if m]}
|
||||||
|
|
||||||
|
# ── Inference Endpoints ───────────────────────────────────────────
|
||||||
|
|
||||||
|
@app.post(f"{route_prefix}/inference/load")
|
||||||
|
async def inference_load(payload: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
"""Load a model for inference using LLaMA-Factory ChatModel.
|
||||||
|
|
||||||
|
Expected payload:
|
||||||
|
model_name_or_path: str (required)
|
||||||
|
adapter_name_or_path: str (optional, for LoRA adapters)
|
||||||
|
template: str (default: "qwen")
|
||||||
|
infer_backend: str (default: "huggingface")
|
||||||
|
infer_dtype: str (default: "auto")
|
||||||
|
"""
|
||||||
|
session = get_inference_session()
|
||||||
|
result = session.load(
|
||||||
|
model_name_or_path=payload.get("model_name_or_path", ""),
|
||||||
|
adapter_name_or_path=payload.get("adapter_name_or_path", ""),
|
||||||
|
template=payload.get("template", "qwen"),
|
||||||
|
infer_backend=payload.get("infer_backend", "huggingface"),
|
||||||
|
infer_dtype=payload.get("infer_dtype", "auto"),
|
||||||
|
)
|
||||||
|
if not result.get("loaded"):
|
||||||
|
raise HTTPException(status_code=500, detail=result.get("error", "model load failed"))
|
||||||
|
return result
|
||||||
|
|
||||||
|
@app.post(f"{route_prefix}/inference/unload")
|
||||||
|
async def inference_unload() -> dict[str, Any]:
|
||||||
|
"""Unload the currently loaded model and free GPU memory."""
|
||||||
|
return get_inference_session().unload()
|
||||||
|
|
||||||
|
@app.get(f"{route_prefix}/inference/status")
|
||||||
|
async def inference_status() -> dict[str, Any]:
|
||||||
|
"""Get the current inference session status."""
|
||||||
|
return get_inference_session().info()
|
||||||
|
|
||||||
|
@app.post(f"{route_prefix}/inference/chat")
|
||||||
|
async def inference_chat(payload: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
"""Chat with the loaded model (non-streaming).
|
||||||
|
|
||||||
|
Expected payload:
|
||||||
|
messages: list[dict] (OpenAI format)
|
||||||
|
temperature: float (default 0.95)
|
||||||
|
top_p: float (default 0.7)
|
||||||
|
max_new_tokens: int (default 1024)
|
||||||
|
"""
|
||||||
|
messages = payload.get("messages") or []
|
||||||
|
if not messages:
|
||||||
|
raise HTTPException(status_code=400, detail="messages is required")
|
||||||
|
result = get_inference_session().chat(
|
||||||
|
messages=messages,
|
||||||
|
temperature=float(payload.get("temperature", 0.95)),
|
||||||
|
top_p=float(payload.get("top_p", 0.7)),
|
||||||
|
max_new_tokens=int(payload.get("max_new_tokens", 1024)),
|
||||||
|
do_sample=bool(payload.get("do_sample", True)),
|
||||||
|
)
|
||||||
|
if result.get("error"):
|
||||||
|
raise HTTPException(status_code=500, detail=result["error"])
|
||||||
|
return {"response": result["response"]}
|
||||||
|
|
||||||
|
@app.post(f"{route_prefix}/inference/chat/stream")
|
||||||
|
async def inference_chat_stream(payload: dict[str, Any]) -> StreamingResponse:
|
||||||
|
"""Chat with streaming response (Server-Sent Events)."""
|
||||||
|
messages = payload.get("messages") or []
|
||||||
|
if not messages:
|
||||||
|
raise HTTPException(status_code=400, detail="messages is required")
|
||||||
|
|
||||||
|
def generate():
|
||||||
|
session = get_inference_session()
|
||||||
|
for chunk in session.chat_stream(
|
||||||
|
messages=messages,
|
||||||
|
temperature=float(payload.get("temperature", 0.95)),
|
||||||
|
top_p=float(payload.get("top_p", 0.7)),
|
||||||
|
max_new_tokens=int(payload.get("max_new_tokens", 1024)),
|
||||||
|
do_sample=bool(payload.get("do_sample", True)),
|
||||||
|
):
|
||||||
|
yield chunk
|
||||||
|
|
||||||
|
return StreamingResponse(generate(), media_type="text/event-stream")
|
||||||
|
|
||||||
@app.post(f"{route_prefix}/compute/files/upload")
|
@app.post(f"{route_prefix}/compute/files/upload")
|
||||||
async def upload_file(
|
async def upload_file(
|
||||||
file: UploadFile | None = File(default=None),
|
file: UploadFile | None = File(default=None),
|
||||||
|
|||||||
125
compute/engines/llama_factory/inference.py
Normal file
125
compute/engines/llama_factory/inference.py
Normal file
@@ -0,0 +1,125 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
|
||||||
|
class InferenceSession:
|
||||||
|
"""Manages a loaded model for inference with LLaMA-Factory ChatModel."""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._model: Any = None
|
||||||
|
self._tokenizer: Any = None
|
||||||
|
self._generating_args: dict[str, Any] = {}
|
||||||
|
self._model_name: str = ""
|
||||||
|
self._adapter_path: str = ""
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
self._loaded_at: float = 0.0
|
||||||
|
self._status: str = "idle"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def status(self) -> str:
|
||||||
|
return self._status
|
||||||
|
|
||||||
|
@property
|
||||||
|
def model_name(self) -> str:
|
||||||
|
return self._model_name
|
||||||
|
|
||||||
|
@property
|
||||||
|
def adapter_path(self) -> str:
|
||||||
|
return self._adapter_path
|
||||||
|
|
||||||
|
@property
|
||||||
|
def loaded_at(self) -> float:
|
||||||
|
return self._loaded_at
|
||||||
|
|
||||||
|
def info(self) -> dict[str, Any]:
|
||||||
|
return {
|
||||||
|
"loaded": self._status == "ready",
|
||||||
|
"status": self._status,
|
||||||
|
"model_name": self._model_name,
|
||||||
|
"adapter_path": self._adapter_path,
|
||||||
|
"loaded_at": self._loaded_at,
|
||||||
|
}
|
||||||
|
|
||||||
|
def load(self, model_name_or_path, adapter_name_or_path="", template="qwen", infer_backend="huggingface", infer_dtype="auto", **kwargs):
|
||||||
|
with self._lock:
|
||||||
|
if self._status == "loading":
|
||||||
|
return {"loaded": False, "error": "model is already loading"}
|
||||||
|
if self._status == "ready":
|
||||||
|
self.unload()
|
||||||
|
self._status = "loading"
|
||||||
|
self._model_name = model_name_or_path
|
||||||
|
self._adapter_path = adapter_name_or_path
|
||||||
|
try:
|
||||||
|
from llamafactory.chat import ChatModel
|
||||||
|
from llamafactory.hparams import get_infer_args
|
||||||
|
args = {"model_name_or_path": model_name_or_path, "template": template, "infer_backend": infer_backend, "infer_dtype": infer_dtype}
|
||||||
|
if adapter_name_or_path:
|
||||||
|
args["adapter_name_or_path"] = adapter_name_or_path
|
||||||
|
args.update(kwargs)
|
||||||
|
model_args, generating_args = get_infer_args(args)
|
||||||
|
self._model = ChatModel(model_args)
|
||||||
|
self._tokenizer = self._model.tokenizer
|
||||||
|
self._generating_args = generating_args
|
||||||
|
self._loaded_at = time.time()
|
||||||
|
self._status = "ready"
|
||||||
|
return {"loaded": True, "status": "ready"}
|
||||||
|
except Exception as exc:
|
||||||
|
self._status = "error"
|
||||||
|
self._model = None
|
||||||
|
return {"loaded": False, "status": "error", "error": str(exc)}
|
||||||
|
|
||||||
|
def unload(self):
|
||||||
|
with self._lock:
|
||||||
|
if self._model is not None:
|
||||||
|
try:
|
||||||
|
del self._model
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
self._model = None
|
||||||
|
self._tokenizer = None
|
||||||
|
self._status = "idle"
|
||||||
|
self._model_name = ""
|
||||||
|
self._adapter_path = ""
|
||||||
|
self._loaded_at = 0.0
|
||||||
|
return {"unloaded": True}
|
||||||
|
|
||||||
|
def chat(self, messages, temperature=0.95, top_p=0.7, max_new_tokens=1024, do_sample=True, **kwargs):
|
||||||
|
with self._lock:
|
||||||
|
if self._status != "ready" or self._model is None:
|
||||||
|
return {"error": "model not loaded", "response": ""}
|
||||||
|
try:
|
||||||
|
generate_kwargs = {**self._generating_args, "temperature": temperature, "top_p": top_p, "max_new_tokens": max_new_tokens, "do_sample": do_sample}
|
||||||
|
generate_kwargs.update(kwargs)
|
||||||
|
formatted = self._model.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||||
|
responses = []
|
||||||
|
for response in self._model.stream_chat(formatted, generate_kwargs):
|
||||||
|
responses.append(response)
|
||||||
|
full_response = "".join(str(r) for r in responses)
|
||||||
|
return {"response": full_response}
|
||||||
|
except Exception as exc:
|
||||||
|
return {"error": str(exc), "response": ""}
|
||||||
|
|
||||||
|
def chat_stream(self, messages, **kwargs):
|
||||||
|
with self._lock:
|
||||||
|
if self._status != "ready" or self._model is None:
|
||||||
|
yield 'data: {"error": "model not loaded"}\n\n'
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
generate_kwargs = {**self._generating_args, **kwargs}
|
||||||
|
formatted = self._model.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||||
|
for new_text in self._model.stream_chat(formatted, generate_kwargs):
|
||||||
|
yield new_text
|
||||||
|
except Exception as exc:
|
||||||
|
yield 'data: {"error": "' + str(exc) + '"}\n\n'
|
||||||
|
|
||||||
|
|
||||||
|
_inference_session = None
|
||||||
|
|
||||||
|
def get_inference_session():
|
||||||
|
global _inference_session
|
||||||
|
if _inference_session is None:
|
||||||
|
_inference_session = InferenceSession()
|
||||||
|
return _inference_session
|
||||||
Reference in New Issue
Block a user