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