diff --git a/astrai/__init__.py b/astrai/__init__.py index 1d223fc..fb1890b 100644 --- a/astrai/__init__.py +++ b/astrai/__init__.py @@ -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", diff --git a/astrai/inference/__init__.py b/astrai/inference/__init__.py index dd33164..5e3edc6 100644 --- a/astrai/inference/__init__.py +++ b/astrai/inference/__init__.py @@ -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", diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index e440cce..17d2d7e 100644 --- a/astrai/inference/engine.py +++ b/astrai/inference/engine.py @@ -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, + ) diff --git a/astrai/inference/network/app.py b/astrai/inference/network/app.py index 4292b93..3a4f3eb 100644 --- a/astrai/inference/network/app.py +++ b/astrai/inference/network/app.py @@ -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 diff --git a/scripts/demo/generate_ar.py b/scripts/demo/generate_ar.py index ccbe49a..917c99e 100644 --- a/scripts/demo/generate_ar.py +++ b/scripts/demo/generate_ar.py @@ -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() diff --git a/scripts/demo/generate_batch.py b/scripts/demo/generate_batch.py index da54922..e0b8d60 100644 --- a/scripts/demo/generate_batch.py +++ b/scripts/demo/generate_batch.py @@ -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, diff --git a/scripts/demo/stream_chat.py b/scripts/demo/stream_chat.py index e09d4bc..2a8b913 100644 --- a/scripts/demo/stream_chat.py +++ b/scripts/demo/stream_chat.py @@ -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(">> ") diff --git a/scripts/eval/evaluate_humaneval.py b/scripts/eval/evaluate_humaneval.py index 3791a47..5d31020 100644 --- a/scripts/eval/evaluate_humaneval.py +++ b/scripts/eval/evaluate_humaneval.py @@ -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) diff --git a/scripts/eval/evaluate_ifeval.py b/scripts/eval/evaluate_ifeval.py index 6a1cdd2..adbbc78 100644 --- a/scripts/eval/evaluate_ifeval.py +++ b/scripts/eval/evaluate_ifeval.py @@ -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, diff --git a/scripts/tools/generate.py b/scripts/tools/generate.py index cd425ca..4bdbe5e 100644 --- a/scripts/tools/generate.py +++ b/scripts/tools/generate.py @@ -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: diff --git a/tests/inference/test_engine.py b/tests/inference/test_engine.py index 6d7a68d..f0c0624 100644 --- a/tests/inference/test_engine.py +++ b/tests/inference/test_engine.py @@ -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) diff --git a/tests/inference/test_server.py b/tests/inference/test_server.py index 926e452..0003635 100644 --- a/tests/inference/test_server.py +++ b/tests/inference/test_server.py @@ -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,