- 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
36 lines
1.1 KiB
Python
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",
|
|
]
|