refactor: assemble inference engines through a shared composition root

- 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

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",

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",

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,
+    )

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

scripts/demo/generate_ar.py

@@ -1,36 +1,68 @@
+from argparse import ArgumentParser
 from pathlib import Path
 
-import torch
+from astrai.inference import build_engine
-
-from astrai.inference import InferenceEngine
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
-def generate_text():
+def parse_args():
-    model = AutoModel.from_pretrained(PARAMETER_ROOT)
+    parser = ArgumentParser(
-    tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
+        description="Autoregressive continuation demo: continue a prompt, "
-    model.to(device="cuda", dtype=torch.bfloat16)
+        "or drop into interactive mode without one"
-
-    query = input(">> ")
-
-    engine = InferenceEngine(
-        model=model,
-        tokenizer=tokenizer,
     )
+    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,
+        max_tokens=args.max_tokens,
-        temperature=0.8,
+        temperature=args.temperature,
-        top_p=0.95,
+        top_p=args.top_p,
-        top_k=50,
+        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

@@ -1,9 +1,6 @@
 from pathlib import Path
 
-import torch
+from astrai.inference import build_engine
-
-from astrai.inference import InferenceEngine
-from astrai.model import AutoModel
 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(
+    engine = build_engine(PARAMETER_ROOT)
-        model=model,
-        tokenizer=tokenizer,
-    )
     responses = engine.generate(
         prompt=prompts,
         stream=False,

scripts/demo/stream_chat.py

@@ -1,11 +1,7 @@
 from argparse import ArgumentParser
 from pathlib import Path
 
-import torch
+from astrai import build_engine
-
-from astrai import InferenceEngine
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
 
 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)
+    engine = build_engine(model_path)
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    tokenizer = engine.tokenizer
-    model.to(device="cuda", dtype=torch.bfloat16)
-    engine = InferenceEngine(model=model, tokenizer=tokenizer)
 
     while True:
         query = input(">> ")

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.inference import build_engine
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
 
 # ---------------------------------------------------------------------------
 # 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

@@ -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.inference import InferenceEngine, build_engine
-from astrai.model import AutoModel
 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)
+    engine = build_engine(
-    tokenizer = AutoTokenizer.from_pretrained(args.param_path)
+        args.param_path,
-    model.to(device="cuda", dtype=torch.bfloat16)
-    model.eval()
-
-    engine = InferenceEngine(
-        model=model,
-        tokenizer=tokenizer,
         max_batch_size=args.batch_size,
         max_seq_len=args.max_seq_len,
     )
+    tokenizer = engine.tokenizer
 
     results = evaluate(
         engine=engine,

scripts/tools/generate.py

@@ -6,9 +6,7 @@ import click
 import torch
 from tqdm import tqdm
 
-from astrai.inference import InferenceEngine
+from astrai.inference import build_engine
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
 
 
 def processor(
@@ -28,17 +26,13 @@ def processor(
 ):
     print(f"Loading model from {param_path} ...")
     t0 = time.time()
-    model = AutoModel.from_pretrained(param_path)
+    engine = build_engine(
-    tokenizer = AutoTokenizer.from_pretrained(param_path)
+        param_path=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,
         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

@@ -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

@@ -5,8 +5,7 @@ from pathlib import Path
 import pytest
 import torch
 
-from astrai.inference import get_app
+from astrai.inference import build_engine, get_app
-from astrai.inference.network.app import _create_engine
 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,