0z5a
90de5bc1bd
- Return structured finish and error reasons for synchronous generation - Reject failed online rollout batches instead of training on empty responses - Verify allocation and extension failures release metrics and KV state
35 lines
1.0 KiB
Python
35 lines
1.0 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
|
|
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",
|
|
"InferenceScheduler",
|
|
"GenerationResult",
|
|
"Executor",
|
|
"STOP",
|
|
"Task",
|
|
"TaskManager",
|
|
"TaskStatus",
|
|
"sample",
|
|
"get_app",
|
|
"run_server",
|
|
]
|