ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: 实现模型动态注册机制

Browse Source
This commit is contained in:
ViperEkura
2026-04-05 19:38:12 +08:00
parent ff43a2fab8
commit fc278d17ab
25 changed files with 686 additions and 651 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/engine.py
+33 -21
View File
@@ -1,11 +1,11 @@
"""Unified inference engine."""
import threading
import torch
import torch.nn as nn
from typing import Any, Dict, Generator, List, Optional, Union
from astrai.config import ModelParameter
from astrai.tokenize.chat_template import build_prompt
from astrai.tokenize.tokenizer import TextTokenizer
from astrai.inference.scheduler import InferenceScheduler
@@ -14,22 +14,18 @@ class GenerationRequest:
def __init__(
self,
query: Union[str, List[str]],
messages: List[Dict[str, str]],
top_k: int = 50,
top_p: float = 1.0,
temperature: float = 1.0,
max_len: int = 1024,
history: Optional[Any] = None,
system_prompt: Optional[str] = None,
stream: bool = False,
):
self.query = query
self.messages = messages
self.top_k = top_k
self.top_p = top_p
self.temperature = temperature
self.max_len = max_len
self.history = history
self.system_prompt = system_prompt
self.stream = stream
self._validate()
@@ -107,26 +103,41 @@ class InferenceEngine:
def __init__(
self,
parameter: ModelParameter,
max_batch_size: int = 16,
model: nn.Module,
tokenizer: TextTokenizer,
max_batch_size: int = 1,
max_seq_len: Optional[int] = None,
):
self.model = parameter.model
self.tokenizer = parameter.tokenizer
self.config = parameter.config
"""
Initialize inference engine with separate model and tokenizer.
model_params = next(self.model.parameters())
self.device = model_params.device
self.dtype = model_params.dtype
Args:
model: The language model for inference (nn.Module, e.g., Transformer)
tokenizer: The tokenizer for encoding/decoding text
config: Model configuration
max_batch_size: Maximum batch size for continuous batching
max_seq_len: Maximum sequence length (defaults to config.max_len)
"""
self.model = model
self.tokenizer = tokenizer
# Get device and dtype from model parameters
try:
first_param = next(model.parameters())
device = first_param.device
dtype = first_param.dtype
except StopIteration:
# Model has no parameters, use default device/dtype
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float32
self.scheduler = InferenceScheduler(
model=self.model,
tokenizer=self.tokenizer,
config=self.config,
max_batch_size=max_batch_size,
max_seq_len=max_seq_len,
device=self.device,
dtype=self.dtype,
device=device,
dtype=dtype,
)
self.kv_cache = self.scheduler.kv_cache
@@ -160,7 +171,8 @@ class InferenceEngine:
self, request: GenerationRequest
) -> Union[Generator[str, None, None], str, List[str]]:
"""Generate with GenerationRequest object."""
prompt = build_prompt(request.query, request.history)
# Use tokenizer's chat template with messages
prompt = self.tokenizer.apply_chat_template(request.messages, tokenize=False)
return self.generate(
prompt=prompt,
astrai/inference/scheduler.py
+17 -16
View File
@@ -8,7 +8,8 @@ from typing import Any, Callable, Dict, List, Optional
import torch
from torch import Tensor
from astrai.config import ModelConfig
from astrai.model.automodel import AutoModel
from astrai.tokenize.tokenizer import TextTokenizer
class TaskStatus:
@@ -98,23 +99,23 @@ class InferenceScheduler:
def __init__(
self,
model,
tokenizer,
config: ModelConfig,
model: AutoModel,
tokenizer: TextTokenizer,
max_batch_size: int = 16,
max_seq_len: Optional[int] = None,
device: str = "cuda",
dtype: torch.dtype = torch.bfloat16,
):
config = model.config
self.model = model
self.tokenizer = tokenizer
self.config = config
self.max_batch_size = max_batch_size
self.max_seq_len = max_seq_len or config.max_len
self.device = device
self.dtype = dtype
self.device = device or next(model.parameters()).device
self.dtype = dtype or next(model.parameters()).dtype
num_heads = config.n_kv_heads
num_kv_heads = config.n_kv_heads
head_dim = config.dim // config.n_heads
n_layers = config.n_layers
@@ -123,26 +124,26 @@ class InferenceScheduler:
max_batch_size,
self.max_seq_len,
n_layers,
num_heads,
num_kv_heads,
head_dim,
),
device=device,
dtype=dtype,
device=self.device,
dtype=self.dtype,
)
v_cache = torch.empty(
(
max_batch_size,
self.max_seq_len,
n_layers,
num_heads,
num_kv_heads,
head_dim,
),
device=device,
dtype=dtype,
device=self.device,
dtype=self.dtype,
)
self.kv_cache = (k_cache, v_cache)
self.seq_mask = torch.ones(
(max_batch_size, self.max_seq_len), device=device, dtype=torch.bool
(max_batch_size, self.max_seq_len), device=self.device, dtype=torch.bool
)
self.waiting_queue: List[Task] = []
@@ -259,7 +260,7 @@ class InferenceScheduler:
)
with torch.inference_mode():
outputs = self.model(
self.model(
input_ids,
input_mask=input_mask,
start_pos=0,
astrai/inference/server.py
+49 -27
View File
@@ -3,8 +3,6 @@ Inference Server with Continuous Batching Support
FastAPI server for inference with continuous batching.
Provides OpenAI-compatible chat completion endpoints.
Author: AstrAI Team
"""
import json
@@ -19,14 +17,15 @@ from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from astrai.config.param_config import ModelParameter
from astrai.inference.engine import GenerationRequest, InferenceEngine
from astrai.inference.engine import InferenceEngine
from astrai.model import AutoModel
from astrai.tokenize import TextTokenizer
logger = logging.getLogger(__name__)
# Global model parameter and engine (loaded once)
_model_param: Optional[ModelParameter] = None
_engine: Optional[InferenceEngine] = None
_model_param: Optional[Any] = None
_project_root = Path(__file__).parent.parent.parent
# Server configuration (set before running server)
@@ -95,13 +94,17 @@ def load_model(
param_path = _project_root / "params"
if not param_path.exists():
raise FileNotFoundError(f"Parameter directory not found: {param_path}")
_model_param = ModelParameter.load(param_path, disable_init=True)
# Load tokenizer separately
tokenizer = TextTokenizer.from_pretrained(param_path)
_model_param = AutoModel.from_pretrained(param_path, tokenizer=tokenizer)
_model_param.to(device=device, dtype=dtype)
logger.info(f"Model loaded on {device} with dtype {dtype}")
# Initialize inference engine with continuous batching
# Initialize inference engine with separate model and tokenizer
_engine = InferenceEngine(
parameter=_model_param,
model=_model_param,
tokenizer=tokenizer,
max_batch_size=max_batch_size,
)
logger.info(f"Inference engine initialized with max_batch_size={max_batch_size}")
@@ -164,27 +167,43 @@ def convert_messages_to_history(
return system_prompt, history if history else None
def convert_messages_to_prompt(messages: List[ChatMessage]) -> str:
def convert_messages_to_prompt(
messages: List[ChatMessage], engine: InferenceEngine = None
) -> str:
"""Convert messages to prompt string.
Args:
messages: List of ChatMessage objects
engine: InferenceEngine instance for accessing tokenizer
Returns:
str: Formatted prompt string
"""
system_prompt, history = convert_messages_to_history(messages)
# Convert to dict format for chat template
msg_dicts = [{"role": m.role, "content": m.content} for m in messages]
# Get the last user message as query
user_messages = [m.content for m in messages if m.role == "user"]
if not user_messages:
raise ValueError("No user message found")
query = user_messages[-1]
# Extract system prompt if present
system_prompt = None
filtered_messages = []
for msg in msg_dicts:
if msg["role"] == "system":
system_prompt = msg["content"]
else:
filtered_messages.append(msg)
# Build prompt using chat template
from astrai.tokenize.chat_template import build_prompt
# Use engine's tokenizer chat template if available
if engine is not None and engine.tokenizer is not None:
return engine.tokenizer.apply_chat_template(
filtered_messages, system_prompt=system_prompt, tokenize=False
)
return build_prompt(query, history)
# Fallback: simple concatenation (deprecated)
prompt_parts = []
for msg in filtered_messages:
prompt_parts.append(
f"<｜im▁start｜>{msg['role']}\n{msg['content']}<｜im▁end｜>"
)
return "\n".join(prompt_parts) + "\n<｜im▁start｜>assistant\n"
@app.get("/health")
@@ -213,8 +232,8 @@ async def chat_completion(request: ChatCompletionRequest):
if _engine is None:
raise HTTPException(status_code=503, detail="Engine not initialized")
# Convert messages to prompt
prompt = convert_messages_to_prompt(request.messages)
# Convert messages to prompt using engine's tokenizer
prompt = convert_messages_to_prompt(request.messages, engine=_engine)
if request.stream:
# Streaming response (use synchronous generator)
@@ -294,15 +313,18 @@ async def generate(
if _engine is None:
raise HTTPException(status_code=503, detail="Engine not initialized")
# Convert history format
hist: Optional[List[Tuple[str, str]]] = None
# Build messages for chat template
messages = []
if history:
hist = [(h[0], h[1]) for h in history]
# Convert history format: List[List[str]] -> List[Dict]
for h in history:
if len(h) >= 2:
messages.append({"role": "user", "content": h[0]})
messages.append({"role": "assistant", "content": h[1]})
messages.append({"role": "user", "content": query})
# Build prompt
from astrai.tokenize.chat_template import build_prompt
prompt = build_prompt(query, hist)
# Use tokenizer's chat template
prompt = _engine.tokenizer.apply_chat_template(messages, tokenize=False)
if stream:
# Synchronous streaming