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
+4 -9
View File
@@ -7,18 +7,13 @@ seq_lens with padding mask), and end-to-end scheduler.run_batch.
import torch
from astrai.extension import ATTN_BACKEND, attn_backend
from astrai.inference.core.cache import PagePool, TaskCacheManager
from astrai.inference.core.workspace import InferenceWorkspace
from astrai.inference.cache import PagePool, TaskCacheManager
from astrai.inference.workspace import InferenceWorkspace
from tests.extension.conftest import D, skip_no_kernel
def _mk_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)
def _ws(pool: PagePool) -> InferenceWorkspace:
@@ -184,7 +179,7 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
@skip_no_kernel
def test_run_batch_cuda_matches_torch_greedy(cuda_model):
"""Greedy decode (temperature=0) should produce identical tokens."""
from astrai.inference.core.scheduler import InferenceScheduler
from astrai.inference.scheduler import InferenceScheduler
from tests.helpers import FakeTokenizer
model, _ = cuda_model
+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,
+1 -1
View File
@@ -3,7 +3,7 @@
import pytest
import torch
from astrai.inference.core.scheduler import InferenceScheduler
from astrai.inference.scheduler import InferenceScheduler
from astrai.trainer.rollout import (
BaseRewardModel,
RawRollout,