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

refactor: assemble inference engines through a shared composition root

Browse Source 
- add build_engine() to astrai.inference.engine as the single load-place-wire path for InferenceEngine, accepting a checkpoint path or live model/tokenizer plus passthrough engine kwargs
- migrate the server lifespan, generate CLI, humaneval/ifeval evals, and all three demos to build_engine; app._create_engine collapses into a direct call
- export build_engine from astrai and astrai.inference
- parameterize the autoregressive demo with --prompt one-shot continuation plus model path and sampling knobs, exiting cleanly on !exit or EOF
- cover the composition root with unit tests for live-object assembly, kwargs passthrough, and argument validation
This commit is contained in:
ViperEkura
2026-09-03 22:16:56 +08:00
parent 9d3ae76683
commit e13fe53475
12 changed files with 199 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/__init__.py
+2 -1
View File
@@ -17,7 +17,7 @@ from astrai.dataset import (
StoreFactory,
)
from astrai.factory import BaseFactory
from astrai.inference import InferenceEngine, get_app, run_server, sample
from astrai.inference import InferenceEngine, build_engine, get_app, run_server, sample
from astrai.inference.network import ProtocolHandler
from astrai.inference.runtime.sample import SamplingPipeline
from astrai.logging import setup_logging
@@ -67,6 +67,7 @@ __all__ = [
"EncoderConfig",
"ExecutorFactory",
"InferenceEngine",
"build_engine",
"LoRAConfig",
"Pipeline",
"PipelineConfig",
astrai/inference/__init__.py
+2 -1
View File
@@ -12,7 +12,7 @@ Modules:
- engine.py: Facade (InferenceEngine)
"""
from astrai.inference.engine import InferenceEngine
from astrai.inference.engine import InferenceEngine, build_engine
from astrai.inference.network import get_app, run_server
from astrai.inference.runtime.executor import Executor
from astrai.inference.runtime.sample import sample
@@ -21,6 +21,7 @@ from astrai.inference.task import STOP, GenerationResult, Task, TaskManager, Tas
__all__ = [
"InferenceEngine",
"build_engine",
"InferenceScheduler",
"GenerationResult",
"Executor",
astrai/inference/engine.py
+55
View File
@@ -2,7 +2,9 @@
import asyncio
import gc
import logging
import threading
from pathlib import Path
from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, Union
import torch
@@ -12,8 +14,11 @@ from astrai.extension import ATTN_BACKEND, AttentionBackend, get_backend
from astrai.inference.cache import PagePool
from astrai.inference.scheduler import InferenceScheduler
from astrai.inference.task import STOP
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer
logger = logging.getLogger(__name__)
class GenerateResult:
"""Thread-safe token accumulator for streaming and non-streaming modes."""
@@ -251,3 +256,53 @@ class InferenceEngine:
if torch.cuda.is_available():
torch.cuda.empty_cache()
gc.collect()
def build_engine(
param_path: Optional[Union[str, Path]] = None,
*,
model: Optional[nn.Module] = None,
tokenizer: Optional[AutoTokenizer] = None,
device: Optional[str] = "cuda",
dtype: Optional[torch.dtype] = torch.bfloat16,
max_batch_size: int = 16,
max_seq_len: Optional[int] = None,
**engine_kwargs: Any,
) -> InferenceEngine:
"""Composition root for inference assembly.
Loads model and tokenizer from *param_path*, or accepts preloaded
objects, places the model, and returns a started InferenceEngine.
Extra *engine_kwargs* (cache, enable_cuda_graph, backend) pass
through to InferenceEngine. Placement parts left as None are skipped.
"""
if param_path is not None:
if model is not None or tokenizer is not None:
raise ValueError("pass either param_path or model+tokenizer, not both")
path = Path(param_path)
if not path.exists():
raise FileNotFoundError(f"Parameter directory not found: {path}")
tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModel.from_pretrained(path)
elif model is None or tokenizer is None:
raise ValueError("build_engine requires param_path or both model and tokenizer")
placement: Dict[str, Any] = {}
if device is not None:
placement["device"] = device
if dtype is not None:
placement["dtype"] = dtype
if placement:
model.to(**placement)
logger.info(
f"Model placed on {placement.get('device')} "
f"with dtype {placement.get('dtype')}"
)
return InferenceEngine(
model=model,
tokenizer=tokenizer,
max_batch_size=max_batch_size,
max_seq_len=max_seq_len,
**engine_kwargs,
)
astrai/inference/network/app.py
+2 -29
View File
@@ -18,12 +18,10 @@ import uvicorn
from fastapi import APIRouter, FastAPI, HTTPException
from pydantic import BaseModel, Field
from astrai.inference.engine import InferenceEngine
from astrai.inference.engine import InferenceEngine, build_engine
from astrai.inference.network.anthropic import AnthropicResponseBuilder
from astrai.inference.network.openai import OpenAIResponseBuilder
from astrai.inference.network.protocol import ProtocolHandler
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer
logger = logging.getLogger(__name__)
@@ -92,7 +90,7 @@ async def lifespan(app: FastAPI):
config = app.state.server_config
if not config.get("_test", False):
try:
app.state.engine = _create_engine(**config)
app.state.engine = build_engine(**config)
except Exception as e:
logger.error(f"Failed to load model: {e}")
raise
@@ -105,31 +103,6 @@ async def lifespan(app: FastAPI):
router = APIRouter()
def _create_engine(
param_path: Path,
device: str = "cuda",
dtype: torch.dtype = torch.bfloat16,
max_batch_size: int = 16,
max_seq_len: Optional[int] = None,
) -> InferenceEngine:
if not param_path.exists():
raise FileNotFoundError(f"Parameter directory not found: {param_path}")
tokenizer = AutoTokenizer.from_pretrained(param_path)
model = AutoModel.from_pretrained(param_path)
model.to(device=device, dtype=dtype)
logger.info(f"Model loaded on {device} with dtype {dtype}")
engine = InferenceEngine(
model=model,
tokenizer=tokenizer,
max_batch_size=max_batch_size,
max_seq_len=max_seq_len,
)
logger.info(f"Inference engine initialized with max_batch_size={max_batch_size}")
return engine
def get_app() -> FastAPI:
"""Return the singleton FastAPI instance (lazily created on first call)."""
global _app_instance
scripts/demo/generate_ar.py
+53 -21
View File
@@ -1,36 +1,68 @@
from argparse import ArgumentParser
from pathlib import Path
import torch
from astrai.inference import InferenceEngine
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer
from astrai.inference import build_engine
PROJECT_ROOT = Path(__file__).resolve().parents[2]
PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
def generate_text():
model = AutoModel.from_pretrained(PARAMETER_ROOT)
tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
model.to(device="cuda", dtype=torch.bfloat16)
query = input(">> ")
engine = InferenceEngine(
model=model,
tokenizer=tokenizer,
def parse_args():
parser = ArgumentParser(
description="Autoregressive continuation demo: continue a prompt, "
"or drop into interactive mode without one"
)
parser.add_argument(
"--prompt",
type=str,
default=None,
help="Text prefix to continue; omit to enter interactive mode",
)
parser.add_argument(
"--model_path",
type=Path,
default=PARAMETER_ROOT,
help="Path to model weights",
)
parser.add_argument("--max_tokens", type=int, default=2048)
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--top_p", type=float, default=0.95)
parser.add_argument("--top_k", type=int, default=50)
return parser.parse_args()
def run_once(engine, prompt, args):
print(prompt, end="", flush=True)
for token in engine.generate(
prompt=query,
prompt=prompt,
stream=True,
max_tokens=2048,
temperature=0.8,
top_p=0.95,
top_k=50,
max_tokens=args.max_tokens,
temperature=args.temperature,
top_p=args.top_p,
top_k=args.top_k,
):
print(token, end="", flush=True)
print()
def main():
args = parse_args()
engine = build_engine(args.model_path)
if args.prompt is not None:
run_once(engine, args.prompt, args)
return
while True:
try:
query = input(">> ")
except EOFError:
break
if query == "!exit":
break
if query:
run_once(engine, query, args)
if __name__ == "__main__":
generate_text()
main()
scripts/demo/generate_batch.py
+2 -11
View File
@@ -1,9 +1,6 @@
from pathlib import Path
import torch
from astrai.inference import InferenceEngine
from astrai.model import AutoModel
from astrai.inference import build_engine
from astrai.tokenize import AutoTokenizer
PROJECT_ROOT = Path(__file__).resolve().parents[2]
@@ -11,10 +8,7 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
def batch_generate():
# Load model using AutoModel
model = AutoModel.from_pretrained(PARAMETER_ROOT)
tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
model.to(device="cuda", dtype=torch.bfloat16)
inputs = [
"你好",
@@ -33,10 +27,7 @@ def batch_generate():
for q in inputs
]
engine = InferenceEngine(
model=model,
tokenizer=tokenizer,
)
engine = build_engine(PARAMETER_ROOT)
responses = engine.generate(
prompt=prompts,
stream=False,
scripts/demo/stream_chat.py
+3 -9
View File
@@ -1,11 +1,7 @@
from argparse import ArgumentParser
from pathlib import Path
import torch
from astrai import InferenceEngine
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer
from astrai import build_engine
PROJECT_ROOT = Path(__file__).resolve().parents[2]
@@ -68,10 +64,8 @@ def chat():
args = parse_args()
model_path = args.model_path
model = AutoModel.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model.to(device="cuda", dtype=torch.bfloat16)
engine = InferenceEngine(model=model, tokenizer=tokenizer)
engine = build_engine(model_path)
tokenizer = engine.tokenizer
while True:
query = input(">> ")
scripts/eval/evaluate_humaneval.py
+6 -19
View File
@@ -18,13 +18,10 @@ from math import prod
from typing import Dict, Iterator, List, Optional, Sequence, Tuple
import numpy as np
import torch
import tqdm
from datasets import load_dataset
from astrai.inference import InferenceEngine
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer
from astrai.inference import build_engine
# ---------------------------------------------------------------------------
# Config
@@ -91,20 +88,6 @@ def save_json(path: str, data):
json.dump(data, f, indent=2, ensure_ascii=False)
def create_engine(
param_path: str, batch_size: int, max_seq_len: int
) -> InferenceEngine:
model = AutoModel.from_pretrained(param_path)
tokenizer = AutoTokenizer.from_pretrained(param_path)
model.to(device="cuda", dtype=torch.bfloat16)
return InferenceEngine(
model=model,
tokenizer=tokenizer,
max_batch_size=batch_size,
max_seq_len=max_seq_len,
)
def trim_stop(text: str) -> str:
for stop in STOP_SEQUENCES:
idx = text.find(stop)
@@ -322,7 +305,11 @@ def run_pipeline(cfg: EvalConfig) -> Dict:
if cfg.problem_indices:
problems = [problems[i] for i in cfg.problem_indices if i < len(problems)]
engine = create_engine(cfg.param_path, cfg.batch_size, cfg.max_seq_len)
engine = build_engine(
cfg.param_path,
max_batch_size=cfg.batch_size,
max_seq_len=cfg.max_seq_len,
)
try:
generated = generate_all(engine, problems, cfg)
scripts/eval/evaluate_ifeval.py
+4 -11
View File
@@ -16,12 +16,10 @@ import os
import re
from typing import Callable, Dict, List, Optional
import torch
import tqdm
from datasets import load_dataset
from astrai.inference import InferenceEngine
from astrai.model import AutoModel
from astrai.inference import InferenceEngine, build_engine
from astrai.tokenize import AutoTokenizer
IFEVAL_HF_DATASET = "google/IFEval"
@@ -536,17 +534,12 @@ def main():
print(f"Loaded {len(problems)} problems")
print(f"Supported constraint types: {len(CONSTRAINT_VERIFIERS)}")
model = AutoModel.from_pretrained(args.param_path)
tokenizer = AutoTokenizer.from_pretrained(args.param_path)
model.to(device="cuda", dtype=torch.bfloat16)
model.eval()
engine = InferenceEngine(
model=model,
tokenizer=tokenizer,
engine = build_engine(
args.param_path,
max_batch_size=args.batch_size,
max_seq_len=args.max_seq_len,
)
tokenizer = engine.tokenizer
results = evaluate(
engine=engine,
scripts/tools/generate.py
+5 -11
View File
@@ -6,9 +6,7 @@ import click
import torch
from tqdm import tqdm
from astrai.inference import InferenceEngine
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer
from astrai.inference import build_engine
def processor(
@@ -28,17 +26,13 @@ def processor(
):
print(f"Loading model from {param_path} ...")
t0 = time.time()
model = AutoModel.from_pretrained(param_path)
tokenizer = AutoTokenizer.from_pretrained(param_path)
model.to(device="cuda", dtype=torch.bfloat16)
print(f" model loaded in {time.time() - t0:.1f}s")
engine = InferenceEngine(
model=model,
tokenizer=tokenizer,
engine = build_engine(
param_path=param_path,
max_batch_size=batch_size * num_samples,
max_seq_len=max_seq_len,
)
tokenizer = engine.tokenizer
print(f" model loaded in {time.time() - t0:.1f}s")
print(f"Reading {input_json_file} ...")
with open(input_json_file, "r", encoding="utf-8") as f:
tests/inference/test_engine.py
+63 -1
View File
@@ -4,9 +4,12 @@ import asyncio
import threading
from unittest.mock import MagicMock, patch
import pytest
from astrai.extension import TorchNativeBackend, attn_backend
from astrai.inference import STOP
from astrai.inference.engine import GenerateResult, InferenceEngine
from astrai.inference.engine import GenerateResult, InferenceEngine, build_engine
from tests.helpers import FakeTokenizer, make_model
def _make_engine_mocks(decode=None):
@@ -301,3 +304,62 @@ def test_generate_captures_calling_backend_context():
assert len(captured) == 1
assert isinstance(captured[0], TorchNativeBackend)
def test_build_engine_from_live_objects_starts_scheduler():
model, _ = make_model("cpu", max_position_embeddings=64)
tokenizer = FakeTokenizer()
engine = build_engine(
model=model,
tokenizer=tokenizer,
device=None,
dtype=None,
max_batch_size=2,
)
try:
assert isinstance(engine, InferenceEngine)
assert engine.tokenizer is tokenizer
assert engine.scheduler._stop_event.is_set() is False
finally:
engine.shutdown()
def test_build_engine_passes_engine_kwargs_through():
model, _ = make_model("cpu", max_position_embeddings=64)
backend = TorchNativeBackend()
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
MockSched.return_value.add_task.side_effect = lambda *args, **k: (
k["stream_callback"](STOP) or "task"
)
engine = build_engine(
model=model,
tokenizer=FakeTokenizer(),
device=None,
dtype=None,
cache=object(),
enable_cuda_graph=False,
backend=backend,
)
engine.generate("hi")
kwargs = MockSched.call_args.kwargs
assert kwargs["cache"] is not None
assert kwargs["enable_cuda_graph"] is False
assert kwargs["backend"] is backend
@pytest.mark.parametrize(
("kwargs", "error", "message"),
[
(
{"param_path": "x", "model": object()},
ValueError,
"not both",
),
({}, ValueError, "requires param_path"),
({"param_path": "/nonexistent-dir-xyz"}, FileNotFoundError, "not found"),
],
)
def test_build_engine_rejects_invalid_arguments(kwargs, error, message):
with pytest.raises(error, match=message):
build_engine(**kwargs)
tests/inference/test_server.py
+2 -3
View File
@@ -5,8 +5,7 @@ from pathlib import Path
import pytest
import torch
from astrai.inference import get_app
from astrai.inference.network.app import _create_engine
from astrai.inference import build_engine, get_app
from astrai.model.transformer import AutoRegressiveLM
from astrai.serialization import save_model
from tests.helpers import CHAT_TEMPLATE, build_test_tokenizer, make_tiny_config
@@ -238,7 +237,7 @@ def test_chat_completions_real_engine(tmp_path, client):
tokenizer = build_test_tokenizer(vocab_size=256, chat_template=CHAT_TEMPLATE)
tokenizer.save_pretrained(str(tmp_path))
engine = _create_engine(
engine = build_engine(
Path(tmp_path),
device="cpu",
dtype=torch.float32,