refactor: split infer core into subpackages by concern

- Eliminate core/ directory into cache/, runtime/, network/ subpackages plus flat modules
- Split cache.py (647 lines) into cache/{buffer,strategy,pool}.py by layer
- Add explicit ContiguousStrategy, make AllocationStrategy a real ABC
- Move TaskCacheState to cache/strategy.py, drop string forward references
- Rename api/ to network/, server.py to app.py
- Move sample.py into runtime/ alongside executor and graph
- Simplify TaskCacheManager.__init__ to single pool param
- Expose pool.strategy and pool.req_pool as public properties
- Fix KVCache import in attention_backend.py (TYPE_CHECKING guard)
- Fix steady-state decode reading uninitialized position_ids on first step
This commit is contained in:
2026-08-08 23:43:05 +08:00
parent 3fa7e66676
commit 0c1b7664c1
37 changed files with 920 additions and 800 deletions
+5 -10
View File
@@ -11,7 +11,7 @@ from astrai.inference import (
TaskCacheManager,
page_hash,
)
from astrai.inference.core.workspace import InferenceWorkspace
from astrai.inference.workspace import InferenceWorkspace
def _ws(pool: PagePool) -> InferenceWorkspace:
@@ -27,12 +27,7 @@ def _ws(pool: PagePool) -> InferenceWorkspace:
def _make_task_cache(pool: PagePool) -> TaskCacheManager:
return TaskCacheManager(
strategy=pool._strategy,
req_pool=pool._req_pool,
max_seq_len=pool.max_seq_len,
pool=pool,
)
return TaskCacheManager(pool)
# ---- page_hash ----
@@ -345,7 +340,7 @@ def test_page_pool_paged_task_alloc():
assert task_cache.task_alloc("t1", list(range(10)))
state = task_cache._states["t1"]
assert len(state.pages) == 10
assert pool._req_pool.req_to_token[state.req_idx, 0].item() == state.pages[0]
assert pool.req_pool.req_to_token[state.req_idx, 0].item() == state.pages[0]
def test_page_pool_paged_task_extend():
@@ -354,7 +349,7 @@ def test_page_pool_paged_task_extend():
task_cache.task_alloc("t1", list(range(4)))
assert task_cache.task_extend("t1", 4)
req_idx = task_cache._states["t1"].req_idx
slot = pool._req_pool.req_to_token[req_idx, 4].item()
slot = pool.req_pool.req_to_token[req_idx, 4].item()
assert slot >= 0
@@ -364,7 +359,7 @@ def test_page_pool_paged_task_free_releases_slots():
task_cache.task_alloc("t1", list(range(8)))
task_cache.task_free("t1")
assert "t1" not in task_cache._states
assert len(pool._req_pool.free_slots) == 4
assert len(pool.req_pool.free_slots) == 4
def test_page_pool_paged_bind_roundtrip():
+3 -3
View File
@@ -5,9 +5,9 @@ from unittest.mock import MagicMock
import pytest
from astrai.inference.api.anthropic import AnthropicResponseBuilder
from astrai.inference.api.openai import OpenAIResponseBuilder
from astrai.inference.api.protocol import GenContext, StopChecker, StopInfo
from astrai.inference.network.anthropic import AnthropicResponseBuilder
from astrai.inference.network.openai import OpenAIResponseBuilder
from astrai.inference.network.protocol import GenContext, StopChecker, StopInfo
def _make_ctx(**kwargs):
+2 -2
View File
@@ -2,7 +2,7 @@
import torch
from astrai.inference.sample import (
from astrai.inference.runtime.sample import (
FrequencyPenaltyStrategy,
SamplingPipeline,
TemperatureStrategy,
@@ -268,7 +268,7 @@ def test_sample_return_logprobs_matches_manual_computation():
logits = torch.randn(2, 30)
tokens, logprobs = sample(logits, temperature=0.7, top_p=0.95, return_logprobs=True)
# Recompute with the same pipeline
from astrai.inference.sample import (
from astrai.inference.runtime.sample import (
SamplingPipeline,
TemperatureStrategy,
TopPStrategy,
+6 -6
View File
@@ -38,8 +38,8 @@ def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
"""Test concurrent add_task operations."""
mock_model, mock_tokenizer = mock_model_and_tokenizer
with patch("astrai.inference.core.scheduler.AutoModel"):
with patch("astrai.inference.core.scheduler.AutoTokenizer"):
with patch("astrai.inference.scheduler.AutoModel"):
with patch("astrai.inference.scheduler.AutoTokenizer"):
scheduler = InferenceScheduler(
model=mock_model,
tokenizer=mock_tokenizer,
@@ -77,8 +77,8 @@ def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
"""Test concurrent add and remove task operations."""
mock_model, mock_tokenizer = mock_model_and_tokenizer
with patch("astrai.inference.core.scheduler.AutoModel"):
with patch("astrai.inference.core.scheduler.AutoTokenizer"):
with patch("astrai.inference.scheduler.AutoModel"):
with patch("astrai.inference.scheduler.AutoTokenizer"):
scheduler = InferenceScheduler(
model=mock_model,
tokenizer=mock_tokenizer,
@@ -126,8 +126,8 @@ def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
"""Test concurrent get_stats operations."""
mock_model, mock_tokenizer = mock_model_and_tokenizer
with patch("astrai.inference.core.scheduler.AutoModel"):
with patch("astrai.inference.core.scheduler.AutoTokenizer"):
with patch("astrai.inference.scheduler.AutoModel"):
with patch("astrai.inference.scheduler.AutoTokenizer"):
scheduler = InferenceScheduler(
model=mock_model,
tokenizer=mock_tokenizer,
+1 -1
View File
@@ -2,7 +2,7 @@
import pytest
from astrai.inference.api.tool_parser import (
from astrai.inference.network.tool_parser import (
_TOOL_CALL_HEAD_RE,
BaseToolParser,
SimpleJsonToolParser,