- avoid constructing model_fn more than once when reading config - keep inference package exports focused on public entry points - rename extra strategy arguments to strategy_kwargs
34 lines
1013 B
Python
34 lines
1013 B
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
|
|
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, Task, TaskManager, TaskStatus
|
|
|
|
__all__ = [
|
|
"InferenceEngine",
|
|
"InferenceScheduler",
|
|
"Executor",
|
|
"STOP",
|
|
"Task",
|
|
"TaskManager",
|
|
"TaskStatus",
|
|
"sample",
|
|
"get_app",
|
|
"run_server",
|
|
]
|