Files
AstrAI/astrai/inference/__init__.py
T
ViperEkura e13fe53475 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
2026-09-03 22:16:56 +08:00

36 lines
1.1 KiB
Python

"""Inference module for continuous batching.
Subpackages:
- cache/: KV cache (buffers, strategies, pool)
- runtime/: Execution + sampling (executor, CUDA graph, sampling strategies)
- task/: Request lifecycle + performance metrics
- network/: HTTP protocol handling (server, protocol, OpenAI/Anthropic builders)
Modules:
- scheduler.py: Continuous batching loop
- workspace.py: Pre-allocated GPU buffers
- engine.py: Facade (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
from astrai.inference.scheduler import InferenceScheduler
from astrai.inference.task import STOP, GenerationResult, Task, TaskManager, TaskStatus
__all__ = [
"InferenceEngine",
"build_engine",
"InferenceScheduler",
"GenerationResult",
"Executor",
"STOP",
"Task",
"TaskManager",
"TaskStatus",
"sample",
"get_app",
"run_server",
]